From cfcab7ccb85c948d5f64dff86d73c755f3ad7775 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 27 Apr 2026 15:40:43 -0700 Subject: [PATCH 01/89] Add --parallel-warmup phase to replay benchmark MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ThreadPoolExecutor fans out (batch, mtp_len, prev_k, dtype, sweep) configs across N threads before the sequential measurement. Triton compile releases the GIL, so threads compile in parallel and populate the on-disk cache. Subsequent timed iterations hit the cache for free. _bench_config gains a warmup_only=True path that calls each kernel once and synchronizes — no timing, no print output. CLI: --parallel-warmup N (default 0 = phase disabled). Wins about 20% wall time on a 7-config sweep with cold Triton cache. Larger sweeps (more batches/mtp/dtype combos) will benefit more since compile time scales with #shapes while wall time is dominated by the slowest single shape compile under N-way parallelism. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 144 ++++++++++++++---- 1 file changed, 115 insertions(+), 29 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index a755579a8600..dd4fc904d570 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -52,6 +52,7 @@ import os import statistics import sys +import time from datetime import datetime from pathlib import Path @@ -415,6 +416,59 @@ def _time_kernel(args, run_fn, reset_fn, tag: str) -> tuple[float, float, float] # Per-config benchmark (consolidated baseline + replay) +def _parallel_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, + baseline_fn, max_workers: int) -> None: + """Run each config once in parallel to compile + cache Triton kernels. + + Triton's ``compile()`` releases the GIL, so a ThreadPoolExecutor + fans out shape compilations across CPU cores in one shared CUDA + context. Compiled binaries land in Triton's on-disk cache (default + ``~/.triton/cache``) and the subsequent sequential measurement + phase loads them with no compile cost. + + Parallel measurement would race for GPU time and skew numbers, so + only the warmup is parallelized; timing stays serial. + """ + from concurrent.futures import ThreadPoolExecutor + + configs = [] + for batch in batch_sizes: + for mtp_len in mtp_lengths: + prev_ks = sorted( + set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs) + ) + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + configs.append((batch, mtp_len, prev_ks, state_dtype, act_dtype)) + + print(f"[parallel-warmup] {len(configs)} configs across {max_workers} threads") + t0 = time.perf_counter() + + def _warm(cfg): + batch, mtp_len, prev_ks, state_dtype, act_dtype = cfg + _bench_config( + args, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, + warmup_only=True, + ) + + errors = [] + with ThreadPoolExecutor(max_workers=max_workers) as ex: + futures = [ex.submit(_warm, cfg) for cfg in configs] + for cfg, fut in zip(configs, futures): + try: + fut.result() + except Exception as e: + errors.append((cfg, e)) + + if errors: + for cfg, e in errors: + print(f"[parallel-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", + file=sys.stderr) + raise errors[0][1] + + print(f"[parallel-warmup] done in {time.perf_counter() - t0:.1f}s") + + def _bench_config( args, batch: int, @@ -423,6 +477,7 @@ def _bench_config( state_dtype: torch.dtype, act_dtype: torch.dtype, baseline_fn, + warmup_only: bool = False, ) -> None: """ Benchmark one (batch, mtp_len, dtype) configuration. @@ -430,6 +485,11 @@ def _bench_config( Runs the baseline kernel (if baseline_fn is not None) followed by the replay kernel for each prev_k value. Tensors are built once and shared across all runs in this config. + + When ``warmup_only`` is True, calls each kernel exactly once instead of + timing it. Used by the parallel-warmup phase to populate Triton's + persistent compile cache across all configs concurrently. No timing + output is produced. """ state_dtype_name = str(state_dtype).split(".")[-1] act_dtype_name = str(act_dtype).split(".")[-1] @@ -600,20 +660,25 @@ def _run_baseline(): ) reset_fn = _reset_conv1d_realistic if with_conv1d else _reset - median_us, p95_us, p99_us = _time_kernel(args, _run_baseline, reset_fn, tag) - - _print_row( - show_kernel_col, - args.baseline, - batch, - mtp_len, - "N/A", - state_dtype_name, - act_dtype_name, - median_us, - p95_us, - p99_us, - ) + if warmup_only: + reset_fn() + _run_baseline() + torch.cuda.synchronize() + else: + median_us, p95_us, p99_us = _time_kernel(args, _run_baseline, reset_fn, tag) + + _print_row( + show_kernel_col, + args.baseline, + batch, + mtp_len, + "N/A", + state_dtype_name, + act_dtype_name, + median_us, + p95_us, + p99_us, + ) # --- Sweep parameter parsing (invariant across prev_k) --- def _parse_sweep(val): @@ -712,21 +777,26 @@ def _run_incr( sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") reset_fn = _reset_conv1d_realistic if with_conv1d else _reset - median_us, p95_us, p99_us = _time_kernel(args, _run_incr, reset_fn, sweep_tag) - - _print_row( - show_kernel_col, - "replay", - batch, - mtp_len, - prev_k, - state_dtype_name, - act_dtype_name, - median_us, - p95_us, - p99_us, - sweep_suffix, - ) + if warmup_only: + reset_fn() + _run_incr() + torch.cuda.synchronize() + else: + median_us, p95_us, p99_us = _time_kernel(args, _run_incr, reset_fn, sweep_tag) + + _print_row( + show_kernel_col, + "replay", + batch, + mtp_len, + prev_k, + state_dtype_name, + act_dtype_name, + median_us, + p95_us, + p99_us, + sweep_suffix, + ) def _print_row( @@ -787,6 +857,12 @@ def _run_benchmark(args) -> None: elif args.l2_flush: _init_l2_flush() + if args.parallel_warmup > 0: + _parallel_warmup_phase( + args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, + baseline_fn, max_workers=args.parallel_warmup, + ) + if args.profile: torch.cuda.cudart().cudaProfilerStart() @@ -877,6 +953,16 @@ def _parse_args() -> argparse.Namespace: ) parser.add_argument("--warmup", type=int, default=20, help="Number of warmup iterations") parser.add_argument("--iters", type=int, default=100, help="Number of timed iterations") + parser.add_argument( + "--parallel-warmup", + type=int, + default=0, + help="Run a parallel-warmup phase that calls each (batch, mtp_len, " + "prev_k, dtype, sweep) config once across N threads before the " + "sequential timed phase. Triton compile releases the GIL so threads " + "compile in parallel, populating the persistent cache for free hits " + "during measurement. 0 disables the phase. Try 8 on a multi-core box.", + ) parser.add_argument( "--profile", action="store_true", From e6694fa1e1da4aca0f99f8d30ab1dc1eb9a04414 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 28 Apr 2026 16:48:14 -0700 Subject: [PATCH 02/89] Fork checkpointing_state_update kernel and test (placeholder = today) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py and the matching test file as forks of replay_selective_state_update. Same two-kernel architecture, same correctness contract. Renames: - replay_selective_state_update -> checkpointing_state_update - _replay_precompute_kernel -> _checkpointing_precompute_kernel - _replay_state_update_kernel -> _checkpointing_main_kernel API change: a new MAX_REPLAY_BUFFER_LENGTH constexpr is plumbed through both kernels and the wrapper. It's the cache buffer T-axis capacity (distinct from T_new = the per-step input token count). In the placeholder build, MAX_REPLAY_BUFFER_LENGTH == T_new (cache buffers sized exactly to T_new), so behavior is byte-identical to the existing replay kernel — verified by 501 unit tests passing on the new file. The BLOCK_SIZE_T heuristic is now keyed on MAX_REPLAY_BUFFER_LENGTH (cache capacity) rather than T (input count); for the placeholder the two are equal so this is a no-op rename. Mixer and cache manager untouched; this is wrapper-level only. Subsequent PRs will: - rectangle-CB Option-2 output factoring (still no checkpointing) - actual checkpoint trigger + double-bank flip on non-x cache - heuristic resweep for the new kernel - lower-precision state with Philox stochastic rounding Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 1060 +++++++++++++++++ .../mamba/test_checkpointing_state_update.py | 764 ++++++++++++ 2 files changed, 1824 insertions(+) create mode 100644 tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py create mode 100644 tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py new file mode 100644 index 000000000000..564c0b203d3a --- /dev/null +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -0,0 +1,1060 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +# +# Adapted from: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +# SPDX-FileCopyrightText: Copyright contributors to the sglang project +# +# Copyright (c) 2024, Tri Dao, Albert Gu. +# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID +from tensorrt_llm._utils import get_sm_version + +from .softplus import softplus + + +@triton.jit +def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 pair → fp16x2 using Philox random bits. + + Uses PTX cvt.rs.f16x2.f32 which rounds each fp32 value to fp16 using + the random bits to break ties, avoiding systematic rounding bias that + accumulates over many decode steps with fp16 state. + + Adapted from flashinfer (Apache-2.0, vLLM/mamba lineage). + """ + return tl.inline_asm_elementwise( + asm="""{ + cvt.rs.f16x2.f32 $0, $2, $1, $3; + }""", + constraints=("=r,r,r,r,r"), + args=(x, rand), + dtype=tl.float16, + is_pure=True, + pack=2, + ) + + +# Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, +# old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. +# Grid: (batch, nheads // HEADS_PER_BLOCK). + + +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)}) +@triton.jit() +def _checkpointing_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache WRITE pointers (write-buffer for next step) + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + # Double-buffer index (per cache slot) + cache_buf_idx_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, # new tokens this step (T_new) + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # cache buffer T-axis capacity + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides: (cache, 2, T, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, +): + pid_b = tl.program_id(axis=0) + pid_hg = tl.program_id(axis=1) # head-group index + first_head = pid_hg * HEADS_PER_BLOCK + + # Resolve cache index for writes + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Signal main kernel to start (internal PDL). Main's replay phase + # reads only from the READ buffer (written by the PREVIOUS step) — + # safe even if conv1d and this kernel are still running. Main's + # gdc_wait() gates the output phase, which reads conv1d outputs + # (x, C) and this kernel's outputs (cb_scaled, decay_vec). + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + # Read buffer index: replay reads from buf_read. We WRITE to 1 - buf_read + # for next step's replay. + buf_read = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + buf_write = 1 - buf_read + + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + t_mask = offs_t < T + n_mask = offs_n < dstate + + # Causal mask is shared across all heads (depends only on offs_t) + causal_mask = offs_t[:, None] >= offs_t[None, :] + valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] + + # --- Loop 1: compute per-head dt/dA_cumsum/decay BEFORE gdc_wait --- + # These only depend on dt (from in_proj, not conv1d) and parameters (A, dt_bias). + # Store to cache; will reload after the wait for CB scaling. + for h_local in range(HEADS_PER_BLOCK): + head_idx = first_head + h_local + + dt_base = dt_ptr + pid_b * stride_dt_batch + head_idx * stride_dt_head + dt = tl.load(dt_base + offs_t * stride_dt_T, mask=t_mask, other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + head_idx * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias + if DT_SOFTPLUS: + dt = softplus(dt) + + A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) + dA_cumsum = tl.cumsum(A * dt, axis=0) + decay_vec = tl.exp(dA_cumsum) + + # Store dt, dA_cumsum, decay_vec to cache + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_write * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + tl.store(old_dt_base + offs_t * stride_old_dt_T, dt, mask=t_mask) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_write * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + tl.store(old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, dA_cumsum, mask=t_mask) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + head_idx * stride_dv_head + tl.store(decay_vec_base + offs_t * stride_dv_t, decay_vec, mask=t_mask) + + # --- Wait for upstream kernel (external PDL) before loading B and C --- + # All dt processing above is independent of conv1d outputs. + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- + group_idx = first_head // nheads_ngroups_ratio + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_all = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_all = tl.load( + B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + # Compute raw CB once — shared across all heads in this block + raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) + + # Store B to cache (once per group, only if this block covers the first heads) + if first_head % nheads_ngroups_ratio == 0: + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + buf_write * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_base + offs_t[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + B_all, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # --- Loop 2: reload per-head dA_cumsum/dt from cache, scale CB --- + # The cache was just written above, so these loads should hit L2. + for h_local in range(HEADS_PER_BLOCK): + head_idx = first_head + h_local + + # Reload dt and dA_cumsum from cache (just written in loop 1) + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_write * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + dt = tl.load(old_dt_base + offs_t * stride_old_dt_T, mask=t_mask, other=0.0).to(tl.float32) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_write * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + dA_cumsum = tl.load( + old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0 + ).to(tl.float32) + + # Scale raw_CB with per-head decay and dt + decay_matrix = tl.exp(dA_cumsum[:, None] - dA_cumsum[None, :]) + CB_scaled = tl.where(valid_mask, raw_CB * decay_matrix * dt[None, :], 0.0) + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head + tl.store( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + CB_scaled, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + ) + + +# Main kernel: tl.dot replay + precomputed CB output. +# Grid: (cdiv(dim, M), batch, nheads). + + +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)}) +@triton.jit() +def _checkpointing_main_kernel( + # Pointers + state_ptr, + # Cache READ pointers (read-buffer from previous step) + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + # Cache WRITE pointer (write-buffer for old_x only; B/dt/dA_cumsum written by precompute) + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + # New input pointers + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + # Precomputed pointers + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + # Stochastic rounding + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, # new tokens this step (T_new) + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # cache buffer T-axis capacity + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # old_x strides: (cache, T, nheads, dim) — single-buffered + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides: (cache, 2, T, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Double-buffer index: buf_read points to the buffer written by LAST step's + # precompute. THIS step's precompute writes to 1-buf_read, which will be + # read by NEXT step's main kernel. Anything not carried between steps is + # single-buffered. + buf_read = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # Load state + state_ptr += cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + + # Phase 1: Replay via tl.dot fast-forward (reads from READ buffer) + group_idx = pid_h // nheads_ngroups_ratio + + # Load precomputed dt and dA_cumsum from READ buffer + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_read * stride_old_dt_dbuf + + pid_h * stride_old_dt_head + ) + old_dt_all = tl.load(old_dt_base + offs_t * stride_old_dt_T, mask=t_mask, other=0.0).to( + tl.float32 + ) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_read * stride_old_dA_cumsum_dbuf + + pid_h * stride_old_dA_cumsum_head + ) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0 + ).to(tl.float32) + + # Load dA_cumsum at prev_k-1 directly via pointer math (avoids masked reduction). + # Clamp to [0, T-1] defensively — out-of-contract PNAT > T would read OOB. + prev_k_idx = tl.minimum(tl.maximum(prev_num_accepted_tokens - 1, 0), T - 1) + total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( + tl.float32 + ) + + # Step 0 invariant: PNAT=0 means `state` is already last step's state (not + # two back). coeff is all-zero (offs_t < 0), total_decay is 1.0, so the + # replay leaves `state` unchanged — cache contents don't matter on step 0. + coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all + coeff = tl.where(offs_t < prev_num_accepted_tokens, coeff, 0.0) + + # Load old_x: (BLOCK_SIZE_T, BLOCK_SIZE_M) — single-buffered + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + old_x_all = tl.load( + old_x_base + offs_t[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + + # Load old_B from READ buffer: (BLOCK_SIZE_T, BLOCK_SIZE_DSTATE) + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + buf_read * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_all = tl.load( + old_B_base + offs_t[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ).to(tl.float32) + + # Scale B by coefficients + dB_scaled = coeff[:, None] * old_B_all + + # Apply total decay to initial state FIRST, then add contributions + total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) + state *= total_decay + + # tl.dot fast-forward: old_x^T @ dB_scaled → (M, dstate) + state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) + + # Write post-replay state + if USE_RS_ROUNDING: + # Stochastic rounding for fp16 state using Philox-4x32 PRNG. + # Each Philox call produces 4 random ints. We call randint4x on + # quarter-sized dstate offsets and join+reshape to get the full + # (M, dstate) random tensor — 4x fewer PRNG rounds. + rand_seed = tl.load(rand_seed_ptr) + base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head + offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // 4) + rand_offsets_q = ( + base_rand + + offs_m[:, None] * stride_state_dim + + offs_n_q[None, :] * (stride_state_dstate * 4) + ) # (M, dstate//4) + if PHILOX_ROUNDS > 0: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + else: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) + # Interleave 4 quarter-sized tensors → full (M, dstate) random tensor + r01 = tl.join(r0, r1) # (M, dstate//4, 2) + r23 = tl.join(r2, r3) # (M, dstate//4, 2) + r0123 = tl.join(r01, r23) # (M, dstate//4, 2, 2) + rand = tl.reshape(r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + tl.store(state_ptrs, _stochastic_round_fp16x2(state, rand), mask=state_mask) + else: + tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=state_mask) + + # Phase 2: Output using precomputed CB_scaled and decay_vec + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + # Wait for precompute kernel (PDL) before reading its outputs. + # With chained PDL (conv1d → precompute → main), gdc_wait() ensures + # precompute has completed — which transitively ensures conv1d has + # completed (precompute waited on conv1d via its own gdc_wait). + # All loads below (x, C from conv1d; CB_scaled, decay_vec from precompute) + # are safe after this point. + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # Load conv1d outputs: C_all and x_all + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + x_all = tl.load( + x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + # Store new x to cache (single-buffered; replay already read the old data) + tl.store( + old_x_base + offs_t[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + x_all, + mask=t_mask[:, None] & m_mask[None, :], + ) + x_all = x_all.to(tl.float32) + + # Load precomputed CB_scaled and decay_vec + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + tl.float32 + ) + + # init_out = C_all @ state^T * decay_vec + init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] + + # cb_out = CB_scaled @ x_all + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) + + out_all = init_out + cb_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + for t in range(T): + z_t = tl.load( + z_ptr + t * stride_z_T + offs_m * stride_z_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + out_t = tl.sum(tl.where((offs_t == t)[:, None], out_all, 0.0), axis=0) + out_t = out_t * z_t * tl.sigmoid(z_t) + tl.store(out_ptr + t * stride_out_T + offs_m * stride_out_dim, out_t, mask=m_mask) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# Python wrapper + + +def checkpointing_state_update( + state: torch.Tensor, + old_x: torch.Tensor, + old_B: torch.Tensor, + old_dt: torch.Tensor, + old_dA_cumsum: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + out: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + pad_slot_id: int = PAD_SLOT_ID, + rand_seed: torch.Tensor | None = None, + philox_rounds: int = 10, + launch_with_pdl=False, + use_internal_pdl=True, + _block_size_m: int | None = None, + _num_warps: int | None = None, + _num_stages: int | None = None, + _precompute_num_warps: int | None = None, + _precompute_num_stages: int | None = None, + _heads_per_block: int | None = None, +): + """ + Replay SSM state update with precomputed CB and tl.dot fast-forward. + + Two-kernel architecture: + 1. Precompute kernel: computes CB_scaled and decay_vec from B, C, dt, A. + Writes processed dt/dA_cumsum/B to double-buffered cache for next step. + 2. Main kernel: replays old tokens via tl.dot fast-forward on cached data, + then computes output using precomputed CB_scaled and new x/C inputs. + + PDL (Programmatic Dependent Launch) chain: + conv1d → (external PDL) → precompute → (internal PDL) → main + External PDL: precompute starts while conv1d is running; gdc_wait() + in precompute blocks until conv1d completes before loading B/C. + Internal PDL: main starts while precompute is running; main's replay + phase uses only cached data from the previous step. gdc_wait() in + main blocks until precompute completes before loading conv1d outputs + (x, C) and precompute outputs (CB_scaled, decay_vec). + + Uses double-buffered cache tensors. cache_buf_idx[slot] indicates which + buffer (0 or 1) to READ from for replay. The WRITE buffer is 1 - read. + Caller must flip cache_buf_idx[slot] after each call. + + Arguments: + state: (cache, nheads, dim, dstate) in-place. After the call, contains + the state after replaying prev_num_accepted_tokens old tokens. + old_x: (cache, T, nheads, dim) bf16 — old x cache (single-buffered). + old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. + old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. + old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. + cache_buf_idx: (cache,) int32 — which buffer to read (0 or 1). + prev_num_accepted_tokens: (cache,) int32. + x: (batch, T, nheads, dim) new token inputs. + dt: (batch, T, nheads, dim) with stride(-1)==0 (tie_hdim). + A: (nheads, dim, dstate) with stride(-1)==0, stride(-2)==0 (tie_hdim). + B: (batch, T, ngroups, dstate). + C: (batch, T, ngroups, dstate). + out: (batch, T, nheads, dim) preallocated output. + D: (nheads, dim) optional feed-through parameter. + z: (batch, T, nheads, dim) optional silu gate. + dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). + state_batch_indices: (batch,) optional cache slot mapping. + rand_seed: optional single-element int64 CUDA tensor for Philox PRNG seed. + When provided, state is stochastically rounded to fp16 on store. + When None, standard deterministic rounding is used. + philox_rounds: number of Philox PRNG rounds (default 10). + launch_with_pdl: enable external PDL (conv1d → precompute chain). + Defaults False; caller opts in when the upstream chain is PDL-safe. + Ignored on hardware that doesn't support PDL (sm < 90). + use_internal_pdl: enable internal PDL (precompute → main overlap). + Defaults True; override for testing only. + Ignored on hardware that doesn't support PDL (sm < 90). + + _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, + _precompute_num_warps, _precompute_num_stages, _heads_per_block) are + benchmark-only overrides; production callers should leave them None + to use the heuristic-tuned defaults. + """ + # PDL needs sm >= 90. + if get_sm_version() < 90: + launch_with_pdl = False + use_internal_pdl = False + + # --- Unsqueeze inputs to canonical shapes --- + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if x.dim() == 3: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if dt.dim() == 3: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if B.dim() == 3: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if C.dim() == 3: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None: + if z.dim() == 2: + z = z.unsqueeze(1) + if z.dim() == 3: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + if out.dim() == 2: + out = out.unsqueeze(1) + if out.dim() == 3: + out = out.unsqueeze(1) + + cache_size, nheads, dim, dstate = state.shape + batch, T, _, _ = x.shape + ngroups = B.shape[2] + assert nheads % ngroups == 0 + + # Cache T-axis is the replay-buffer capacity (committed-history slots + + # provisional-draft slots). In the placeholder build, this matches T_new. + max_replay_buffer_length = old_x.shape[1] + assert T <= max_replay_buffer_length, ( + f"x has T_new={T} > max_replay_buffer_length={max_replay_buffer_length}" + ) + + assert x.shape == (batch, T, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + assert B.shape == (batch, T, ngroups, dstate) + assert C.shape == B.shape + assert old_x.shape == (cache_size, max_replay_buffer_length, nheads, dim) + assert old_B.shape == (cache_size, 2, max_replay_buffer_length, ngroups, dstate) + assert old_dt.shape == (cache_size, 2, nheads, max_replay_buffer_length) + assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_replay_buffer_length) + assert cache_buf_idx.shape == (cache_size,) + assert prev_num_accepted_tokens.shape == (cache_size,) + + tie_hdim = ( + A.stride(-1) == 0 + and A.stride(-2) == 0 + and dt.stride(-1) == 0 + and (dt_bias is None or dt_bias.stride(-1) == 0) + ) + assert tie_hdim + + device = x.device + BLOCK_SIZE_T = max(triton.next_power_of_2(max_replay_buffer_length), 16) + + # Allocate precomputed intermediates (per-call, not cached) + cb_scaled = torch.empty( + batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_T, device=device, dtype=torch.float32 + ) + decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) + + z_strides = ( + (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) if z is not None else (0, 0, 0, 0) + ) + + # Kernel tuning: BLOCK_SIZE_M, num_warps, HEADS_PER_BLOCK, precompute_num_warps. + # Dtype-aware heuristic from B200 sweeps (batch 1-512, T=6/32, TP=8, conv1d + + # chained PDL). Keyed on total_heads, BLOCK_SIZE_T, and state dtype; 16-bit + # states prefer different tiles from fp32 due to lower bandwidth. Philox + # gets its own branch — stochastic rounding shifts compute toward CUDA cores, + # so small-batch configs want more warps to hide the extra work. + total_heads = batch * nheads + heads_per_group = nheads // ngroups + state_is_16bit = state.dtype in (torch.float16, torch.bfloat16) + use_philox = rand_seed is not None + if BLOCK_SIZE_T <= 16: + if use_philox and state_is_16bit: + # Philox: more warps at small batch to hide CUDA core work. + # At large batch, converges to non-Philox fp16 config. + if total_heads <= 16: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 4, 4, 4, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + elif state_is_16bit: + if total_heads <= 16: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 4, 1 + elif total_heads <= 64: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 2, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 1, + min(2, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + else: # fp32 state (no Philox — fp32 doesn't need stochastic rounding) + if total_heads <= 32: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 4, 1 + elif total_heads <= 64: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + elif total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 2, 2, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 1, 2, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 2, + min(2, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + else: # T > 16 + if state_is_16bit: + if total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 16, + 1, + 4, + min(2, heads_per_group), + ) + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 1, + min(4, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 4, + min(2, heads_per_group), + ) + else: # fp32 state + if total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 2, + 4, + min(2, heads_per_group), + ) + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 2, + min(4, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 4, + min(2, heads_per_group), + ) + if _block_size_m is not None: + BLOCK_SIZE_M = _block_size_m + if _num_warps is not None: + num_warps = _num_warps + if _heads_per_block is not None: + heads_per_block = _heads_per_block + if _precompute_num_warps is not None: + precompute_num_warps = _precompute_num_warps + + HAS_CACHE_BATCH_INDICES = state_batch_indices is not None + + with torch.cuda.device(device.index): + # --- Precompute kernel --- + assert nheads % heads_per_block == 0, ( + f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" + ) + assert heads_per_block <= heads_per_group, ( + f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" + ) + _checkpointing_precompute_kernel[(batch, nheads // heads_per_block)]( + dt, + dt_bias, + A, + B, + C, + cb_scaled, + decay_vec, + old_B, + old_dt, + old_dA_cumsum, + cache_buf_idx, + state_batch_indices, + pad_slot_id, + T, + max_replay_buffer_length, + dstate, + nheads // ngroups, + # dt strides + dt.stride(0), + dt.stride(1), + dt.stride(2), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + # B strides + B.stride(0), + B.stride(1), + B.stride(2), + B.stride(3), + # C strides + C.stride(0), + C.stride(1), + C.stride(2), + C.stride(3), + # cb_scaled strides + cb_scaled.stride(0), + cb_scaled.stride(1), + cb_scaled.stride(2), + cb_scaled.stride(3), + # decay_vec strides + decay_vec.stride(0), + decay_vec.stride(1), + decay_vec.stride(2), + # old_B strides + old_B.stride(0), + old_B.stride(1), + old_B.stride(2), + old_B.stride(3), + old_B.stride(4), + # old_dt strides + old_dt.stride(0), + old_dt.stride(1), + old_dt.stride(2), + old_dt.stride(3), + # old_dA_cumsum strides + old_dA_cumsum.stride(0), + old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), + old_dA_cumsum.stride(3), + dt_softplus, + HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, + LAUNCH_WITH_PDL=launch_with_pdl, + LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, + HEADS_PER_BLOCK=heads_per_block, + num_warps=precompute_num_warps, + **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), + launch_pdl=launch_with_pdl, + ) + + # --- Main kernel --- + def grid(META): + return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + + _checkpointing_main_kernel[grid]( + state, + old_x, + old_B, + old_dt, + old_dA_cumsum, + prev_num_accepted_tokens, + cache_buf_idx, + x, + C, + D, + z, + out, + cb_scaled, + decay_vec, + state_batch_indices, + rand_seed, + pad_slot_id, + T, + max_replay_buffer_length, + dim, + dstate, + nheads // ngroups, + # state strides + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + # old_x strides (single-buffered: cache, T, nheads, dim) + old_x.stride(0), + old_x.stride(1), + old_x.stride(2), + old_x.stride(3), + # old_B strides + old_B.stride(0), + old_B.stride(1), + old_B.stride(2), + old_B.stride(3), + old_B.stride(4), + # old_dt strides + old_dt.stride(0), + old_dt.stride(1), + old_dt.stride(2), + old_dt.stride(3), + # old_dA_cumsum strides + old_dA_cumsum.stride(0), + old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), + old_dA_cumsum.stride(3), + # x strides + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + # C strides + C.stride(0), + C.stride(1), + C.stride(2), + C.stride(3), + # D strides + *(D.stride(0), D.stride(1)) if D is not None else (0, 0), + # z strides + z_strides[0], + z_strides[1], + z_strides[2], + z_strides[3], + # out strides + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + # cb_scaled strides + cb_scaled.stride(0), + cb_scaled.stride(1), + cb_scaled.stride(2), + cb_scaled.stride(3), + # decay_vec strides + decay_vec.stride(0), + decay_vec.stride(1), + decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + launch_pdl=use_internal_pdl, + ) diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py new file mode 100644 index 000000000000..c6327252a6aa --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -0,0 +1,764 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import pytest +import torch +import torch.nn.functional as F +from einops import repeat + +from tensorrt_llm._torch.modules.mamba.checkpointing_state_update import ( + checkpointing_state_update, +) +from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update +from tensorrt_llm._utils import get_sm_version + +# Philox stochastic rounding uses PTX cvt.rs.f16x2.f32 which requires sm >= 100. +_skip_pre_sm100 = pytest.mark.skipif( + get_sm_version() < 100, reason="Philox stochastic rounding needs sm >= 100" +) + +# Configs derived from NVIDIA-Nemotron-3-Super-120B-A12B Mamba2 parameters +# (nheads=128, headdim=64, d_state=128, ngroups=8) with TP split applied: +# TP=8: nheads=16, ngroups=1 — primary production config +# TP=4: nheads=32, ngroups=2 — exercises ngroups>1 (grouped B/C path) +_CONFIGS = [ + # (nheads, head_dim, d_state, ngroups) + (16, 64, 128, 1), # TP=8 production config + (32, 64, 128, 2), # TP=4, ngroups>1 (more heads than B/C groups) +] + + +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize("state_dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) +@pytest.mark.parametrize( + "T", [6, 10, 16, 27, 32, 55], ids=["T6", "T10", "T16", "T27", "T32", "T55"] +) +def test_checkpointing_state_update( + nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T +): + """ + Verify that: + checkpointing_state_update(state0, old_caches, k, new_x, ...) + produces the same output as: + selective_state_update(state_after_k_old_tokens, new_x, ...) + and writes state_after_k_old_tokens back to the state tensor. + """ + batch = 2 + device = "cuda" + dtype = torch.bfloat16 # input activations are bf16 + assert nheads % ngroups == 0 + + if paged_cache: + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + else: + cache_size = batch + state_batch_indices = None + + torch.manual_seed(42) + + # A: (nheads, head_dim, d_state) with stride(-2)=0, stride(-1)=0 [tie_hdim] + A_base = -torch.rand(nheads, device=device) - 0.5 # float32, negative + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + + # dt_bias: (nheads, head_dim) with stride(-1)=0 [tie_hdim] + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + + # D: (nheads, head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + # Initial SSM state (cache_size slots) + state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + + # Old inputs: T tokens per batch request + x1 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) # stride(-1)=0 + B1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + # Capture intermediate SSM states using selective_state_update. + states_buffer_f32 = torch.zeros( + cache_size, T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = ( + state_batch_indices + if paged_cache + else torch.arange(batch, device=device, dtype=torch.int32) + ) + out1 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + state0.clone(), + x1, + dt1, + A, + B1, + C1, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=T, + out=out1, + disable_state_update=True, + ) + + # Build cache tensors for the replay kernel. + # old_x: (cache, T, nheads, dim) bf16 — single-buffered + # old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered + # old_dt: (cache, 2, nheads, T) fp32 — double-buffered, T contiguous + # old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered, T contiguous + # cache_buf_idx: random 0s and 1s to verify indexing correctness + old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) + + # Fill each slot's READ buffer (indexed by cache_buf_idx) with step 1's data. + # The OTHER buffer has random garbage to catch indexing bugs. + slots = state_batch_indices if paged_cache else slice(None) + old_x[slots] = x1 + + # Compute processed dt and dA_cumsum for step 1 + dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) + + # Write to each slot's read buffer based on its cache_buf_idx + slot_indices = state_batch_indices.tolist() if paged_cache else list(range(cache_size)) + for i, slot in enumerate(slot_indices): + buf = cache_buf_idx[slot].item() + batch_idx = i # maps slot back to the batch index + old_B[slot, buf] = B1[batch_idx] + old_dt[slot, buf] = dt1[batch_idx].T # (T, nheads) → (nheads, T) + old_dA_cumsum[slot, buf] = dA_cumsum1[batch_idx].T # (T, nheads) → (nheads, T) + + # Main loop: test each k (number of old tokens replayed) + for k in range(T + 1): + torch.manual_seed(k + 100) + + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + # Reference + ref_state_f32 = state0.float().clone() + if k > 0: + ref_state_f32[slots] = states_buffer_f32[slots, k - 1] + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=(state_batch_indices if paged_cache else None), + out=ref_out, + ) + + # Replay kernel + test_state = state0.clone() + prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + # cache_buf_idx stays at its random values — each slot reads from its own buffer + + checkpointing_state_update( + test_state, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + x=x2, + dt=dt2, + A=A, + B=B2, + C=C2, + out=test_out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + ) + + # Tolerance rationale: the replay kernel uses bf16 tl.dot for four + # matmuls (dB_scaled @ old_x, C @ state, CB_scaled @ x, and C @ B in + # precompute). The reference (selective_state_update) and flashinfer + # baseline use fp32 element-wise MACs. The bf16 input casts lose the + # dt_bias/A-derived bits that the baselines keep — per-element rounding, + # not accumulating. Prefill (ssd_chunk_scan) does identical bf16 tl.dot + # casts, so we match prefill precision exactly. Empirical: max ~1.0 at + # T<=16, ~2.0 at T=32-55; mean ~0.014; <0.02% of elements exceed 0.5. + # State dtype (fp16/bf16/fp32) doesn't shift the error — bf16 dot + # inputs dominate, not state storage. + torch.testing.assert_close( + test_out, ref_out, rtol=2e-2, atol=1.0, msg=f"Output mismatch at k={k}" + ) + + expected_state = ( + state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) + ) + torch.testing.assert_close( + test_state[slots], expected_state, rtol=2e-2, atol=1.0, msg=f"State mismatch at k={k}" + ) + + +@_skip_pre_sm100 +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) +@pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) +def test_checkpointing_state_update_philox(nheads, head_dim, d_state, ngroups, paged_cache, T): + """ + Verify that Philox stochastic rounding produces correct results. + + Runs our kernel twice with identical inputs: once without rounding + (fp16 state, deterministic), once with rounding (fp16 state, Philox). + The outputs should be nearly identical — stochastic rounding only + perturbs the state by ±1 fp16 ULP, which barely affects output. + Also verifies the state dtype remains fp16. + """ + batch = 2 + device = "cuda" + dtype = torch.bfloat16 + state_dtype = torch.float16 + assert nheads % ngroups == 0 + + if paged_cache: + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + else: + cache_size = batch + state_batch_indices = None + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + + # Cache tensors + old_x = torch.randn(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) + + # New token inputs + x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) + B = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + prev_tokens = torch.full((cache_size,), T // 2, device=device, dtype=torch.int32) + + common_kwargs = dict( + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + ) + + # --- Run without rounding (deterministic fp16 state store) --- + state_no_round = state0.clone() + out_no_round = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_no_round, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + out=out_no_round, + **common_kwargs, + ) + + # --- Run with Philox rounding --- + rand_seed = torch.tensor([12345], device=device, dtype=torch.int64) + state_rounded = state0.clone() + out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_rounded, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + out=out_rounded, + rand_seed=rand_seed, + philox_rounds=10, + **common_kwargs, + ) + + # Outputs should be nearly identical — rounding only perturbs the + # post-replay state by ±1 ULP before the output phase reads it. + torch.testing.assert_close( + out_rounded, out_no_round, rtol=2e-2, atol=1.0, msg="Output diverged with Philox rounding" + ) + + # State should remain fp16 + assert state_rounded.dtype == torch.float16 + + # States should differ by at most 1 fp16 ULP per element. + # fp16 ULP depends on magnitude: up to 0.5 for values near 512. + # Use rtol to account for magnitude-dependent ULP. + slots = state_batch_indices if paged_cache else slice(None) + torch.testing.assert_close( + state_rounded[slots], + state_no_round[slots], + rtol=2e-3, + atol=0.2, + msg="State diverged with Philox rounding", + ) + + +@_skip_pre_sm100 +def test_philox_rounding_unbiased(): + """ + Verify that Philox stochastic rounding is unbiased. + + Runs the replay kernel with fp32 state (capturing the true fp32 + post-replay state) and with fp16 state + Philox rounding. Compares the + rounding residual (fp16_state.float() - fp32_state) against deterministic + rounding (fp32_state.to(fp16).float() - fp32_state). + + Deterministic round-to-nearest-even has a systematic positive bias on + the residual. Philox stochastic rounding should be unbiased: the mean + residual should be near zero. + + Uses a large batch (16) for ~2M state elements — plenty of statistics. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + batch, T = 16, 6 + device = "cuda" + dtype = torch.bfloat16 + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + # Use fp32 initial state so replay produces non-fp16-representable values + state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=torch.float32) + + old_x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) + + x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt_val = repeat(dt_base, "b t h -> b t h p", p=head_dim) + B = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + prev_tokens = torch.full((batch,), T, device=device, dtype=torch.int32) + + common_kwargs = dict( + x=x, + dt=dt_val, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + ) + + # 1. fp32 state — captures true post-replay state + state_fp32 = state0.clone() + out_fp32 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_fp32, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + out=out_fp32, + **common_kwargs, + ) + + # 2. fp16 state with Philox rounding + rand_seed = torch.tensor([99999], device=device, dtype=torch.int64) + state_rounded = state0.to(torch.float16).clone() + out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_rounded, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + out=out_rounded, + rand_seed=rand_seed, + philox_rounds=10, + **common_kwargs, + ) + + # Compute rounding residuals where fp32 state has non-zero values + fp32_vals = state_fp32.flatten() + stochastic_residual = state_rounded.float().flatten() - fp32_vals + deterministic_residual = fp32_vals.to(torch.float16).float() - fp32_vals + + # Only consider elements where rounding matters (non-zero residual possible) + nonzero_mask = deterministic_residual.abs() > 0 + num_nonzero = nonzero_mask.sum().item() + assert num_nonzero > 1000, f"Too few roundable elements: {num_nonzero}" + + stochastic_mean = stochastic_residual[nonzero_mask].mean().item() + deterministic_mean = deterministic_residual[nonzero_mask].mean().item() + + # Stochastic rounding should be less biased than deterministic. + # With ~millions of elements, the stochastic mean should be very close to 0. + # Deterministic round-to-nearest-even has a small but systematic bias. + assert abs(stochastic_mean) < abs(deterministic_mean) or abs(stochastic_mean) < 1e-5, ( + f"Stochastic rounding appears biased: stochastic_mean={stochastic_mean:.6f}, " + f"deterministic_mean={deterministic_mean:.6f}, n_elements={num_nonzero}" + ) + + +# HEADS_PER_BLOCK > 1 test. The default heuristic only picks HPB > 1 at large +# total_heads (>= 256-512), which the main test with batch=2 never reaches. +# This test overrides _heads_per_block to exercise the two-loop structure in +# the precompute kernel (store-then-reload of per-head dt/dA_cumsum). +# Configs: (nheads=16, ngroups=1) and (nheads=32, ngroups=2) both have +# heads_per_group=16. The heuristic caps HPB at min(2|4, hpg), so HPB=2, 4. +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) +@pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) +@pytest.mark.parametrize("launch_with_pdl", [False, True], ids=["no_ext_pdl", "ext_pdl"]) +@pytest.mark.parametrize("use_internal_pdl", [False, True], ids=["no_int_pdl", "int_pdl"]) +@pytest.mark.parametrize("batch", [1, 2, 8, 16], ids=["B1", "B2", "B8", "B16"]) +def test_checkpointing_heads_per_block( + nheads, + head_dim, + d_state, + ngroups, + state_dtype, + T, + heads_per_block, + launch_with_pdl, + use_internal_pdl, + batch, +): + """ + Verify checkpointing_state_update produces correct results when + _heads_per_block > 1, exercising the precompute kernel's two-loop + structure (store per-head dt/dA_cumsum in loop 1, reload in loop 2). + """ + device = "cuda" + dtype = torch.bfloat16 + + if nheads % heads_per_block != 0: + pytest.skip(f"nheads ({nheads}) not divisible by heads_per_block ({heads_per_block})") + if heads_per_block > nheads // ngroups: + pytest.skip( + f"heads_per_block ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})" + ) + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + cache_size = batch + state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + + x1 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + cache_size, T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + state0.clone(), + x1, + dt1, + A, + B1, + C1, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) + + old_x[:] = x1 + dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) + + for slot in range(cache_size): + buf = cache_buf_idx[slot].item() + old_B[slot, buf] = B1[slot] + old_dt[slot, buf] = dt1[slot].T + old_dA_cumsum[slot, buf] = dA_cumsum1[slot].T + + k = T + torch.manual_seed(123) + + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = state0.float().clone() + ref_state_f32[:] = states_buffer_f32[:, k - 1] + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + out=ref_out, + ) + + test_state = state0.clone() + prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + + checkpointing_state_update( + test_state, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + x=x2, + dt=dt2, + A=A, + B=B2, + C=C2, + out=test_out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + _heads_per_block=heads_per_block, + launch_with_pdl=launch_with_pdl, + use_internal_pdl=use_internal_pdl, + ) + + torch.testing.assert_close( + test_out, + ref_out, + rtol=2e-2, + atol=1.0, + msg=f"Output mismatch with HPB={heads_per_block}, T={T}, " + f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", + ) + + expected_state = states_buffer_f32[:, k - 1].to(state_dtype) + torch.testing.assert_close( + test_state, + expected_state, + rtol=2e-2, + atol=1.0, + msg=f"State mismatch with HPB={heads_per_block}, T={T}, " + f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", + ) + + +# HPB > 1 multi-step test. Production chains decode steps; bugs in +# buffer ordering or stale cache values accumulate across steps and can +# be invisible in a single-step test. +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("T", [6, 16], ids=["T6", "T16"]) +@pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) +@pytest.mark.parametrize("paged_cache", [False, True], ids=["contig", "paged"]) +def test_checkpointing_heads_per_block_multistep( + nheads, head_dim, d_state, ngroups, state_dtype, T, heads_per_block, paged_cache +): + """ + Chain N decode steps with HPB > 1 and verify each step's output matches + a fresh reference. A bug that mixes up WRITE/READ buffers, writes wrong + data to cache, or races in the two-loop structure would accumulate + across steps. + """ + batch = 2 + device = "cuda" + dtype = torch.bfloat16 + n_steps = 8 + + if nheads % heads_per_block != 0: + pytest.skip(f"nheads ({nheads}) not divisible by HPB ({heads_per_block})") + if heads_per_block > nheads // ngroups: + pytest.skip(f"HPB ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})") + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + if paged_cache: + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + slots = state_batch_indices + else: + cache_size = batch + state_batch_indices = None + slots = slice(None) + + all_x = [] + all_dt = [] + all_B = [] + all_C = [] + for step in range(n_steps): + torch.manual_seed(1000 + step) + all_x.append(torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype)) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + all_dt.append(repeat(dt_base, "b t h -> b t h p", p=head_dim)) + all_B.append(torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype)) + all_C.append(torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype)) + + torch.manual_seed(999) + state_init = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + + ref_state = state_init.float().clone() + ref_outs = [] + ref_slots = ( + state_batch_indices + if paged_cache + else torch.arange(batch, device=device, dtype=torch.int32) + ) + for step in range(n_steps): + out_step = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state, + all_x[step], + all_dt[step], + A, + all_B[step], + all_C[step], + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=ref_slots, + out=out_step, + ) + ref_outs.append(out_step) + + test_state = state_init.clone() + old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.zeros(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) + + for step in range(n_steps): + k = T if step > 0 else 0 + prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + + checkpointing_state_update( + test_state, + old_x, + old_B, + old_dt, + old_dA_cumsum, + cache_buf_idx, + prev_tokens, + x=all_x[step], + dt=all_dt[step], + A=A, + B=all_B[step], + C=all_C[step], + out=test_out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + _heads_per_block=heads_per_block, + ) + + if paged_cache: + cache_buf_idx[slots] = 1 - cache_buf_idx[slots] + else: + cache_buf_idx[:] = 1 - cache_buf_idx + + torch.testing.assert_close( + test_out, + ref_outs[step], + rtol=2e-2, + atol=2.0, + msg=f"Output mismatch at step {step} with HPB={heads_per_block}, " + f"T={T}, nheads={nheads}, ngroups={ngroups}, " + f"state_dtype={state_dtype}, paged_cache={paged_cache}", + ) From de30678fea40951767c78ae1b7608842d14d19a2 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 28 Apr 2026 16:53:03 -0700 Subject: [PATCH 03/89] Trim heads_per_block tests: drop PDL parametrize, pin batch=8 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both test_replay_heads_per_block and test_checkpointing_heads_per_block were sweeping launch_with_pdl × use_internal_pdl × batch={1,2,8,16} on top of the actual params, blowing up the matrix ~16x for no extra correctness coverage — PDL flag combinations are already exercised by the dedicated correctness tests above (test_replay_selective_state_update and test_checkpointing_state_update). Drop the PDL parametrizes (use wrapper defaults), pin batch=8. Both suites combined: 73s -> 45s. 282 tests pass across both files. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../modules/mamba/test_checkpointing_state_update.py | 12 ++++-------- .../mamba/test_replay_selective_state_update.py | 12 ++++-------- 2 files changed, 8 insertions(+), 16 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index c6327252a6aa..703da1b502fd 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -470,9 +470,6 @@ def test_philox_rounding_unbiased(): @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) @pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) -@pytest.mark.parametrize("launch_with_pdl", [False, True], ids=["no_ext_pdl", "ext_pdl"]) -@pytest.mark.parametrize("use_internal_pdl", [False, True], ids=["no_int_pdl", "int_pdl"]) -@pytest.mark.parametrize("batch", [1, 2, 8, 16], ids=["B1", "B2", "B8", "B16"]) def test_checkpointing_heads_per_block( nheads, head_dim, @@ -481,10 +478,11 @@ def test_checkpointing_heads_per_block( state_dtype, T, heads_per_block, - launch_with_pdl, - use_internal_pdl, - batch, ): + # PDL flags use wrapper defaults; trimming the parametrize keeps this + # suite fast. Coverage of {launch_with_pdl, use_internal_pdl} variations + # lives in the dedicated correctness tests above (test_checkpointing_state_update). + batch = 8 """ Verify checkpointing_state_update produces correct results when _heads_per_block > 1, exercising the precompute kernel's two-loop @@ -605,8 +603,6 @@ def test_checkpointing_heads_per_block( dt_softplus=True, state_batch_indices=None, _heads_per_block=heads_per_block, - launch_with_pdl=launch_with_pdl, - use_internal_pdl=use_internal_pdl, ) torch.testing.assert_close( diff --git a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py index 7757f83ab2ad..89737937fbfb 100644 --- a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py @@ -470,9 +470,6 @@ def test_philox_rounding_unbiased(): @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) @pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) -@pytest.mark.parametrize("launch_with_pdl", [False, True], ids=["no_ext_pdl", "ext_pdl"]) -@pytest.mark.parametrize("use_internal_pdl", [False, True], ids=["no_int_pdl", "int_pdl"]) -@pytest.mark.parametrize("batch", [1, 2, 8, 16], ids=["B1", "B2", "B8", "B16"]) def test_replay_heads_per_block( nheads, head_dim, @@ -481,10 +478,11 @@ def test_replay_heads_per_block( state_dtype, T, heads_per_block, - launch_with_pdl, - use_internal_pdl, - batch, ): + # PDL flags use wrapper defaults; trimming the parametrize keeps this + # suite fast. Coverage of {launch_with_pdl, use_internal_pdl} variations + # lives in the dedicated correctness tests above (test_replay_selective_state_update). + batch = 8 """ Verify replay_selective_state_update produces correct results when _heads_per_block > 1, exercising the precompute kernel's two-loop @@ -605,8 +603,6 @@ def test_replay_heads_per_block( dt_softplus=True, state_batch_indices=None, _heads_per_block=heads_per_block, - launch_with_pdl=launch_with_pdl, - use_internal_pdl=use_internal_pdl, ) torch.testing.assert_close( From 1819cb453142ee5bb33b8f9acf7bcce4e441cfb2 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 28 Apr 2026 17:23:51 -0700 Subject: [PATCH 04/89] Add --variant {replay,checkpointing} flag to benchmark Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 49 +++++++++++++++---- 1 file changed, 40 insertions(+), 9 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index dd4fc904d570..aec41c318228 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -114,11 +114,13 @@ def _load(mod_name: str, file_name: str): _load("softplus", "softplus.py") # 3. The actual kernels replay_mod = _load("replay_selective_state_update", "replay_selective_state_update.py") + checkpoint_mod = _load("checkpointing_state_update", "checkpointing_state_update.py") base_mod = _load("selective_state_update", "selective_state_update.py") conv1d_mod = _load("causal_conv1d_triton", "causal_conv1d_triton.py") return ( replay_mod.replay_selective_state_update, + checkpoint_mod.checkpointing_state_update, base_mod.selective_state_update, conv1d_mod.causal_conv1d_update, ) @@ -127,25 +129,39 @@ def _load(mod_name: str, file_name: str): def _import_mamba_kernels_full(): """Import via the standard tensorrt_llm package (slow but safe).""" from tensorrt_llm._torch.modules.mamba.causal_conv1d_triton import causal_conv1d_update + from tensorrt_llm._torch.modules.mamba.checkpointing_state_update import ( + checkpointing_state_update, + ) from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( replay_selective_state_update, ) from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update - return replay_selective_state_update, selective_state_update, causal_conv1d_update + return ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) # Use fast import by default; --full-import parsed later but we need the # functions at module level. Check sys.argv early. if "--full-import" in sys.argv: - replay_selective_state_update, selective_state_update, causal_conv1d_update = ( - _import_mamba_kernels_full() - ) + ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) = _import_mamba_kernels_full() else: try: - replay_selective_state_update, selective_state_update, causal_conv1d_update = ( - _import_mamba_kernels_fast() - ) + ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) = _import_mamba_kernels_fast() except Exception as e: # noqa: BLE001 - exit loudly; don't hide a fast-import regression print( f"ERROR: fast import failed ({type(e).__name__}: {e})\n" @@ -155,6 +171,12 @@ def _import_mamba_kernels_full(): ) sys.exit(1) + +_VARIANT_FNS = { + "replay": lambda: replay_selective_state_update, + "checkpointing": lambda: checkpointing_state_update, +} + # Model config defaults (Nemotron-3-Super-120B full model). # --tp-size divides nheads and ngroups to get the per-GPU slice. # TP=1: nheads=128, ngroups=8 @@ -535,6 +557,7 @@ def _bench_config( d_state = args.d_state with_conv1d = getattr(args, "with_conv1d", False) use_philox = getattr(args, "philox_rounding", False) + variant_fn = _VARIANT_FNS[args.variant]() # Philox rounding: allocate rand_seed tensor rand_seed = None @@ -731,7 +754,7 @@ def _run_incr( else: x_call, B_call, C_call = x, B, C extra_kwargs = {} - replay_selective_state_update( + variant_fn( state_work, old_x_work, old_B_work, @@ -786,7 +809,7 @@ def _run_incr( _print_row( show_kernel_col, - "replay", + args.variant, batch, mtp_len, prev_k, @@ -1069,6 +1092,14 @@ def _parse_args() -> argparse.Namespace: help="Enable Philox stochastic rounding for fp16 state " "(rand_seed generated per iteration, philox_rounds=10).", ) + parser.add_argument( + "--variant", + choices=["replay", "checkpointing"], + default="replay", + help="Which kernel to time as the 'replay' row. 'replay' = today's " + "kernel (selective_state_update.py:replay). 'checkpointing' = " + "checkpointing_state_update.py. Both share the same wrapper signature.", + ) parser.add_argument( "--full-import", action="store_true", From 4dfd3000dc0b45bb96af4ed3c75e6646456bfd0c Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 28 Apr 2026 17:28:19 -0700 Subject: [PATCH 05/89] Print mismatch magnitude in checkpointing tests on assertion failure Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/test_checkpointing_state_update.py | 38 ++++++++++++++++--- 1 file changed, 32 insertions(+), 6 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index 703da1b502fd..674717048230 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -215,16 +215,42 @@ def test_checkpointing_state_update( # T<=16, ~2.0 at T=32-55; mean ~0.014; <0.02% of elements exceed 0.5. # State dtype (fp16/bf16/fp32) doesn't shift the error — bf16 dot # inputs dominate, not state storage. - torch.testing.assert_close( - test_out, ref_out, rtol=2e-2, atol=1.0, msg=f"Output mismatch at k={k}" - ) + out_diff = (test_out.float() - ref_out.float()).abs() + out_max = out_diff.max().item() + out_mean = out_diff.mean().item() + try: + torch.testing.assert_close( + test_out, ref_out, rtol=2e-2, atol=1.0, msg=f"Output mismatch at k={k}" + ) + except AssertionError: + print( + f"k={k} out: max={out_max:.4f} mean={out_mean:.4f} " + f"nan={torch.isnan(test_out).any().item()} " + f"inf={torch.isinf(test_out).any().item()}" + ) + raise expected_state = ( state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) ) - torch.testing.assert_close( - test_state[slots], expected_state, rtol=2e-2, atol=1.0, msg=f"State mismatch at k={k}" - ) + state_diff = (test_state[slots].float() - expected_state.float()).abs() + state_max = state_diff.max().item() + state_mean = state_diff.mean().item() + try: + torch.testing.assert_close( + test_state[slots], + expected_state, + rtol=2e-2, + atol=1.0, + msg=f"State mismatch at k={k}", + ) + except AssertionError: + print( + f"k={k} state: max={state_max:.4f} mean={state_mean:.4f} " + f"nan={torch.isnan(test_state).any().item()} " + f"inf={torch.isinf(test_state).any().item()}" + ) + raise @_skip_pre_sm100 From 8ff8e94dabe5352f21e4cf2aaf3c72ab9ac3abbd Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 28 Apr 2026 17:28:19 -0700 Subject: [PATCH 06/89] Rectangle CB factoring in checkpointing_state_update kernel Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 298 ++++++++++++++---- 1 file changed, 230 insertions(+), 68 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 564c0b203d3a..399cf006ac1d 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -59,7 +59,9 @@ def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: @triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) @triton.heuristics({"BLOCK_SIZE_T": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)}) + triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics({"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"] + args["T"]), 16)}) @triton.jit() def _checkpointing_precompute_kernel( # Input pointers @@ -71,12 +73,15 @@ def _checkpointing_precompute_kernel( # Output pointers cb_scaled_ptr, decay_vec_ptr, - # Cache WRITE pointers (write-buffer for next step) + # Cache pointers — both READ (for rectangle CB factoring) and WRITE + # (next step's replay) buffers are accessed via stride_*_dbuf offsets. old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, # Double-buffer index (per cache slot) cache_buf_idx_ptr, + # Per-request: number of committed tokens since last checkpoint + prev_num_accepted_tokens_ptr, state_batch_indices_ptr, pad_slot_id, # Dimensions @@ -131,6 +136,7 @@ def _checkpointing_precompute_kernel( HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, LAUNCH_DEPENDENT_KERNELS: tl.constexpr, HEADS_PER_BLOCK: tl.constexpr, @@ -160,15 +166,12 @@ def _checkpointing_precompute_kernel( buf_read = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) buf_write = 1 - buf_read - offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis (output rows): T_new positions + offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis (rectangle input cols): up to PNAT+T_new offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) t_mask = offs_t < T n_mask = offs_n < dstate - # Causal mask is shared across all heads (depends only on offs_t) - causal_mask = offs_t[:, None] >= offs_t[None, :] - valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] - # --- Loop 1: compute per-head dt/dA_cumsum/decay BEFORE gdc_wait --- # These only depend on dt (from in_proj, not conv1d) and parameters (A, dt_bias). # Store to cache; will reload after the wait for CB scaling. @@ -212,72 +215,182 @@ def _checkpointing_precompute_kernel( if LAUNCH_WITH_PDL: tl.extra.cuda.gdc_wait() - # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- + # --- Rectangle setup: combined k-axis [old | new | pad] of length up to --- + # MAX_REPLAY_BUFFER_LENGTH + T = BLOCK_SIZE_K. is_old_k / is_new_k are + # disjoint masks covering the two halves; safe_k_new clamps the new half's + # source row to [0, T) and safe_old_k clamps the old half's cache row to + # [0, MAX_REPLAY_BUFFER_LENGTH) so masked-off load addresses stay in-bounds. + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # --- Load C and B (per-group, shared across HEADS_PER_BLOCK heads) --- group_idx = first_head // nheads_ngroups_ratio C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group - B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + B_new_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group C_all = tl.load( C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, mask=t_mask[:, None] & n_mask[None, :], other=0.0, ) - B_all = tl.load( - B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + # New B at unshifted rows [0, T) — used to populate the next step's cache. + B_new_orig = tl.load( + B_new_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, mask=t_mask[:, None] & n_mask[None, :], other=0.0, ) + # New B at shifted rows [PNAT, PNAT+T) along the k-axis — populates the new + # half of B_combined. Same memory range as B_new_orig; second load hits L1. + B_new_shifted = tl.load( + B_new_base + safe_k_new[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=is_new_k[:, None] & n_mask[None, :], + other=0.0, + ) + # Old B from READ buffer at k-axis rows [0, PNAT) — populates the old half. + old_B_read_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + buf_read * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_load = tl.load( + old_B_read_base + + safe_old_k[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + mask=is_old_k[:, None] & n_mask[None, :], + other=0.0, + ) + # B_combined (BLOCK_SIZE_K × dstate): rows [0, PNAT) from old, rows + # [PNAT, PNAT+T) from new, else 0. Disjoint masks → addition gives union. + # Single rectangle matmul replaces today's square C @ B.T. + B_combined = old_B_load + B_new_shifted + raw_rect_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_combined).to(tl.bfloat16)) - # Compute raw CB once — shared across all heads in this block - raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) - - # Store B to cache (once per group, only if this block covers the first heads) + # Store new B to cache write buffer (once per group) if first_head % nheads_ngroups_ratio == 0: - old_B_base = ( + old_B_write_base = ( old_B_ptr + cache_batch_idx * stride_old_B_cache + buf_write * stride_old_B_dbuf + group_idx * stride_old_B_group ) tl.store( - old_B_base + offs_t[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - B_all, + old_B_write_base + + offs_t[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_new_orig, mask=t_mask[:, None] & n_mask[None, :], ) - # --- Loop 2: reload per-head dA_cumsum/dt from cache, scale CB --- - # The cache was just written above, so these loads should hit L2. + # Causal mask for the rectangle (BLOCK_SIZE_T × BLOCK_SIZE_K, same across heads): + # k < PNAT (old): always passes — causal in the combined timeline. + # k_new ∈ [0, T) (new): passes iff k_new ≤ t. + t_idx_2d = offs_t[:, None] + k_idx_2d = offs_k[None, :] + is_old_k_2d = k_idx_2d < prev_num_accepted_tokens + k_new_idx_2d = k_idx_2d - prev_num_accepted_tokens + is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) + causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] + + # --- Loop 2: per-head post-scaling of the rectangle CB --- + # Numerical guard: factoring dt[k] * exp(cumAdt_new[t] - cumAdt_at_k[k]) + # into factor_dt[k] * exp_diff[t,k] keeps both sides bounded. Splitting + # further into dt[k]*exp(-cumAdt_at_k[k]) * exp(cumAdt_new[t]) would let + # the new-side `exp(-cumAdt_new[k_new])` overflow at large T (cumAdt grows + # ~T*A*dt) while the t-side `exp(cumAdt_new[t])` underflows, producing + # inf*0 = NaN at causal-valid positions. Computing the SUM of the two + # cumsum offsets before exp keeps the argument ≤ 0 for all valid (t,k). for h_local in range(HEADS_PER_BLOCK): head_idx = first_head + h_local - # Reload dt and dA_cumsum from cache (just written in loop 1) - old_dt_base = ( + # Reload new dt and dA_cumsum from WRITE buffer (just stored in loop 1). + old_dt_write_base = ( old_dt_ptr + cache_batch_idx * stride_old_dt_cache + buf_write * stride_old_dt_dbuf + head_idx * stride_old_dt_head ) - dt = tl.load(old_dt_base + offs_t * stride_old_dt_T, mask=t_mask, other=0.0).to(tl.float32) - - old_dA_cumsum_base = ( + old_dA_cumsum_write_base = ( old_dA_cumsum_ptr + cache_batch_idx * stride_old_dA_cumsum_cache + buf_write * stride_old_dA_cumsum_dbuf + head_idx * stride_old_dA_cumsum_head ) - dA_cumsum = tl.load( - old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0 + dA_cumsum_new = tl.load( + old_dA_cumsum_write_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0 + ).to(tl.float32) + + # Old-token side: load old dt and dA_cumsum from READ buffer. + old_dt_read_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_read * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + old_dA_cumsum_read_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_read * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + # Old-token per-k load: address indexed by safe_old_k (clamped to in-bounds + # for masked-off positions); cache T-axis size = MAX_REPLAY_BUFFER_LENGTH. + old_dt_all = tl.load( + old_dt_read_base + safe_old_k * stride_old_dt_T, mask=is_old_k, other=0.0 + ).to(tl.float32) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_read_base + safe_old_k * stride_old_dA_cumsum_T, + mask=is_old_k, + other=0.0, + ).to(tl.float32) + total_dA_cumsum = tl.load( + old_dA_cumsum_read_base + prev_k_idx * stride_old_dA_cumsum_T + ).to(tl.float32) + + # New-token per-k loads from the WRITE buffer (cache hits L1 from loop 1). + dt_at_kn = tl.load( + old_dt_write_base + safe_k_new * stride_old_dt_T, mask=is_new_k, other=0.0 + ).to(tl.float32) + dA_cumsum_at_kn = tl.load( + old_dA_cumsum_write_base + safe_k_new * stride_old_dA_cumsum_T, + mask=is_new_k, + other=0.0, ).to(tl.float32) - # Scale raw_CB with per-head decay and dt - decay_matrix = tl.exp(dA_cumsum[:, None] - dA_cumsum[None, :]) - CB_scaled = tl.where(valid_mask, raw_CB * decay_matrix * dt[None, :], 0.0) + # Per-k dt factor (no exp). Disjoint region select. + factor_dt = tl.where(is_old_k, old_dt_all, dt_at_kn) + + # Per-k cumsum offset s_k. For valid k, s_k ≤ 0 for old (cumAdt + # non-increasing → total ≤ old_cumAdt[k]) and s_k ≥ 0 for new + # (-cumAdt_new[k_new] flips the negative cumsum). Bounded ≤ |cumAdt|. + s_k = tl.where( + is_old_k, total_dA_cumsum - old_dA_cumsum_all, -dA_cumsum_at_kn + ) + + # exp_diff[t, k] = exp(cumAdt_new[t] + s_k[k]) + # For valid (t, k_new) with k_new ≤ t (causal): exp arg ≤ 0 → bounded. + # For valid (t, k_old) with k < PNAT: exp arg ≤ 0 → bounded. + # For invalid positions: causal mask zeroes the result downstream. + exp_diff = tl.exp(s_k[None, :] + dA_cumsum_new[:, None]) + + rect_CB_scaled = tl.where( + causal_combined, + raw_rect_CB * factor_dt[None, :] * exp_diff, + 0.0, + ) cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head tl.store( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, - CB_scaled, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + rect_CB_scaled, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), ) @@ -293,7 +406,9 @@ def _checkpointing_precompute_kernel( @triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) @triton.heuristics({"BLOCK_SIZE_T": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)}) + triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics({"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"] + args["T"]), 16)}) @triton.jit() def _checkpointing_main_kernel( # Pointers @@ -390,6 +505,7 @@ def _checkpointing_main_kernel( HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, USE_RS_ROUNDING: tl.constexpr, PHILOX_ROUNDS: tl.constexpr, @@ -413,7 +529,8 @@ def _checkpointing_main_kernel( offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis: T_new positions + offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis: combined [0, PNAT+T_new) m_mask = offs_m < dim n_mask = offs_n < dstate t_mask = offs_t < T @@ -427,8 +544,13 @@ def _checkpointing_main_kernel( state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - # Phase 1: Replay via tl.dot fast-forward (reads from READ buffer) + # Phase 1: Replay via tl.dot fast-forward (reads from READ buffer). + # Old-token loads are indexed along the K-axis (offs_k, BLOCK_SIZE_K rows + # of which only [0, PNAT) carry data); safe_old_k clamps the masked-off + # row addresses to in-bounds cache positions. group_idx = pid_h // nheads_ngroups_ratio + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) # Load precomputed dt and dA_cumsum from READ buffer old_dt_base = ( @@ -437,9 +559,9 @@ def _checkpointing_main_kernel( + buf_read * stride_old_dt_dbuf + pid_h * stride_old_dt_head ) - old_dt_all = tl.load(old_dt_base + offs_t * stride_old_dt_T, mask=t_mask, other=0.0).to( - tl.float32 - ) + old_dt_all = tl.load( + old_dt_base + safe_old_k * stride_old_dt_T, mask=is_old_k, other=0.0 + ).to(tl.float32) old_dA_cumsum_base = ( old_dA_cumsum_ptr @@ -448,31 +570,32 @@ def _checkpointing_main_kernel( + pid_h * stride_old_dA_cumsum_head ) old_dA_cumsum_all = tl.load( - old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0 + old_dA_cumsum_base + safe_old_k * stride_old_dA_cumsum_T, mask=is_old_k, other=0.0 ).to(tl.float32) - # Load dA_cumsum at prev_k-1 directly via pointer math (avoids masked reduction). - # Clamp to [0, T-1] defensively — out-of-contract PNAT > T would read OOB. - prev_k_idx = tl.minimum(tl.maximum(prev_num_accepted_tokens - 1, 0), T - 1) + # Total cumsum at end of old: scalar load at PNAT-1, clamped against the + # cache T-axis (= MAX_REPLAY_BUFFER_LENGTH) for OOB safety on PNAT=0. + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( tl.float32 ) # Step 0 invariant: PNAT=0 means `state` is already last step's state (not - # two back). coeff is all-zero (offs_t < 0), total_decay is 1.0, so the - # replay leaves `state` unchanged — cache contents don't matter on step 0. + # two back). coeff is all-zero (is_old_k all-false), total_decay is 1.0, + # so the replay leaves `state` unchanged — cache contents don't matter. coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all - coeff = tl.where(offs_t < prev_num_accepted_tokens, coeff, 0.0) - # Load old_x: (BLOCK_SIZE_T, BLOCK_SIZE_M) — single-buffered + # Load old_x: (BLOCK_SIZE_K, BLOCK_SIZE_M) — single-buffered cache old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head old_x_all = tl.load( - old_x_base + offs_t[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - mask=t_mask[:, None] & m_mask[None, :], + old_x_base + safe_old_k[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + mask=is_old_k[:, None] & m_mask[None, :], other=0.0, ) - # Load old_B from READ buffer: (BLOCK_SIZE_T, BLOCK_SIZE_DSTATE) + # Load old_B from READ buffer: (BLOCK_SIZE_K, BLOCK_SIZE_DSTATE) old_B_base = ( old_B_ptr + cache_batch_idx * stride_old_B_cache @@ -480,20 +603,28 @@ def _checkpointing_main_kernel( + group_idx * stride_old_B_group ) old_B_all = tl.load( - old_B_base + offs_t[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - mask=t_mask[:, None] & n_mask[None, :], + old_B_base + + safe_old_k[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + mask=is_old_k[:, None] & n_mask[None, :], other=0.0, ).to(tl.float32) - # Scale B by coefficients + # Scale B by coefficients (k-axis) dB_scaled = coeff[:, None] * old_B_all - # Apply total decay to initial state FIRST, then add contributions + # Apply total decay to initial state. Save the decayed-but-pre-replay + # state for the output (state_out uses C @ state_prev_decayed.T because + # the rectangle CB carries the old-tokens contribution to output through + # the token_out matmul). total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) - state *= total_decay + state_prev_decayed = state * total_decay - # tl.dot fast-forward: old_x^T @ dB_scaled → (M, dstate) - state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) + # tl.dot fast-forward: old_x^T @ dB_scaled → (M, dstate); add to decayed + # prior state to get the post-replay (checkpointed) state. + state = state_prev_decayed + tl.dot( + tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16) + ) # Write post-replay state if USE_RS_ROUNDING: @@ -563,26 +694,52 @@ def _checkpointing_main_kernel( ) x_all = x_all.to(tl.float32) - # Load precomputed CB_scaled and decay_vec + # --- Build x_combined for the rectangle token_out matmul --- + # K-axis layout: rows [0, PNAT) from old_x (already loaded; is_old_k mask), + # rows [PNAT, PNAT+T) from new_x via a shifted load over the same memory + # as `x_all` (hits L1). Disjoint masks → safe to add. + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + old_x_in_combined = tl.where( + is_old_k[:, None] & m_mask[None, :], + old_x_all.to(tl.float32), + 0.0, + ) + new_x_shifted = tl.load( + x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=is_new_k[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + x_combined = old_x_in_combined + new_x_shifted # (BLOCK_SIZE_K, BLOCK_SIZE_M) + + # Load precomputed rectangle CB_scaled (BLOCK_SIZE_T × BLOCK_SIZE_K) and + # decay_vec_new (= exp(cumAdt_new), BLOCK_SIZE_T). cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), other=0.0, ).to(tl.float32) decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + decay_vec_new = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( tl.float32 ) - # init_out = C_all @ state^T * decay_vec - init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] + # state_out: pre-replay state contribution to output. + # (C[t] · state_prev) * total_decay * decay_vec_new[t] + # Folded: state_prev_decayed already = state_prev * total_decay. + state_out = ( + tl.dot(C_all.to(tl.bfloat16), tl.trans(state_prev_decayed).to(tl.bfloat16)) + * decay_vec_new[:, None] + ) - # cb_out = CB_scaled @ x_all - cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) + # token_out: combined old+new tokens contribution via the rectangle CB. + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) - out_all = init_out + cb_out + out_all = state_out + token_out if HAS_D: out_all = out_all + x_all * D[None, :] @@ -764,11 +921,15 @@ def checkpointing_state_update( assert tie_hdim device = x.device - BLOCK_SIZE_T = max(triton.next_power_of_2(max_replay_buffer_length), 16) + # BLOCK_SIZE_T sizes the T-axis (T_new positions); BLOCK_SIZE_K sizes the + # rectangle K-axis (combined [old | new] up to PNAT+T_new ≤ MAX+T_new). + # Decoupling them keeps T-axis matmuls compact at np2(T) when MAX > T. + BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) + BLOCK_SIZE_K = max(triton.next_power_of_2(max_replay_buffer_length + T), 16) # Allocate precomputed intermediates (per-call, not cached) cb_scaled = torch.empty( - batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_T, device=device, dtype=torch.float32 + batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 ) decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) @@ -910,6 +1071,7 @@ def checkpointing_state_update( old_dt, old_dA_cumsum, cache_buf_idx, + prev_num_accepted_tokens, state_batch_indices, pad_slot_id, T, From ed569d38d891d3db256cc87aedad4a9865da9050 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Thu, 14 May 2026 11:19:27 -0700 Subject: [PATCH 07/89] Track Mamba replay history window in cache manager Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../_torch/pyexecutor/mamba_cache_manager.py | 102 ++++++++++++------ .../executor/test_mamba_cache_manager.py | 45 +++++++- 2 files changed, 112 insertions(+), 35 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index da682bebc858..a40ea37f472a 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -263,11 +263,11 @@ class SpeculativeState(State): # 0 means temporal saved state is actually the last state, not two back. prev_num_accepted_tokens: torch.Tensor | None = None # (cache,) int — shared across layers cache_buf_idx: torch.Tensor | None = None # (cache,) int32 — shared across layers - old_x: torch.Tensor | None = None # (layers, cache, T, nheads, dim) - old_B: torch.Tensor | None = None # (layers, cache, 2, T, ngroups, dstate) + old_x: torch.Tensor | None = None # (layers, cache, history, nheads, dim) + old_B: torch.Tensor | None = None # (layers, cache, 2, history, ngroups, dstate) # Processed dt: softplus(raw_dt + dt_bias), clamped to dt_limit. - old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, T) fp32 - old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, T) fp32 + old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 + old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 def __init__( self, @@ -292,6 +292,8 @@ def __init__( self.speculative_num_draft_tokens = speculative_num_draft_tokens self.spec_state_size = spec_state_size self._use_replay_state_update = use_replay_state_update + self.replay_history_size: Optional[int] = None + self.replay_step_width: Optional[int] = None # get tp size tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size @@ -355,6 +357,7 @@ def __init__( # create state container if speculative_num_draft_tokens is not None: T = speculative_num_draft_tokens + 1 + self.replay_step_width = T # Conv intermediate cache — same for both paths intermediate_conv_window_cache = torch.zeros( @@ -370,6 +373,7 @@ def __init__( assert n_groups % tp_size == 0, \ "replay state update requires n_groups divisible by tp_size" n_groups_per_rank = n_groups // tp_size + self.replay_history_size = max(16, T) # Compact replay cache. # old_x is single-buffered (written by main kernel after replay). @@ -382,7 +386,7 @@ def __init__( device=device) spec_kwargs['old_x'] = torch.zeros(num_local_layers, max_batch_size, - T, + self.replay_history_size, nheads, head_dim, dtype=dtype, @@ -390,7 +394,7 @@ def __init__( spec_kwargs['old_B'] = torch.zeros(num_local_layers, max_batch_size, 2, - T, + self.replay_history_size, n_groups_per_rank, d_state, dtype=dtype, @@ -399,16 +403,17 @@ def __init__( max_batch_size, 2, nheads, - T, + self.replay_history_size, dtype=torch.float32, device=device) - spec_kwargs['old_dA_cumsum'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - nheads, - T, - dtype=torch.float32, - device=device) + spec_kwargs['old_dA_cumsum'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + nheads, + self.replay_history_size, + dtype=torch.float32, + device=device) ssm_spec_cache = [ spec_kwargs['old_x'], spec_kwargs['old_B'], spec_kwargs['old_dt'], spec_kwargs['old_dA_cumsum'] @@ -638,14 +643,26 @@ def update_mamba_states(self, attn_metadata: "AttentionMetadata", src_state_indices = self.intermediate_state_indices[:num_gens] if self._use_replay_state_update: - # SSM state is handled incrementally by the kernel. Update the - # number of accepted tokens and flip the double-buffer index so the - # next step's replay reads from the buffer that was just written by - # the precompute kernel. + # SSM state is handled incrementally by the kernel. Mirror the + # kernel's per-slot checkpoint predicate from the previous PNAT and + # fixed replay step width: checkpoint steps write a fresh history + # buffer and flip, while no-checkpoint steps append to the active + # buffer and keep reading from it next step. + accepted_tokens = num_accepted_tokens[num_contexts:num_contexts + + num_gens] + prev_num_accepted_tokens = \ + self.mamba_cache.prev_num_accepted_tokens[state_indices_d] + wrote_checkpoint = (prev_num_accepted_tokens + + self.replay_step_width + > self.replay_history_size) + next_num_accepted_tokens = torch.where( + wrote_checkpoint, accepted_tokens, + prev_num_accepted_tokens + accepted_tokens) + cache_buf_idx = self.mamba_cache.cache_buf_idx[state_indices_d] self.mamba_cache.prev_num_accepted_tokens[state_indices_d] = \ - num_accepted_tokens[num_contexts:num_contexts + num_gens] + next_num_accepted_tokens self.mamba_cache.cache_buf_idx[state_indices_d] = \ - 1 - self.mamba_cache.cache_buf_idx[state_indices_d] + torch.where(wrote_checkpoint, 1 - cache_buf_idx, cache_buf_idx) else: # Legacy: copy accepted SSM state from intermediate cache. ssm_states = self.mamba_cache.temporal @@ -1039,6 +1056,12 @@ def __init__( # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work # on ranks with no local mamba layers. self._use_replay_state_update = use_replay_state_update + self.replay_step_width: Optional[int] = ( + spec_config.max_draft_len + + 1 if spec_config is not None and use_replay_state_update else None) + self.replay_history_size: Optional[int] = (max( + 16, self.replay_step_width) if self.replay_step_width is not None + else None) self.ssm_state_dtype = mamba_ssm_cache_dtype if self.local_num_mamba_layers == 0: @@ -1351,15 +1374,26 @@ def update_mamba_states(self, src_state_indices = self.intermediate_state_indices[:num_gens] if self._use_replay_state_update: - # SSM state is handled incrementally by the replay kernel. Update - # the per-slot accepted-token counter and flip the double-buffer - # index so the next step reads from the buffer that was just - # written by the precompute kernel. - accepted = num_accepted_tokens[num_contexts:num_contexts + num_gens] - self.prev_num_accepted_tokens[state_indices_d] = accepted.to( - self.prev_num_accepted_tokens.dtype) - self.cache_buf_idx[state_indices_d] = \ - 1 - self.cache_buf_idx[state_indices_d] + # SSM state is handled incrementally by the kernel. Mirror the + # kernel's checkpoint predicate from the previous PNAT and fixed + # replay step width: checkpoint steps flip buffers, while no-write + # steps append to the active history. + accepted = num_accepted_tokens[num_contexts:num_contexts + + num_gens].to( + self.prev_num_accepted_tokens. + dtype) + prev_num_accepted_tokens = \ + self.prev_num_accepted_tokens[state_indices_d] + wrote_checkpoint = (prev_num_accepted_tokens + + self.replay_step_width + > self.replay_history_size) + next_num_accepted_tokens = torch.where( + wrote_checkpoint, accepted, prev_num_accepted_tokens + accepted) + cache_buf_idx = self.cache_buf_idx[state_indices_d] + self.prev_num_accepted_tokens[state_indices_d] = \ + next_num_accepted_tokens + self.cache_buf_idx[state_indices_d] = torch.where( + wrote_checkpoint, 1 - cache_buf_idx, cache_buf_idx) else: # Legacy: copy accepted SSM states from intermediate buffer back to pool accepted_ssm = self.intermediate_ssm_states[:, src_state_indices, @@ -1584,7 +1618,7 @@ def _setup_replay_buffers(self, spec_config) -> None: self.old_dA_cumsum = None return - T = spec_config.max_draft_len + 1 + history_size = self.replay_history_size num_local_mamba_layers = self.local_num_mamba_layers # all_ssm_states: [num_local_mamba_layers, num_blocks_in_pool, ...] cache_size = self.all_ssm_states.shape[1] @@ -1602,7 +1636,7 @@ def _setup_replay_buffers(self, spec_config) -> None: # x is not double-buffered self.old_x = torch.zeros(num_local_mamba_layers, cache_size, - T, + history_size, nheads, head_dim, dtype=self.conv_state_dtype, @@ -1611,7 +1645,7 @@ def _setup_replay_buffers(self, spec_config) -> None: self.old_B = torch.zeros(num_local_mamba_layers, cache_size, 2, - T, + history_size, n_groups_per_rank, d_state, dtype=self.conv_state_dtype, @@ -1620,14 +1654,14 @@ def _setup_replay_buffers(self, spec_config) -> None: cache_size, 2, nheads, - T, + history_size, dtype=torch.float32, device=device) self.old_dA_cumsum = torch.zeros(num_local_mamba_layers, cache_size, 2, nheads, - T, + history_size, dtype=torch.float32, device=device) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index d959cdfc2547..f7a4cd14e271 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -23,7 +23,7 @@ skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def _make_mgr(max_batch_size=4, max_draft_len=2): +def _make_mgr(max_batch_size=4, max_draft_len=2, use_replay_state_update=False): # +1 headroom matches MixedMambaHybridCacheManager.pool_size. pool = max_batch_size + 1 return PythonMambaCacheManager( @@ -39,6 +39,7 @@ def _make_mgr(max_batch_size=4, max_draft_len=2): dtype=torch.float16, ssm_cache_dtype=torch.float16, speculative_num_draft_tokens=max_draft_len, + use_replay_state_update=use_replay_state_update, ) @@ -113,6 +114,48 @@ def _fake(rid): assert mgr._padding_slot == shared +@skip_no_cuda +def test_replay_update_mamba_states_uses_history_window(): + """Replay path appends PNAT until the layer kernels checkpointed.""" + mgr = _make_mgr(max_batch_size=4, max_draft_len=5, use_replay_state_update=True) + assert mgr.replay_step_width == 6 + assert mgr.replay_history_size == 16 + assert mgr.mamba_cache.old_x.shape[2] == 16 + assert mgr.mamba_cache.old_B.shape[3] == 16 + assert mgr.mamba_cache.old_dt.shape[4] == 16 + assert mgr.mamba_cache.old_dA_cumsum.shape[4] == 16 + + mgr._prepare_mamba_cache_blocks([100, 101]) + slot_appended = mgr.mamba_cache_index[100] + slot_checkpointed = mgr.mamba_cache_index[101] + + mgr.mamba_cache.prev_num_accepted_tokens[slot_appended] = 7 + mgr.mamba_cache.prev_num_accepted_tokens[slot_checkpointed] = 13 + mgr.mamba_cache.cache_buf_idx[slot_appended] = 0 + mgr.mamba_cache.cache_buf_idx[slot_checkpointed] = 1 + mgr.mamba_cache.conv.zero_() + mgr.mamba_cache.intermediate_conv_window.zero_() + mgr.mamba_cache.intermediate_conv_window[:, 0, 2] = 11.0 + mgr.mamba_cache.intermediate_conv_window[:, 1, 2] = 13.0 + + state_indices = torch.tensor( + [slot_appended, slot_checkpointed], dtype=torch.int32, device="cuda" + ) + attn = SimpleNamespace(num_seqs=2, num_contexts=0) + mgr.update_mamba_states( + attn, + torch.tensor([3, 3], dtype=torch.int32, device="cuda"), + state_indices=state_indices, + ) + + assert mgr.mamba_cache.prev_num_accepted_tokens[slot_appended].item() == 10 + assert mgr.mamba_cache.prev_num_accepted_tokens[slot_checkpointed].item() == 3 + assert mgr.mamba_cache.cache_buf_idx[slot_appended].item() == 0 + assert mgr.mamba_cache.cache_buf_idx[slot_checkpointed].item() == 0 + assert torch.all(mgr.mamba_cache.conv[:, slot_appended] == 11.0) + assert torch.all(mgr.mamba_cache.conv[:, slot_checkpointed] == 13.0) + + @skip_no_cuda def test_update_mamba_states_mtp_path(): """MTP forward path: update_mamba_states must scatter using the From b4c53c542044abc796e043353567b9123bba5059 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 29 Apr 2026 15:22:51 -0700 Subject: [PATCH 08/89] Cache-write semantics: PNAT-aware writes + WRITE_CHECKPOINT/RECTANGLE constexprs Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 512 ++++++++---------- ...benchmark_replay_selective_state_update.py | 111 +++- .../mamba/test_checkpointing_state_update.py | 158 +++++- 3 files changed, 434 insertions(+), 347 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 399cf006ac1d..9a13554aceb7 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -58,10 +58,7 @@ def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: @triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max( - triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics({"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"] + args["T"]), 16)}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) @triton.jit() def _checkpointing_precompute_kernel( # Input pointers @@ -73,20 +70,24 @@ def _checkpointing_precompute_kernel( # Output pointers cb_scaled_ptr, decay_vec_ptr, - # Cache pointers — both READ (for rectangle CB factoring) and WRITE - # (next step's replay) buffers are accessed via stride_*_dbuf offsets. + # Cache pointers (both buffers reachable via stride_*_dbuf). This + # kernel writes to either the active (= cache_buf_idx) or inactive + # (= 1 - cache_buf_idx) buffer depending on WRITE_CHECKPOINT — see + # comment block at top of kernel body. old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - # Double-buffer index (per cache slot) + # Double-buffer index (per cache slot) — selects this step's "active" + # buffer (= where the historical inputs for this step live). cache_buf_idx_ptr, - # Per-request: number of committed tokens since last checkpoint + # Per-request accepted-tokens count (already-cached old tokens at + # [0, PNAT) of the active buffer; new tokens this step go after them + # on no-checkpoint steps). prev_num_accepted_tokens_ptr, state_batch_indices_ptr, pad_slot_id, # Dimensions - T: tl.constexpr, # new tokens this step (T_new) - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # cache buffer T-axis capacity + T: tl.constexpr, dstate: tl.constexpr, nheads_ngroups_ratio: tl.constexpr, # dt strides @@ -136,10 +137,12 @@ def _checkpointing_precompute_kernel( HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, LAUNCH_DEPENDENT_KERNELS: tl.constexpr, HEADS_PER_BLOCK: tl.constexpr, + # Checkpointing flag — selects target buffer + offset for new-token + # cache writes. See "Cache write semantics" block below. + WRITE_CHECKPOINT: tl.constexpr, ): pid_b = tl.program_id(axis=0) pid_hg = tl.program_id(axis=1) # head-group index @@ -161,17 +164,38 @@ def _checkpointing_precompute_kernel( if LAUNCH_DEPENDENT_KERNELS: tl.extra.cuda.gdc_launch_dependents() - # Read buffer index: replay reads from buf_read. We WRITE to 1 - buf_read - # for next step's replay. - buf_read = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - buf_write = 1 - buf_read + # --- Cache write semantics --- + # cache_buf_idx names this step's "active" buffer — the one with the + # historical inputs at [0, PNAT). The other buffer is "staging". + # + # Where do we write new tokens this step? + # WRITE_CHECKPOINT=False (no overflow): append to ACTIVE buffer at + # offset [PNAT : PNAT+T). Caller does NOT flip cache_buf_idx + # afterward; PNAT_next = PNAT + accepted. [0, PNAT) preserved. + # WRITE_CHECKPOINT=True (would overflow): write to STAGING buffer at + # [0, T). Caller flips cache_buf_idx afterward; next step's + # active = the one we just wrote. PNAT_next = accepted. Old + # data in the previous active buffer is folded into state via + # the replay update and discarded. This matches today's replay + # kernel behavior exactly. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if WRITE_CHECKPOINT: + write_buf = 1 - buf_active + write_offset = 0 + else: + write_buf = buf_active + write_offset = prev_num_accepted_tokens - offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis (output rows): T_new positions - offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis (rectangle input cols): up to PNAT+T_new + offs_t = tl.arange(0, BLOCK_SIZE_T) offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) t_mask = offs_t < T n_mask = offs_n < dstate + # Causal mask is shared across all heads (depends only on offs_t) + causal_mask = offs_t[:, None] >= offs_t[None, :] + valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] + # --- Loop 1: compute per-head dt/dA_cumsum/decay BEFORE gdc_wait --- # These only depend on dt (from in_proj, not conv1d) and parameters (A, dt_bias). # Store to cache; will reload after the wait for CB scaling. @@ -190,23 +214,33 @@ def _checkpointing_precompute_kernel( dA_cumsum = tl.cumsum(A * dt, axis=0) decay_vec = tl.exp(dA_cumsum) - # Store dt, dA_cumsum, decay_vec to cache + # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of + # write_buf (selected by WRITE_CHECKPOINT — see top of kernel). old_dt_base = ( old_dt_ptr + cache_batch_idx * stride_old_dt_cache - + buf_write * stride_old_dt_dbuf + + write_buf * stride_old_dt_dbuf + head_idx * stride_old_dt_head ) - tl.store(old_dt_base + offs_t * stride_old_dt_T, dt, mask=t_mask) + tl.store( + old_dt_base + (write_offset + offs_t) * stride_old_dt_T, + dt, + mask=t_mask, + ) old_dA_cumsum_base = ( old_dA_cumsum_ptr + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_write * stride_old_dA_cumsum_dbuf + + write_buf * stride_old_dA_cumsum_dbuf + head_idx * stride_old_dA_cumsum_head ) - tl.store(old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, dA_cumsum, mask=t_mask) + tl.store( + old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, + dA_cumsum, + mask=t_mask, + ) + # decay_vec is per-call scratch (not cached); always write at offs_t. decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + head_idx * stride_dv_head tl.store(decay_vec_base + offs_t * stride_dv_t, decay_vec, mask=t_mask) @@ -215,182 +249,80 @@ def _checkpointing_precompute_kernel( if LAUNCH_WITH_PDL: tl.extra.cuda.gdc_wait() - # --- Rectangle setup: combined k-axis [old | new | pad] of length up to --- - # MAX_REPLAY_BUFFER_LENGTH + T = BLOCK_SIZE_K. is_old_k / is_new_k are - # disjoint masks covering the two halves; safe_k_new clamps the new half's - # source row to [0, T) and safe_old_k clamps the old half's cache row to - # [0, MAX_REPLAY_BUFFER_LENGTH) so masked-off load addresses stay in-bounds. - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - prev_k_idx = tl.minimum( - tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 - ) - is_old_k = offs_k < prev_num_accepted_tokens - safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - prev_num_accepted_tokens - is_new_k = (k_new_idx >= 0) & (k_new_idx < T) - safe_k_new = tl.where(is_new_k, k_new_idx, 0) - - # --- Load C and B (per-group, shared across HEADS_PER_BLOCK heads) --- + # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- group_idx = first_head // nheads_ngroups_ratio C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group - B_new_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group C_all = tl.load( C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, mask=t_mask[:, None] & n_mask[None, :], other=0.0, ) - # New B at unshifted rows [0, T) — used to populate the next step's cache. - B_new_orig = tl.load( - B_new_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + B_all = tl.load( + B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, mask=t_mask[:, None] & n_mask[None, :], other=0.0, ) - # New B at shifted rows [PNAT, PNAT+T) along the k-axis — populates the new - # half of B_combined. Same memory range as B_new_orig; second load hits L1. - B_new_shifted = tl.load( - B_new_base + safe_k_new[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, - mask=is_new_k[:, None] & n_mask[None, :], - other=0.0, - ) - # Old B from READ buffer at k-axis rows [0, PNAT) — populates the old half. - old_B_read_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + buf_read * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - old_B_load = tl.load( - old_B_read_base - + safe_old_k[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - mask=is_old_k[:, None] & n_mask[None, :], - other=0.0, - ) - # B_combined (BLOCK_SIZE_K × dstate): rows [0, PNAT) from old, rows - # [PNAT, PNAT+T) from new, else 0. Disjoint masks → addition gives union. - # Single rectangle matmul replaces today's square C @ B.T. - B_combined = old_B_load + B_new_shifted - raw_rect_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_combined).to(tl.bfloat16)) - # Store new B to cache write buffer (once per group) + # Compute raw CB once — shared across all heads in this block + raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) + + # Store B to cache at [write_offset : write_offset+T) of write_buf. if first_head % nheads_ngroups_ratio == 0: - old_B_write_base = ( + old_B_base = ( old_B_ptr + cache_batch_idx * stride_old_B_cache - + buf_write * stride_old_B_dbuf + + write_buf * stride_old_B_dbuf + group_idx * stride_old_B_group ) tl.store( - old_B_write_base - + offs_t[:, None] * stride_old_B_T + old_B_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - B_new_orig, + B_all, mask=t_mask[:, None] & n_mask[None, :], ) - # Causal mask for the rectangle (BLOCK_SIZE_T × BLOCK_SIZE_K, same across heads): - # k < PNAT (old): always passes — causal in the combined timeline. - # k_new ∈ [0, T) (new): passes iff k_new ≤ t. - t_idx_2d = offs_t[:, None] - k_idx_2d = offs_k[None, :] - is_old_k_2d = k_idx_2d < prev_num_accepted_tokens - k_new_idx_2d = k_idx_2d - prev_num_accepted_tokens - is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) - causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] - - # --- Loop 2: per-head post-scaling of the rectangle CB --- - # Numerical guard: factoring dt[k] * exp(cumAdt_new[t] - cumAdt_at_k[k]) - # into factor_dt[k] * exp_diff[t,k] keeps both sides bounded. Splitting - # further into dt[k]*exp(-cumAdt_at_k[k]) * exp(cumAdt_new[t]) would let - # the new-side `exp(-cumAdt_new[k_new])` overflow at large T (cumAdt grows - # ~T*A*dt) while the t-side `exp(cumAdt_new[t])` underflows, producing - # inf*0 = NaN at causal-valid positions. Computing the SUM of the two - # cumsum offsets before exp keeps the argument ≤ 0 for all valid (t,k). + # --- Loop 2: reload per-head dA_cumsum/dt from cache, scale CB --- + # Reload from where loop 1 stored: write_buf at [write_offset, write_offset+T). for h_local in range(HEADS_PER_BLOCK): head_idx = first_head + h_local - # Reload new dt and dA_cumsum from WRITE buffer (just stored in loop 1). - old_dt_write_base = ( + # Reload dt and dA_cumsum from cache (just written in loop 1) + old_dt_base = ( old_dt_ptr + cache_batch_idx * stride_old_dt_cache - + buf_write * stride_old_dt_dbuf + + write_buf * stride_old_dt_dbuf + head_idx * stride_old_dt_head ) - old_dA_cumsum_write_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_write * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - dA_cumsum_new = tl.load( - old_dA_cumsum_write_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0 + dt = tl.load( + old_dt_base + (write_offset + offs_t) * stride_old_dt_T, + mask=t_mask, + other=0.0, ).to(tl.float32) - # Old-token side: load old dt and dA_cumsum from READ buffer. - old_dt_read_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + buf_read * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - old_dA_cumsum_read_base = ( + old_dA_cumsum_base = ( old_dA_cumsum_ptr + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_read * stride_old_dA_cumsum_dbuf + + write_buf * stride_old_dA_cumsum_dbuf + head_idx * stride_old_dA_cumsum_head ) - # Old-token per-k load: address indexed by safe_old_k (clamped to in-bounds - # for masked-off positions); cache T-axis size = MAX_REPLAY_BUFFER_LENGTH. - old_dt_all = tl.load( - old_dt_read_base + safe_old_k * stride_old_dt_T, mask=is_old_k, other=0.0 - ).to(tl.float32) - old_dA_cumsum_all = tl.load( - old_dA_cumsum_read_base + safe_old_k * stride_old_dA_cumsum_T, - mask=is_old_k, - other=0.0, - ).to(tl.float32) - total_dA_cumsum = tl.load( - old_dA_cumsum_read_base + prev_k_idx * stride_old_dA_cumsum_T - ).to(tl.float32) - - # New-token per-k loads from the WRITE buffer (cache hits L1 from loop 1). - dt_at_kn = tl.load( - old_dt_write_base + safe_k_new * stride_old_dt_T, mask=is_new_k, other=0.0 - ).to(tl.float32) - dA_cumsum_at_kn = tl.load( - old_dA_cumsum_write_base + safe_k_new * stride_old_dA_cumsum_T, - mask=is_new_k, + dA_cumsum = tl.load( + old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, + mask=t_mask, other=0.0, ).to(tl.float32) - # Per-k dt factor (no exp). Disjoint region select. - factor_dt = tl.where(is_old_k, old_dt_all, dt_at_kn) - - # Per-k cumsum offset s_k. For valid k, s_k ≤ 0 for old (cumAdt - # non-increasing → total ≤ old_cumAdt[k]) and s_k ≥ 0 for new - # (-cumAdt_new[k_new] flips the negative cumsum). Bounded ≤ |cumAdt|. - s_k = tl.where( - is_old_k, total_dA_cumsum - old_dA_cumsum_all, -dA_cumsum_at_kn - ) - - # exp_diff[t, k] = exp(cumAdt_new[t] + s_k[k]) - # For valid (t, k_new) with k_new ≤ t (causal): exp arg ≤ 0 → bounded. - # For valid (t, k_old) with k < PNAT: exp arg ≤ 0 → bounded. - # For invalid positions: causal mask zeroes the result downstream. - exp_diff = tl.exp(s_k[None, :] + dA_cumsum_new[:, None]) - - rect_CB_scaled = tl.where( - causal_combined, - raw_rect_CB * factor_dt[None, :] * exp_diff, - 0.0, - ) + # Scale raw_CB with per-head decay and dt + decay_matrix = tl.exp(dA_cumsum[:, None] - dA_cumsum[None, :]) + CB_scaled = tl.where(valid_mask, raw_CB * decay_matrix * dt[None, :], 0.0) cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head tl.store( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, - rect_CB_scaled, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + CB_scaled, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), ) @@ -405,10 +337,7 @@ def _checkpointing_precompute_kernel( ) @triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max( - triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics({"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"] + args["T"]), 16)}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) @triton.jit() def _checkpointing_main_kernel( # Pointers @@ -435,8 +364,7 @@ def _checkpointing_main_kernel( rand_seed_ptr, pad_slot_id, # Dimensions - T: tl.constexpr, # new tokens this step (T_new) - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # cache buffer T-axis capacity + T: tl.constexpr, dim: tl.constexpr, dstate: tl.constexpr, nheads_ngroups_ratio: tl.constexpr, @@ -505,10 +433,15 @@ def _checkpointing_main_kernel( HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, USE_RS_ROUNDING: tl.constexpr, PHILOX_ROUNDS: tl.constexpr, + # Checkpointing flags + WRITE_CHECKPOINT: tl.constexpr, # When True: quantize+write post-replay state to HBM (checkpoint step). + # When False: skip state write entirely (non-checkpoint step). + RECTANGLE: tl.constexpr, # Reserved for the rectangle non-checkpoint optimization path. + # Currently asserted False at the wrapper; kernel takes the + # replay-style code path unconditionally. ): pid_m = tl.program_id(axis=0) pid_b = tl.program_id(axis=1) @@ -521,16 +454,22 @@ def _checkpointing_main_kernel( else: cache_batch_idx = pid_b.to(tl.int64) - # Double-buffer index: buf_read points to the buffer written by LAST step's - # precompute. THIS step's precompute writes to 1-buf_read, which will be - # read by NEXT step's main kernel. Anything not carried between steps is - # single-buffered. - buf_read = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + # Active buffer (= cache_buf_idx) holds the historical inputs for this + # step at [0, PNAT). The replay phase reads from there. The new-tokens + # write target depends on WRITE_CHECKPOINT — see Cache write semantics + # block in the precompute kernel for the full rationale. + active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if WRITE_CHECKPOINT: + write_buf = 1 - active_buf # noqa: F841 — old_x is single-buffered (no use here) + write_offset = 0 + else: + write_buf = active_buf # noqa: F841 — old_x is single-buffered (no use here) + write_offset = prev_num_accepted_tokens offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis: T_new positions - offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis: combined [0, PNAT+T_new) + offs_t = tl.arange(0, BLOCK_SIZE_T) m_mask = offs_m < dim n_mask = offs_n < dstate t_mask = offs_t < T @@ -542,116 +481,105 @@ def _checkpointing_main_kernel( ) state_mask = m_mask[:, None] & n_mask[None, :] state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - # Phase 1: Replay via tl.dot fast-forward (reads from READ buffer). - # Old-token loads are indexed along the K-axis (offs_k, BLOCK_SIZE_K rows - # of which only [0, PNAT) carry data); safe_old_k clamps the masked-off - # row addresses to in-bounds cache positions. + # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) group_idx = pid_h // nheads_ngroups_ratio - is_old_k = offs_k < prev_num_accepted_tokens - safe_old_k = tl.where(is_old_k, offs_k, 0) # Load precomputed dt and dA_cumsum from READ buffer old_dt_base = ( old_dt_ptr + cache_batch_idx * stride_old_dt_cache - + buf_read * stride_old_dt_dbuf + + active_buf * stride_old_dt_dbuf + pid_h * stride_old_dt_head ) - old_dt_all = tl.load( - old_dt_base + safe_old_k * stride_old_dt_T, mask=is_old_k, other=0.0 - ).to(tl.float32) + old_dt_all = tl.load(old_dt_base + offs_t * stride_old_dt_T, mask=t_mask, other=0.0).to( + tl.float32 + ) old_dA_cumsum_base = ( old_dA_cumsum_ptr + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_read * stride_old_dA_cumsum_dbuf + + active_buf * stride_old_dA_cumsum_dbuf + pid_h * stride_old_dA_cumsum_head ) old_dA_cumsum_all = tl.load( - old_dA_cumsum_base + safe_old_k * stride_old_dA_cumsum_T, mask=is_old_k, other=0.0 + old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0 ).to(tl.float32) - # Total cumsum at end of old: scalar load at PNAT-1, clamped against the - # cache T-axis (= MAX_REPLAY_BUFFER_LENGTH) for OOB safety on PNAT=0. - prev_k_idx = tl.minimum( - tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 - ) + # Load dA_cumsum at prev_k-1 directly via pointer math (avoids masked reduction). + # Clamp to [0, T-1] defensively — out-of-contract PNAT > T would read OOB. + prev_k_idx = tl.minimum(tl.maximum(prev_num_accepted_tokens - 1, 0), T - 1) total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( tl.float32 ) # Step 0 invariant: PNAT=0 means `state` is already last step's state (not - # two back). coeff is all-zero (is_old_k all-false), total_decay is 1.0, - # so the replay leaves `state` unchanged — cache contents don't matter. + # two back). coeff is all-zero (offs_t < 0), total_decay is 1.0, so the + # replay leaves `state` unchanged — cache contents don't matter on step 0. coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all + coeff = tl.where(offs_t < prev_num_accepted_tokens, coeff, 0.0) - # Load old_x: (BLOCK_SIZE_K, BLOCK_SIZE_M) — single-buffered cache + # Load old_x: (BLOCK_SIZE_T, BLOCK_SIZE_M) — single-buffered old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head old_x_all = tl.load( - old_x_base + safe_old_k[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - mask=is_old_k[:, None] & m_mask[None, :], + old_x_base + offs_t[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, ) - # Load old_B from READ buffer: (BLOCK_SIZE_K, BLOCK_SIZE_DSTATE) + # Load old_B from READ buffer: (BLOCK_SIZE_T, BLOCK_SIZE_DSTATE) old_B_base = ( old_B_ptr + cache_batch_idx * stride_old_B_cache - + buf_read * stride_old_B_dbuf + + active_buf * stride_old_B_dbuf + group_idx * stride_old_B_group ) old_B_all = tl.load( - old_B_base - + safe_old_k[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - mask=is_old_k[:, None] & n_mask[None, :], + old_B_base + offs_t[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], other=0.0, ).to(tl.float32) - # Scale B by coefficients (k-axis) + # Scale B by coefficients dB_scaled = coeff[:, None] * old_B_all - # Apply total decay to initial state. Save the decayed-but-pre-replay - # state for the output (state_out uses C @ state_prev_decayed.T because - # the rectangle CB carries the old-tokens contribution to output through - # the token_out matmul). + # Apply total decay to initial state FIRST, then add contributions total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) - state_prev_decayed = state * total_decay - - # tl.dot fast-forward: old_x^T @ dB_scaled → (M, dstate); add to decayed - # prior state to get the post-replay (checkpointed) state. - state = state_prev_decayed + tl.dot( - tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16) - ) - - # Write post-replay state - if USE_RS_ROUNDING: - # Stochastic rounding for fp16 state using Philox-4x32 PRNG. - # Each Philox call produces 4 random ints. We call randint4x on - # quarter-sized dstate offsets and join+reshape to get the full - # (M, dstate) random tensor — 4x fewer PRNG rounds. - rand_seed = tl.load(rand_seed_ptr) - base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head - offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // 4) - rand_offsets_q = ( - base_rand - + offs_m[:, None] * stride_state_dim - + offs_n_q[None, :] * (stride_state_dstate * 4) - ) # (M, dstate//4) - if PHILOX_ROUNDS > 0: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + state *= total_decay + + # tl.dot fast-forward: old_x^T @ dB_scaled → (M, dstate) + state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) + + # Write post-replay state — only on checkpoint steps. When + # WRITE_CHECKPOINT is False, the replay computed `state` is local-only and + # discarded; skipping the HBM store + Philox path is the main performance + # win of replay-style checkpointing on the common (non-checkpoint) step. + if WRITE_CHECKPOINT: + if USE_RS_ROUNDING: + # Stochastic rounding for fp16 state using Philox-4x32 PRNG. + # Each Philox call produces 4 random ints. We call randint4x on + # quarter-sized dstate offsets and join+reshape to get the full + # (M, dstate) random tensor — 4x fewer PRNG rounds. + rand_seed = tl.load(rand_seed_ptr) + base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head + offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // 4) + rand_offsets_q = ( + base_rand + + offs_m[:, None] * stride_state_dim + + offs_n_q[None, :] * (stride_state_dstate * 4) + ) # (M, dstate//4) + if PHILOX_ROUNDS > 0: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + else: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) + # Interleave 4 quarter-sized tensors → full (M, dstate) random tensor + r01 = tl.join(r0, r1) # (M, dstate//4, 2) + r23 = tl.join(r2, r3) # (M, dstate//4, 2) + r0123 = tl.join(r01, r23) # (M, dstate//4, 2, 2) + rand = tl.reshape(r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + tl.store(state_ptrs, _stochastic_round_fp16x2(state, rand), mask=state_mask) else: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) - # Interleave 4 quarter-sized tensors → full (M, dstate) random tensor - r01 = tl.join(r0, r1) # (M, dstate//4, 2) - r23 = tl.join(r2, r3) # (M, dstate//4, 2) - r0123 = tl.join(r01, r23) # (M, dstate//4, 2, 2) - rand = tl.reshape(r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) - tl.store(state_ptrs, _stochastic_round_fp16x2(state, rand), mask=state_mask) - else: - tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=state_mask) + tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=state_mask) # Phase 2: Output using precomputed CB_scaled and decay_vec x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head @@ -686,60 +614,40 @@ def _checkpointing_main_kernel( mask=t_mask[:, None] & m_mask[None, :], other=0.0, ) - # Store new x to cache (single-buffered; replay already read the old data) + # Store new x to old_x cache at [write_offset : write_offset+T). + # old_x is single-buffered: write goes to the active buffer regardless; + # replay already read positions [0, PNAT) so write_offset = PNAT (no + # overlap) on no-checkpoint steps. On checkpoint steps write_offset = 0 + # (cache reset; old data folded into state via replay update). tl.store( - old_x_base + offs_t[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + old_x_base + + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, x_all, mask=t_mask[:, None] & m_mask[None, :], ) x_all = x_all.to(tl.float32) - # --- Build x_combined for the rectangle token_out matmul --- - # K-axis layout: rows [0, PNAT) from old_x (already loaded; is_old_k mask), - # rows [PNAT, PNAT+T) from new_x via a shifted load over the same memory - # as `x_all` (hits L1). Disjoint masks → safe to add. - k_new_idx = offs_k - prev_num_accepted_tokens - is_new_k = (k_new_idx >= 0) & (k_new_idx < T) - safe_k_new = tl.where(is_new_k, k_new_idx, 0) - - old_x_in_combined = tl.where( - is_old_k[:, None] & m_mask[None, :], - old_x_all.to(tl.float32), - 0.0, - ) - new_x_shifted = tl.load( - x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=is_new_k[:, None] & m_mask[None, :], - other=0.0, - ).to(tl.float32) - x_combined = old_x_in_combined + new_x_shifted # (BLOCK_SIZE_K, BLOCK_SIZE_M) - - # Load precomputed rectangle CB_scaled (BLOCK_SIZE_T × BLOCK_SIZE_K) and - # decay_vec_new (= exp(cumAdt_new), BLOCK_SIZE_T). + # Load precomputed CB_scaled and decay_vec cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), other=0.0, ).to(tl.float32) decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec_new = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( tl.float32 ) - # state_out: pre-replay state contribution to output. - # (C[t] · state_prev) * total_decay * decay_vec_new[t] - # Folded: state_prev_decayed already = state_prev * total_decay. - state_out = ( - tl.dot(C_all.to(tl.bfloat16), tl.trans(state_prev_decayed).to(tl.bfloat16)) - * decay_vec_new[:, None] - ) + # init_out = C_all @ state^T * decay_vec + init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] - # token_out: combined old+new tokens contribution via the rectangle CB. - token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) + # cb_out = CB_scaled @ x_all + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) - out_all = state_out + token_out + out_all = init_out + cb_out if HAS_D: out_all = out_all + x_all * D[None, :] @@ -784,6 +692,8 @@ def checkpointing_state_update( philox_rounds: int = 10, launch_with_pdl=False, use_internal_pdl=True, + write_checkpoint: bool = True, + rectangle: bool = False, _block_size_m: int | None = None, _num_warps: int | None = None, _num_stages: int | None = None, @@ -853,6 +763,20 @@ def checkpointing_state_update( launch_with_pdl = False use_internal_pdl = False + # Constexpr modes: + # write_checkpoint=True, rectangle=False → checkpoint step (default). + # write_checkpoint=False, rectangle=False → non-checkpoint step (skip state HBM write). + # write_checkpoint=False, rectangle=True → reserved for the rectangle non-checkpoint + # optimization path (not wired yet). + # write_checkpoint=True, rectangle=True → not supported. + if rectangle: + raise NotImplementedError( + "RECTANGLE path is not wired yet; pass rectangle=False." + ) + assert not (write_checkpoint and rectangle), ( + "WRITE_CHECKPOINT and RECTANGLE are mutually exclusive." + ) + # --- Unsqueeze inputs to canonical shapes --- if state.dim() == 3: state = state.unsqueeze(1) @@ -893,11 +817,20 @@ def checkpointing_state_update( ngroups = B.shape[2] assert nheads % ngroups == 0 - # Cache T-axis is the replay-buffer capacity (committed-history slots + - # provisional-draft slots). In the placeholder build, this matches T_new. - max_replay_buffer_length = old_x.shape[1] - assert T <= max_replay_buffer_length, ( - f"x has T_new={T} > max_replay_buffer_length={max_replay_buffer_length}" + # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the + # placeholder degenerate case max_window = T (every step is a checkpoint + # step). For real replay-style checkpointing, max_window > T and + # `prev_num_accepted_tokens` can be 0..max_window. + max_window = old_x.shape[1] + assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" + # Replay-style code path uses BLOCK_SIZE_T = max(np2(T), 16) for the + # combined T-axis (T_new tile size) and reuses it for window loads. Until + # the heuristic is generalized to track max_window separately, require + # max_window to fit within that tile. + block_size_t = max(triton.next_power_of_2(T), 16) + assert max_window <= block_size_t, ( + f"max_window={max_window} exceeds BLOCK_SIZE_T={block_size_t} " + f"derived from T={T}; extend the heuristic to include max_window." ) assert x.shape == (batch, T, nheads, dim) @@ -905,10 +838,10 @@ def checkpointing_state_update( assert A.shape == (nheads, dim, dstate) assert B.shape == (batch, T, ngroups, dstate) assert C.shape == B.shape - assert old_x.shape == (cache_size, max_replay_buffer_length, nheads, dim) - assert old_B.shape == (cache_size, 2, max_replay_buffer_length, ngroups, dstate) - assert old_dt.shape == (cache_size, 2, nheads, max_replay_buffer_length) - assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_replay_buffer_length) + assert old_x.shape == (cache_size, max_window, nheads, dim) + assert old_B.shape == (cache_size, 2, max_window, ngroups, dstate) + assert old_dt.shape == (cache_size, 2, nheads, max_window) + assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_window) assert cache_buf_idx.shape == (cache_size,) assert prev_num_accepted_tokens.shape == (cache_size,) @@ -921,15 +854,11 @@ def checkpointing_state_update( assert tie_hdim device = x.device - # BLOCK_SIZE_T sizes the T-axis (T_new positions); BLOCK_SIZE_K sizes the - # rectangle K-axis (combined [old | new] up to PNAT+T_new ≤ MAX+T_new). - # Decoupling them keeps T-axis matmuls compact at np2(T) when MAX > T. BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) - BLOCK_SIZE_K = max(triton.next_power_of_2(max_replay_buffer_length + T), 16) # Allocate precomputed intermediates (per-call, not cached) cb_scaled = torch.empty( - batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 + batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_T, device=device, dtype=torch.float32 ) decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) @@ -1075,7 +1004,6 @@ def checkpointing_state_update( state_batch_indices, pad_slot_id, T, - max_replay_buffer_length, dstate, nheads // ngroups, # dt strides @@ -1124,6 +1052,7 @@ def checkpointing_state_update( LAUNCH_WITH_PDL=launch_with_pdl, LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, HEADS_PER_BLOCK=heads_per_block, + WRITE_CHECKPOINT=write_checkpoint, num_warps=precompute_num_warps, **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), launch_pdl=launch_with_pdl, @@ -1152,7 +1081,6 @@ def grid(META): rand_seed, pad_slot_id, T, - max_replay_buffer_length, dim, dstate, nheads // ngroups, @@ -1216,6 +1144,8 @@ def grid(META): BLOCK_SIZE_M, LAUNCH_WITH_PDL=use_internal_pdl, PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + WRITE_CHECKPOINT=write_checkpoint, + RECTANGLE=rectangle, num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), launch_pdl=use_internal_pdl, diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index aec41c318228..d6ba608cd4d9 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -205,6 +205,27 @@ def _flush_l2() -> None: torch.cuda.synchronize() +def _resolve_prev_ks(args, mtp_len: int) -> list[int]: + """Resolve prev_k values for one mtp_len cell. + + Two input modes (mutually exclusive in spirit; absolute wins if both given): + --prev-tokens-int "0,10,11,16" → use literal integers, clamped to + [0, max_window] (where max_window is the cache T-axis capacity). + --prev-tokens-fracs "0,0.5,1.0" → fractions of mtp_len, clamped to + [0, mtp_len] (current behavior). + + For replay-style checkpointing the cache holds up to max_window old + tokens, so absolute integers are the right knob. Fractions are kept + for back-compat with prior placeholder runs. + """ + upper = getattr(args, "max_window", 0) or mtp_len + if getattr(args, "prev_tokens_int", None): + return sorted(set(max(0, min(upper, int(v))) for v in args.prev_tokens_int)) + return sorted( + set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs) + ) + + # Tensor construction helpers @@ -217,6 +238,7 @@ def _build_tensors( head_dim: int, d_state: int, ngroups: int, + max_window: int | None = None, ): """ Build all tensors for one benchmark configuration. @@ -250,14 +272,18 @@ def _build_tensors( state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=state_dtype) # --- Cache tensors for replay kernel --- - # old_x: single-buffered (cache, T, nheads, dim) - old_x = torch.randn(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - # old_B: double-buffered (cache, 2, T, ngroups, dstate) - old_B = torch.randn(batch, 2, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) - # old_dt: double-buffered (cache, 2, nheads, T) fp32 — T contiguous - old_dt = torch.randn(batch, 2, nheads, mtp_len, device=device, dtype=torch.float32) - # old_dA_cumsum: double-buffered (cache, 2, nheads, T) fp32 — T contiguous - old_dA_cumsum = torch.randn(batch, 2, nheads, mtp_len, device=device, dtype=torch.float32) + # max_window is the cache T-axis capacity; defaults to mtp_len (the + # placeholder/degenerate case where every step is a checkpoint step). + # For real replay-style checkpointing, max_window > mtp_len. + cache_T = max_window if max_window is not None else mtp_len + # old_x: single-buffered (cache, max_window, nheads, dim) + old_x = torch.randn(batch, cache_T, nheads, head_dim, device=device, dtype=act_dtype) + # old_B: double-buffered (cache, 2, max_window, ngroups, dstate) + old_B = torch.randn(batch, 2, cache_T, ngroups, d_state, device=device, dtype=act_dtype) + # old_dt: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous + old_dt = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) + # old_dA_cumsum: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous + old_dA_cumsum = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) # cache_buf_idx: which buffer to read (0 or 1) cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) @@ -438,8 +464,8 @@ def _time_kernel(args, run_fn, reset_fn, tag: str) -> tuple[float, float, float] # Per-config benchmark (consolidated baseline + replay) -def _parallel_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, - baseline_fn, max_workers: int) -> None: +def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, + baseline_fn, max_workers: int) -> None: """Run each config once in parallel to compile + cache Triton kernels. Triton's ``compile()`` releases the GIL, so a ThreadPoolExecutor @@ -456,14 +482,12 @@ def _parallel_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dty configs = [] for batch in batch_sizes: for mtp_len in mtp_lengths: - prev_ks = sorted( - set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs) - ) + prev_ks = _resolve_prev_ks(args, mtp_len) for state_dtype in state_dtypes: for act_dtype in act_dtypes: configs.append((batch, mtp_len, prev_ks, state_dtype, act_dtype)) - print(f"[parallel-warmup] {len(configs)} configs across {max_workers} threads") + print(f"[compile-warmup] {len(configs)} configs across {max_workers} threads") t0 = time.perf_counter() def _warm(cfg): @@ -484,11 +508,11 @@ def _warm(cfg): if errors: for cfg, e in errors: - print(f"[parallel-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", + print(f"[compile-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", file=sys.stderr) raise errors[0][1] - print(f"[parallel-warmup] done in {time.perf_counter() - t0:.1f}s") + print(f"[compile-warmup] done in {time.perf_counter() - t0:.1f}s") def _bench_config( @@ -549,6 +573,7 @@ def _bench_config( args.head_dim, args.d_state, args.tp_ngroups, + max_window=getattr(args, "max_window", None) or None, ) nheads = args.tp_nheads @@ -754,6 +779,10 @@ def _run_incr( else: x_call, B_call, C_call = x, B, C extra_kwargs = {} + # write_checkpoint is only meaningful for the checkpointing + # variant; replay variant ignores the kwarg. + if args.variant == "checkpointing": + extra_kwargs["write_checkpoint"] = args.write_checkpoint variant_fn( state_work, old_x_work, @@ -880,10 +909,10 @@ def _run_benchmark(args) -> None: elif args.l2_flush: _init_l2_flush() - if args.parallel_warmup > 0: - _parallel_warmup_phase( + if args.compile_threads > 0: + _compile_warmup_phase( args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, - baseline_fn, max_workers=args.parallel_warmup, + baseline_fn, max_workers=args.compile_threads, ) if args.profile: @@ -913,9 +942,7 @@ def _run_benchmark(args) -> None: for batch in batch_sizes: for mtp_len in mtp_lengths: # Resolve prev_k fractions → clamped integers in [0, mtp_len] - prev_ks = sorted( - set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs) - ) + prev_ks = _resolve_prev_ks(args, mtp_len) for state_dtype in state_dtypes: for act_dtype in act_dtypes: _bench_config( @@ -977,14 +1004,14 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--warmup", type=int, default=20, help="Number of warmup iterations") parser.add_argument("--iters", type=int, default=100, help="Number of timed iterations") parser.add_argument( - "--parallel-warmup", + "--compile-threads", type=int, - default=0, - help="Run a parallel-warmup phase that calls each (batch, mtp_len, " - "prev_k, dtype, sweep) config once across N threads before the " - "sequential timed phase. Triton compile releases the GIL so threads " - "compile in parallel, populating the persistent cache for free hits " - "during measurement. 0 disables the phase. Try 8 on a multi-core box.", + default=64, + help="Number of THREADS used in the compile-warmup phase (one call " + "per (batch, mtp_len, prev_k, dtype, sweep) cell, parallelized over " + "N threads). Triton compile releases the GIL, so threads compile " + "in parallel and populate the persistent cache for free hits during " + "the sequential timed phase. 0 disables the phase. Default 64.", ) parser.add_argument( "--profile", @@ -1066,6 +1093,32 @@ def _parse_args() -> argparse.Namespace: default=None, help="Override num_stages for precompute kernel (comma-separated sweep).", ) + parser.add_argument( + "--max-window", + type=int, + default=0, + help="Cache T-axis capacity (max replay buffer length). 0 (default) " + "= use mtp_len, the placeholder/degenerate every-step-checkpoint " + "case. Set to e.g. 16 for real replay-style checkpointing on " + "Nemotron-3-Super-120B.", + ) + parser.add_argument( + "--prev-tokens-int", + type=lambda s: [int(x) for x in s.split(",")] if s else None, + default=None, + help="Absolute prev_num_accepted_tokens values to test, comma-separated " + "(e.g. '0,10,11,16'). Clamped to [0, max_window]. When set, " + "overrides --prev-tokens-fracs.", + ) + parser.add_argument( + "--write-checkpoint", + action=argparse.BooleanOptionalAction, + default=True, + help="Whether the checkpointing kernel should write the post-replay " + "state to HBM. True = checkpoint step (default). False = " + "non-checkpoint step (skip state HBM write + Philox). No effect on " + "the replay variant.", + ) parser.add_argument( "--with-conv1d", action="store_true", diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index 674717048230..259f29103ae3 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -16,6 +16,7 @@ import pytest import torch import torch.nn.functional as F +import triton from einops import repeat from tensorrt_llm._torch.modules.mamba.checkpointing_state_update import ( @@ -46,8 +47,11 @@ @pytest.mark.parametrize( "T", [6, 10, 16, 27, 32, 55], ids=["T6", "T10", "T16", "T27", "T32", "T55"] ) +@pytest.mark.parametrize( + "write_checkpoint", [True, False], ids=["write", "no_write"] +) def test_checkpointing_state_update( - nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T + nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T, write_checkpoint ): """ Verify that: @@ -61,6 +65,12 @@ def test_checkpointing_state_update( dtype = torch.bfloat16 # input activations are bf16 assert nheads % ngroups == 0 + # Cache T-axis size (max_window). Use the kernel's BLOCK_SIZE_T as the + # ceiling — this is what the wrapper allows and enables PNAT-aware writes + # at [PNAT, PNAT+T) for no-checkpoint mode. For T=6 that's 16 (production + # max_window); for larger T it scales with np2(T). + max_window = max(triton.next_power_of_2(T), 16) + if paged_cache: cache_size = 4 state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) @@ -120,37 +130,43 @@ def test_checkpointing_state_update( ) # Build cache tensors for the replay kernel. - # old_x: (cache, T, nheads, dim) bf16 — single-buffered - # old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered - # old_dt: (cache, 2, nheads, T) fp32 — double-buffered, T contiguous - # old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered, T contiguous + # old_x: (cache, max_window, nheads, dim) bf16 — single-buffered + # old_B: (cache, 2, max_window, ngroups, dstate) bf16 — double-buffered + # old_dt: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous + # old_dA_cumsum: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous # cache_buf_idx: random 0s and 1s to verify indexing correctness - old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_x = torch.zeros(cache_size, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) - # Fill each slot's READ buffer (indexed by cache_buf_idx) with step 1's data. - # The OTHER buffer has random garbage to catch indexing bugs. + # Fill each slot's active buffer (= cache_buf_idx) with step 1's data at + # positions [0:T). Positions [T:max_window) stay as torch.randn garbage — + # they're outside the test's PNAT range, kernel doesn't read them. + # The OTHER (inactive) buffer has random garbage to catch indexing bugs. slots = state_batch_indices if paged_cache else slice(None) - old_x[slots] = x1 + old_x[slots, :T] = x1 # Compute processed dt and dA_cumsum for step 1 dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) - # Write to each slot's read buffer based on its cache_buf_idx + # Write to each slot's active buffer based on its cache_buf_idx slot_indices = state_batch_indices.tolist() if paged_cache else list(range(cache_size)) for i, slot in enumerate(slot_indices): buf = cache_buf_idx[slot].item() batch_idx = i # maps slot back to the batch index - old_B[slot, buf] = B1[batch_idx] - old_dt[slot, buf] = dt1[batch_idx].T # (T, nheads) → (nheads, T) - old_dA_cumsum[slot, buf] = dA_cumsum1[batch_idx].T # (T, nheads) → (nheads, T) + old_B[slot, buf, :T] = B1[batch_idx] + old_dt[slot, buf, :, :T] = dt1[batch_idx].T # (T, nheads) → (nheads, T) + old_dA_cumsum[slot, buf, :, :T] = dA_cumsum1[batch_idx].T # (T, nheads) → (nheads, T) - # Main loop: test each k (number of old tokens replayed) + # Main loop: test each k (number of old tokens replayed). For + # write_checkpoint=False, the kernel writes new tokens at [k:k+T) of the + # active buffer — skip k where this would exceed max_window. for k in range(T + 1): + if not write_checkpoint and k + T > max_window: + continue torch.manual_seed(k + 100) x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) @@ -179,18 +195,23 @@ def test_checkpointing_state_update( out=ref_out, ) - # Replay kernel + # Replay kernel — clone caches into mutable working copies that we + # can inspect AFTER the call to verify cache postconditions. test_state = state0.clone() prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_x_w = old_x.clone() + old_B_w = old_B.clone() + old_dt_w = old_dt.clone() + old_dA_cumsum_w = old_dA_cumsum.clone() # cache_buf_idx stays at its random values — each slot reads from its own buffer checkpointing_state_update( test_state, - old_x.clone(), - old_B.clone(), - old_dt.clone(), - old_dA_cumsum.clone(), + old_x_w, + old_B_w, + old_dt_w, + old_dA_cumsum_w, cache_buf_idx.clone(), prev_tokens, x=x2, @@ -203,6 +224,7 @@ def test_checkpointing_state_update( dt_bias=dt_bias, dt_softplus=True, state_batch_indices=state_batch_indices, + write_checkpoint=write_checkpoint, ) # Tolerance rationale: the replay kernel uses bf16 tl.dot for four @@ -230,9 +252,17 @@ def test_checkpointing_state_update( ) raise - expected_state = ( - state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) - ) + # State expectation depends on write_checkpoint: + # True → kernel writes the post-replay state; expect the + # selective_state_update reference's state at step k-1. + # False → kernel skips the HBM store; state must be UNCHANGED + # from the input (state0). + if write_checkpoint: + expected_state = ( + state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) + ) + else: + expected_state = state0[slots] state_diff = (test_state[slots].float() - expected_state.float()).abs() state_max = state_diff.max().item() state_mean = state_diff.mean().item() @@ -241,8 +271,8 @@ def test_checkpointing_state_update( test_state[slots], expected_state, rtol=2e-2, - atol=1.0, - msg=f"State mismatch at k={k}", + atol=1.0 if write_checkpoint else 0.0, + msg=f"State mismatch at k={k} (write_checkpoint={write_checkpoint})", ) except AssertionError: print( @@ -252,6 +282,80 @@ def test_checkpointing_state_update( ) raise + # --- Cache postconditions --- + # Compute step 2's processed values (what the kernel should have + # stored at [write_offset : write_offset+T) of write_buf): + # write_buf = (1 - active_buf) if write_checkpoint else active_buf + # write_offset = 0 if write_checkpoint else k + # Untouched cache regions must equal their pre-call snapshots + # (old_x / old_B / old_dt / old_dA_cumsum captured before the call). + dt2_proc = F.softplus(dt2_base.float() + dt_bias_base.float()[None, None, :]) # (B,T,H) + dA_cumsum2 = torch.cumsum(A_base.float()[None, None, :] * dt2_proc, dim=1) + write_offset = 0 if write_checkpoint else k + + for batch_idx, slot in enumerate(slot_indices): + active = cache_buf_idx[slot].item() + wb = (1 - active) if write_checkpoint else active + + # --- old_x (single-buffered): write at [write_offset : +T) of slot --- + written_x = old_x_w[slot, write_offset : write_offset + T] + torch.testing.assert_close( + written_x, x2[batch_idx], rtol=0, atol=0, + msg=f"old_x written region wrong at k={k} write={write_checkpoint}", + ) + # Untouched ranges of old_x[slot] + if write_offset > 0: + torch.testing.assert_close( + old_x_w[slot, :write_offset], old_x[slot, :write_offset], + rtol=0, atol=0, + msg=f"old_x [0:{write_offset}) modified at k={k} write={write_checkpoint}", + ) + if write_offset + T < max_window: + torch.testing.assert_close( + old_x_w[slot, write_offset + T:], old_x[slot, write_offset + T:], + rtol=0, atol=0, + msg=f"old_x [{write_offset+T}:) modified at k={k} write={write_checkpoint}", + ) + + # --- old_B (double-buffered): write at write_buf, [write_offset:+T) --- + torch.testing.assert_close( + old_B_w[slot, wb, write_offset : write_offset + T], + B2[batch_idx], rtol=0, atol=0, + msg=f"old_B written region wrong at k={k} write={write_checkpoint}", + ) + # Other-buffer (= 1-wb) untouched + torch.testing.assert_close( + old_B_w[slot, 1 - wb], old_B[slot, 1 - wb], + rtol=0, atol=0, + msg=f"old_B inactive buffer modified at k={k} write={write_checkpoint}", + ) + + # --- old_dt (double-buffered, fp32, layout (heads, T)): --- + torch.testing.assert_close( + old_dt_w[slot, wb, :, write_offset : write_offset + T], + dt2_proc[batch_idx].T, + rtol=1e-4, atol=1e-4, + msg=f"old_dt written region wrong at k={k} write={write_checkpoint}", + ) + torch.testing.assert_close( + old_dt_w[slot, 1 - wb], old_dt[slot, 1 - wb], + rtol=0, atol=0, + msg=f"old_dt inactive buffer modified at k={k} write={write_checkpoint}", + ) + + # --- old_dA_cumsum (double-buffered, fp32, layout (heads, T)): --- + torch.testing.assert_close( + old_dA_cumsum_w[slot, wb, :, write_offset : write_offset + T], + dA_cumsum2[batch_idx].T, + rtol=1e-4, atol=1e-4, + msg=f"old_dA_cumsum written region wrong at k={k} write={write_checkpoint}", + ) + torch.testing.assert_close( + old_dA_cumsum_w[slot, 1 - wb], old_dA_cumsum[slot, 1 - wb], + rtol=0, atol=0, + msg=f"old_dA_cumsum inactive buf modified at k={k} write={write_checkpoint}", + ) + @_skip_pre_sm100 @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) From a9dff37d3a6787d48d7e082a461c9178f56f7100 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 18 May 2026 12:26:19 -0700 Subject: [PATCH 09/89] Add computation of n_writes and the nwrite-partition metadata vector to trtllm frame work. still need to update mamba2_mixer to pass them, and the kernels to use them. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../_torch/modules/mamba/mamba2_metadata.py | 67 +++++++++++++++++++ .../_torch/modules/mamba/mamba2_mixer.py | 6 ++ .../_torch/pyexecutor/mamba_cache_manager.py | 41 ++++++++++++ .../modules/mamba/test_mamba2_metadata.py | 59 ++++++++++++++++ 4 files changed, 173 insertions(+) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index beedeccab165..677492b74eed 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -25,6 +25,12 @@ CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm._utils import prefer_pinned +REPLAY_WORK_POSITION_IN_DECODE_BATCH = 0 +REPLAY_WORK_CACHE_SLOT = 1 +REPLAY_WORK_PNAT = 2 +REPLAY_WORK_CACHE_BUF_IDX = 3 +REPLAY_WORK_ITEM_WIDTH = 4 + @triton.jit def _cu_seqlens_triton_kernel( @@ -214,6 +220,16 @@ def __init__(self, max_batch_size: int, chunk_size: int): dtype=torch.int32, device="cuda") + self.replay_work_items = torch.zeros( + max_batch_size, + REPLAY_WORK_ITEM_WIDTH, + dtype=torch.int32, + device="cuda") + self.replay_n_writes = torch.zeros(1, + dtype=torch.int32, + device="cuda") + self.replay_num_decodes = 0 + # Pre-allocated buffers. self._arange_buffer = torch.arange(max_batch_size + 1, dtype=torch.int, @@ -223,6 +239,54 @@ def __init__(self, max_batch_size: int, chunk_size: int): dtype=torch.long, device="cuda") + def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, + num_contexts: int): + self.replay_num_decodes = 0 + if not getattr(kv_cache_manager, 'use_replay_state_update', False): + return + if not hasattr(kv_cache_manager, 'get_replay_state_update_metadata'): + return + + replay_metadata = kv_cache_manager.get_replay_state_update_metadata() + if replay_metadata is None: + return + self.replay_n_writes.zero_() + + prev_num_accepted_tokens, cache_buf_idx, replay_step_width, \ + replay_history_size = replay_metadata + num_decodes = batch_size - num_contexts + self.replay_num_decodes = num_decodes + if num_decodes == 0: + return + + position_in_decode_batch = torch.arange(num_decodes, + dtype=torch.int32, + device=self.state_indices.device) + cache_slot = self.state_indices[num_contexts:batch_size] + cache_slot_idx = cache_slot.to(torch.long) + pnat = prev_num_accepted_tokens[cache_slot_idx].to(torch.int32) + active_cache_buf_idx = cache_buf_idx[cache_slot_idx].to(torch.int32) + + writes = (pnat + replay_step_width > replay_history_size) + writes_i32 = writes.to(torch.int32) + write_offsets = torch.cumsum(writes_i32, dim=0) - writes_i32 + n_writes = torch.sum(writes_i32, dim=0, + keepdim=True).to(torch.int32) + no_write_offsets = position_in_decode_batch - write_offsets + output_offsets = torch.where(writes, write_offsets, + n_writes + no_write_offsets) + output_offsets = output_offsets.to(torch.long) + + work_items = self.replay_work_items[:num_decodes] + work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH].scatter_( + 0, output_offsets, position_in_decode_batch) + work_items[:, REPLAY_WORK_CACHE_SLOT].scatter_(0, output_offsets, + cache_slot) + work_items[:, REPLAY_WORK_PNAT].scatter_(0, output_offsets, pnat) + work_items[:, REPLAY_WORK_CACHE_BUF_IDX].scatter_( + 0, output_offsets, active_cache_buf_idx) + self.replay_n_writes.copy_(n_writes) + def prepare(self, attn_metadata: AttentionMetadata): batch_size = attn_metadata.seq_lens.shape[0] num_contexts = attn_metadata.num_contexts @@ -247,6 +311,9 @@ def prepare(self, attn_metadata: AttentionMetadata): self.state_indices[:batch_size].copy_( self.state_indices_cpu[:batch_size], non_blocking=True) + self._prepare_replay_work_items(kv_cache_manager, batch_size, + num_contexts) + if num_contexts > 0: torch.cumsum(context_lens, dim=0, diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 9f43720a98c7..12626bdbfc13 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -514,6 +514,12 @@ def convert_dt(): philox_kwargs['philox_rounds'] = self._philox_rounds if use_replay: + # TODO: persistent-main replay should consume + # attn_metadata.mamba_metadata.replay_work_items and + # replay_n_writes. Each work item carries + # position_in_decode_batch, cache_slot, PNAT, and + # cache_buf_idx so the kernel can avoid repeated scalar + # pointer chasing. replay_selective_state_update( ssm_states, layer_cache.old_x, diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index a40ea37f472a..9efda8abd536 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -74,6 +74,12 @@ def get_state_indices(self, *args, **kwargs) -> torch.Tensor: """ ... + def get_replay_state_update_metadata( + self + ) -> Optional[tuple[torch.Tensor, torch.Tensor, int, int]]: + """Return replay metadata tensors and fixed replay sizes.""" + return None + @abstractmethod def get_conv_states(self, layer_idx: int) -> torch.Tensor: """Return conv states for specific layer. @@ -603,6 +609,20 @@ def get_mamba_ssm_cache_dtype(self) -> torch.dtype: def use_replay_state_update(self) -> bool: return self._use_replay_state_update + def get_replay_state_update_metadata( + self + ) -> Optional[tuple[torch.Tensor, torch.Tensor, int, int]]: + if (not self._use_replay_state_update + or not isinstance(self.mamba_cache, self.SpeculativeState) + or self.mamba_cache.prev_num_accepted_tokens is None + or self.mamba_cache.cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return (self.mamba_cache.prev_num_accepted_tokens, + self.mamba_cache.cache_buf_idx, self.replay_step_width, + self.replay_history_size) + def shutdown(self): """Release tensor memory.""" # Clear mamba cache states @@ -794,6 +814,15 @@ def get_mamba_ssm_cache_dtype(self) -> torch.dtype: def use_replay_state_update(self) -> bool: return getattr(self._impl, 'use_replay_state_update', False) + def get_replay_state_update_metadata( + self + ) -> Optional[tuple[torch.Tensor, torch.Tensor, int, int]]: + get_metadata = getattr(self._impl, 'get_replay_state_update_metadata', + None) + if get_metadata is None: + return None + return get_metadata() + def get_intermediate_ssm_states(self, layer_idx: int) -> Optional[torch.Tensor]: assert not self._use_cpp, "get_intermediate_ssm_states is not supported in CppMambaCacheManager" @@ -1669,5 +1698,17 @@ def _setup_replay_buffers(self, spec_config) -> None: def use_replay_state_update(self) -> bool: return self._use_replay_state_update + def get_replay_state_update_metadata( + self + ) -> Optional[tuple[torch.Tensor, torch.Tensor, int, int]]: + if (not self._use_replay_state_update + or self.prev_num_accepted_tokens is None + or self.cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return (self.prev_num_accepted_tokens, self.cache_buf_idx, + self.replay_step_width, self.replay_history_size) + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: return self.ssm_state_dtype diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index c44e0d7f9f87..9e88699be875 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -21,6 +21,10 @@ from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( Mamba2Metadata, + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, cu_seqlens_to_chunk_indices_offsets, cu_seqlens_to_chunk_indices_offsets_triton, ) @@ -89,6 +93,61 @@ def test_prepare_handles_tensor_cached_tokens(self): assert metadata.chunk_indices is not None assert metadata.chunk_offsets is not None + def test_prepare_replay_work_items_write_first(self): + class ReplayCacheManager: + use_replay_state_update = True + + def __init__(self): + self.state_indices = [0, 3, 1, 4, 2] + self.prev_num_accepted_tokens = torch.tensor( + [0, 4, 10, 11, 20], dtype=torch.int32, device="cuda") + self.cache_buf_idx = torch.tensor([0, 1, 0, 1, 0], + dtype=torch.int32, + device="cuda") + + def get_state_indices(self, request_ids, is_padding): + return self.state_indices[:len(request_ids)] + + def get_replay_state_update_metadata(self): + return (self.prev_num_accepted_tokens, self.cache_buf_idx, 6, + 16) + + metadata = Mamba2Metadata(max_batch_size=5, chunk_size=8) + seq_lens = torch.tensor([2, 7, 7, 7, 7], dtype=torch.int) + attn_metadata = SimpleNamespace( + seq_lens=seq_lens, + seq_lens_cuda=seq_lens.cuda(), + num_contexts=1, + num_ctx_tokens=2, + kv_cache_manager=ReplayCacheManager(), + request_ids=[10, 11, 12, 13, 14], + kv_cache_params=SimpleNamespace( + num_cached_tokens_per_seq=torch.tensor([0], + dtype=torch.int), + ), + ) + + metadata.prepare(attn_metadata) + + expected = torch.tensor( + [ + [0, 3, 11, 1], + [2, 4, 20, 0], + [1, 1, 4, 1], + [3, 2, 10, 0], + ], + dtype=torch.int32, + device="cuda", + ) + actual = metadata.replay_work_items[:4] + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(metadata.replay_n_writes.cpu(), + torch.tensor([2], dtype=torch.int32)) + assert actual[0, REPLAY_WORK_POSITION_IN_DECODE_BATCH] == 0 + assert actual[0, REPLAY_WORK_CACHE_SLOT] == 3 + assert actual[0, REPLAY_WORK_PNAT] == 11 + assert actual[0, REPLAY_WORK_CACHE_BUF_IDX] == 1 + def test_single_sequence_unaligned(self): """Test with a single sequence that doesn't align with chunk size.""" cu_seqlens = torch.tensor([0, 10], dtype=torch.int, device="cuda") From 93bf1d5ca23dc7a055e42317a6afd5b270514f6a Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Thu, 30 Apr 2026 09:57:20 -0700 Subject: [PATCH 10/89] Add int8/int16/fp8 quant state to checkpointing kernel + benchmark/test plumbing Kernel (checkpointing_state_update.py): - Quant state path with QUANT_MAX constexpr (0=non-quant, 127/32767/448 for int8/int16/fp8_e4m3fn). Per-(cache, head, dim) decode_scale tensor (fp32, broadcast over dstate) for dequant on load and rescale on store. - fp8 stochastic-rounding PTX helper: cvt.rs.satfinite.e4m3x4.f32 with reversed source-register order (load-bearing per pack=4 little-endian). - state_q rebind: separate variable for the quantized-on-store value so the dequantized fp32 state stays alive for the output phase (fixes the output-uses-quantized-state bug). - SM gating: fp8_e4m3fn requires SM 89+ (fp32 <-> fp8 cvt PTX); fp16 SR and fp8 SR additionally require SM 100+ (cvt.rs.* family). - maxnreg and num_ctas Triton config args plumbed through wrapper as _-prefixed override kwargs. Tests (test_checkpointing_state_update.py): - int8/int16/fp8 folded into the normal test_checkpointing_state_update parametrize, across all T, k, write_checkpoint, paged_cache cells. - test_checkpointing_state_update_philox: quant dtypes added with per-channel-ULP-aware comparison (replaces flat atol that flaked on per-channel scale variation). - test_philox_rounding_unbiased: K*SE statistic auto-calibrates per dtype residual noise (replaces fixed 1e-5 threshold which was below SE for int8/fp8 and would always fail by chance). - test_sr_grid_bracket: parametrized over fp8/fp16; catches the PTX source-register byte-order trap. - Tolerances empirically derived (audit agent + 30-seed stress test); bf16-baseline rationale verified to the digit. Benchmark (benchmark_replay_selective_state_update.py): - --philox-rounds CLI flag, default 5 (matches Nemotron-3-Super-120B production configs at examples/models/core/nemotron and tests/integration). - --maxnreg and --num-ctas sweep flags plumbed to the kernel call. - Silent baseline-skip for incompatible (baseline, dtype, SR) combos instead of failing the whole sweep with ValueError. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 236 ++++++- ...benchmark_replay_selective_state_update.py | 151 ++++- .../mamba/test_checkpointing_state_update.py | 621 +++++++++++++++--- 3 files changed, 876 insertions(+), 132 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 9a13554aceb7..648d0442a446 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -51,6 +51,32 @@ def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: ) +@triton.jit +def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 quad → fp8 e4m3 using Philox random bits. + + Uses PTX cvt.rs.satfinite.e4m3x4.f32 which combines stochastic rounding + and saturating cast in a single op (output is final fp8, no separate + clamp needed). The reversed source-register order {$4,$3,$2,$1} is + load-bearing — PTX packs leftmost source into the high byte but Triton's + pack=4 is little-endian, so the natural {$1,$2,$3,$4} order would + silently shuffle every group of 4 contiguous outputs. + + Requires SM_100a+ (Blackwell B200). Caller must gate at the wrapper + level — this kernel does not check. + + Adapted from vLLM PR #40012 (Apache-2.0). + """ + return tl.inline_asm_elementwise( + asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", + constraints="=r,r,r,r,r,r,r,r,r", + args=(x, rand), + dtype=tl.float8e4nv, + is_pure=True, + pack=4, + ) + + # Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, # old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. # Grid: (batch, nheads // HEADS_PER_BLOCK). @@ -342,6 +368,9 @@ def _checkpointing_precompute_kernel( def _checkpointing_main_kernel( # Pointers state_ptr, + # Per-(cache, head, dim) decode scale, fp32, only consulted when QUANT_MAX>0. + # Layout (cache, nheads, dim) — broadcast over dstate at load/store. + state_scales_ptr, # Cache READ pointers (read-buffer from previous step) old_x_ptr, old_B_ptr, @@ -373,6 +402,10 @@ def _checkpointing_main_kernel( stride_state_head, stride_state_dim, stride_state_dstate, + # state_scales strides: (cache, nheads, dim) — only used when QUANT_MAX>0 + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, # old_x strides: (cache, T, nheads, dim) — single-buffered stride_old_x_cache, stride_old_x_T, @@ -436,6 +469,12 @@ def _checkpointing_main_kernel( LAUNCH_WITH_PDL: tl.constexpr, USE_RS_ROUNDING: tl.constexpr, PHILOX_ROUNDS: tl.constexpr, + # State quantization: 0.0 means non-quantized (fp16/bf16/fp32); >0 means + # quantized (int8=127, int16=32767, fp8_e4m3fn=448). Single in-kernel + # switch for the dequant-on-load and encode-on-store paths. Wrapper sets + # this from state.dtype; kernel-entry static_assert below pins the + # invariant that it must coincide with int8/int16/float8e4nv state dtype. + QUANT_MAX: tl.constexpr, # Checkpointing flags WRITE_CHECKPOINT: tl.constexpr, # When True: quantize+write post-replay state to HBM (checkpoint step). # When False: skip state write entirely (non-checkpoint step). @@ -443,6 +482,19 @@ def _checkpointing_main_kernel( # Currently asserted False at the wrapper; kernel takes the # replay-style code path unconditionally. ): + # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized + # state dtype (int8 / int16 / float8e4nv) and only those. Cheap + # insurance against a wrapper bug that desynchronizes the two. + tl.static_assert( + (QUANT_MAX > 0.0) + == ( + (state_ptr.dtype.element_ty == tl.int8) + or (state_ptr.dtype.element_ty == tl.int16) + or (state_ptr.dtype.element_ty == tl.float8e4nv) + ), + "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", + ) + pid_m = tl.program_id(axis=0) pid_b = tl.program_id(axis=1) pid_h = tl.program_id(axis=2) @@ -481,6 +533,20 @@ def _checkpointing_main_kernel( ) state_mask = m_mask[:, None] & n_mask[None, :] state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + # Dequantize on load (per-(head, dim) decode scale, broadcast over dstate). + # Only consulted when QUANT_MAX>0 — non-quantized paths skip entirely. + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, + other=1.0, + ).to(tl.float32) + state = state * decode_scale[:, None] # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) group_idx = pid_h // nheads_ngroups_ratio @@ -556,10 +622,10 @@ def _checkpointing_main_kernel( # win of replay-style checkpointing on the common (non-checkpoint) step. if WRITE_CHECKPOINT: if USE_RS_ROUNDING: - # Stochastic rounding for fp16 state using Philox-4x32 PRNG. - # Each Philox call produces 4 random ints. We call randint4x on - # quarter-sized dstate offsets and join+reshape to get the full - # (M, dstate) random tensor — 4x fewer PRNG rounds. + # Generate (M, dstate) random tensor for stochastic rounding. + # Used by fp16 / int8 / int16 / fp8 SR paths below. Quarter-sized + # randint4x calls produce 4 u32s each; we interleave them into the + # full (M, dstate) tensor — 4x fewer PRNG rounds vs per-element. rand_seed = tl.load(rand_seed_ptr) base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // 4) @@ -577,8 +643,78 @@ def _checkpointing_main_kernel( r23 = tl.join(r2, r3) # (M, dstate//4, 2) r0123 = tl.join(r01, r23) # (M, dstate//4, 2, 2) rand = tl.reshape(r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + + if QUANT_MAX > 0.0: + # Quantized state path: int8 / int16 / fp8_e4m3fn (RN or SR). + # 1) Per-(head, dim) channel scale via amax over dstate. + amax = tl.max(tl.abs(state), axis=1) # (M,) + encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) # (M,) + decode_scale = 1.0 / encode_scale # (M,) + # 2) Store decode_scale (1/encode) so reads do a single multiply. + state_scales_ptrs = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + + offs_m * stride_state_scales_dim + ) + tl.store(state_scales_ptrs, decode_scale, mask=m_mask) + # 3) Scale state into quant range — into a NEW variable so the + # downstream output phase still sees the dequantized fp32 state. + state_q = state * encode_scale[:, None] + # 4) Round per dtype. Order matters: handle fp8 SR first (PTX + # combines round + saturating cast in one op, output is final fp8 + # so we store and finish on that branch). Other branches share + # the clamp + cast tail below. + if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): + # fp8_e4m3fn + SR — PTX cvt.rs.satfinite.e4m3x4.f32. Output + # is final fp8 (saturate included); store directly. + tl.store( + state_ptrs, + _stochastic_round_fp8x4_e4m3(state_q, rand), + mask=state_mask, + ) + else: + if USE_RS_ROUNDING: + # int8 / int16 + SR — uniform-noise + floor. + # (fp8 SR was handled by the early branch above.) + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized SR fall-through expects int8 or int16; " + "fp8 SR is handled by the prior branch.", + ) + rand01 = (rand & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) + state_q = tl.extra.cuda.libdevice.floor(state_q + rand01) + elif state_ptrs.dtype.element_ty != tl.float8e4nv: + # int8 / int16 + RN — explicit round before clamp. + # fp8 + RN deliberately skips this — explicit round() would + # destroy fp8 sub-integer precision; native cast at store + # does RN at the fp8 grid resolution. + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized RN with explicit round() expects int8 or int16.", + ) + state_q = tl.extra.cuda.libdevice.round(state_q) + # Clamp + cast tail: int8/int16 (RN+SR) and fp8 RN. + # fp8 RN reaches here without prior round() — .to(float8e4nv) + # does native RN at the fp8 grid. + state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) + tl.store( + state_ptrs, + state_q.to(state_ptrs.dtype.element_ty), + mask=state_mask, + ) + elif USE_RS_ROUNDING: + # Non-quantized + SR: only fp16 (bf16 has no PTX SR cast; fp32 + # doesn't need rounding). + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "Non-quantized SR only supports fp16 state.", + ) tl.store(state_ptrs, _stochastic_round_fp16x2(state, rand), mask=state_mask) else: + # Non-quantized + RN: fp16 / bf16 / fp32 native cast. tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=state_mask) # Phase 2: Output using precomputed CB_scaled and decay_vec @@ -668,6 +804,13 @@ def _checkpointing_main_kernel( # Python wrapper +_QUANT_MAX_BY_DTYPE = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, +} + + def checkpointing_state_update( state: torch.Tensor, old_x: torch.Tensor, @@ -690,6 +833,7 @@ def checkpointing_state_update( pad_slot_id: int = PAD_SLOT_ID, rand_seed: torch.Tensor | None = None, philox_rounds: int = 10, + state_scales: torch.Tensor | None = None, launch_with_pdl=False, use_internal_pdl=True, write_checkpoint: bool = True, @@ -700,6 +844,8 @@ def checkpointing_state_update( _precompute_num_warps: int | None = None, _precompute_num_stages: int | None = None, _heads_per_block: int | None = None, + _maxnreg: int | None = None, + _num_ctas: int | None = None, ): """ Replay SSM state update with precomputed CB and tl.dot fast-forward. @@ -743,9 +889,16 @@ def checkpointing_state_update( dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). state_batch_indices: (batch,) optional cache slot mapping. rand_seed: optional single-element int64 CUDA tensor for Philox PRNG seed. - When provided, state is stochastically rounded to fp16 on store. - When None, standard deterministic rounding is used. + When provided, state is stochastically rounded on store. Supported + for state.dtype in (fp16, int8, int16, fp8_e4m3fn); other dtypes + silently use deterministic rounding. fp16+SR and fp8+SR both + require sm_100a (Blackwell B200+) — wrapper asserts this loudly. philox_rounds: number of Philox PRNG rounds (default 10). + state_scales: required when state.dtype in (int8, int16, fp8_e4m3fn). + Shape (cache_size, nheads, dim), fp32. Per-(head, dim) channel + decode scale (= 1 / encode_scale). The kernel writes scales on + checkpoint steps and reads them on load (broadcast over dstate). + Ignored for non-quantized state dtypes. launch_with_pdl: enable external PDL (conv1d → precompute chain). Defaults False; caller opts in when the upstream chain is PDL-safe. Ignored on hardware that doesn't support PDL (sm < 90). @@ -754,9 +907,9 @@ def checkpointing_state_update( Ignored on hardware that doesn't support PDL (sm < 90). _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, - _precompute_num_warps, _precompute_num_stages, _heads_per_block) are - benchmark-only overrides; production callers should leave them None - to use the heuristic-tuned defaults. + _precompute_num_warps, _precompute_num_stages, _heads_per_block, + _maxnreg, _num_ctas) are benchmark-only overrides; production callers + should leave them None to use the heuristic-tuned defaults. """ # PDL needs sm >= 90. if get_sm_version() < 90: @@ -777,6 +930,31 @@ def checkpointing_state_update( "WRITE_CHECKPOINT and RECTANGLE are mutually exclusive." ) + # --- Hardware support gates --- + # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX + # instructions (Ada Lovelace introduced them; Hopper/Blackwell carry them). + if state.dtype == torch.float8_e4m3fn: + assert get_sm_version() >= 89, ( + "fp8_e4m3fn state requires SM 89+ (Ada Lovelace / Hopper / Blackwell) " + f"for fp32↔fp8 cvt PTX instructions; current SM is {get_sm_version()}." + ) + + # PTX cvt.rs.* (stochastic rounding) family lands on Blackwell only. + # Wrapper fails loud; framework decides fall-back (e.g. drop SR, use RN). + # int8 / int16 SR uses pure-Triton libdevice.floor + uniform noise — no + # PTX SR instruction needed, runs anywhere. + if rand_seed is not None: + if state.dtype == torch.float16: + assert get_sm_version() >= 100, ( + "fp16 stochastic rounding (PTX cvt.rs.f16x2.f32) requires " + f"sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + ) + elif state.dtype == torch.float8_e4m3fn: + assert get_sm_version() >= 100, ( + "fp8 stochastic rounding (PTX cvt.rs.satfinite.e4m3x4.f32) " + f"requires sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + ) + # --- Unsqueeze inputs to canonical shapes --- if state.dim() == 3: state = state.unsqueeze(1) @@ -817,6 +995,25 @@ def checkpointing_state_update( ngroups = B.shape[2] assert nheads % ngroups == 0 + # --- Quantization plumbing --- + # QUANT_MAX > 0 ⇔ state is int8 / int16 / fp8_e4m3fn. Kernel-entry + # static_assert on the Triton side mirrors this invariant. + quant_max = _QUANT_MAX_BY_DTYPE.get(state.dtype, 0.0) + is_quantized = quant_max > 0.0 + if is_quantized: + assert state_scales is not None, ( + f"state.dtype={state.dtype} requires state_scales tensor " + "(shape (cache_size, nheads, dim), fp32)." + ) + assert state_scales.shape == (cache_size, nheads, dim), ( + f"state_scales shape mismatch: expected {(cache_size, nheads, dim)}, " + f"got {state_scales.shape}." + ) + assert state_scales.dtype == torch.float32, ( + f"state_scales must be fp32, got {state_scales.dtype}." + ) + assert state_scales.device == state.device + # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the # placeholder degenerate case max_window = T (every step is a checkpoint # step). For real replay-style checkpointing, max_window > T and @@ -1062,8 +1259,22 @@ def checkpointing_state_update( def grid(META): return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + # state_scales pointer + strides: real tensor when quantized, otherwise + # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). + if is_quantized: + state_scales_arg = state_scales + state_scales_strides = ( + state_scales.stride(0), + state_scales.stride(1), + state_scales.stride(2), + ) + else: + state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 + state_scales_strides = (0, 0, 0) + _checkpointing_main_kernel[grid]( state, + state_scales_arg, old_x, old_B, old_dt, @@ -1089,6 +1300,10 @@ def grid(META): state.stride(1), state.stride(2), state.stride(3), + # state_scales strides (cache, head, dim) + state_scales_strides[0], + state_scales_strides[1], + state_scales_strides[2], # old_x strides (single-buffered: cache, T, nheads, dim) old_x.stride(0), old_x.stride(1), @@ -1144,9 +1359,12 @@ def grid(META): BLOCK_SIZE_M, LAUNCH_WITH_PDL=use_internal_pdl, PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, WRITE_CHECKPOINT=write_checkpoint, RECTANGLE=rectangle, num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), launch_pdl=use_internal_pdl, ) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index d6ba608cd4d9..3321aa90674b 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -269,7 +269,34 @@ def _build_tensors( D = repeat(D_base, "h -> h p", p=head_dim) # --- SSM state --- - state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=state_dtype) + # Quantized dtypes need their own initializer (torch.randn doesn't accept + # int) and a parallel fp32 scales tensor (per-(head, dim) channel decode + # scale, broadcast over dstate). Quant state is filled with realistic- + # range values via fp32 → quant; scales are derived consistently so the + # initial state isn't garbage on dequant. + _QUANT_BENCH = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, + } + if state_dtype in _QUANT_BENCH: + quant_max = _QUANT_BENCH[state_dtype] + state_fp32 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + amax = state_fp32.abs().amax(dim=-1) # (batch, nheads, head_dim) + encode_scale = quant_max / amax.clamp(min=1e-30) + state_scales0 = (1.0 / encode_scale).to(torch.float32) # decode scale + scaled = state_fp32 * encode_scale.unsqueeze(-1) + if state_dtype == torch.float8_e4m3fn: + state0 = scaled.clamp(-quant_max, quant_max).to(state_dtype) + else: + state0 = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) + else: + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state_scales0 = None # --- Cache tensors for replay kernel --- # max_window is the cache T-axis capacity; defaults to mtp_len (the @@ -301,8 +328,12 @@ def _build_tensors( out_incr = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) out_base = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) + # intermediate_states_buffer is only consumed by the fp/baseline path; + # for quantized state dtypes we'll skip baselines entirely, so the buffer + # dtype falls back to fp32 to keep selective_state_update happy. + int_buffer_dtype = state_dtype if state_dtype not in _QUANT_BENCH else torch.float32 intermediate_states_buffer = torch.zeros( - batch, mtp_len, nheads, head_dim, d_state, device=device, dtype=state_dtype + batch, mtp_len, nheads, head_dim, d_state, device=device, dtype=int_buffer_dtype ) # --- Conv1d tensors (for --with-conv1d mode) --- @@ -328,6 +359,7 @@ def _build_tensors( return ( state0, + state_scales0, old_x, old_B, old_dt, @@ -542,6 +574,7 @@ def _bench_config( ( state0, + state_scales0, old_x0, old_B0, old_dt0, @@ -584,19 +617,26 @@ def _bench_config( use_philox = getattr(args, "philox_rounding", False) variant_fn = _VARIANT_FNS[args.variant]() - # Philox rounding: allocate rand_seed tensor + # Philox rounding: allow fp16 and the quantized dtypes (int8/int16/fp8). + # bf16/fp32 SR is not supported (no PTX path for bf16; fp32 doesn't need + # rounding). This validates the user's input — it doesn't apply to the + # baseline (incompatible baselines are silently skipped below). rand_seed = None + _SR_SUPPORTED = ( + torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn, + ) if use_philox: - if state_dtype != torch.float16: - raise ValueError(f"--philox-rounding requires --state-dtypes fp16, got {state_dtype}") - if args.baseline == "triton": + if state_dtype not in _SR_SUPPORTED: raise ValueError( - "--philox-rounding not supported with --baseline triton " - "(only flashinfer and replay support it)" + f"--philox-rounding requires state dtype in {{fp16, int8, int16, fp8}}, " + f"got {state_dtype}" ) rand_seed = torch.randint(0, 2**62, (1,), device="cuda", dtype=torch.int64) + is_quantized = state_dtype in (torch.int8, torch.int16, torch.float8_e4m3fn) + state_work = state0.clone() + state_scales_work = state_scales0.clone() if state_scales0 is not None else None old_x_work = old_x0.clone() old_B_work = old_B0.clone() old_dt_work = old_dt0.clone() @@ -607,6 +647,8 @@ def _bench_config( def _reset(): state_work.copy_(state0) + if state_scales_work is not None: + state_scales_work.copy_(state_scales0) old_x_work.copy_(old_x0) old_B_work.copy_(old_B0) old_dt_work.copy_(old_dt0) @@ -619,6 +661,8 @@ def _reset_conv1d_realistic(): """Realistic reset: cold cache, L2 flush, then hot in_proj output.""" # 1. Reset cold state (cache tensors, SSM state) state_work.copy_(state0) + if state_scales_work is not None: + state_scales_work.copy_(state_scales0) old_x_work.copy_(old_x0) old_B_work.copy_(old_B0) old_dt_work.copy_(old_dt0) @@ -631,6 +675,33 @@ def _reset_conv1d_realistic(): # 3. Write hot tensors (simulates in_proj output landing in L2) xbc_input_work.copy_(xbc_input0) + # Silently skip the baseline row for any (baseline, state_dtype, SR) + # combo it can't run. Better than erroring on a partial sweep — our + # kernel rows still print. Compatibility: + # * Quantized states (int8 / int16 / fp8): no baseline supports them. + # * Triton baseline (selective_state_update): no rand_seed kwarg. + # * flashinfer baseline: rand_seed only on fp16 state. + def _baseline_supports() -> bool: + if baseline_fn is None: + return False + if is_quantized: + return False + if use_philox: + if args.baseline == "triton": + return False + if args.baseline == "flashinfer" and state_dtype != torch.float16: + return False + return True + + if baseline_fn is not None and not _baseline_supports(): + if not warmup_only: + sr_tag = " + SR" if use_philox else "" + print( + f"# Skipping {args.baseline} baseline for " + f"state_dtype={state_dtype_name}{sr_tag} (unsupported)." + ) + baseline_fn = None + show_kernel_col = baseline_fn is not None def _conv1d_split(xbc_in, conv_st, launch_dependent_kernels=False): @@ -665,7 +736,7 @@ def _conv1d_split(xbc_in, conv_st, launch_dependent_kernels=False): philox_kwargs = {} if rand_seed is not None and args.baseline == "flashinfer": - philox_kwargs = {"rand_seed": rand_seed, "philox_rounds": 10} + philox_kwargs = {"rand_seed": rand_seed, "philox_rounds": args.philox_rounds} if with_conv1d: @@ -740,6 +811,8 @@ def _parse_sweep(val): precompute_num_warps_values = _parse_sweep(args.precompute_num_warps) precompute_num_stages_values = _parse_sweep(args.precompute_num_stages) heads_per_block_values = _parse_sweep(args.heads_per_block) + maxnreg_values = _parse_sweep(args.maxnreg) + num_ctas_values = _parse_sweep(args.num_ctas) # --- Replay kernel, one row per prev_k --- for prev_k in prev_ks: @@ -753,6 +826,8 @@ def _parse_sweep(val): precompute_num_warps, precompute_num_stages, heads_per_block, + maxnreg, + num_ctas, ) in itertools.product( block_size_m_values, num_warps_values, @@ -760,6 +835,8 @@ def _parse_sweep(val): precompute_num_warps_values, precompute_num_stages_values, heads_per_block_values, + maxnreg_values, + num_ctas_values, ): def _run_incr( @@ -770,6 +847,8 @@ def _run_incr( precompute_num_warps=precompute_num_warps, precompute_num_stages=precompute_num_stages, heads_per_block=heads_per_block, + maxnreg=maxnreg, + num_ctas=num_ctas, ): if with_conv1d: x_call, B_call, C_call = _conv1d_split( @@ -780,9 +859,12 @@ def _run_incr( x_call, B_call, C_call = x, B, C extra_kwargs = {} # write_checkpoint is only meaningful for the checkpointing - # variant; replay variant ignores the kwarg. + # variant; replay variant ignores the kwarg. state_scales + # is also checkpointing-only (replay kernel doesn't quantize). if args.variant == "checkpointing": extra_kwargs["write_checkpoint"] = args.write_checkpoint + if state_scales_work is not None: + extra_kwargs["state_scales"] = state_scales_work variant_fn( state_work, old_x_work, @@ -802,6 +884,7 @@ def _run_incr( dt_softplus=True, state_batch_indices=None, rand_seed=rand_seed, + philox_rounds=args.philox_rounds, use_internal_pdl=args.internal_pdl, _block_size_m=block_size_m, _num_warps=num_warps, @@ -809,6 +892,8 @@ def _run_incr( _precompute_num_warps=precompute_num_warps, _precompute_num_stages=precompute_num_stages, _heads_per_block=heads_per_block, + _maxnreg=maxnreg, + _num_ctas=num_ctas, **extra_kwargs, ) @@ -825,6 +910,10 @@ def _run_incr( parts.append(f"pS={precompute_num_stages}") if heads_per_block is not None: parts.append(f"H={heads_per_block}") + if maxnreg is not None: + parts.append(f"R={maxnreg}") + if num_ctas is not None: + parts.append(f"CT={num_ctas}") sweep_suffix = (" " + ",".join(parts)) if parts else "" sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") @@ -889,7 +978,14 @@ def _run_benchmark(args) -> None: batch_sizes = [int(x) for x in args.batch_sizes.split(",")] mtp_lengths = [int(x) for x in args.mtp_lengths.split(",")] - dtype_map = {"bf16": torch.bfloat16, "fp32": torch.float32, "fp16": torch.float16} + dtype_map = { + "bf16": torch.bfloat16, + "fp32": torch.float32, + "fp16": torch.float16, + "int8": torch.int8, + "int16": torch.int16, + "fp8": torch.float8_e4m3fn, + } state_dtypes = [dtype_map[s] for s in args.state_dtypes.split(",")] act_dtypes = [dtype_map[s] for s in args.act_dtypes.split(",")] @@ -994,7 +1090,11 @@ def _parse_args() -> argparse.Namespace: help="Comma-separated per-request sequence lengths (num_draft_tokens + 1 target)", ) parser.add_argument( - "--state-dtypes", default="fp32", help="Comma-separated state dtypes: fp16,bf16,fp32" + "--state-dtypes", + default="fp32", + help="Comma-separated state dtypes: fp16,bf16,fp32,int8,int16,fp8. " + "Quantized dtypes (int8/int16/fp8) require the checkpointing variant " + "and skip baselines (selective_state_update doesn't accept them).", ) parser.add_argument( "--act-dtypes", @@ -1139,11 +1239,34 @@ def _parse_args() -> argparse.Namespace: default=None, help="Override HEADS_PER_BLOCK for precompute kernel (comma-separated sweep).", ) + parser.add_argument( + "--maxnreg", + type=str, + default=None, + help="Override maxnreg for the main kernel (comma-separated sweep).", + ) + parser.add_argument( + "--num-ctas", + type=str, + default=None, + help="Override num_ctas for the main kernel (comma-separated sweep).", + ) parser.add_argument( "--philox-rounding", action="store_true", - help="Enable Philox stochastic rounding for fp16 state " - "(rand_seed generated per iteration, philox_rounds=10).", + help="Enable Philox stochastic rounding (rand_seed generated per " + "iteration). Supported state dtypes: fp16, int8, int16, fp8. " + "fp16 SR and fp8 SR require sm_100a (Blackwell B200+).", + ) + parser.add_argument( + "--philox-rounds", + type=int, + default=5, + help="Number of Philox PRNG rounds. Default 5 matches the " + "Nemotron-3-Super-120B production config (mamba_ssm_philox_rounds=5 " + "in examples/configs and tests/integration/perf configs). The " + "wrapper's generic fallback default is 10; callers without explicit " + "config see 10. Only consulted when --philox-rounding is enabled.", ) parser.add_argument( "--variant", diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index 259f29103ae3..a3871ee6d4f4 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -17,6 +17,7 @@ import torch import torch.nn.functional as F import triton +import triton.language as tl from einops import repeat from tensorrt_llm._torch.modules.mamba.checkpointing_state_update import ( @@ -40,9 +41,60 @@ (32, 64, 128, 2), # TP=4, ngroups>1 (more heads than B/C groups) ] +# Quantized state dtypes and their representable-magnitude limits (== QUANT_MAX +# in the kernel). fp8_e4m3fn cells require SM 89+ for the fp32↔fp8 cvt PTX +# instructions; SR variants of fp16/fp8 additionally need SM 100+. +_QUANT_MAX_BY_DTYPE = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, +} + + +def _quantize_state(state_fp32: torch.Tensor, state_dtype: torch.dtype, quant_max: float): + """Quantize fp32 state to (state_quant, decode_scale) using the same + per-(head, dim) channel scheme the kernel does on store. decode_scale = + max_abs_per_channel / quant_max (= 1/encode_scale). + """ + amax = state_fp32.abs().amax(dim=-1) # (cache, nheads, head_dim) + encode_scale = quant_max / amax.clamp(min=1e-30) + decode_scale = 1.0 / encode_scale + scaled = state_fp32 * encode_scale.unsqueeze(-1) + if state_dtype == torch.float8_e4m3fn: + # Native cast does RN at the fp8 grid; explicit round() would destroy + # sub-integer precision (matches the kernel's fp8 RN path). + state_quant = scaled.clamp(-quant_max, quant_max).to(state_dtype) + else: + state_quant = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) + return state_quant, decode_scale + + +def _dequantize_state(state_quant: torch.Tensor, decode_scale: torch.Tensor): + return state_quant.to(torch.float32) * decode_scale.unsqueeze(-1) + + +def _maybe_skip_dtype(state_dtype, use_sr): + """Skip on insufficient SM. fp8 e4m3fn (any) needs SM 89+; fp16/fp8 SR + needs SM 100+; int8/int16 (RN or SR) runs anywhere.""" + if state_dtype == torch.float8_e4m3fn and get_sm_version() < 89: + pytest.skip("fp8_e4m3fn requires SM 89+ (Ada Lovelace / Hopper / Blackwell)") + if use_sr and state_dtype in (torch.float16, torch.float8_e4m3fn) and get_sm_version() < 100: + pytest.skip(f"{state_dtype} stochastic rounding requires SM 100+ (Blackwell B200+)") + @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) -@pytest.mark.parametrize("state_dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize( + "state_dtype", + [ + torch.float16, + torch.bfloat16, + torch.float32, + torch.int8, + torch.int16, + torch.float8_e4m3fn, + ], + ids=["fp16", "bf16", "fp32", "int8", "int16", "fp8"], +) @pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) @pytest.mark.parametrize( "T", [6, 10, 16, 27, 32, 55], ids=["T6", "T10", "T16", "T27", "T32", "T55"] @@ -59,7 +111,16 @@ def test_checkpointing_state_update( produces the same output as: selective_state_update(state_after_k_old_tokens, new_x, ...) and writes state_after_k_old_tokens back to the state tensor. + + Quantized state dtypes (int8/int16/fp8) follow the same flow with + a per-(head, dim) channel decode-scale tensor; comparison is done + via dequant(state, scales) against the fp32 reference. """ + _maybe_skip_dtype(state_dtype, use_sr=False) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + batch = 2 device = "cuda" dtype = torch.bfloat16 # input activations are bf16 @@ -92,8 +153,24 @@ def test_checkpointing_state_update( D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) - # Initial SSM state (cache_size slots) - state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + # Initial SSM state (cache_size slots). Quantized dtypes need a separate + # init: derive scales from a fp32 source so the quantized state isn't + # garbage on dequant. ref_input_state is what the fp32 reference run + # sees — for non-quant it's state0 (cast to fp32 inside reference); for + # quant it's the lossy dequant of state0 (matches what the kernel sees + # internally on load). + if is_quantized: + state0_fp32 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + state0, state0_scales = _quantize_state(state0_fp32, state_dtype, quant_max) + ref_input_state = _dequantize_state(state0, state0_scales) + else: + state0 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state0_scales = None + ref_input_state = state0.float() # Old inputs: T tokens per batch request x1 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) @@ -113,7 +190,7 @@ def test_checkpointing_state_update( ) out1 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) selective_state_update( - state0.clone(), + ref_input_state.clone(), x1, dt1, A, @@ -175,8 +252,9 @@ def test_checkpointing_state_update( B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - # Reference - ref_state_f32 = state0.float().clone() + # Reference (fp32, starting from the same lossy-or-not state the + # kernel sees). + ref_state_f32 = ref_input_state.clone() if k > 0: ref_state_f32[slots] = states_buffer_f32[slots, k - 1] @@ -198,6 +276,7 @@ def test_checkpointing_state_update( # Replay kernel — clone caches into mutable working copies that we # can inspect AFTER the call to verify cache postconditions. test_state = state0.clone() + test_scales = state0_scales.clone() if is_quantized else None prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) old_x_w = old_x.clone() @@ -224,6 +303,7 @@ def test_checkpointing_state_update( dt_bias=dt_bias, dt_softplus=True, state_batch_indices=state_batch_indices, + state_scales=test_scales, write_checkpoint=write_checkpoint, ) @@ -237,12 +317,36 @@ def test_checkpointing_state_update( # T<=16, ~2.0 at T=32-55; mean ~0.014; <0.02% of elements exceed 0.5. # State dtype (fp16/bf16/fp32) doesn't shift the error — bf16 dot # inputs dominate, not state storage. + # + # Quantized states add a per-element state quant error eps that + # propagates through C @ state in the output dot. With dstate=128 + # and C ~ N(0,1), the output channel std from this noise is roughly + # eps * sqrt(128/3) ≈ 6.5 * eps. Stack with the bf16 baseline: + # out_atol = bf16_atol + 6.5 * eps_max + # where eps_max is the worst-case per-element error at the + # post-replay state magnitude (T=55 → amax ≈ 23). + # + # Per-element error (eps_max for T=55): + # int8 (uniform grid): amax/(2*127) ≈ 0.091 + # int16 (uniform grid): amax/(2*32767) ≈ 3.5e-4 + # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 (worst-case + # cell at top of channel; smaller for + # smaller-magnitude elements) + out_atol = ( + {torch.int8: 1.6, torch.int16: 1.05, torch.float8_e4m3fn: 4.0}[state_dtype] + if is_quantized else 1.0 + ) + out_rtol = ( + {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] + if is_quantized else 2e-2 + ) out_diff = (test_out.float() - ref_out.float()).abs() out_max = out_diff.max().item() out_mean = out_diff.mean().item() try: torch.testing.assert_close( - test_out, ref_out, rtol=2e-2, atol=1.0, msg=f"Output mismatch at k={k}" + test_out, ref_out, rtol=out_rtol, atol=out_atol, + msg=f"Output mismatch at k={k}", ) except AssertionError: print( @@ -256,31 +360,85 @@ def test_checkpointing_state_update( # True → kernel writes the post-replay state; expect the # selective_state_update reference's state at step k-1. # False → kernel skips the HBM store; state must be UNCHANGED - # from the input (state0). - if write_checkpoint: - expected_state = ( - state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) - ) + # from the input (state0; for quant, scales also unchanged). + if is_quantized: + if write_checkpoint: + # Compare via dequant against the fp32 reference state. + expected_fp32 = ( + ref_input_state[slots] if k == 0 else states_buffer_f32[slots, k - 1] + ) + actual_fp32 = _dequantize_state(test_state[slots], test_scales[slots]) + # State diff = bf16_replay_error + quant_error (per element). + # The bf16 component is the SAME error source the non-quant + # test absorbs in its atol=1.0 baseline (replay's tl.dot is + # bf16-input fp32-accum; per-element error ~ 2^-7 * amax, + # empirically ≤ ~0.2 at T=55 amax≈23). Quant adds: + # int8: amax/(2*127) ≈ 0.091 worst-case + # int16: amax/(2*32767) ≈ 3.5e-4 (negligible vs bf16) + # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 worst-case + # Atol = bf16_baseline (1.0) + quant_eps_max. + state_atol = { + torch.int8: 1.1, torch.int16: 1.0, torch.float8_e4m3fn: 2.5, + }[state_dtype] + state_rtol = { + torch.int8: 5e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 1e-1, + }[state_dtype] + try: + torch.testing.assert_close( + actual_fp32, expected_fp32, + rtol=state_rtol, atol=state_atol, + msg=f"State mismatch at k={k} dtype={state_dtype}", + ) + except AssertionError: + diff = (actual_fp32 - expected_fp32).abs() + print( + f"k={k} state(dequant): max={diff.max().item():.4f} " + f"mean={diff.mean().item():.4f}" + ) + raise + # Scales sanity (fp32, finite, positive). + assert test_scales.dtype == torch.float32 + assert torch.isfinite(test_scales[slots]).all(), ( + f"state_scales has non-finite values at k={k}" + ) + assert (test_scales[slots] > 0).all(), ( + f"state_scales has non-positive values at k={k}" + ) + else: + # No write: raw quant state and scales unchanged. Use + # torch.equal for byte-level equality (dtype-agnostic; works + # for int8 / int16 / fp8 alike). + assert torch.equal(test_state[slots], state0[slots]), ( + f"Quant state changed at k={k} write_checkpoint=False" + ) + assert torch.equal(test_scales[slots], state0_scales[slots]), ( + f"State scales changed at k={k} write_checkpoint=False" + ) else: - expected_state = state0[slots] - state_diff = (test_state[slots].float() - expected_state.float()).abs() - state_max = state_diff.max().item() - state_mean = state_diff.mean().item() - try: - torch.testing.assert_close( - test_state[slots], - expected_state, - rtol=2e-2, - atol=1.0 if write_checkpoint else 0.0, - msg=f"State mismatch at k={k} (write_checkpoint={write_checkpoint})", - ) - except AssertionError: - print( - f"k={k} state: max={state_max:.4f} mean={state_mean:.4f} " - f"nan={torch.isnan(test_state).any().item()} " - f"inf={torch.isinf(test_state).any().item()}" - ) - raise + if write_checkpoint: + expected_state = ( + state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) + ) + else: + expected_state = state0[slots] + state_diff = (test_state[slots].float() - expected_state.float()).abs() + state_max = state_diff.max().item() + state_mean = state_diff.mean().item() + try: + torch.testing.assert_close( + test_state[slots], + expected_state, + rtol=2e-2, + atol=1.0 if write_checkpoint else 0.0, + msg=f"State mismatch at k={k} (write_checkpoint={write_checkpoint})", + ) + except AssertionError: + print( + f"k={k} state: max={state_max:.4f} mean={state_mean:.4f} " + f"nan={torch.isnan(test_state).any().item()} " + f"inf={torch.isinf(test_state).any().item()}" + ) + raise # --- Cache postconditions --- # Compute step 2's processed values (what the kernel should have @@ -357,24 +515,33 @@ def test_checkpointing_state_update( ) -@_skip_pre_sm100 @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize( + "state_dtype", + [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], + ids=["fp16", "int8", "int16", "fp8"], +) @pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) @pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) -def test_checkpointing_state_update_philox(nheads, head_dim, d_state, ngroups, paged_cache, T): +def test_checkpointing_state_update_philox(state_dtype, nheads, head_dim, d_state, ngroups, paged_cache, T): """ - Verify that Philox stochastic rounding produces correct results. - - Runs our kernel twice with identical inputs: once without rounding - (fp16 state, deterministic), once with rounding (fp16 state, Philox). - The outputs should be nearly identical — stochastic rounding only - perturbs the state by ±1 fp16 ULP, which barely affects output. - Also verifies the state dtype remains fp16. + Verify that Philox stochastic rounding produces correct results across + all SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). + + Runs our kernel twice with identical inputs — once without rand_seed + (deterministic RN), once with rand_seed (Philox SR) — and confirms: + - Outputs are within bf16-dot tolerance (state perturbation ≤ 1 ULP). + - State dtype is preserved. + - State difference is bounded by ~1 ULP of the chosen grid. """ + _maybe_skip_dtype(state_dtype, use_sr=True) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + batch = 2 device = "cuda" dtype = torch.bfloat16 - state_dtype = torch.float16 assert nheads % ngroups == 0 if paged_cache: @@ -393,7 +560,16 @@ def test_checkpointing_state_update_philox(nheads, head_dim, d_state, ngroups, p D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) - state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + if is_quantized: + state0_fp32 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + state0, state0_scales = _quantize_state(state0_fp32, state_dtype, quant_max) + else: + state0 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state0_scales = None # Cache tensors old_x = torch.randn(cache_size, T, nheads, head_dim, device=device, dtype=dtype) @@ -423,8 +599,9 @@ def test_checkpointing_state_update_philox(nheads, head_dim, d_state, ngroups, p state_batch_indices=state_batch_indices, ) - # --- Run without rounding (deterministic fp16 state store) --- + # --- Run without rounding (deterministic RN store) --- state_no_round = state0.clone() + scales_no_round = state0_scales.clone() if is_quantized else None out_no_round = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) checkpointing_state_update( state_no_round, @@ -435,12 +612,14 @@ def test_checkpointing_state_update_philox(nheads, head_dim, d_state, ngroups, p cache_buf_idx.clone(), prev_tokens, out=out_no_round, + state_scales=scales_no_round, **common_kwargs, ) # --- Run with Philox rounding --- rand_seed = torch.tensor([12345], device=device, dtype=torch.int64) state_rounded = state0.clone() + scales_rounded = state0_scales.clone() if is_quantized else None out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) checkpointing_state_update( state_rounded, @@ -453,47 +632,100 @@ def test_checkpointing_state_update_philox(nheads, head_dim, d_state, ngroups, p out=out_rounded, rand_seed=rand_seed, philox_rounds=10, + state_scales=scales_rounded, **common_kwargs, ) # Outputs should be nearly identical — rounding only perturbs the # post-replay state by ±1 ULP before the output phase reads it. + # Out_atol = bf16_baseline + 6.5 * per_elem_ULP_after_dequant: + # non-quant fp16: fp16 ULP at typical magnitude is tiny → 1.0 + # int8: amax/127 ≈ 23/127 → 6.5*0.18 ≈ 1.2 + bf16_baseline + # int16: amax/32767 ≈ 7e-4 → ~bf16_baseline only + # fp8: amax/14 ≈ 23/14 → 6.5*1.6 ≈ 10.7 + bf16_baseline + out_atol = ( + {torch.int8: 1.5, torch.int16: 1.0, torch.float8_e4m3fn: 6.0}[state_dtype] + if is_quantized else 1.0 + ) + out_rtol = ( + {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] + if is_quantized else 2e-2 + ) torch.testing.assert_close( - out_rounded, out_no_round, rtol=2e-2, atol=1.0, msg="Output diverged with Philox rounding" + out_rounded, out_no_round, rtol=out_rtol, atol=out_atol, + msg=f"Output diverged with Philox rounding ({state_dtype})", ) - # State should remain fp16 - assert state_rounded.dtype == torch.float16 + # State dtype preserved. + assert state_rounded.dtype == state_dtype - # States should differ by at most 1 fp16 ULP per element. - # fp16 ULP depends on magnitude: up to 0.5 for values near 512. - # Use rtol to account for magnitude-dependent ULP. + # State diff between RN and SR is bounded by 1 quant cell per element. + # Per-channel decode_scale varies by 10x+ across channels (amax depends + # on randn extremes), so a single flat atol can't bound it accurately — + # use per-channel ULP-aware comparison. slots = state_batch_indices if paged_cache else slice(None) - torch.testing.assert_close( - state_rounded[slots], - state_no_round[slots], - rtol=2e-3, - atol=0.2, - msg="State diverged with Philox rounding", - ) + if is_quantized: + rounded_fp32 = _dequantize_state(state_rounded[slots], scales_rounded[slots]) + no_round_fp32 = _dequantize_state(state_no_round[slots], scales_no_round[slots]) + diff = (rounded_fp32 - no_round_fp32).abs() + # Per-element bound = max(decode_scale_no_round, decode_scale_rounded). + # decode_scale is shape (cache, nheads, dim); broadcast over dstate. + scale_bound = torch.maximum( + scales_no_round[slots], scales_rounded[slots] + ).unsqueeze(-1) + # int8 / int16: 1 cell after dequant = decode_scale exactly. + # fp8_e4m3: variable grid; the largest cell within a channel scaled + # to fit ±448 is at the channel's max-magnitude element, where the + # cell is ~32x larger than the average. Bound = decode_scale * 32. + # Apply a 1.5x slack pad for floating-point compare quirks at the + # exact-cell boundary. + cell_pad = ( + 32.0 if state_dtype == torch.float8_e4m3fn else 1.0 + ) + bound = scale_bound * (cell_pad * 1.5) + if not (diff <= bound).all(): + offenders = (diff > bound).sum().item() + n_total = diff.numel() + pytest.fail( + f"State RN-SR diff exceeds 1 cell per element for " + f"{offenders}/{n_total} elements ({state_dtype}). " + f"max_diff={diff.max().item():.4g}, " + f"max_bound={bound.max().item():.4g}." + ) + else: + # fp16 ULP depends on magnitude — rtol absorbs that. + torch.testing.assert_close( + state_rounded[slots], + state_no_round[slots], + rtol=2e-3, + atol=0.2, + msg=f"State diverged with Philox rounding ({state_dtype})", + ) -@_skip_pre_sm100 -def test_philox_rounding_unbiased(): +@pytest.mark.parametrize( + "state_dtype", + [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], + ids=["fp16", "int8", "int16", "fp8"], +) +def test_philox_rounding_unbiased(state_dtype): """ - Verify that Philox stochastic rounding is unbiased. + Verify that Philox stochastic rounding is unbiased across all + SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). - Runs the replay kernel with fp32 state (capturing the true fp32 - post-replay state) and with fp16 state + Philox rounding. Compares the - rounding residual (fp16_state.float() - fp32_state) against deterministic - rounding (fp32_state.to(fp16).float() - fp32_state). - - Deterministic round-to-nearest-even has a systematic positive bias on - the residual. Philox stochastic rounding should be unbiased: the mean - residual should be near zero. + Captures the true fp32 post-replay state by running with fp32 storage, + then runs the kernel with the target dtype + Philox SR. Compares the + SR rounding residual against the deterministic-RN residual: SR should + have mean residual closer to zero than RN, since RN has a systematic + round-to-nearest-even bias and SR is unbiased by construction. Uses a large batch (16) for ~2M state elements — plenty of statistics. """ + _maybe_skip_dtype(state_dtype, use_sr=True) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 batch, T = 16, 6 device = "cuda" @@ -507,8 +739,11 @@ def test_philox_rounding_unbiased(): D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) - # Use fp32 initial state so replay produces non-fp16-representable values - state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=torch.float32) + # fp32 reference state — replay produces values that don't fit cleanly + # in the target dtype's grid, exposing the rounding bias. + state0_fp32 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) old_x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(batch, 2, T, ngroups, d_state, device=device, dtype=dtype) @@ -525,68 +760,83 @@ def test_philox_rounding_unbiased(): prev_tokens = torch.full((batch,), T, device=device, dtype=torch.int32) common_kwargs = dict( - x=x, - dt=dt_val, - A=A, - B=B, - C=C, - D=D, - dt_bias=dt_bias, - dt_softplus=True, + x=x, dt=dt_val, A=A, B=B, C=C, D=D, dt_bias=dt_bias, dt_softplus=True, ) - # 1. fp32 state — captures true post-replay state - state_fp32 = state0.clone() + # 1. fp32 state — captures true post-replay fp32 state. + state_fp32 = state0_fp32.clone() out_fp32 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) checkpointing_state_update( state_fp32, - old_x.clone(), - old_B.clone(), - old_dt.clone(), - old_dA_cumsum.clone(), - cache_buf_idx.clone(), - prev_tokens, - out=out_fp32, - **common_kwargs, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), prev_tokens, out=out_fp32, **common_kwargs, ) - # 2. fp16 state with Philox rounding + # 2. Target dtype + Philox SR. For quant we also need scales (derived + # from the same per-channel amax used by the kernel on store). rand_seed = torch.tensor([99999], device=device, dtype=torch.int64) - state_rounded = state0.to(torch.float16).clone() + if is_quantized: + state_rounded, scales_rounded = _quantize_state(state0_fp32, state_dtype, quant_max) + else: + state_rounded = state0_fp32.to(state_dtype) + scales_rounded = None out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) checkpointing_state_update( state_rounded, - old_x.clone(), - old_B.clone(), - old_dt.clone(), - old_dA_cumsum.clone(), - cache_buf_idx.clone(), - prev_tokens, - out=out_rounded, - rand_seed=rand_seed, - philox_rounds=10, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), prev_tokens, out=out_rounded, + rand_seed=rand_seed, philox_rounds=10, + state_scales=scales_rounded, **common_kwargs, ) - # Compute rounding residuals where fp32 state has non-zero values - fp32_vals = state_fp32.flatten() - stochastic_residual = state_rounded.float().flatten() - fp32_vals - deterministic_residual = fp32_vals.to(torch.float16).float() - fp32_vals + # Compute residuals. For non-quant: stochastic_residual = SR(fp32) - + # fp32, deterministic_residual = RN(fp32) - fp32. For quant: dequant + # both, comparing in fp32. + if is_quantized: + fp32_vals = state_fp32.flatten() + stochastic_residual = ( + _dequantize_state(state_rounded, scales_rounded).flatten() - fp32_vals + ) + # Deterministic reference: do the same per-channel quant on the + # captured fp32 state, then dequant. This is what the kernel would + # have produced with rand_seed=None. + det_quant, det_scales = _quantize_state(state_fp32, state_dtype, quant_max) + deterministic_residual = ( + _dequantize_state(det_quant, det_scales).flatten() - fp32_vals + ) + else: + fp32_vals = state_fp32.flatten() + stochastic_residual = state_rounded.float().flatten() - fp32_vals + deterministic_residual = fp32_vals.to(state_dtype).float() - fp32_vals - # Only consider elements where rounding matters (non-zero residual possible) + # Only consider elements where rounding matters (non-zero residual possible). nonzero_mask = deterministic_residual.abs() > 0 num_nonzero = nonzero_mask.sum().item() assert num_nonzero > 1000, f"Too few roundable elements: {num_nonzero}" stochastic_mean = stochastic_residual[nonzero_mask].mean().item() + stochastic_std = stochastic_residual[nonzero_mask].std().item() deterministic_mean = deterministic_residual[nonzero_mask].mean().item() - # Stochastic rounding should be less biased than deterministic. - # With ~millions of elements, the stochastic mean should be very close to 0. - # Deterministic round-to-nearest-even has a small but systematic bias. - assert abs(stochastic_mean) < abs(deterministic_mean) or abs(stochastic_mean) < 1e-5, ( - f"Stochastic rounding appears biased: stochastic_mean={stochastic_mean:.6f}, " - f"deterministic_mean={deterministic_mean:.6f}, n_elements={num_nonzero}" + # SE-based bias check. An unbiased estimator's sample mean has standard + # error SE = std / sqrt(n). We require |sr_mean| < K*SE (K=4 ≈ ~3.2e-5 + # one-sided false-positive rate). This auto-calibrates per dtype: + # * int16: residual std ~1e-4 → SE ~9e-8 (very tight bound) + # * int8: residual std ~3e-2 → SE ~2e-5 + # * fp8: residual std ~1e-1 → SE ~9e-5 (loosest, magnitude-driven) + # The previous fixed-1e-5 threshold was below SE for int8/fp8 and would + # always fail by chance. Note the |sr|<|det| fallback was also dropped: + # on Gaussian (symmetric) inputs RN's bias is ~0 by symmetry, so SR vs RN + # is just two unbiased estimators racing — unreliable as a unbias test. + se_sr = stochastic_std / (num_nonzero ** 0.5) + K = 4 + assert abs(stochastic_mean) < K * se_sr, ( + f"SR mean exceeds {K}*SE (likely biased) ({state_dtype}): " + f"stochastic_mean={stochastic_mean:.3e}, " + f"SE={se_sr:.3e} (K*SE={K * se_sr:.3e}), " + f"deterministic_mean={deterministic_mean:.3e} (for reference), " + f"n_elements={num_nonzero}" ) @@ -888,3 +1138,156 @@ def test_checkpointing_heads_per_block_multistep( f"T={T}, nheads={nheads}, ngroups={ngroups}, " f"state_dtype={state_dtype}, paged_cache={paged_cache}", ) + + +# ----- SR grid-bracket tests (fp8 and fp16) ----- +# +# Verify that each PTX SR output lands on the destination dtype's grid as +# a bracket neighbour of the fp32 input. Catches byte-order traps in the +# inline-asm source-register specifier: +# * fp8: cvt.rs.satfinite.e4m3x4.f32 with pack=4, asm "{$4,$3,$2,$1}" +# * fp16: cvt.rs.f16x2.f32 with pack=2, asm "$0, $2, $1, $3" +# The unbiased test (test_philox_rounding_unbiased) wouldn't catch a +# shuffle: outputs that are still on-grid but swapped within a pack still +# average correctly. Only the per-element bracket check exposes it. +# +# Both kernels are inline copies of the production helpers — kept here so +# the test exercises the exact PTX form independent of wrapper changes. + + +@triton.jit +def _bracket_kernel_fp8(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + offs) + y = tl.inline_asm_elementwise( + asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", + constraints="=r,r,r,r,r,r,r,r,r", + args=(x, rand), + dtype=tl.float8e4nv, + is_pure=True, + pack=4, + ) + tl.store(out_ptr + offs, y) + + +@triton.jit +def _bracket_kernel_fp16(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + offs) + y = tl.inline_asm_elementwise( + asm="""{ + cvt.rs.f16x2.f32 $0, $2, $1, $3; + }""", + constraints=("=r,r,r,r,r"), + args=(x, rand), + dtype=tl.float16, + is_pure=True, + pack=2, + ) + tl.store(out_ptr + offs, y) + + +_BRACKET_KERNEL = { + torch.float8_e4m3fn: _bracket_kernel_fp8, + torch.float16: _bracket_kernel_fp16, +} + + +def _build_finite_grid(dtype: torch.dtype, device: str) -> torch.Tensor: + """Reinterpret all bit patterns of ``dtype`` as floats; return sorted + unique finite values (drops ±inf, NaNs).""" + if dtype == torch.float8_e4m3fn: + ints = torch.arange(256, dtype=torch.uint8, device=device) + full = ints.view(torch.float8_e4m3fn).to(torch.float32) + elif dtype == torch.float16: + # int16 view of all 65536 patterns (covers fp16 normals + subnormals + # + ±inf + NaN; we filter to finite below). + ints = torch.arange(65536, dtype=torch.int32, device=device).to(torch.int16) + full = ints.view(torch.float16).to(torch.float32) + else: + raise ValueError(f"Unsupported bracket-test dtype: {dtype}") + return full[torch.isfinite(full)].sort()[0].unique() + + +def _build_bracket_inputs(dtype: torch.dtype, n: int, device: str) -> torch.Tensor: + """Test inputs spanning the dtype's grid range. Includes on-grid points + so we exercise the no-rounding case; for fp8 also includes overflow to + test saturation (PTX `cvt.rs.satfinite.e4m3x4.f32` clamps in-op). + + fp16 inputs are kept inside the finite range — `cvt.rs.f16x2.f32` does + NOT have a `satfinite` modifier and produces ±inf for OOR inputs (not + a saturate-to-±max). The kernel only ever sees in-range fp32 state in + practice (state_amax is always ≪ fp16_max), so the test mirrors that. + """ + grid = _build_finite_grid(dtype, device) + g_min, g_max = grid[0].item(), grid[-1].item() + x = torch.empty(n, device=device, dtype=torch.float32) + if dtype == torch.float8_e4m3fn: + # 1.5x range exercises saturation; satfinite handles it in-op. + x.uniform_(g_min * 1.5, g_max * 1.5) + else: # fp16: four magnitude bands, all within finite range. + x[: n // 4].uniform_(-1.0, 1.0) + x[n // 4 : n // 2].uniform_(-100, 100) + x[n // 2 : 3 * n // 4].uniform_(-1000, 1000) + x[3 * n // 4 :].uniform_(g_min * 0.99, g_max * 0.99) + return x, grid + + +@_skip_pre_sm100 +@pytest.mark.parametrize( + "state_dtype", + [torch.float8_e4m3fn, torch.float16], + ids=["fp8", "fp16"], +) +def test_sr_grid_bracket(state_dtype): + """Verify SR PTX outputs each lie on the destination grid as a bracket + neighbour of the fp32 input.""" + device = "cuda" + n = 1024 # multiple of both pack=4 (fp8) and pack=2 (fp16) + + torch.manual_seed(42) + x, grid_finite = _build_bracket_inputs(state_dtype, n, device) + g_min, g_max = grid_finite[0].item(), grid_finite[-1].item() + + # Bracket [lo, hi] in the destination grid for each input. For + # out-of-range inputs the bracket is the saturating endpoint pair. + x_clamped = x.clamp(g_min, g_max) + idx = torch.searchsorted(grid_finite, x_clamped, right=False).clamp( + min=1, max=len(grid_finite) - 1 + ) + lo = grid_finite[idx - 1] + hi = grid_finite[idx] + # For x exactly on grid, idx points at it; lo = grid[i-1], hi = x — the + # bracket allows out==hi (=x) which is what RN-on-grid produces. + + kernel = _BRACKET_KERNEL[state_dtype] + + for seed in range(4): + torch.manual_seed(seed) + # int32 for raw random bits — PTX takes the bit pattern, sign + # interpretation doesn't matter. + rand = torch.randint(-(2**31), 2**31, (n,), device=device, dtype=torch.int32) + out = torch.empty(n, device=device, dtype=state_dtype) + kernel[(1,)](x, rand, out, BLOCK=n) + out_fp32 = out.to(torch.float32) + + on_grid = (out_fp32 == lo) | (out_fp32 == hi) + if not on_grid.all(): + offenders = ~on_grid + n_off = offenders.sum().item() + sample = ( + x[offenders][:5].tolist(), + lo[offenders][:5].tolist(), + hi[offenders][:5].tolist(), + out_fp32[offenders][:5].tolist(), + ) + pytest.fail( + f"{state_dtype} SR output not on grid bracket for {n_off}/{n} " + f"elements (seed={seed}). x={sample[0]} lo={sample[1]} " + f"hi={sample[2]} out={sample[3]}. Likely the PTX byte-order " + "bug (cvt.rs source-register order)." + ) + + From 93f5d2ae79dbd3eed26f1189be213b2e2bff2539 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Thu, 30 Apr 2026 16:52:54 -0700 Subject: [PATCH 11/89] Benchmark: --sr-modes RN,SR sweep param replaces --philox-rounding flag Lets one invocation sweep both rounding modes (e.g. for A/B kernel comparisons across RN and SR), avoiding the prior need for two separate runs. Skips SR silently for unsupported state dtypes (bf16, fp32). Backward compat: --philox-rounding still works as an alias for --sr-modes SR; setting both is an error. NVTX tag gains _SR=0 / _SR=1 sweep tag. The companion collect.py regex was updated separately (in the scripts repo) to capture this axis. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 72 ++++++++++++++----- 1 file changed, 54 insertions(+), 18 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 3321aa90674b..9db6d3cb2a65 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -511,22 +511,26 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp """ from concurrent.futures import ThreadPoolExecutor + sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) + configs = [] for batch in batch_sizes: for mtp_len in mtp_lengths: prev_ks = _resolve_prev_ks(args, mtp_len) for state_dtype in state_dtypes: for act_dtype in act_dtypes: - configs.append((batch, mtp_len, prev_ks, state_dtype, act_dtype)) + for sr_mode in sr_modes_list: + configs.append( + (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode)) print(f"[compile-warmup] {len(configs)} configs across {max_workers} threads") t0 = time.perf_counter() def _warm(cfg): - batch, mtp_len, prev_ks, state_dtype, act_dtype = cfg + batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode = cfg _bench_config( args, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, - warmup_only=True, + sr_mode=sr_mode, warmup_only=True, ) errors = [] @@ -555,6 +559,7 @@ def _bench_config( state_dtype: torch.dtype, act_dtype: torch.dtype, baseline_fn, + sr_mode: str = "RN", warmup_only: bool = False, ) -> None: """ @@ -614,23 +619,21 @@ def _bench_config( head_dim = args.head_dim d_state = args.d_state with_conv1d = getattr(args, "with_conv1d", False) - use_philox = getattr(args, "philox_rounding", False) + use_philox = (sr_mode == "SR") variant_fn = _VARIANT_FNS[args.variant]() - # Philox rounding: allow fp16 and the quantized dtypes (int8/int16/fp8). + # SR rounding: allow fp16 and the quantized dtypes (int8/int16/fp8). # bf16/fp32 SR is not supported (no PTX path for bf16; fp32 doesn't need - # rounding). This validates the user's input — it doesn't apply to the - # baseline (incompatible baselines are silently skipped below). + # rounding). When sweeping --sr-modes RN,SR over a mixed dtype set, + # silently skip the SR cell for unsupported dtypes — the RN cell still + # prints, and other dtypes still get their SR row. rand_seed = None _SR_SUPPORTED = ( torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn, ) if use_philox: if state_dtype not in _SR_SUPPORTED: - raise ValueError( - f"--philox-rounding requires state dtype in {{fp16, int8, int16, fp8}}, " - f"got {state_dtype}" - ) + return rand_seed = torch.randint(0, 2**62, (1,), device="cuda", dtype=torch.int64) is_quantized = state_dtype in (torch.int8, torch.int16, torch.float8_e4m3fn) @@ -914,6 +917,7 @@ def _run_incr( parts.append(f"R={maxnreg}") if num_ctas is not None: parts.append(f"CT={num_ctas}") + parts.append(f"SR={1 if use_philox else 0}") sweep_suffix = (" " + ",".join(parts)) if parts else "" sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") @@ -1035,15 +1039,19 @@ def _run_benchmark(args) -> None: f"|{'-' * 7}|{'-' * 9}|{'-' * 8}|{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" ) + sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) + for batch in batch_sizes: for mtp_len in mtp_lengths: # Resolve prev_k fractions → clamped integers in [0, mtp_len] prev_ks = _resolve_prev_ks(args, mtp_len) for state_dtype in state_dtypes: for act_dtype in act_dtypes: - _bench_config( - args, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn - ) + for sr_mode in sr_modes_list: + _bench_config( + args, batch, mtp_len, prev_ks, state_dtype, act_dtype, + baseline_fn, sr_mode=sr_mode, + ) if args.profile: torch.cuda.cudart().cudaProfilerStop() @@ -1251,12 +1259,21 @@ def _parse_args() -> argparse.Namespace: default=None, help="Override num_ctas for the main kernel (comma-separated sweep).", ) + parser.add_argument( + "--sr-modes", + type=str, + default="RN", + help="Comma-separated rounding modes to sweep: any combination of " + "{RN, SR}. SR (stochastic rounding) is silently skipped for state " + "dtypes that don't support it (bf16, fp32). Default 'RN' matches " + "legacy --philox-rounding=False behavior.", + ) parser.add_argument( "--philox-rounding", action="store_true", - help="Enable Philox stochastic rounding (rand_seed generated per " - "iteration). Supported state dtypes: fp16, int8, int16, fp8. " - "fp16 SR and fp8 SR require sm_100a (Blackwell B200+).", + help="DEPRECATED — equivalent to --sr-modes SR. Retained for " + "backward compatibility; use --sr-modes for new scripts. fp16 SR " + "and fp8 SR require sm_100a (Blackwell B200+).", ) parser.add_argument( "--philox-rounds", @@ -1283,7 +1300,26 @@ def _parse_args() -> argparse.Namespace: "module loading. Slower (~40s startup) but guaranteed correct " "if the fast path breaks due to package changes.", ) - return parser.parse_args() + args = parser.parse_args() + + # Backward-compat: --philox-rounding implies --sr-modes SR if --sr-modes + # was left at the default. If both are set explicitly, error. + sr_modes_default = (args.sr_modes == "RN") + if args.philox_rounding: + if not sr_modes_default and args.sr_modes != "SR": + parser.error( + "--philox-rounding (deprecated) is incompatible with explicit " + f"--sr-modes={args.sr_modes!r}. Use --sr-modes SR (or " + "RN,SR) instead and drop --philox-rounding." + ) + args.sr_modes = "SR" + + sr_modes = [m.strip() for m in args.sr_modes.split(",") if m.strip()] + for m in sr_modes: + if m not in ("RN", "SR"): + parser.error(f"--sr-modes value must be RN or SR, got {m!r}") + args.sr_modes_list = sr_modes + return args class _Tee: From f83793546327ae7a2efd13088ed0c39f74f1816f Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 5 May 2026 22:13:10 -0700 Subject: [PATCH 12/89] =?UTF-8?q?rect=5Fprecompute=20opt=20B:=20pre-wait?= =?UTF-8?q?=20factor=5Fdt=20=C3=97=20exp=5Fdiff=20via=20cb=5Fscaled=20scra?= =?UTF-8?q?tch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move the post-wait per-head loop's factor_dt/s_k/exp_diff math pre-gdc_wait, storing the (T, K) per-head combo tile (= factor_dt[None, :] * exp_diff) to cb_scaled as scratch. Post-wait per-head loop becomes: combo = tl.load(cb_scaled[...]) rect_CB_scaled = tl.where(causal, raw_rect_CB * combo, 0) tl.store(cb_scaled[...], rect_CB_scaled) Drops the L1-warmer trick (warm_old_dt/warm_old_dA/warm_dt_at_kn/warm_dA_at_kn + warm_keep wraparound) — those loads now feed real factor_dt/s_k/exp_diff computation pre-wait. Same registers (48/thread) — change shifts cycles from post-PDL critical path to pre-PDL phase, letting precompute finish earlier so main's gdc_wait returns sooner. Wins concentrated at b=16/32 (-3 to -15% across all 4 dtypes); flat or small loss at b=64+; small regressions at b=1/b=4 (replay already won there) and b=128 fp8/int8 (rect still wins by 13-25%). Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 135 ++++++------------ 1 file changed, 45 insertions(+), 90 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index d9c9fe278def..465ac23f5579 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -555,13 +555,12 @@ def _rectangle_precompute_kernel( other=0.0, ) - # Per-head: pre-compute decay_vec_full (= total_decay * exp(cumAdt_new[t])) - # and store to scratch. Also pre-load the rest of the per-head cache - # data (old_dt, old_dA_cumsum, dt_at_kn, dA_cumsum_at_kn) as L1 warmers - # — Triton can't span their values across the gdc_wait barrier without a - # DRAM round-trip, but L1 is per-SM-persistent so the post-wait loop's - # reload hits the cache cheaply. Folding the warmer loads into - # decay_vec_full's `tl.where(False, ...)` keeps Triton from DCE-ing them. + # Per-head pre-wait: compute decay_vec_full AND combo = factor_dt * exp_diff, + # storing both to scratch (combo overlays cb_scaled — overwritten post-wait + # with rect_CB_scaled). This pulls all factor_dt / exp_diff math out of + # the post-wait loop, which becomes a single load + multiply + masked store + # per head. Loads of old_dt/old_dA_cumsum/dt_at_kn/dA_cumsum_at_kn that + # used to be L1 warmers are now real consumers (no fake-keep wraparound). prev_k_idx = tl.minimum( tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 ) @@ -591,7 +590,14 @@ def _rectangle_precompute_kernel( + write_buf * stride_old_dA_cumsum_dbuf + head_idx * stride_old_dA_cumsum_head ) - # Required for decay_vec_full + old_dt_all_h = tl.load( + old_dt_read_base_h + safe_old_k * stride_old_dt_T, + mask=is_old_k, other=0.0, + ).to(tl.float32) + old_dA_cumsum_all_h = tl.load( + old_dA_cumsum_read_base_h + safe_old_k * stride_old_dA_cumsum_T, + mask=is_old_k, other=0.0, + ).to(tl.float32) total_dA_cumsum_h = tl.load( old_dA_cumsum_read_base_h + prev_k_idx * stride_old_dA_cumsum_T ).to(tl.float32) @@ -599,44 +605,42 @@ def _rectangle_precompute_kernel( old_dA_cumsum_write_base_h + (write_offset + offs_t) * stride_old_dA_cumsum_T, mask=t_mask, other=0.0, ).to(tl.float32) - # L1 warmers — same addresses the post-wait loop will reload. - warm_old_dt = tl.load( - old_dt_read_base_h + safe_old_k * stride_old_dt_T, - mask=is_old_k, other=0.0, - ).to(tl.float32) - warm_old_dA = tl.load( - old_dA_cumsum_read_base_h + safe_old_k * stride_old_dA_cumsum_T, - mask=is_old_k, other=0.0, - ).to(tl.float32) - warm_dt_at_kn = tl.load( + dt_at_kn_h = tl.load( old_dt_write_base_h + (write_offset + safe_k_new) * stride_old_dt_T, mask=is_new_k, other=0.0, ).to(tl.float32) - warm_dA_at_kn = tl.load( + dA_cumsum_at_kn_h = tl.load( old_dA_cumsum_write_base_h + (write_offset + safe_k_new) * stride_old_dA_cumsum_T, mask=is_new_k, other=0.0, ).to(tl.float32) + # decay_vec_full = total_decay * exp(cumAdt_new) total_decay_h = tl.where( prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum_h), 1.0 ) decay_vec_full_h = total_decay_h * tl.exp(dA_cumsum_new_h) - # Force the warmer loads to participate via a runtime-False guard. - # prev_num_accepted_tokens is dynamic, so Triton can't DCE this. - warm_keep = ( - tl.sum(warm_old_dt) + tl.sum(warm_old_dA) - + tl.sum(warm_dt_at_kn) + tl.sum(warm_dA_at_kn) - ) - decay_vec_full_h = tl.where( - prev_num_accepted_tokens < 0, - decay_vec_full_h + warm_keep, - decay_vec_full_h, - ) decay_vec_base_h = decay_vec_ptr + pid_b * stride_dv_batch + head_idx * stride_dv_head tl.store( decay_vec_base_h + offs_t * stride_dv_t, decay_vec_full_h, mask=t_mask ) + # combo = factor_dt * exp_diff — stored to cb_scaled scratch (overwritten + # post-wait with rect_CB_scaled = where(causal, raw_rect_CB * combo, 0)). + factor_dt_h = tl.where(is_old_k, old_dt_all_h, dt_at_kn_h) + s_k_h = tl.where( + is_old_k, total_dA_cumsum_h - old_dA_cumsum_all_h, -dA_cumsum_at_kn_h + ) + exp_diff_h = tl.exp(s_k_h[None, :] + dA_cumsum_new_h[:, None]) + combo_h = factor_dt_h[None, :] * exp_diff_h + cb_scaled_base_h = ( + cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head + ) + tl.store( + cb_scaled_base_h + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + combo_h, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + ) + # ---- gdc_wait: from here on we depend on conv1d's outputs ---- if LAUNCH_WITH_PDL: tl.extra.cuda.gdc_wait() @@ -688,76 +692,27 @@ def _rectangle_precompute_kernel( is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] - # Per-head: factor_dt + exp_diff + rect_CB_scaled. Loads are post-wait - # but cache-resident from this kernel's earlier writes (loop 1 + the - # hoisted decay_vec_full block above), so they hit L1. Recomputing - # total_dA_cumsum / dA_cumsum_new here is cheaper than spanning the - # gdc_wait via a DRAM round-trip. + # Per-head post-wait: load pre-computed combo (factor_dt * exp_diff) from + # cb_scaled scratch, apply causal mask × raw_rect_CB, write rect_CB_scaled + # back to the same address. All factor_dt / exp_diff math happened + # pre-wait above; this loop is pure load + multiply + masked store. for h_local in range(HEADS_PER_BLOCK): head_idx = first_head + h_local - old_dt_read_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + buf_active * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - old_dA_cumsum_read_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_active * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - old_dt_write_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - old_dA_cumsum_write_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - old_dt_all = tl.load( - old_dt_read_base + safe_old_k * stride_old_dt_T, mask=is_old_k, other=0.0 - ).to(tl.float32) - old_dA_cumsum_all = tl.load( - old_dA_cumsum_read_base + safe_old_k * stride_old_dA_cumsum_T, - mask=is_old_k, other=0.0, - ).to(tl.float32) - total_dA_cumsum = tl.load( - old_dA_cumsum_read_base + prev_k_idx * stride_old_dA_cumsum_T - ).to(tl.float32) - dA_cumsum_new = tl.load( - old_dA_cumsum_write_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, - mask=t_mask, other=0.0, - ).to(tl.float32) - dt_at_kn = tl.load( - old_dt_write_base + (write_offset + safe_k_new) * stride_old_dt_T, - mask=is_new_k, other=0.0, - ).to(tl.float32) - dA_cumsum_at_kn = tl.load( - old_dA_cumsum_write_base + (write_offset + safe_k_new) * stride_old_dA_cumsum_T, - mask=is_new_k, other=0.0, - ).to(tl.float32) - - factor_dt = tl.where(is_old_k, old_dt_all, dt_at_kn) - s_k = tl.where( - is_old_k, total_dA_cumsum - old_dA_cumsum_all, -dA_cumsum_at_kn + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head + cb_load_mask = (offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K) + combo_loaded = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + mask=cb_load_mask, other=0.0, ) - exp_diff = tl.exp(s_k[None, :] + dA_cumsum_new[:, None]) - rect_CB_scaled = tl.where( causal_combined, - raw_rect_CB * factor_dt[None, :] * exp_diff, + raw_rect_CB * combo_loaded, 0.0, ) - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head tl.store( cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, rect_CB_scaled, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + mask=cb_load_mask, ) From d2bcdc5dca67251d724141977e1f8397f6284dc1 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 4 May 2026 19:17:53 -0700 Subject: [PATCH 13/89] Fix _checkpointing_main_kernel: handle PNAT > T-1 when max_window > T Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 55 ++++++++++++----- .../mamba/test_checkpointing_state_update.py | 59 +++++++++++-------- 2 files changed, 75 insertions(+), 39 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 648d0442a446..6d40df528030 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -364,6 +364,16 @@ def _checkpointing_precompute_kernel( @triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) @triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +# Replay axis: separate from the T-output axis. Sized to MAX_REPLAY_BUFFER_LENGTH +# so the kernel can iterate over up to PNAT old-cache positions independent of +# the new-token T axis. In production T=6, MAX=16 → BLOCK_SIZE_T=16, +# BLOCK_SIZE_WINDOW=16 (coincidentally equal); separating them is the +# semantically correct fix and avoids the prior `offs_t < T` masking bug +# that under-read old data when PNAT > T-1. +@triton.heuristics( + {"BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) @triton.jit() def _checkpointing_main_kernel( # Pointers @@ -394,6 +404,7 @@ def _checkpointing_main_kernel( pad_slot_id, # Dimensions T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # cache T-axis capacity (= max_window) dim: tl.constexpr, dstate: tl.constexpr, nheads_ngroups_ratio: tl.constexpr, @@ -466,6 +477,7 @@ def _checkpointing_main_kernel( HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, USE_RS_ROUNDING: tl.constexpr, PHILOX_ROUNDS: tl.constexpr, @@ -522,6 +534,9 @@ def _checkpointing_main_kernel( offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) offs_t = tl.arange(0, BLOCK_SIZE_T) + # Replay axis: separate from offs_t. Spans [0, BLOCK_SIZE_WINDOW) ⊇ + # [0, MAX_REPLAY_BUFFER_LENGTH); used for old-token loads (mask: offs_window < PNAT). + offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) m_mask = offs_m < dim n_mask = offs_n < dstate t_mask = offs_t < T @@ -551,16 +566,21 @@ def _checkpointing_main_kernel( # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) group_idx = pid_h // nheads_ngroups_ratio - # Load precomputed dt and dA_cumsum from READ buffer + # Old-token mask along the WINDOW (replay) axis. PNAT ≤ MAX ≤ + # BLOCK_SIZE_WINDOW, so this enables all valid old-cache positions. + # (Distinct from t_mask = offs_t < T which gates output T-rows only.) + old_window_mask = offs_window < prev_num_accepted_tokens + + # Load precomputed dt and dA_cumsum from READ buffer at [0, PNAT). old_dt_base = ( old_dt_ptr + cache_batch_idx * stride_old_dt_cache + active_buf * stride_old_dt_dbuf + pid_h * stride_old_dt_head ) - old_dt_all = tl.load(old_dt_base + offs_t * stride_old_dt_T, mask=t_mask, other=0.0).to( - tl.float32 - ) + old_dt_all = tl.load( + old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 + ).to(tl.float32) old_dA_cumsum_base = ( old_dA_cumsum_ptr @@ -569,31 +589,33 @@ def _checkpointing_main_kernel( + pid_h * stride_old_dA_cumsum_head ) old_dA_cumsum_all = tl.load( - old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0 + old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, + mask=old_window_mask, other=0.0, ).to(tl.float32) # Load dA_cumsum at prev_k-1 directly via pointer math (avoids masked reduction). - # Clamp to [0, T-1] defensively — out-of-contract PNAT > T would read OOB. - prev_k_idx = tl.minimum(tl.maximum(prev_num_accepted_tokens - 1, 0), T - 1) + # Clamp to [0, MAX-1] defensively — caller contract gives PNAT ≤ MAX. + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( tl.float32 ) # Step 0 invariant: PNAT=0 means `state` is already last step's state (not - # two back). coeff is all-zero (offs_t < 0), total_decay is 1.0, so the - # replay leaves `state` unchanged — cache contents don't matter on step 0. + # two back). coeff is all-zero (old_window_mask all-false), total_decay + # is 1.0, so the replay leaves `state` unchanged — cache contents don't matter. coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all - coeff = tl.where(offs_t < prev_num_accepted_tokens, coeff, 0.0) - # Load old_x: (BLOCK_SIZE_T, BLOCK_SIZE_M) — single-buffered + # Load old_x at [0, PNAT) of the WINDOW axis (single-buffered cache). old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head old_x_all = tl.load( - old_x_base + offs_t[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - mask=t_mask[:, None] & m_mask[None, :], + old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + mask=old_window_mask[:, None] & m_mask[None, :], other=0.0, ) - # Load old_B from READ buffer: (BLOCK_SIZE_T, BLOCK_SIZE_DSTATE) + # Load old_B from READ buffer at [0, PNAT) of the WINDOW axis. old_B_base = ( old_B_ptr + cache_batch_idx * stride_old_B_cache @@ -601,8 +623,8 @@ def _checkpointing_main_kernel( + group_idx * stride_old_B_group ) old_B_all = tl.load( - old_B_base + offs_t[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - mask=t_mask[:, None] & n_mask[None, :], + old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + mask=old_window_mask[:, None] & n_mask[None, :], other=0.0, ).to(tl.float32) @@ -1292,6 +1314,7 @@ def grid(META): rand_seed, pad_slot_id, T, + max_window, # MAX_REPLAY_BUFFER_LENGTH dim, dstate, nheads // ngroups, diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index a3871ee6d4f4..fa39ade6d3bb 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -172,23 +172,27 @@ def test_checkpointing_state_update( state0_scales = None ref_input_state = state0.float() - # Old inputs: T tokens per batch request - x1 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + # Old inputs: up to `max_window` tokens per batch request, so the test + # loop can probe PNAT > T-1 (which the prior T-token setup couldn't + # reach). step1_T = max_window covers the full PNAT range we sweep. + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) # stride(-1)=0 - B1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - # Capture intermediate SSM states using selective_state_update. + # Capture intermediate SSM states using selective_state_update across + # all step1_T positions — gives us reference states for k ∈ [0, step1_T]. states_buffer_f32 = torch.zeros( - cache_size, T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + cache_size, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 ) cache_idx_for_capture = ( state_batch_indices if paged_cache else torch.arange(batch, device=device, dtype=torch.int32) ) - out1 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) selective_state_update( ref_input_state.clone(), x1, @@ -201,7 +205,7 @@ def test_checkpointing_state_update( dt_softplus=True, state_batch_indices=cache_idx_for_capture, intermediate_states_buffer=states_buffer_f32, - cache_steps=T, + cache_steps=step1_T, out=out1, disable_state_update=True, ) @@ -219,11 +223,11 @@ def test_checkpointing_state_update( cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) # Fill each slot's active buffer (= cache_buf_idx) with step 1's data at - # positions [0:T). Positions [T:max_window) stay as torch.randn garbage — - # they're outside the test's PNAT range, kernel doesn't read them. - # The OTHER (inactive) buffer has random garbage to catch indexing bugs. + # positions [0:step1_T) = [0:max_window). Whole buffer covered so PNAT + # values up to max_window are exercised. Inactive buffer has random + # garbage to catch indexing bugs. slots = state_batch_indices if paged_cache else slice(None) - old_x[slots, :T] = x1 + old_x[slots, :step1_T] = x1 # Compute processed dt and dA_cumsum for step 1 dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) @@ -234,16 +238,25 @@ def test_checkpointing_state_update( for i, slot in enumerate(slot_indices): buf = cache_buf_idx[slot].item() batch_idx = i # maps slot back to the batch index - old_B[slot, buf, :T] = B1[batch_idx] - old_dt[slot, buf, :, :T] = dt1[batch_idx].T # (T, nheads) → (nheads, T) - old_dA_cumsum[slot, buf, :, :T] = dA_cumsum1[batch_idx].T # (T, nheads) → (nheads, T) - - # Main loop: test each k (number of old tokens replayed). For - # write_checkpoint=False, the kernel writes new tokens at [k:k+T) of the - # active buffer — skip k where this would exceed max_window. - for k in range(T + 1): - if not write_checkpoint and k + T > max_window: - continue + old_B[slot, buf, :step1_T] = B1[batch_idx] + old_dt[slot, buf, :, :step1_T] = dt1[batch_idx].T # (step1_T, nheads) → (nheads, step1_T) + old_dA_cumsum[slot, buf, :, :step1_T] = dA_cumsum1[batch_idx].T + + # Main loop: test each k (number of old tokens replayed). + # write_checkpoint=False (nowrite): k ∈ [0, max_window-T] — new tokens + # append at [k, k+T) of the active buffer; need k+T ≤ max_window. + # write_checkpoint=True (write): k ∈ [max_window-T+1, max_window] — + # new tokens land in the staging buffer at [0, T); k > max_window-T + # captures the overflow case that triggers a checkpoint in production. + # Combined sweep covers the full k ∈ [0, max_window] with the + # appropriate boundary handling per mode. + if write_checkpoint: + k_lo = max(0, max_window - T + 1) + k_hi = max_window + 1 # exclusive + else: + k_lo = 0 + k_hi = max_window - T + 1 # exclusive + for k in range(k_lo, k_hi): torch.manual_seed(k + 100) x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) From aeab6a765d7b9d7c92aa6bb8f8431e45ad5e4ac2 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 6 May 2026 00:48:42 -0700 Subject: [PATCH 14/89] rect_precompute opt C: vectorize per-head loop, keep combo (H,T,K) live across gdc_wait MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace both the pre-wait and post-wait per-head loops with single vectorized 3D tensor ops. Pre-wait builds combo_block of shape (HEADS_PER_BLOCK, T, K) in one pass; the tile stays in registers across gdc_wait (Triton handles the cross-barrier liveness when the value naturally flows to a post-wait use). Post-wait computes: rect_CB_scaled_block = where(causal[None,:,:], raw_rect_CB[None,:,:] * combo_block, 0) tl.store(cb_scaled[H,T,K]) # one 3D store vs optB which stored combo to cb_scaled scratch pre-wait and reloaded post-wait. Eliminates the global memory roundtrip and the two per-head loops; pre-wait and post-wait are each a single tensor op. Same registers as optB. Net wins everywhere vs optB (-1 to -8%) and vs baseline (-2 to -8%, peak at b=16/32 across all 4 dtypes). Closes 7 cells that previously favored replay (b=1 fp16/int16/int8, b=4 int16/int8, b=16 int8). Only 3 of 64 cells still favor replay (b=1 int8 SR, b=4 int16 SR, b=16 int16 SR — all by ≤1.5%). Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 211 +++++++++--------- 1 file changed, 109 insertions(+), 102 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 465ac23f5579..5beb16bf3fc2 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -555,91 +555,100 @@ def _rectangle_precompute_kernel( other=0.0, ) - # Per-head pre-wait: compute decay_vec_full AND combo = factor_dt * exp_diff, - # storing both to scratch (combo overlays cb_scaled — overwritten post-wait - # with rect_CB_scaled). This pulls all factor_dt / exp_diff math out of - # the post-wait loop, which becomes a single load + multiply + masked store - # per head. Loads of old_dt/old_dA_cumsum/dt_at_kn/dA_cumsum_at_kn that - # used to be L1 warmers are now real consumers (no fake-keep wraparound). + # Pre-wait: vectorized across HEADS_PER_BLOCK heads. Compute decay_vec_full + # (H, T) and combo = factor_dt * exp_diff (H, T, K). Store decay_vec_full; + # combo_block stays in registers across gdc_wait — used directly post-wait + # to compute rect_CB_scaled without a global memory roundtrip. prev_k_idx = tl.minimum( tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 ) - for h_local in range(HEADS_PER_BLOCK): - head_idx = first_head + h_local - old_dt_read_base_h = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + buf_active * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - old_dA_cumsum_read_base_h = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_active * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - old_dt_write_base_h = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - old_dA_cumsum_write_base_h = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - old_dt_all_h = tl.load( - old_dt_read_base_h + safe_old_k * stride_old_dt_T, - mask=is_old_k, other=0.0, - ).to(tl.float32) - old_dA_cumsum_all_h = tl.load( - old_dA_cumsum_read_base_h + safe_old_k * stride_old_dA_cumsum_T, - mask=is_old_k, other=0.0, - ).to(tl.float32) - total_dA_cumsum_h = tl.load( - old_dA_cumsum_read_base_h + prev_k_idx * stride_old_dA_cumsum_T - ).to(tl.float32) - dA_cumsum_new_h = tl.load( - old_dA_cumsum_write_base_h + (write_offset + offs_t) * stride_old_dA_cumsum_T, - mask=t_mask, other=0.0, - ).to(tl.float32) - dt_at_kn_h = tl.load( - old_dt_write_base_h + (write_offset + safe_k_new) * stride_old_dt_T, - mask=is_new_k, other=0.0, - ).to(tl.float32) - dA_cumsum_at_kn_h = tl.load( - old_dA_cumsum_write_base_h + (write_offset + safe_k_new) * stride_old_dA_cumsum_T, - mask=is_new_k, other=0.0, - ).to(tl.float32) + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) - # decay_vec_full = total_decay * exp(cumAdt_new) - total_decay_h = tl.where( - prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum_h), 1.0 - ) - decay_vec_full_h = total_decay_h * tl.exp(dA_cumsum_new_h) - decay_vec_base_h = decay_vec_ptr + pid_b * stride_dv_batch + head_idx * stride_dv_head - tl.store( - decay_vec_base_h + offs_t * stride_dv_t, decay_vec_full_h, mask=t_mask - ) + # Per-head bases (H,) — broadcast with offs_k or offs_t for 2D loads. + old_dt_read_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_active * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_read_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + ) + old_dt_write_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_write_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + ) - # combo = factor_dt * exp_diff — stored to cb_scaled scratch (overwritten - # post-wait with rect_CB_scaled = where(causal, raw_rect_CB * combo, 0)). - factor_dt_h = tl.where(is_old_k, old_dt_all_h, dt_at_kn_h) - s_k_h = tl.where( - is_old_k, total_dA_cumsum_h - old_dA_cumsum_all_h, -dA_cumsum_at_kn_h - ) - exp_diff_h = tl.exp(s_k_h[None, :] + dA_cumsum_new_h[:, None]) - combo_h = factor_dt_h[None, :] * exp_diff_h - cb_scaled_base_h = ( - cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head - ) - tl.store( - cb_scaled_base_h + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, - combo_h, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), - ) + # (H, K) loads at [0, PNAT) — old data from previous step. + hk_mask = is_old_k[None, :] # (1, K) + old_dt_all = tl.load( + old_dt_read_h[:, None] + safe_old_k[None, :] * stride_old_dt_T, + mask=hk_mask, other=0.0, + ).to(tl.float32) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, + mask=hk_mask, other=0.0, + ).to(tl.float32) + # (H,) scalar-per-head: total_dA_cumsum at prev_k_idx. + total_dA_cumsum = tl.load( + old_dA_cumsum_read_h + prev_k_idx * stride_old_dA_cumsum_T + ).to(tl.float32) + # (H, T) loads at [PNAT, PNAT+T) — this step's dA_cumsum_new from loop 1. + ht_mask = t_mask[None, :] # (1, T) + dA_cumsum_new = tl.load( + old_dA_cumsum_write_h[:, None] + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T, + mask=ht_mask, other=0.0, + ).to(tl.float32) + # (H, K) loads at K_NEW_SHIFT-shifted positions for new tokens. + hkn_mask = is_new_k[None, :] + dt_at_kn = tl.load( + old_dt_write_h[:, None] + + (write_offset + safe_k_new)[None, :] * stride_old_dt_T, + mask=hkn_mask, other=0.0, + ).to(tl.float32) + dA_cumsum_at_kn = tl.load( + old_dA_cumsum_write_h[:, None] + + (write_offset + safe_k_new)[None, :] * stride_old_dA_cumsum_T, + mask=hkn_mask, other=0.0, + ).to(tl.float32) + + # decay_vec_full = total_decay * exp(cumAdt_new). (H, T). + total_decay = tl.where( + prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0 + ) # (H,) + decay_vec_full_block = total_decay[:, None] * tl.exp(dA_cumsum_new) # (H, T) + decay_vec_addrs = ( + decay_vec_ptr + + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) # (H, T) + tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) + + # combo_block = factor_dt * exp_diff — (H, T, K). Stays in registers + # across gdc_wait. + factor_dt = tl.where(is_old_k[None, :], old_dt_all, dt_at_kn) # (H, K) + s_k = tl.where( + is_old_k[None, :], + total_dA_cumsum[:, None] - old_dA_cumsum_all, + -dA_cumsum_at_kn, + ) # (H, K) + # exp_diff (H, T, K) = exp(s_k (H, 1, K) + dA_cumsum_new (H, T, 1)). + exp_diff = tl.exp(s_k[:, None, :] + dA_cumsum_new[:, :, None]) + combo_block = factor_dt[:, None, :] * exp_diff # (H, T, K) # ---- gdc_wait: from here on we depend on conv1d's outputs ---- if LAUNCH_WITH_PDL: @@ -692,28 +701,26 @@ def _rectangle_precompute_kernel( is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] - # Per-head post-wait: load pre-computed combo (factor_dt * exp_diff) from - # cb_scaled scratch, apply causal mask × raw_rect_CB, write rect_CB_scaled - # back to the same address. All factor_dt / exp_diff math happened - # pre-wait above; this loop is pure load + multiply + masked store. - for h_local in range(HEADS_PER_BLOCK): - head_idx = first_head + h_local - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head - cb_load_mask = (offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K) - combo_loaded = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, - mask=cb_load_mask, other=0.0, - ) - rect_CB_scaled = tl.where( - causal_combined, - raw_rect_CB * combo_loaded, - 0.0, - ) - tl.store( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, - rect_CB_scaled, - mask=cb_load_mask, - ) + # Post-wait vectorized: combo_block (H, T, K) is still live in registers. + # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store as + # one (H, T, K) tile. + rect_CB_scaled_block = tl.where( + causal_combined[None, :, :], + raw_rect_CB[None, :, :] * combo_block, + 0.0, + ) # (H, T, K) + cb_scaled_addrs = ( + cb_scaled_ptr + + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_k[None, None, :] * stride_cb_j + ) # (H, T, K) + cb_store_mask_3d = ( + (offs_t[None, :, None] < BLOCK_SIZE_T) + & (offs_k[None, None, :] < BLOCK_SIZE_K) + ) # (1, T, K) → broadcasts to (H, T, K) + tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) # Main kernel: tl.dot replay + precomputed CB output. From 7740397e44b897b13b19b024b4b22e82f73c78f7 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 4 May 2026 16:13:08 -0700 Subject: [PATCH 15/89] Rectangle nowrite kernel + benchmark axes (--rectangle-for-nowrite, --write-modes) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds dedicated _rectangle_precompute_kernel and _rectangle_main_kernel for the nowrite path. Rectangle CB combines old (cache) and new (input) B contributions in a single (T, BLOCK_K=16) matmul; main kernel skips replay update and HBM state write entirely. Wrapper picks via rectangle_for_nowrite kwarg (no-op on write_checkpoint=True; replay is unaffected). Hoisted optimizations in rectangle precompute: - Group-level old_B from cache loaded above gdc_wait so HBM latency overlaps with conv1d - Per-head decay_vec_full = total_decay * exp(cumAdt_new) computed and stored above gdc_wait (folds total_decay into precomp output so main can apply post-matmul without materializing a state_prev_decayed (M, dstate) tile) - Per-head L1 warmer trick: old_dt, old_dA_cumsum, dt_at_kn, dA_cumsum_at_kn pre-loaded above gdc_wait via tl.where(False, ..., 0.0) sentinel so Triton can't DCE; post-wait loop's reload hits L1 cheaply. Hoisted in rectangle main: old_x_load above gdc_wait (cache read, no conv1d dep). Wrapper renamed rectangle (placeholder) to rectangle_for_nowrite; existing RECTANGLE constexpr removed from _checkpointing_main_kernel signature. Tests: test_checkpointing_state_update parametrized over {write, no_write_replay, no_write_rectangle}; 542 pass (398 base + 144 rectangle). Rectangle correctness verified across all dtypes (bf16/fp16/fp32/int8/int16/fp8) and full PNAT range up to MAX-T. Benchmark plumbing: - --rectangle-for-nowrite 0,1 sweeps both replay-nowrite and rectangle-nowrite in one nsys process (silent no-op on write cells) - --write-modes 0,1 sweeps write+nowrite in one nsys process for apples-to-apples comparison; auto-skips invalid (write=0, prev_k+T>MAX) - NVTX tag carries _RECT={0,1} and _WC={0,1} Production sweep summary (--with-conv1d, full M/W/pW/HPB at all dtypes/SR): rectangle wins replay-nowrite at b≥32 (-6 to -32%); replay-nowrite wins or ties at b≤16 on int variants. Full results in CHECKPOINTING_DESIGN.md 'Overnight sweep (2026-05-05)' section. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 883 +++++++++++++++++- ...benchmark_replay_selective_state_update.py | 94 +- .../mamba/test_checkpointing_state_update.py | 12 +- 3 files changed, 943 insertions(+), 46 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 6d40df528030..d9c9fe278def 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -352,6 +352,415 @@ def _checkpointing_precompute_kernel( ) +# Rectangle precompute kernel: produces a (T, K) CB rectangle that combines +# old-token (B from cache, k ∈ [0, PNAT)) and new-token (B from input, k ∈ +# [MAX-T, MAX) at compile-time-static shift) contributions in a single matmul. +# Used only on no-checkpoint steps (nowrite path); pairs with +# `_rectangle_main_kernel`. K-axis size = max(np2(MAX_REPLAY_BUFFER_LENGTH), +# 16); the static layout is sound because nowrite implies PNAT + T <= +# MAX_REPLAY_BUFFER_LENGTH, so old [0, PNAT) and new [MAX-T, MAX) never +# overlap. Also folds total_decay into decay_vec at precomp time so main +# can skip materializing a state_prev_decayed (M, dstate) tile. + + +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _rectangle_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle + decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) + # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite + # path: read from buf_active at [0, PNAT), write new tokens at + # [PNAT, PNAT+T) of buf_active (same buffer). + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # rectangle K-axis bound + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides (rectangle: (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides: (cache, 2, T_max, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T_max) + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T_max) + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, +): + pid_b = tl.program_id(axis=0) + pid_hg = tl.program_id(axis=1) + first_head = pid_hg * HEADS_PER_BLOCK + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_buf = buf_active + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout: old at [0, PNAT) (mask is_old_k); + # new at [MAX-T, MAX) at compile-time shift K_NEW_SHIFT = MAX - T. + # PNAT + T <= MAX is guaranteed on the nowrite path → no overlap. + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis (output rows) + offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis (rectangle input cols) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + t_mask = offs_t < T + n_mask = offs_n < dstate + + # K-axis masks + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - K_NEW_SHIFT + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Loop 1: per-head dt processing. dt → dt_processed → dA_cumsum → + # decay_vec_new (= exp(cumAdt_new)). Stored to write_buf for next step. + # decay_vec_full (= total_decay * decay_vec_new) is finalized in loop 2 + # once total_decay is loaded; loop 1 stores raw decay_vec_new to scratch. + for h_local in range(HEADS_PER_BLOCK): + head_idx = first_head + h_local + + dt_base = dt_ptr + pid_b * stride_dt_batch + head_idx * stride_dt_head + dt = tl.load(dt_base + offs_t * stride_dt_T, mask=t_mask, other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + head_idx * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias + if DT_SOFTPLUS: + dt = softplus(dt) + + A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) + dA_cumsum = tl.cumsum(A * dt, axis=0) + + # Store dt and dA_cumsum to write_buf at [write_offset, write_offset+T) + # for next step's replay/rectangle use. + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + tl.store( + old_dt_base + (write_offset + offs_t) * stride_old_dt_T, + dt, + mask=t_mask, + ) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + tl.store( + old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, + dA_cumsum, + mask=t_mask, + ) + + # ---- Hoisted: cache-only loads independent of conv1d ---- + # old_B (group-level, BLOCK_K × BLOCK_DSTATE = ~8KB tile) and the + # decay_vec_full per-head pre-compute (which writes to DRAM and doesn't + # need cross-gdc_wait variables) are issued BEFORE gdc_wait so their + # HBM latency overlaps with conv1d. Per-head factor_dt/exp_diff stay + # below gdc_wait — they need cross-iteration spans, which Triton can't + # express without a DRAM round-trip; the per-head LOADS in the post- + # wait loop are small and cheap, so leave them. + group_idx = first_head // nheads_ngroups_ratio + + # Group-level: old B from active buffer at [0, PNAT) of the K-axis. + old_B_read_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + buf_active * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_load = tl.load( + old_B_read_base + + safe_old_k[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + mask=is_old_k[:, None] & n_mask[None, :], + other=0.0, + ) + + # Per-head: pre-compute decay_vec_full (= total_decay * exp(cumAdt_new[t])) + # and store to scratch. Also pre-load the rest of the per-head cache + # data (old_dt, old_dA_cumsum, dt_at_kn, dA_cumsum_at_kn) as L1 warmers + # — Triton can't span their values across the gdc_wait barrier without a + # DRAM round-trip, but L1 is per-SM-persistent so the post-wait loop's + # reload hits the cache cheaply. Folding the warmer loads into + # decay_vec_full's `tl.where(False, ...)` keeps Triton from DCE-ing them. + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + for h_local in range(HEADS_PER_BLOCK): + head_idx = first_head + h_local + old_dt_read_base_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_active * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + old_dA_cumsum_read_base_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + old_dt_write_base_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + old_dA_cumsum_write_base_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + # Required for decay_vec_full + total_dA_cumsum_h = tl.load( + old_dA_cumsum_read_base_h + prev_k_idx * stride_old_dA_cumsum_T + ).to(tl.float32) + dA_cumsum_new_h = tl.load( + old_dA_cumsum_write_base_h + (write_offset + offs_t) * stride_old_dA_cumsum_T, + mask=t_mask, other=0.0, + ).to(tl.float32) + # L1 warmers — same addresses the post-wait loop will reload. + warm_old_dt = tl.load( + old_dt_read_base_h + safe_old_k * stride_old_dt_T, + mask=is_old_k, other=0.0, + ).to(tl.float32) + warm_old_dA = tl.load( + old_dA_cumsum_read_base_h + safe_old_k * stride_old_dA_cumsum_T, + mask=is_old_k, other=0.0, + ).to(tl.float32) + warm_dt_at_kn = tl.load( + old_dt_write_base_h + (write_offset + safe_k_new) * stride_old_dt_T, + mask=is_new_k, other=0.0, + ).to(tl.float32) + warm_dA_at_kn = tl.load( + old_dA_cumsum_write_base_h + (write_offset + safe_k_new) * stride_old_dA_cumsum_T, + mask=is_new_k, other=0.0, + ).to(tl.float32) + + total_decay_h = tl.where( + prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum_h), 1.0 + ) + decay_vec_full_h = total_decay_h * tl.exp(dA_cumsum_new_h) + # Force the warmer loads to participate via a runtime-False guard. + # prev_num_accepted_tokens is dynamic, so Triton can't DCE this. + warm_keep = ( + tl.sum(warm_old_dt) + tl.sum(warm_old_dA) + + tl.sum(warm_dt_at_kn) + tl.sum(warm_dA_at_kn) + ) + decay_vec_full_h = tl.where( + prev_num_accepted_tokens < 0, + decay_vec_full_h + warm_keep, + decay_vec_full_h, + ) + decay_vec_base_h = decay_vec_ptr + pid_b * stride_dv_batch + head_idx * stride_dv_head + tl.store( + decay_vec_base_h + offs_t * stride_dv_t, decay_vec_full_h, mask=t_mask + ) + + # ---- gdc_wait: from here on we depend on conv1d's outputs ---- + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # Conv1d outputs: B and C + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + B_new_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_all = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_new_orig = tl.load( + B_new_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_new_shifted = tl.load( + B_new_base + safe_k_new[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=is_new_k[:, None] & n_mask[None, :], + other=0.0, + ) + # Disjoint masks: old at [0, PNAT), new at [K_NEW_SHIFT, K_NEW_SHIFT+T). + B_combined = old_B_load + B_new_shifted + raw_rect_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_combined).to(tl.bfloat16)) + + # Append new B to cache at [PNAT, PNAT+T) of write_buf (once per group). + if first_head % nheads_ngroups_ratio == 0: + old_B_write_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + write_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_write_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_new_orig, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # Causal mask (BLOCK_SIZE_T × BLOCK_SIZE_K, shared across heads). + t_idx_2d = offs_t[:, None] + k_idx_2d = offs_k[None, :] + is_old_k_2d = k_idx_2d < prev_num_accepted_tokens + k_new_idx_2d = k_idx_2d - K_NEW_SHIFT + is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) + causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] + + # Per-head: factor_dt + exp_diff + rect_CB_scaled. Loads are post-wait + # but cache-resident from this kernel's earlier writes (loop 1 + the + # hoisted decay_vec_full block above), so they hit L1. Recomputing + # total_dA_cumsum / dA_cumsum_new here is cheaper than spanning the + # gdc_wait via a DRAM round-trip. + for h_local in range(HEADS_PER_BLOCK): + head_idx = first_head + h_local + old_dt_read_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_active * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + old_dA_cumsum_read_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + old_dt_write_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + old_dA_cumsum_write_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + old_dt_all = tl.load( + old_dt_read_base + safe_old_k * stride_old_dt_T, mask=is_old_k, other=0.0 + ).to(tl.float32) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_read_base + safe_old_k * stride_old_dA_cumsum_T, + mask=is_old_k, other=0.0, + ).to(tl.float32) + total_dA_cumsum = tl.load( + old_dA_cumsum_read_base + prev_k_idx * stride_old_dA_cumsum_T + ).to(tl.float32) + dA_cumsum_new = tl.load( + old_dA_cumsum_write_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, + mask=t_mask, other=0.0, + ).to(tl.float32) + dt_at_kn = tl.load( + old_dt_write_base + (write_offset + safe_k_new) * stride_old_dt_T, + mask=is_new_k, other=0.0, + ).to(tl.float32) + dA_cumsum_at_kn = tl.load( + old_dA_cumsum_write_base + (write_offset + safe_k_new) * stride_old_dA_cumsum_T, + mask=is_new_k, other=0.0, + ).to(tl.float32) + + factor_dt = tl.where(is_old_k, old_dt_all, dt_at_kn) + s_k = tl.where( + is_old_k, total_dA_cumsum - old_dA_cumsum_all, -dA_cumsum_at_kn + ) + exp_diff = tl.exp(s_k[None, :] + dA_cumsum_new[:, None]) + + rect_CB_scaled = tl.where( + causal_combined, + raw_rect_CB * factor_dt[None, :] * exp_diff, + 0.0, + ) + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head + tl.store( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + rect_CB_scaled, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + ) + + # Main kernel: tl.dot replay + precomputed CB output. # Grid: (cdiv(dim, M), batch, nheads). @@ -487,12 +896,12 @@ def _checkpointing_main_kernel( # this from state.dtype; kernel-entry static_assert below pins the # invariant that it must coincide with int8/int16/float8e4nv state dtype. QUANT_MAX: tl.constexpr, - # Checkpointing flags + # Checkpointing flag WRITE_CHECKPOINT: tl.constexpr, # When True: quantize+write post-replay state to HBM (checkpoint step). # When False: skip state write entirely (non-checkpoint step). - RECTANGLE: tl.constexpr, # Reserved for the rectangle non-checkpoint optimization path. - # Currently asserted False at the wrapper; kernel takes the - # replay-style code path unconditionally. + # The rectangle non-checkpoint path is implemented in + # _rectangle_main_kernel (separate kernel pair, picked by + # the wrapper via rectangle_for_nowrite=True). ): # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized # state dtype (int8 / int16 / float8e4nv) and only those. Cheap @@ -823,6 +1232,261 @@ def _checkpointing_main_kernel( tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) +# Rectangle main kernel (nowrite-only): no replay step, no state HBM write, +# no SR codegen. state_out is computed from state_prev directly using the +# precomp-folded decay_vec_full; token_out is a single rectangle matmul over +# the (T, K) CB rectangle and (K, M) x_combined. Pairs with +# `_rectangle_precompute_kernel`. + + +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _rectangle_main_kernel( + # Pointers + state_ptr, + state_scales_ptr, # only consulted when QUANT_MAX > 0 + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, # rectangle (batch, nheads, T, K) + decay_vec_ptr, # folded (batch, nheads, T) — total_decay * exp(cumAdt_new[t]) + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides (no quant-store path; state read-only for state_out) + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides: (cache, nheads, dim) — fp32, broadcast over dstate + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides: (cache, T_max, nheads, dim) — single-buffered + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides (rectangle (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + QUANT_MAX: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout (matches precompute). + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_k = tl.arange(0, BLOCK_SIZE_K) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # K-axis masks + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - K_NEW_SHIFT + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Load state (with dequant on load if quantized). Read-only — no HBM + # write on the nowrite path. + state_ptr += cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, other=1.0, + ).to(tl.float32) + state = state * decode_scale[:, None] + + # Group / pointer offset setup + group_idx = pid_h // nheads_ngroups_ratio + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + # Hoist: old_x from cache doesn't depend on conv1d/precompute, so issue + # the load BEFORE gdc_wait so its HBM latency overlaps with conv1d. + old_x_load = tl.load( + old_x_base + + safe_old_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + mask=is_old_k[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + + # PDL gate: precompute outputs (cb_scaled, decay_vec_full) become safe + # after gdc_wait. conv1d outputs (x, C) also gated by the chained PDL. + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # Load C and x (conv1d outputs after PDL wait) + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + x_all = tl.load( + x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + # Append new x to cache at [PNAT, PNAT+T) of the active buffer. + tl.store( + old_x_base + + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_all, + mask=t_mask[:, None] & m_mask[None, :], + ) + x_all = x_all.to(tl.float32) + + # Build x_combined for the rectangle token_out matmul (STATIC layout). + # old_x_load is from cache (loaded above gdc_wait); new_x_shifted reads + # x (conv1d output) at compile-time-shifted [K_NEW_SHIFT, K_NEW_SHIFT+T). + # Disjoint masks → safe to add. + new_x_shifted = tl.load( + x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=is_new_k[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + x_combined = old_x_load + new_x_shifted # (BLOCK_SIZE_K, BLOCK_SIZE_M) + + # Load precomputed rectangle CB and folded decay_vec. + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec_full = tl.load( + decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 + ).to(tl.float32) + + # state_out: state_prev contribution to output, with decay folded post-matmul. + # No state_prev_decayed (M, dstate) materialization — state is consumed + # directly by the matmul, then decay_vec_full multiplies the (T, M) result. + state_out = ( + tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) + * decay_vec_full[:, None] + ) + + # token_out: combined old + new tokens contribution via the rectangle. + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) + + out_all = state_out + token_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + for t in range(T): + z_t = tl.load( + z_ptr + t * stride_z_T + offs_m * stride_z_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + out_t = tl.sum(tl.where((offs_t == t)[:, None], out_all, 0.0), axis=0) + out_t = out_t * z_t * tl.sigmoid(z_t) + tl.store(out_ptr + t * stride_out_T + offs_m * stride_out_dim, out_t, mask=m_mask) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + # Python wrapper @@ -859,7 +1523,7 @@ def checkpointing_state_update( launch_with_pdl=False, use_internal_pdl=True, write_checkpoint: bool = True, - rectangle: bool = False, + rectangle_for_nowrite: bool = False, _block_size_m: int | None = None, _num_warps: int | None = None, _num_stages: int | None = None, @@ -938,19 +1602,14 @@ def checkpointing_state_update( launch_with_pdl = False use_internal_pdl = False - # Constexpr modes: - # write_checkpoint=True, rectangle=False → checkpoint step (default). - # write_checkpoint=False, rectangle=False → non-checkpoint step (skip state HBM write). - # write_checkpoint=False, rectangle=True → reserved for the rectangle non-checkpoint - # optimization path (not wired yet). - # write_checkpoint=True, rectangle=True → not supported. - if rectangle: - raise NotImplementedError( - "RECTANGLE path is not wired yet; pass rectangle=False." - ) - assert not (write_checkpoint and rectangle), ( - "WRITE_CHECKPOINT and RECTANGLE are mutually exclusive." - ) + # Path selection: + # write_checkpoint=True, rectangle_for_nowrite=* → checkpoint step + # (replay-style write kernel; rectangle_for_nowrite ignored). + # write_checkpoint=False, rectangle_for_nowrite=False → non-checkpoint step + # via replay-style nowrite kernel. + # write_checkpoint=False, rectangle_for_nowrite=True → non-checkpoint step + # via dedicated rectangle nowrite kernel (single fused old+new matmul). + use_rectangle = rectangle_for_nowrite and not write_checkpoint # --- Hardware support gates --- # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX @@ -1074,10 +1733,15 @@ def checkpointing_state_update( device = x.device BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) + # Rectangle K-axis bound = window (max_window). Computed unconditionally + # so the launch sites can refer to it; only used on the rectangle path. + BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) - # Allocate precomputed intermediates (per-call, not cached) + # Allocate precomputed intermediates (per-call, not cached). CB shape + # depends on path: replay-style is square (T, T); rectangle is (T, K). + cb_T_dim = BLOCK_SIZE_K if use_rectangle else BLOCK_SIZE_T cb_scaled = torch.empty( - batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_T, device=device, dtype=torch.float32 + batch, nheads, BLOCK_SIZE_T, cb_T_dim, device=device, dtype=torch.float32 ) decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) @@ -1207,6 +1871,171 @@ def checkpointing_state_update( assert heads_per_block <= heads_per_group, ( f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" ) + + # state_scales pointer + strides: real tensor when quantized, otherwise + # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). + # Hoisted out of the branch so both kernels can share. + if is_quantized: + state_scales_arg = state_scales + state_scales_strides = ( + state_scales.stride(0), + state_scales.stride(1), + state_scales.stride(2), + ) + else: + state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 + state_scales_strides = (0, 0, 0) + + if use_rectangle: + _rectangle_precompute_kernel[(batch, nheads // heads_per_block)]( + dt, + dt_bias, + A, + B, + C, + cb_scaled, + decay_vec, + old_B, + old_dt, + old_dA_cumsum, + cache_buf_idx, + prev_num_accepted_tokens, + state_batch_indices, + pad_slot_id, + T, + max_window, # MAX_REPLAY_BUFFER_LENGTH + dstate, + nheads // ngroups, + # dt strides + dt.stride(0), + dt.stride(1), + dt.stride(2), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + # B strides + B.stride(0), + B.stride(1), + B.stride(2), + B.stride(3), + # C strides + C.stride(0), + C.stride(1), + C.stride(2), + C.stride(3), + # cb_scaled strides (rectangle) + cb_scaled.stride(0), + cb_scaled.stride(1), + cb_scaled.stride(2), + cb_scaled.stride(3), + # decay_vec strides + decay_vec.stride(0), + decay_vec.stride(1), + decay_vec.stride(2), + # old_B strides + old_B.stride(0), + old_B.stride(1), + old_B.stride(2), + old_B.stride(3), + old_B.stride(4), + # old_dt strides + old_dt.stride(0), + old_dt.stride(1), + old_dt.stride(2), + old_dt.stride(3), + # old_dA_cumsum strides + old_dA_cumsum.stride(0), + old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), + old_dA_cumsum.stride(3), + dt_softplus, + HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, + LAUNCH_WITH_PDL=launch_with_pdl, + LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, + HEADS_PER_BLOCK=heads_per_block, + num_warps=precompute_num_warps, + **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), + launch_pdl=launch_with_pdl, + ) + + def grid(META): + return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + + _rectangle_main_kernel[grid]( + state, + state_scales_arg, + old_x, + prev_num_accepted_tokens, + cache_buf_idx, + x, + C, + D, + z, + out, + cb_scaled, + decay_vec, + state_batch_indices, + pad_slot_id, + T, + max_window, # MAX_REPLAY_BUFFER_LENGTH + dim, + dstate, + nheads // ngroups, + # state strides + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + # state_scales strides (cache, head, dim) + state_scales_strides[0], + state_scales_strides[1], + state_scales_strides[2], + # old_x strides + old_x.stride(0), + old_x.stride(1), + old_x.stride(2), + old_x.stride(3), + # x strides + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + # C strides + C.stride(0), + C.stride(1), + C.stride(2), + C.stride(3), + # D strides + *(D.stride(0), D.stride(1)) if D is not None else (0, 0), + # z strides + z_strides[0], + z_strides[1], + z_strides[2], + z_strides[3], + # out strides + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + # cb_scaled strides + cb_scaled.stride(0), + cb_scaled.stride(1), + cb_scaled.stride(2), + cb_scaled.stride(3), + # decay_vec strides + decay_vec.stride(0), + decay_vec.stride(1), + decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + QUANT_MAX=quant_max, + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + return # rectangle path done; skip the replay-style launches below + _checkpointing_precompute_kernel[(batch, nheads // heads_per_block)]( dt, dt_bias, @@ -1281,19 +2110,6 @@ def checkpointing_state_update( def grid(META): return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) - # state_scales pointer + strides: real tensor when quantized, otherwise - # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). - if is_quantized: - state_scales_arg = state_scales - state_scales_strides = ( - state_scales.stride(0), - state_scales.stride(1), - state_scales.stride(2), - ) - else: - state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 - state_scales_strides = (0, 0, 0) - _checkpointing_main_kernel[grid]( state, state_scales_arg, @@ -1384,7 +2200,6 @@ def grid(META): PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, QUANT_MAX=quant_max, WRITE_CHECKPOINT=write_checkpoint, - RECTANGLE=rectangle, num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 9db6d3cb2a65..62ce155086ef 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -513,6 +513,9 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) + rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) + write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) + configs = [] for batch in batch_sizes: for mtp_len in mtp_lengths: @@ -520,17 +523,26 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp for state_dtype in state_dtypes: for act_dtype in act_dtypes: for sr_mode in sr_modes_list: - configs.append( - (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode)) + for write_ckpt in write_modes_list: + # Rectangle is only meaningful for nowrite cells. + effective_rect_list = ( + [False] if write_ckpt else rect_list + ) + for rect in effective_rect_list: + configs.append(( + batch, mtp_len, prev_ks, state_dtype, act_dtype, + sr_mode, rect, write_ckpt, + )) print(f"[compile-warmup] {len(configs)} configs across {max_workers} threads") t0 = time.perf_counter() def _warm(cfg): - batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode = cfg + batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, rect, write_ckpt = cfg _bench_config( args, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, - sr_mode=sr_mode, warmup_only=True, + sr_mode=sr_mode, rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, warmup_only=True, ) errors = [] @@ -560,6 +572,8 @@ def _bench_config( act_dtype: torch.dtype, baseline_fn, sr_mode: str = "RN", + rectangle_for_nowrite: bool = False, + write_checkpoint: bool = True, warmup_only: bool = False, ) -> None: """ @@ -818,7 +832,15 @@ def _parse_sweep(val): num_ctas_values = _parse_sweep(args.num_ctas) # --- Replay kernel, one row per prev_k --- + # Cache T-axis capacity (for prev_k validity check on the nowrite path). + max_window = getattr(args, "max_window", 0) or mtp_len for prev_k in prev_ks: + # On the nowrite path, new tokens append at [prev_k, prev_k+T) of the + # active buffer, so prev_k+T must fit within max_window. Skip + # silently for combinations that don't satisfy this — lets a single + # nsys run sweep both write modes against a shared prev_k list. + if not write_checkpoint and prev_k + mtp_len > max_window: + continue prev_tokens.fill_(prev_k) tag = f"incr_b{batch}_mtp{mtp_len}_k{prev_k}_s{state_dtype_name}_a{act_dtype_name}" @@ -865,7 +887,8 @@ def _run_incr( # variant; replay variant ignores the kwarg. state_scales # is also checkpointing-only (replay kernel doesn't quantize). if args.variant == "checkpointing": - extra_kwargs["write_checkpoint"] = args.write_checkpoint + extra_kwargs["write_checkpoint"] = write_checkpoint + extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite if state_scales_work is not None: extra_kwargs["state_scales"] = state_scales_work variant_fn( @@ -918,6 +941,8 @@ def _run_incr( if num_ctas is not None: parts.append(f"CT={num_ctas}") parts.append(f"SR={1 if use_philox else 0}") + parts.append(f"RECT={1 if rectangle_for_nowrite else 0}") + parts.append(f"WC={1 if write_checkpoint else 0}") sweep_suffix = (" " + ",".join(parts)) if parts else "" sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") @@ -1040,6 +1065,8 @@ def _run_benchmark(args) -> None: ) sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) + rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) + write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) for batch in batch_sizes: for mtp_len in mtp_lengths: @@ -1048,10 +1075,18 @@ def _run_benchmark(args) -> None: for state_dtype in state_dtypes: for act_dtype in act_dtypes: for sr_mode in sr_modes_list: - _bench_config( - args, batch, mtp_len, prev_ks, state_dtype, act_dtype, - baseline_fn, sr_mode=sr_mode, - ) + for write_ckpt in write_modes_list: + # Rectangle only meaningful for nowrite cells. + effective_rect_list = ( + [False] if write_ckpt else rect_list + ) + for rect in effective_rect_list: + _bench_config( + args, batch, mtp_len, prev_ks, state_dtype, act_dtype, + baseline_fn, sr_mode=sr_mode, + rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, + ) if args.profile: torch.cuda.cudart().cudaProfilerStop() @@ -1225,7 +1260,17 @@ def _parse_args() -> argparse.Namespace: help="Whether the checkpointing kernel should write the post-replay " "state to HBM. True = checkpoint step (default). False = " "non-checkpoint step (skip state HBM write + Philox). No effect on " - "the replay variant.", + "the replay variant. Ignored if --write-modes is set.", + ) + parser.add_argument( + "--write-modes", + type=str, + default=None, + help="Comma-separated 0/1 values to sweep both write modes in a " + "single nsys process — for apples-to-apples comparison of write " + "vs nowrite (replay) vs nowrite (rectangle) within one timeline. " + "Skips silently for (write=False, prev_k+T>max_window) combos. " + "When set, overrides --write-checkpoint.", ) parser.add_argument( "--with-conv1d", @@ -1268,6 +1313,16 @@ def _parse_args() -> argparse.Namespace: "dtypes that don't support it (bf16, fp32). Default 'RN' matches " "legacy --philox-rounding=False behavior.", ) + parser.add_argument( + "--rectangle-for-nowrite", + type=str, + default="0", + help="Comma-separated 0/1 values: 0 = replay-style nowrite kernel, " + "1 = dedicated rectangle nowrite kernel. Sweep both with '0,1' to " + "compare in one invocation. Silently no-op for write cells (the " + "write path always uses replay-style). Only applies to the " + "checkpointing variant.", + ) parser.add_argument( "--philox-rounding", action="store_true", @@ -1319,6 +1374,25 @@ def _parse_args() -> argparse.Namespace: if m not in ("RN", "SR"): parser.error(f"--sr-modes value must be RN or SR, got {m!r}") args.sr_modes_list = sr_modes + + rect_modes = [v.strip() for v in args.rectangle_for_nowrite.split(",") if v.strip()] + rect_list = [] + for v in rect_modes: + if v not in ("0", "1"): + parser.error(f"--rectangle-for-nowrite value must be 0 or 1, got {v!r}") + rect_list.append(v == "1") + args.rectangle_for_nowrite_list = rect_list + + if args.write_modes is not None: + wm = [v.strip() for v in args.write_modes.split(",") if v.strip()] + write_list = [] + for v in wm: + if v not in ("0", "1"): + parser.error(f"--write-modes value must be 0 or 1, got {v!r}") + write_list.append(v == "1") + args.write_modes_list = write_list + else: + args.write_modes_list = [args.write_checkpoint] return args diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index fa39ade6d3bb..215e357a0888 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -100,10 +100,17 @@ def _maybe_skip_dtype(state_dtype, use_sr): "T", [6, 10, 16, 27, 32, 55], ids=["T6", "T10", "T16", "T27", "T32", "T55"] ) @pytest.mark.parametrize( - "write_checkpoint", [True, False], ids=["write", "no_write"] + "write_checkpoint,rectangle_for_nowrite", + [ + (True, False), # write path (rectangle_for_nowrite is ignored) + (False, False), # nowrite path via replay-style kernels + (False, True), # nowrite path via dedicated rectangle kernels + ], + ids=["write", "no_write_replay", "no_write_rectangle"], ) def test_checkpointing_state_update( - nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T, write_checkpoint + nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T, + write_checkpoint, rectangle_for_nowrite, ): """ Verify that: @@ -318,6 +325,7 @@ def test_checkpointing_state_update( state_batch_indices=state_batch_indices, state_scales=test_scales, write_checkpoint=write_checkpoint, + rectangle_for_nowrite=rectangle_for_nowrite, ) # Tolerance rationale: the replay kernel uses bf16 tl.dot for four From 736276747263b7ba16f04dc1a7c0b28be0979cc4 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 6 May 2026 10:38:39 -0700 Subject: [PATCH 16/89] checkpointing_precompute (replay) opt: vectorize loops 1+2 via (H,T,T) tile across gdc_wait MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Apply the same pattern as rect_precompute optC. Replace both per-head loops in _checkpointing_precompute_kernel with vectorized (H, T, T) tensor ops. Pre-wait builds scale_combo (H, T, T) = decay_matrix * dt[:, None, :]; the tile stays in registers across gdc_wait. Post-wait collapses to a single 3D op: CB_scaled_block = where(valid_mask[None,:,:], raw_CB[None,:,:] * scale_combo, 0) tl.store(cb_scaled[H,T,T]) vs baseline which stored dt + dA_cumsum to cache pre-wait, RELOADED them post-wait, then per-head computed decay_matrix * dt and CB_scaled. The post-wait cache reload was the obvious optC analog — it disappears here. Affects both write-checkpoint and replay-nowrite paths (both use _checkpointing_precompute_kernel). Same-node sweep at b=1-64 shows 67 better, 1 worse (b=16 int16 RN +1.1% at the magic config we've been investigating), 12 same. Best wins at b=32 fp8 SR -17.1% and b=16 fp16 SR (WC=1) -11.6%. Both write-checkpoint and nowrite-replay benefit. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 154 ++++++++---------- 1 file changed, 71 insertions(+), 83 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 5beb16bf3fc2..8594c13b9004 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -222,53 +222,62 @@ def _checkpointing_precompute_kernel( causal_mask = offs_t[:, None] >= offs_t[None, :] valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] - # --- Loop 1: compute per-head dt/dA_cumsum/decay BEFORE gdc_wait --- - # These only depend on dt (from in_proj, not conv1d) and parameters (A, dt_bias). - # Store to cache; will reload after the wait for CB scaling. - for h_local in range(HEADS_PER_BLOCK): - head_idx = first_head + h_local - - dt_base = dt_ptr + pid_b * stride_dt_batch + head_idx * stride_dt_head - dt = tl.load(dt_base + offs_t * stride_dt_T, mask=t_mask, other=0.0).to(tl.float32) - if HAS_DT_BIAS: - dt_bias = tl.load(dt_bias_ptr + head_idx * stride_dt_bias_head).to(tl.float32) - dt = dt + dt_bias - if DT_SOFTPLUS: - dt = softplus(dt) + # --- Vectorized pre-wait phase across HEADS_PER_BLOCK heads --- + # Compute dt, dA_cumsum, decay_vec as (H, T) tiles. Pre-compute + # scale_combo = decay_matrix * dt[:, None, :] as an (H, T, T) tile that + # stays in registers across gdc_wait — eliminates the post-wait reload + # of dt + dA_cumsum and the per-head loop. + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) - A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) - dA_cumsum = tl.cumsum(A * dt, axis=0) - decay_vec = tl.exp(dA_cumsum) + # Load dt (H, T) + dt_addrs = ( + dt_ptr + pid_b * stride_dt_batch + + heads_block[:, None] * stride_dt_head + + offs_t[None, :] * stride_dt_T + ) + dt = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias[:, None] + if DT_SOFTPLUS: + dt = softplus(dt) + + A = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) # (H,) + dA_cumsum = tl.cumsum(A[:, None] * dt, axis=1) # (H, T) + decay_vec = tl.exp(dA_cumsum) # (H, T) + + # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of write_buf. + old_dt_addrs = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block[:, None] * stride_old_dt_head + + (write_offset + offs_t)[None, :] * stride_old_dt_T + ) + tl.store(old_dt_addrs, dt, mask=t_mask[None, :]) - # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of - # write_buf (selected by WRITE_CHECKPOINT — see top of kernel). - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - tl.store( - old_dt_base + (write_offset + offs_t) * stride_old_dt_T, - dt, - mask=t_mask, - ) + old_dA_cumsum_addrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block[:, None] * stride_old_dA_cumsum_head + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T + ) + tl.store(old_dA_cumsum_addrs, dA_cumsum, mask=t_mask[None, :]) - old_dA_cumsum_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - tl.store( - old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, - dA_cumsum, - mask=t_mask, - ) + # decay_vec scratch — always at offs_t. + decay_vec_addrs = ( + decay_vec_ptr + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) + tl.store(decay_vec_addrs, decay_vec, mask=t_mask[None, :]) - # decay_vec is per-call scratch (not cached); always write at offs_t. - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + head_idx * stride_dv_head - tl.store(decay_vec_base + offs_t * stride_dv_t, decay_vec, mask=t_mask) + # scale_combo (H, T, T) = exp(dA_cumsum[h, t1] - dA_cumsum[h, t2]) * dt[h, t2] + # Stays live across gdc_wait — used post-wait to compute CB_scaled. + decay_matrix = tl.exp(dA_cumsum[:, :, None] - dA_cumsum[:, None, :]) # (H, T, T) + scale_combo = decay_matrix * dt[:, None, :] # (H, T, T) # --- Wait for upstream kernel (external PDL) before loading B and C --- # All dt processing above is independent of conv1d outputs. @@ -310,46 +319,25 @@ def _checkpointing_precompute_kernel( mask=t_mask[:, None] & n_mask[None, :], ) - # --- Loop 2: reload per-head dA_cumsum/dt from cache, scale CB --- - # Reload from where loop 1 stored: write_buf at [write_offset, write_offset+T). - for h_local in range(HEADS_PER_BLOCK): - head_idx = first_head + h_local - - # Reload dt and dA_cumsum from cache (just written in loop 1) - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - dt = tl.load( - old_dt_base + (write_offset + offs_t) * stride_old_dt_T, - mask=t_mask, - other=0.0, - ).to(tl.float32) - - old_dA_cumsum_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - dA_cumsum = tl.load( - old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, - mask=t_mask, - other=0.0, - ).to(tl.float32) - - # Scale raw_CB with per-head decay and dt - decay_matrix = tl.exp(dA_cumsum[:, None] - dA_cumsum[None, :]) - CB_scaled = tl.where(valid_mask, raw_CB * decay_matrix * dt[None, :], 0.0) - - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head - tl.store( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, - CB_scaled, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), - ) + # --- Vectorized post-wait phase: scale_combo (H, T, T) is still live in + # registers from pre-wait; multiply by raw_CB (T, T), apply causal mask, + # store as one (H, T, T) tile. --- + CB_scaled_block = tl.where( + valid_mask[None, :, :], + raw_CB[None, :, :] * scale_combo, + 0.0, + ) # (H, T, T) + cb_scaled_addrs = ( + cb_scaled_ptr + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_t[None, None, :] * stride_cb_j + ) # (H, T, T) + cb_store_mask = ( + (offs_t[None, :, None] < BLOCK_SIZE_T) + & (offs_t[None, None, :] < BLOCK_SIZE_T) + ) + tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) # Rectangle precompute kernel: produces a (T, K) CB rectangle that combines From 39f08544fee2bca1d4a2d2ff9a36ba3e9a27550e Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 5 May 2026 14:31:30 -0700 Subject: [PATCH 17/89] Add maindl + dlgrouped modes + PDL-out-of-main for dl-family chains Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 2577 +++++++++++++---- ...benchmark_replay_selective_state_update.py | 278 +- .../modules/mamba/checkpoint_mix_sim.py | 331 +++ .../mamba/test_checkpointing_state_update.py | 195 +- 4 files changed, 2713 insertions(+), 668 deletions(-) create mode 100644 tests/unittest/_torch/modules/mamba/checkpoint_mix_sim.py diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index d9c9fe278def..0f0134a6aa62 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -82,11 +82,8 @@ def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: # Grid: (batch, nheads // HEADS_PER_BLOCK). -@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) @triton.jit() -def _checkpointing_precompute_kernel( +def _replay_precompute_impl( # Input pointers dt_ptr, dt_bias_ptr, @@ -164,11 +161,15 @@ def _checkpointing_precompute_kernel( BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, HEADS_PER_BLOCK: tl.constexpr, # Checkpointing flag — selects target buffer + offset for new-token # cache writes. See "Cache write semantics" block below. - WRITE_CHECKPOINT: tl.constexpr, + # Runtime (not constexpr): the only WRITE_CHECKPOINT-dependent code in + # this body is the write_buf/write_offset selection, which is plain + # arithmetic — no constexpr-shaped tile or whole-block gate. Letting + # it be runtime lets the dynamic dispatch kernel call us once with the + # per-slot needs_write flag instead of inlining two specializations. + write_checkpoint, ): pid_b = tl.program_id(axis=0) pid_hg = tl.program_id(axis=1) # head-group index @@ -182,13 +183,6 @@ def _checkpointing_precompute_kernel( else: cache_batch_idx = pid_b.to(tl.int64) - # Signal main kernel to start (internal PDL). Main's replay phase - # reads only from the READ buffer (written by the PREVIOUS step) — - # safe even if conv1d and this kernel are still running. Main's - # gdc_wait() gates the output phase, which reads conv1d outputs - # (x, C) and this kernel's outputs (cb_scaled, decay_vec). - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() # --- Cache write semantics --- # cache_buf_idx names this step's "active" buffer — the one with the @@ -206,7 +200,7 @@ def _checkpointing_precompute_kernel( # kernel behavior exactly. buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - if WRITE_CHECKPOINT: + if write_checkpoint: write_buf = 1 - buf_active write_offset = 0 else: @@ -352,6 +346,173 @@ def _checkpointing_precompute_kernel( ) +# Replay-style precompute kernel. Thin wrapper around _replay_precompute_impl +# that carries the @triton.heuristics for constexpr derivation; called from +# the Python wrapper on the replay-style path (write or replay-nowrite). +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.jit() +def _checkpointing_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + EARLY_OUT: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does, so + # main can start its setup regardless of how this program ends (pad, + # early-out, or full body). PDL signals are idempotent; main's + # gdc_wait still gates on prerequisite-kernel completion for + # correctness. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + # Per-program early-out gate (option-2 double-launch). When EARLY_OUT + # is False the entire block is constexpr-folded out and the wrapper is + # just an impl call. When True, this kernel only runs for slots whose + # (PNAT + T > MAX) status matches WRITE_CHECKPOINT. + if EARLY_OUT: + pid_b_eo = tl.program_id(axis=0) + if HAS_CACHE_BATCH_INDICES: + cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) + if cbi_eo == pad_slot_id: + return + else: + cbi_eo = pid_b_eo.to(tl.int64) + pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) + if (pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH) != WRITE_CHECKPOINT: + return + _replay_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + WRITE_CHECKPOINT, + ) + + # Rectangle precompute kernel: produces a (T, K) CB rectangle that combines # old-token (B from cache, k ∈ [0, PNAT)) and new-token (B from input, k ∈ # [MAX-T, MAX) at compile-time-static shift) contributions in a single matmul. @@ -363,15 +524,8 @@ def _checkpointing_precompute_kernel( # can skip materializing a state_prev_decayed (M, dstate) tile. -@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) @triton.jit() -def _rectangle_precompute_kernel( +def _rectangle_precompute_impl( # Input pointers dt_ptr, dt_bias_ptr, @@ -445,7 +599,6 @@ def _rectangle_precompute_kernel( BLOCK_SIZE_T: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, HEADS_PER_BLOCK: tl.constexpr, ): pid_b = tl.program_id(axis=0) @@ -459,9 +612,6 @@ def _rectangle_precompute_kernel( else: cache_batch_idx = pid_b.to(tl.int64) - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) @@ -761,115 +911,56 @@ def _rectangle_precompute_kernel( ) -# Main kernel: tl.dot replay + precomputed CB output. -# Grid: (cdiv(dim, M), batch, nheads). - - -@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) -@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics( - {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} -) -@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +# Rectangle precompute kernel. Thin wrapper around _rectangle_precompute_impl +# that carries the @triton.heuristics for constexpr derivation; called from +# the Python wrapper on the rectangle nowrite path. +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) @triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -# Replay axis: separate from the T-output axis. Sized to MAX_REPLAY_BUFFER_LENGTH -# so the kernel can iterate over up to PNAT old-cache positions independent of -# the new-token T axis. In production T=6, MAX=16 → BLOCK_SIZE_T=16, -# BLOCK_SIZE_WINDOW=16 (coincidentally equal); separating them is the -# semantically correct fix and avoids the prior `offs_t < T` masking bug -# that under-read old data when PNAT > T-1. @triton.heuristics( - {"BLOCK_SIZE_WINDOW": lambda args: max( + {"BLOCK_SIZE_K": lambda args: max( triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} ) @triton.jit() -def _checkpointing_main_kernel( - # Pointers - state_ptr, - # Per-(cache, head, dim) decode scale, fp32, only consulted when QUANT_MAX>0. - # Layout (cache, nheads, dim) — broadcast over dstate at load/store. - state_scales_ptr, - # Cache READ pointers (read-buffer from previous step) - old_x_ptr, +def _rectangle_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - # Cache WRITE pointer (write-buffer for old_x only; B/dt/dA_cumsum written by precompute) - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - # New input pointers - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - # Precomputed pointers - cb_scaled_ptr, - decay_vec_ptr, + prev_num_accepted_tokens_ptr, state_batch_indices_ptr, - # Stochastic rounding - rand_seed_ptr, pad_slot_id, # Dimensions T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # cache T-axis capacity (= max_window) - dim: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, dstate: tl.constexpr, nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides: (cache, nheads, dim) — only used when QUANT_MAX>0 - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides: (cache, T, nheads, dim) — single-buffered - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides: (cache, 2, T, ngroups, dstate) - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, # C strides stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, # cb_scaled strides stride_cb_batch, stride_cb_head, @@ -879,38 +970,498 @@ def _checkpointing_main_kernel( stride_dv_batch, stride_dv_head, stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_WINDOW: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, - # State quantization: 0.0 means non-quantized (fp16/bf16/fp32); >0 means - # quantized (int8=127, int16=32767, fp8_e4m3fn=448). Single in-kernel - # switch for the dequant-on-load and encode-on-store paths. Wrapper sets - # this from state.dtype; kernel-entry static_assert below pins the - # invariant that it must coincide with int8/int16/float8e4nv state dtype. - QUANT_MAX: tl.constexpr, - # Checkpointing flag - WRITE_CHECKPOINT: tl.constexpr, # When True: quantize+write post-replay state to HBM (checkpoint step). - # When False: skip state write entirely (non-checkpoint step). - # The rectangle non-checkpoint path is implemented in - # _rectangle_main_kernel (separate kernel pair, picked by - # the wrapper via rectangle_for_nowrite=True). + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + EARLY_OUT: tl.constexpr, ): - # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized - # state dtype (int8 / int16 / float8e4nv) and only those. Cheap - # insurance against a wrapper bug that desynchronizes the two. - tl.static_assert( - (QUANT_MAX > 0.0) - == ( - (state_ptr.dtype.element_ty == tl.int8) - or (state_ptr.dtype.element_ty == tl.int16) + # Hoisted PDL signal: fire as the first thing every program does, so + # main can start its setup regardless of how this program ends. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + # Per-program early-out gate. Rectangle is nowrite-only, so EARLY_OUT + # skips slots whose PNAT + T > MAX (slots that would need write). + if EARLY_OUT: + pid_b_eo = tl.program_id(axis=0) + if HAS_CACHE_BATCH_INDICES: + cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) + if cbi_eo == pad_slot_id: + return + else: + cbi_eo = pid_b_eo.to(tl.int64) + pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) + if pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH: + return + _rectangle_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + ) + + +# Dynamic precompute kernel. Single launchable kernel that, per program, +# reads PNAT and dispatches to one of the existing impls: +# +# if pnat + T > MAX: replay_precompute_impl(WRITE_CHECKPOINT=True) +# else if RECTANGLE: rectangle_precompute_impl +# else: replay_precompute_impl(WRITE_CHECKPOINT=False) +# +# RECTANGLE is constexpr (compile-time tuning param); the inner branch +# is folded so only one of the two nowrite paths is emitted per +# specialization. Reg envelope = max(replay_write, X) where X depends +# on RECTANGLE. cb_scaled is allocated (T, K) by the wrapper regardless; +# replay paths write to the first T columns, rectangle writes the full K. +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _dynamic_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. + RECTANGLE: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + pid_b = tl.program_id(axis=0) + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH + # write_checkpoint is now runtime in replay precompute, so a single + # call site handles both write and nowrite for the replay branch. + # Take rectangle only when RECTANGLE is True AND this slot doesn't + # need write; everything else funnels into replay. + if needs_write_runtime or not RECTANGLE: + _replay_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + needs_write_runtime, + ) + else: + _rectangle_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + ) + + +# Main kernel: tl.dot replay + precomputed CB output. +# Grid: (cdiv(dim, M), batch, nheads). + + +@triton.jit() +def _replay_main_impl( + # Pointers + state_ptr, + # Per-(cache, head, dim) decode scale, fp32, only consulted when QUANT_MAX>0. + # Layout (cache, nheads, dim) — broadcast over dstate at load/store. + state_scales_ptr, + # Cache READ pointers (read-buffer from previous step) + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + # Cache WRITE pointer (write-buffer for old_x only; B/dt/dA_cumsum written by precompute) + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + # New input pointers + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + # Precomputed pointers + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + # Stochastic rounding + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # cache T-axis capacity (= max_window) + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides: (cache, nheads, dim) — only used when QUANT_MAX>0 + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides: (cache, T, nheads, dim) — single-buffered + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides: (cache, 2, T, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + # State quantization: 0.0 means non-quantized (fp16/bf16/fp32); >0 means + # quantized (int8=127, int16=32767, fp8_e4m3fn=448). Single in-kernel + # switch for the dequant-on-load and encode-on-store paths. Wrapper sets + # this from state.dtype; kernel-entry static_assert below pins the + # invariant that it must coincide with int8/int16/float8e4nv state dtype. + QUANT_MAX: tl.constexpr, + # Checkpointing flag + WRITE_CHECKPOINT: tl.constexpr, # When True: quantize+write post-replay state to HBM (checkpoint step). + # When False: skip state write entirely (non-checkpoint step). + # The rectangle non-checkpoint path is implemented in + # _rectangle_main_kernel (separate kernel pair, picked by + # the wrapper via rectangle_for_nowrite=True). + # When True: signal PDL dependents at the very top of every program + # (including pad/early-out programs). Used by doublelaunch and maindl + # so the next kernel (the second main, or the second precompute) can + # start its setup while this main is still computing. Default False + # for monolithic / dynamic / the LAST main in dl/maindl chains. + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does, so + # downstream kernels can start setup regardless of how this program + # ends (pad / early-out / full body). + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized + # state dtype (int8 / int16 / float8e4nv) and only those. Cheap + # insurance against a wrapper bug that desynchronizes the two. + tl.static_assert( + (QUANT_MAX > 0.0) + == ( + (state_ptr.dtype.element_ty == tl.int8) + or (state_ptr.dtype.element_ty == tl.int16) or (state_ptr.dtype.element_ty == tl.float8e4nv) ), "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", @@ -1232,29 +1783,243 @@ def _checkpointing_main_kernel( tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) -# Rectangle main kernel (nowrite-only): no replay step, no state HBM write, -# no SR codegen. state_out is computed from state_prev directly using the -# precomp-folded decay_vec_full; token_out is a single rectangle matmul over -# the (T, K) CB rectangle and (K, M) x_combined. Pairs with -# `_rectangle_precompute_kernel`. - - +# Replay-style main kernel. Thin wrapper around _replay_main_impl that carries +# the @triton.heuristics for constexpr derivation; called from the Python +# wrapper on the replay-style path (write or replay-nowrite). @triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) @triton.heuristics( {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} ) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) @triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) @triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( + {"BLOCK_SIZE_WINDOW": lambda args: max( triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} ) @triton.jit() -def _rectangle_main_kernel( +def _checkpointing_main_kernel( # Pointers state_ptr, - state_scales_ptr, # only consulted when QUANT_MAX > 0 + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + EARLY_OUT: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + # Per-program early-out gate. Signal-then-skip lets the next kernel + # in dl/maindl chains start regardless of early-out outcome. + if EARLY_OUT: + pid_b_eo = tl.program_id(axis=1) + if HAS_CACHE_BATCH_INDICES: + cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) + if cbi_eo == pad_slot_id: + return + else: + cbi_eo = pid_b_eo.to(tl.int64) + pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) + if (pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH) != WRITE_CHECKPOINT: + return + _replay_main_impl( + state_ptr, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + WRITE_CHECKPOINT, + LAUNCH_DEPENDENT_KERNELS, + ) + + +# Rectangle main kernel (nowrite-only): no replay step, no state HBM write, +# no SR codegen. state_out is computed from state_prev directly using the +# precomp-folded decay_vec_full; token_out is a single rectangle matmul over +# the (T, K) CB rectangle and (K, M) x_combined. Pairs with +# `_rectangle_precompute_kernel`. + + +@triton.jit() +def _rectangle_main_impl( + # Pointers + state_ptr, + state_scales_ptr, # only consulted when QUANT_MAX > 0 old_x_ptr, prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, @@ -1310,7 +2075,478 @@ def _rectangle_main_kernel( stride_out_T, stride_out_head, stride_out_dim, - # cb_scaled strides (rectangle (batch, nheads, T, K)) + # cb_scaled strides (rectangle (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + QUANT_MAX: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, +): + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + pid_m = tl.program_id(axis=0) + pid_b = tl.program_id(axis=1) + pid_h = tl.program_id(axis=2) + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout (matches precompute). + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_k = tl.arange(0, BLOCK_SIZE_K) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # K-axis masks + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - K_NEW_SHIFT + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Load state (with dequant on load if quantized). Read-only — no HBM + # write on the nowrite path. + state_ptr += cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, other=1.0, + ).to(tl.float32) + state = state * decode_scale[:, None] + + # Group / pointer offset setup + group_idx = pid_h // nheads_ngroups_ratio + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + # Hoist: old_x from cache doesn't depend on conv1d/precompute, so issue + # the load BEFORE gdc_wait so its HBM latency overlaps with conv1d. + old_x_load = tl.load( + old_x_base + + safe_old_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + mask=is_old_k[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + + # PDL gate: precompute outputs (cb_scaled, decay_vec_full) become safe + # after gdc_wait. conv1d outputs (x, C) also gated by the chained PDL. + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # Load C and x (conv1d outputs after PDL wait) + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + x_all = tl.load( + x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + # Append new x to cache at [PNAT, PNAT+T) of the active buffer. + tl.store( + old_x_base + + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_all, + mask=t_mask[:, None] & m_mask[None, :], + ) + x_all = x_all.to(tl.float32) + + # Build x_combined for the rectangle token_out matmul (STATIC layout). + # old_x_load is from cache (loaded above gdc_wait); new_x_shifted reads + # x (conv1d output) at compile-time-shifted [K_NEW_SHIFT, K_NEW_SHIFT+T). + # Disjoint masks → safe to add. + new_x_shifted = tl.load( + x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=is_new_k[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + x_combined = old_x_load + new_x_shifted # (BLOCK_SIZE_K, BLOCK_SIZE_M) + + # Load precomputed rectangle CB and folded decay_vec. + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec_full = tl.load( + decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 + ).to(tl.float32) + + # state_out: state_prev contribution to output, with decay folded post-matmul. + # No state_prev_decayed (M, dstate) materialization — state is consumed + # directly by the matmul, then decay_vec_full multiplies the (T, M) result. + state_out = ( + tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) + * decay_vec_full[:, None] + ) + + # token_out: combined old + new tokens contribution via the rectangle. + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) + + out_all = state_out + token_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + for t in range(T): + z_t = tl.load( + z_ptr + t * stride_z_T + offs_m * stride_z_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + out_t = tl.sum(tl.where((offs_t == t)[:, None], out_all, 0.0), axis=0) + out_t = out_t * z_t * tl.sigmoid(z_t) + tl.store(out_ptr + t * stride_out_T + offs_m * stride_out_dim, out_t, mask=m_mask) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# Rectangle main kernel. Thin wrapper around _rectangle_main_impl that carries +# the @triton.heuristics for constexpr derivation; called from the Python +# wrapper on the rectangle nowrite path. +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _rectangle_main_kernel( + # Pointers + state_ptr, + state_scales_ptr, + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + QUANT_MAX: tl.constexpr, + EARLY_OUT: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, +): + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + # Per-program early-out gate. Rectangle is nowrite-only. + if EARLY_OUT: + pid_b_eo = tl.program_id(axis=1) + if HAS_CACHE_BATCH_INDICES: + cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) + if cbi_eo == pad_slot_id: + return + else: + cbi_eo = pid_b_eo.to(tl.int64) + pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) + if pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH: + return + _rectangle_main_impl( + state_ptr, + state_scales_ptr, + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + QUANT_MAX, + LAUNCH_DEPENDENT_KERNELS, + ) + + +# Dynamic main kernel. Single launchable kernel that, per program, reads +# PNAT and dispatches to one of the existing impls: +# +# if pnat + T > MAX: replay_main_impl(WRITE_CHECKPOINT=True) +# elif RECTANGLE (constexpr): rectangle_main_impl +# else: replay_main_impl(WRITE_CHECKPOINT=False) +# +# Unlike precompute, WRITE_CHECKPOINT stays constexpr in the main impl — +# the body has constexpr-gated state-write code (quant + Philox + HBM +# store) where folding meaningfully shrinks the codegen. So this kernel +# has TWO replay call sites (one per WRITE_CHECKPOINT specialization) +# both inlined, with a runtime branch picking which runs. Reg envelope = +# max(replay_write, X) where X = rectangle_nowrite (RECTANGLE=True) or +# replay_nowrite (RECTANGLE=False). cb_scaled is allocated (T, K) by the +# wrapper regardless. +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _dynamic_main_kernel( + # Pointers — union of replay-main and rectangle-main pointer args. + state_ptr, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides (replay only; passed but unused on rectangle path) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides (replay only) + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides (replay only) + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K stride_cb_batch, stride_cb_head, stride_cb_t, @@ -1326,14 +2562,20 @@ def _rectangle_main_kernel( HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, # for replay path + BLOCK_SIZE_K: tl.constexpr, # for rectangle path LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, QUANT_MAX: tl.constexpr, + # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. + RECTANGLE: tl.constexpr, + # Default False — dynamic main is normally terminal in its chain. + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, ): - pid_m = tl.program_id(axis=0) + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() pid_b = tl.program_id(axis=1) - pid_h = tl.program_id(axis=2) - if HAS_CACHE_BATCH_INDICES: cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) if cache_batch_idx == pad_slot_id: @@ -1341,150 +2583,254 @@ def _rectangle_main_kernel( else: cache_batch_idx = pid_b.to(tl.int64) - # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - write_offset = prev_num_accepted_tokens - - # Static rectangle K-axis layout (matches precompute). - K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_k = tl.arange(0, BLOCK_SIZE_K) - m_mask = offs_m < dim - n_mask = offs_n < dstate - t_mask = offs_t < T - - # K-axis masks - is_old_k = offs_k < prev_num_accepted_tokens - safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - K_NEW_SHIFT - is_new_k = (k_new_idx >= 0) & (k_new_idx < T) - safe_k_new = tl.where(is_new_k, k_new_idx, 0) - - # Load state (with dequant on load if quantized). Read-only — no HBM - # write on the nowrite path. - state_ptr += cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - ) - state_mask = m_mask[:, None] & n_mask[None, :] - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) - if QUANT_MAX > 0.0: - state_scales_base = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head + pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if pnat_local + T > MAX_REPLAY_BUFFER_LENGTH: + # Write slot — replay-style write (WRITE_CHECKPOINT=True constexpr). + _replay_main_impl( + state_ptr, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + True, # WRITE_CHECKPOINT (constexpr) + False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top ) - decode_scale = tl.load( - state_scales_base + offs_m * stride_state_scales_dim, - mask=m_mask, other=1.0, - ).to(tl.float32) - state = state * decode_scale[:, None] - - # Group / pointer offset setup - group_idx = pid_h // nheads_ngroups_ratio - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head - C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group - if HAS_Z: - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head - - if HAS_D: - D = tl.load( - D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - - # Hoist: old_x from cache doesn't depend on conv1d/precompute, so issue - # the load BEFORE gdc_wait so its HBM latency overlaps with conv1d. - old_x_load = tl.load( - old_x_base - + safe_old_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - mask=is_old_k[:, None] & m_mask[None, :], - other=0.0, - ).to(tl.float32) - - # PDL gate: precompute outputs (cb_scaled, decay_vec_full) become safe - # after gdc_wait. conv1d outputs (x, C) also gated by the chained PDL. - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - # Load C and x (conv1d outputs after PDL wait) - C_all = tl.load( - C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - x_all = tl.load( - x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=t_mask[:, None] & m_mask[None, :], - other=0.0, - ) - # Append new x to cache at [PNAT, PNAT+T) of the active buffer. - tl.store( - old_x_base - + (write_offset + offs_t)[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - x_all, - mask=t_mask[:, None] & m_mask[None, :], - ) - x_all = x_all.to(tl.float32) - - # Build x_combined for the rectangle token_out matmul (STATIC layout). - # old_x_load is from cache (loaded above gdc_wait); new_x_shifted reads - # x (conv1d output) at compile-time-shifted [K_NEW_SHIFT, K_NEW_SHIFT+T). - # Disjoint masks → safe to add. - new_x_shifted = tl.load( - x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=is_new_k[:, None] & m_mask[None, :], - other=0.0, - ).to(tl.float32) - x_combined = old_x_load + new_x_shifted # (BLOCK_SIZE_K, BLOCK_SIZE_M) - - # Load precomputed rectangle CB and folded decay_vec. - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head - CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), - other=0.0, - ).to(tl.float32) - - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec_full = tl.load( - decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 - ).to(tl.float32) - - # state_out: state_prev contribution to output, with decay folded post-matmul. - # No state_prev_decayed (M, dstate) materialization — state is consumed - # directly by the matmul, then decay_vec_full multiplies the (T, M) result. - state_out = ( - tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) - * decay_vec_full[:, None] - ) - - # token_out: combined old + new tokens contribution via the rectangle. - token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) - - out_all = state_out + token_out - - if HAS_D: - out_all = out_all + x_all * D[None, :] - - if HAS_Z: - for t in range(T): - z_t = tl.load( - z_ptr + t * stride_z_T + offs_m * stride_z_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - out_t = tl.sum(tl.where((offs_t == t)[:, None], out_all, 0.0), axis=0) - out_t = out_t * z_t * tl.sigmoid(z_t) - tl.store(out_ptr + t * stride_out_T + offs_m * stride_out_dim, out_t, mask=m_mask) else: - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + if RECTANGLE: + _rectangle_main_impl( + state_ptr, + state_scales_ptr, + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + QUANT_MAX, + False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top + ) + else: + # Replay-style nowrite (WRITE_CHECKPOINT=False constexpr). + _replay_main_impl( + state_ptr, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + False, # WRITE_CHECKPOINT (constexpr) + False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top + ) # Python wrapper @@ -1524,6 +2870,7 @@ def checkpointing_state_update( use_internal_pdl=True, write_checkpoint: bool = True, rectangle_for_nowrite: bool = False, + mode: str = "monolithic", _block_size_m: int | None = None, _num_warps: int | None = None, _num_stages: int | None = None, @@ -1602,13 +2949,32 @@ def checkpointing_state_update( launch_with_pdl = False use_internal_pdl = False - # Path selection: - # write_checkpoint=True, rectangle_for_nowrite=* → checkpoint step - # (replay-style write kernel; rectangle_for_nowrite ignored). - # write_checkpoint=False, rectangle_for_nowrite=False → non-checkpoint step - # via replay-style nowrite kernel. - # write_checkpoint=False, rectangle_for_nowrite=True → non-checkpoint step - # via dedicated rectangle nowrite kernel (single fused old+new matmul). + # Mode selection: + # mode="monolithic" (default): today's behavior. write_checkpoint and + # rectangle_for_nowrite together pick a single kernel pair for the + # whole batch. Calls the corresponding kernel pair with EARLY_OUT=False. + # mode="dynamic": single kernel pair (_dynamic_*_kernel) that dispatches + # per-slot at runtime based on PNAT. RECTANGLE constexpr (= + # rectangle_for_nowrite) picks whether the nowrite path is rectangle + # or replay-nowrite. write_checkpoint is ignored (per-slot from PNAT). + # mode="doublelaunch": two kernel pairs launched in sequence, each with + # EARLY_OUT=True, partitioning the batch by PNAT-derived mode. + # Write half: replay-write. Nowrite half: rectangle if + # rectangle_for_nowrite else replay-nowrite. write_checkpoint ignored. + # mode="dlgrouped": same 4 kernels as doublelaunch, but reordered to + # launch both precomputes first, then both mains. Lets the GPU + # run precomp1 || precomp2 in parallel before the mains start. + # write_checkpoint ignored. + # mode="maindl": shared (dynamic) precompute + doublelaunched main. + # One precompute call (_dynamic_precompute_kernel) handles per-slot + # dispatch, then two main kernels with EARLY_OUT=True for the write + # and nowrite halves. Strictly fewer kernel launches than + # doublelaunch (3 vs 4) at the cost of dispatch precompute's wider + # reg envelope. write_checkpoint ignored. + assert mode in ("monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl"), ( + f"unknown mode {mode!r}; expected one of " + "'monolithic', 'dynamic', 'doublelaunch', 'dlgrouped', or 'maindl'" + ) use_rectangle = rectangle_for_nowrite and not write_checkpoint # --- Hardware support gates --- @@ -1737,11 +3103,13 @@ def checkpointing_state_update( # so the launch sites can refer to it; only used on the rectangle path. BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) - # Allocate precomputed intermediates (per-call, not cached). CB shape - # depends on path: replay-style is square (T, T); rectangle is (T, K). - cb_T_dim = BLOCK_SIZE_K if use_rectangle else BLOCK_SIZE_T + # Allocate precomputed intermediates (per-call, not cached). Always + # allocate (T, K) — the largest layout that any path uses. Replay-style + # paths only touch the first T columns; rectangle/dynamic use the full K. + # The few extra unused columns per row are negligible (~6KB per layer at + # production sizes) and let the dispatch helpers share one buffer. cb_scaled = torch.empty( - batch, nheads, BLOCK_SIZE_T, cb_T_dim, device=device, dtype=torch.float32 + batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 ) decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) @@ -1863,346 +3231,305 @@ def checkpointing_state_update( HAS_CACHE_BATCH_INDICES = state_batch_indices is not None - with torch.cuda.device(device.index): - # --- Precompute kernel --- - assert nheads % heads_per_block == 0, ( - f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" + assert nheads % heads_per_block == 0, ( + f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" + ) + assert heads_per_block <= heads_per_group, ( + f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" + ) + + # state_scales pointer + strides: real tensor when quantized, otherwise + # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). + if is_quantized: + state_scales_arg = state_scales + state_scales_strides = ( + state_scales.stride(0), + state_scales.stride(1), + state_scales.stride(2), ) - assert heads_per_block <= heads_per_group, ( - f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" + else: + state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 + state_scales_strides = (0, 0, 0) + + # Grid for main kernels (M tiling × batch × nheads). + def main_grid(META): + return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + + precomp_grid = (batch, nheads // heads_per_block) + d_strides = (D.stride(0), D.stride(1)) if D is not None else (0, 0) + + # ---- Launch helpers (close over locals) ------------------------------- + # Each helper is a thin closure that calls one Triton kernel with the + # full positional + kwarg argument list. Mode-dependent constexprs + # (write_checkpoint, early_out, rectangle) are passed in. + + def launch_replay_precompute(write_checkpoint: bool, early_out: bool): + _checkpointing_precompute_kernel[precomp_grid]( + dt, dt_bias, A, B, C, + cb_scaled, decay_vec, + old_B, old_dt, old_dA_cumsum, + cache_buf_idx, prev_num_accepted_tokens, + state_batch_indices, pad_slot_id, + T, max_window, dstate, nheads // ngroups, + dt.stride(0), dt.stride(1), dt.stride(2), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), B.stride(1), B.stride(2), B.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + dt_softplus, + HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, + LAUNCH_WITH_PDL=launch_with_pdl, + LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, + HEADS_PER_BLOCK=heads_per_block, + WRITE_CHECKPOINT=write_checkpoint, + EARLY_OUT=early_out, + num_warps=precompute_num_warps, + **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), + launch_pdl=launch_with_pdl, ) - # state_scales pointer + strides: real tensor when quantized, otherwise - # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). - # Hoisted out of the branch so both kernels can share. - if is_quantized: - state_scales_arg = state_scales - state_scales_strides = ( - state_scales.stride(0), - state_scales.stride(1), - state_scales.stride(2), - ) - else: - state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 - state_scales_strides = (0, 0, 0) - - if use_rectangle: - _rectangle_precompute_kernel[(batch, nheads // heads_per_block)]( - dt, - dt_bias, - A, - B, - C, - cb_scaled, - decay_vec, - old_B, - old_dt, - old_dA_cumsum, - cache_buf_idx, - prev_num_accepted_tokens, - state_batch_indices, - pad_slot_id, - T, - max_window, # MAX_REPLAY_BUFFER_LENGTH - dstate, - nheads // ngroups, - # dt strides - dt.stride(0), - dt.stride(1), - dt.stride(2), - dt_bias.stride(0) if dt_bias is not None else 0, - A.stride(0), - # B strides - B.stride(0), - B.stride(1), - B.stride(2), - B.stride(3), - # C strides - C.stride(0), - C.stride(1), - C.stride(2), - C.stride(3), - # cb_scaled strides (rectangle) - cb_scaled.stride(0), - cb_scaled.stride(1), - cb_scaled.stride(2), - cb_scaled.stride(3), - # decay_vec strides - decay_vec.stride(0), - decay_vec.stride(1), - decay_vec.stride(2), - # old_B strides - old_B.stride(0), - old_B.stride(1), - old_B.stride(2), - old_B.stride(3), - old_B.stride(4), - # old_dt strides - old_dt.stride(0), - old_dt.stride(1), - old_dt.stride(2), - old_dt.stride(3), - # old_dA_cumsum strides - old_dA_cumsum.stride(0), - old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), - old_dA_cumsum.stride(3), - dt_softplus, - HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, - LAUNCH_WITH_PDL=launch_with_pdl, - LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, - HEADS_PER_BLOCK=heads_per_block, - num_warps=precompute_num_warps, - **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), - launch_pdl=launch_with_pdl, - ) - - def grid(META): - return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) - - _rectangle_main_kernel[grid]( - state, - state_scales_arg, - old_x, - prev_num_accepted_tokens, - cache_buf_idx, - x, - C, - D, - z, - out, - cb_scaled, - decay_vec, - state_batch_indices, - pad_slot_id, - T, - max_window, # MAX_REPLAY_BUFFER_LENGTH - dim, - dstate, - nheads // ngroups, - # state strides - state.stride(0), - state.stride(1), - state.stride(2), - state.stride(3), - # state_scales strides (cache, head, dim) - state_scales_strides[0], - state_scales_strides[1], - state_scales_strides[2], - # old_x strides - old_x.stride(0), - old_x.stride(1), - old_x.stride(2), - old_x.stride(3), - # x strides - x.stride(0), - x.stride(1), - x.stride(2), - x.stride(3), - # C strides - C.stride(0), - C.stride(1), - C.stride(2), - C.stride(3), - # D strides - *(D.stride(0), D.stride(1)) if D is not None else (0, 0), - # z strides - z_strides[0], - z_strides[1], - z_strides[2], - z_strides[3], - # out strides - out.stride(0), - out.stride(1), - out.stride(2), - out.stride(3), - # cb_scaled strides - cb_scaled.stride(0), - cb_scaled.stride(1), - cb_scaled.stride(2), - cb_scaled.stride(3), - # decay_vec strides - decay_vec.stride(0), - decay_vec.stride(1), - decay_vec.stride(2), - BLOCK_SIZE_M, - LAUNCH_WITH_PDL=use_internal_pdl, - QUANT_MAX=quant_max, - num_warps=num_warps, - **({"num_stages": _num_stages} if _num_stages else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), - launch_pdl=use_internal_pdl, - ) - return # rectangle path done; skip the replay-style launches below + def launch_rectangle_precompute(early_out: bool): + _rectangle_precompute_kernel[precomp_grid]( + dt, dt_bias, A, B, C, + cb_scaled, decay_vec, + old_B, old_dt, old_dA_cumsum, + cache_buf_idx, prev_num_accepted_tokens, + state_batch_indices, pad_slot_id, + T, max_window, dstate, nheads // ngroups, + dt.stride(0), dt.stride(1), dt.stride(2), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), B.stride(1), B.stride(2), B.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + dt_softplus, + HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, + LAUNCH_WITH_PDL=launch_with_pdl, + LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, + HEADS_PER_BLOCK=heads_per_block, + EARLY_OUT=early_out, + num_warps=precompute_num_warps, + **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), + launch_pdl=launch_with_pdl, + ) - _checkpointing_precompute_kernel[(batch, nheads // heads_per_block)]( - dt, - dt_bias, - A, - B, - C, - cb_scaled, - decay_vec, - old_B, - old_dt, - old_dA_cumsum, - cache_buf_idx, - prev_num_accepted_tokens, - state_batch_indices, - pad_slot_id, - T, - dstate, - nheads // ngroups, - # dt strides - dt.stride(0), - dt.stride(1), - dt.stride(2), + def launch_dynamic_precompute(rectangle: bool): + _dynamic_precompute_kernel[precomp_grid]( + dt, dt_bias, A, B, C, + cb_scaled, decay_vec, + old_B, old_dt, old_dA_cumsum, + cache_buf_idx, prev_num_accepted_tokens, + state_batch_indices, pad_slot_id, + T, max_window, dstate, nheads // ngroups, + dt.stride(0), dt.stride(1), dt.stride(2), dt_bias.stride(0) if dt_bias is not None else 0, A.stride(0), - # B strides - B.stride(0), - B.stride(1), - B.stride(2), - B.stride(3), - # C strides - C.stride(0), - C.stride(1), - C.stride(2), - C.stride(3), - # cb_scaled strides - cb_scaled.stride(0), - cb_scaled.stride(1), - cb_scaled.stride(2), - cb_scaled.stride(3), - # decay_vec strides - decay_vec.stride(0), - decay_vec.stride(1), - decay_vec.stride(2), - # old_B strides - old_B.stride(0), - old_B.stride(1), - old_B.stride(2), - old_B.stride(3), - old_B.stride(4), - # old_dt strides - old_dt.stride(0), - old_dt.stride(1), - old_dt.stride(2), - old_dt.stride(3), - # old_dA_cumsum strides - old_dA_cumsum.stride(0), - old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), - old_dA_cumsum.stride(3), + B.stride(0), B.stride(1), B.stride(2), B.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), dt_softplus, HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, LAUNCH_WITH_PDL=launch_with_pdl, LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, HEADS_PER_BLOCK=heads_per_block, - WRITE_CHECKPOINT=write_checkpoint, + RECTANGLE=rectangle, num_warps=precompute_num_warps, **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), launch_pdl=launch_with_pdl, ) - # --- Main kernel --- - def grid(META): - return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) - - _checkpointing_main_kernel[grid]( - state, - state_scales_arg, - old_x, - old_B, - old_dt, - old_dA_cumsum, - prev_num_accepted_tokens, - cache_buf_idx, - x, - C, - D, - z, - out, - cb_scaled, - decay_vec, - state_batch_indices, - rand_seed, - pad_slot_id, - T, - max_window, # MAX_REPLAY_BUFFER_LENGTH - dim, - dstate, - nheads // ngroups, - # state strides - state.stride(0), - state.stride(1), - state.stride(2), - state.stride(3), - # state_scales strides (cache, head, dim) - state_scales_strides[0], - state_scales_strides[1], - state_scales_strides[2], - # old_x strides (single-buffered: cache, T, nheads, dim) - old_x.stride(0), - old_x.stride(1), - old_x.stride(2), - old_x.stride(3), - # old_B strides - old_B.stride(0), - old_B.stride(1), - old_B.stride(2), - old_B.stride(3), - old_B.stride(4), - # old_dt strides - old_dt.stride(0), - old_dt.stride(1), - old_dt.stride(2), - old_dt.stride(3), - # old_dA_cumsum strides - old_dA_cumsum.stride(0), - old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), - old_dA_cumsum.stride(3), - # x strides - x.stride(0), - x.stride(1), - x.stride(2), - x.stride(3), - # C strides - C.stride(0), - C.stride(1), - C.stride(2), - C.stride(3), - # D strides - *(D.stride(0), D.stride(1)) if D is not None else (0, 0), - # z strides - z_strides[0], - z_strides[1], - z_strides[2], - z_strides[3], - # out strides - out.stride(0), - out.stride(1), - out.stride(2), - out.stride(3), - # cb_scaled strides - cb_scaled.stride(0), - cb_scaled.stride(1), - cb_scaled.stride(2), - cb_scaled.stride(3), - # decay_vec strides - decay_vec.stride(0), - decay_vec.stride(1), - decay_vec.stride(2), + def launch_replay_main(write_checkpoint: bool, early_out: bool, + launch_dependent_kernels: bool = False): + _checkpointing_main_kernel[main_grid]( + state, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, rand_seed, pad_slot_id, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), BLOCK_SIZE_M, LAUNCH_WITH_PDL=use_internal_pdl, PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, QUANT_MAX=quant_max, WRITE_CHECKPOINT=write_checkpoint, + EARLY_OUT=early_out, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + def launch_rectangle_main(early_out: bool, + launch_dependent_kernels: bool = False): + _rectangle_main_kernel[main_grid]( + state, state_scales_arg, old_x, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, pad_slot_id, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + QUANT_MAX=quant_max, + EARLY_OUT=early_out, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + def launch_dynamic_main(rectangle: bool, + launch_dependent_kernels: bool = False): + _dynamic_main_kernel[main_grid]( + state, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, rand_seed, pad_slot_id, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + RECTANGLE=rectangle, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), **({"maxnreg": _maxnreg} if _maxnreg else {}), launch_pdl=use_internal_pdl, ) + + # ---- Mode dispatch ---------------------------------------------------- + with torch.cuda.device(device.index): + if mode == "monolithic": + if use_rectangle: + launch_rectangle_precompute(early_out=False) + launch_rectangle_main(early_out=False) + else: + launch_replay_precompute(write_checkpoint=write_checkpoint, early_out=False) + launch_replay_main(write_checkpoint=write_checkpoint, early_out=False) + elif mode == "dynamic": + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_dynamic_main(rectangle=rectangle_for_nowrite) + elif mode == "maindl": + # Shared dispatch precompute, doublelaunched main. Precompute + # runs once with per-slot dispatch (saves the second precomp + # empty-grid tax of doublelaunch). Mains stay split with + # EARLY_OUT so each retains its constexpr-specialized reg + # envelope. First main signals PDL dependents so the second + # main can start its setup while the first is still computing. + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_replay_main(write_checkpoint=True, early_out=True, + launch_dependent_kernels=True) + if rectangle_for_nowrite: + launch_rectangle_main(early_out=True) + else: + launch_replay_main(write_checkpoint=False, early_out=True) + elif mode == "dlgrouped": + # Same 4 kernels as doublelaunch but reordered: both precomputes + # first, then both mains. Lets the GPU run precomp1 || precomp2 + # in parallel (they're tiny grids) before the mains start, vs + # doublelaunch's interleaved precomp1→main1→precomp2→main2. + # First main signals PDL so the second main's setup overlaps. + launch_replay_precompute(write_checkpoint=True, early_out=True) + if rectangle_for_nowrite: + launch_rectangle_precompute(early_out=True) + else: + launch_replay_precompute(write_checkpoint=False, early_out=True) + launch_replay_main(write_checkpoint=True, early_out=True, + launch_dependent_kernels=True) + if rectangle_for_nowrite: + launch_rectangle_main(early_out=True) + else: + launch_replay_main(write_checkpoint=False, early_out=True) + else: # mode == "doublelaunch" + # Write half: always replay-style write. First main signals + # PDL dependents so the second precompute can start its setup + # while the first main is still computing. + launch_replay_precompute(write_checkpoint=True, early_out=True) + launch_replay_main(write_checkpoint=True, early_out=True, + launch_dependent_kernels=True) + # Nowrite half: rectangle if asked, else replay-nowrite. + if rectangle_for_nowrite: + launch_rectangle_precompute(early_out=True) + launch_rectangle_main(early_out=True) + else: + launch_replay_precompute(write_checkpoint=False, early_out=True) + launch_replay_main(write_checkpoint=False, early_out=True) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 62ce155086ef..5ae190c6aebb 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -402,6 +402,8 @@ def _time_kernel_cuda_graph( run_fn, reset_fn, tag: str, + pre_iter_fn=None, + iters_override: int | None = None, ) -> tuple[float, float, float]: """ All-in-one CUDA graph timing. @@ -409,15 +411,26 @@ def _time_kernel_cuda_graph( Captures a single graph containing warmup iterations followed by timed iterations with per-iteration event pairs recorded inside the graph. One replay, one sync, then all timings are read. + + ``pre_iter_fn(i)`` (if not None) is invoked inside the captured graph + before each iteration's run_fn(). Used by mix-mode benchmarking to + inject per-iter prev_tokens copies from a pre-baked samples tensor + (each call captures a copy from samples[i] into the graph). + + ``iters_override`` (if not None) overrides ``args.iters`` for this + call. Used to give mix scenarios a higher iter count than pure + (more iters = more independent mix draws averaged in). """ warmup = args.warmup - iters = args.iters + iters = iters_override if iters_override is not None else args.iters start_events = [torch.cuda.Event(enable_timing=True, external=True) for _ in range(iters)] end_events = [torch.cuda.Event(enable_timing=True, external=True) for _ in range(iters)] # Eager warmup before graph capture (triggers Triton autotune if active) reset_fn() + if pre_iter_fn is not None: + pre_iter_fn(0) run_fn() torch.cuda.synchronize() @@ -426,18 +439,24 @@ def _time_kernel_cuda_graph( g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): - # Warmup iterations (unrolled into the graph) - for _ in range(warmup): + # Warmup iterations (unrolled into the graph) — use the FIRST + # `warmup` samples so they don't share data with timed iters. + for i in range(warmup): reset_fn() if args.l2_flush: _l2_flush.fill_(0.0) + if pre_iter_fn is not None: + pre_iter_fn(i) run_fn() - # Timed iterations with events inside the graph + # Timed iterations with events inside the graph — use samples + # [warmup, warmup+iters), distinct from the warmup samples. for i in range(iters): reset_fn() if args.l2_flush: _l2_flush.fill_(0.0) + if pre_iter_fn is not None: + pre_iter_fn(warmup + i) start_events[i].record() run_fn() end_events[i].record() @@ -459,11 +478,16 @@ def _time_kernel_eager( run_fn, reset_fn, tag: str, + pre_iter_fn=None, + iters_override: int | None = None, ) -> tuple[float, float, float]: """Non-CUDA-graph timing path (for debugging, ncu, etc.).""" + iters = iters_override if iters_override is not None else args.iters # Warmup - for _ in range(args.warmup): + for i in range(args.warmup): reset_fn() + if pre_iter_fn is not None: + pre_iter_fn(i) run_fn() torch.cuda.synchronize() @@ -472,10 +496,12 @@ def _time_kernel_eager( latencies_us: list[float] = [] torch.cuda.nvtx.range_push(tag) - for _ in range(args.iters): + for i in range(iters): reset_fn() if args.l2_flush: _flush_l2() # includes synchronize + if pre_iter_fn is not None: + pre_iter_fn(args.warmup + i) start_event.record() run_fn() end_event.record() @@ -486,11 +512,16 @@ def _time_kernel_eager( return _compute_stats(latencies_us) -def _time_kernel(args, run_fn, reset_fn, tag: str) -> tuple[float, float, float]: +def _time_kernel(args, run_fn, reset_fn, tag: str, pre_iter_fn=None, + iters_override: int | None = None) -> tuple[float, float, float]: """Dispatch to CUDA-graph or eager timing path.""" if args.cuda_graph: - return _time_kernel_cuda_graph(args, run_fn, reset_fn, tag) - return _time_kernel_eager(args, run_fn, reset_fn, tag) + return _time_kernel_cuda_graph(args, run_fn, reset_fn, tag, + pre_iter_fn=pre_iter_fn, + iters_override=iters_override) + return _time_kernel_eager(args, run_fn, reset_fn, tag, + pre_iter_fn=pre_iter_fn, + iters_override=iters_override) # Per-config benchmark (consolidated baseline + replay) @@ -515,6 +546,7 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) + modes_list = getattr(args, "modes_list", ["monolithic"]) configs = [] for batch in batch_sizes: @@ -523,26 +555,33 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp for state_dtype in state_dtypes: for act_dtype in act_dtypes: for sr_mode in sr_modes_list: - for write_ckpt in write_modes_list: - # Rectangle is only meaningful for nowrite cells. - effective_rect_list = ( - [False] if write_ckpt else rect_list + for mode in modes_list: + effective_write_modes = ( + write_modes_list if mode == "monolithic" else [True] ) - for rect in effective_rect_list: - configs.append(( - batch, mtp_len, prev_ks, state_dtype, act_dtype, - sr_mode, rect, write_ckpt, - )) + for write_ckpt in effective_write_modes: + if mode == "monolithic": + effective_rect_list = ( + [False] if write_ckpt else rect_list + ) + else: + effective_rect_list = rect_list + for rect in effective_rect_list: + configs.append(( + batch, mtp_len, prev_ks, state_dtype, act_dtype, + sr_mode, rect, write_ckpt, mode, + )) print(f"[compile-warmup] {len(configs)} configs across {max_workers} threads") t0 = time.perf_counter() def _warm(cfg): - batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, rect, write_ckpt = cfg + (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, + rect, write_ckpt, mode) = cfg _bench_config( args, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, sr_mode=sr_mode, rectangle_for_nowrite=rect, - write_checkpoint=write_ckpt, warmup_only=True, + write_checkpoint=write_ckpt, mode=mode, warmup_only=True, ) errors = [] @@ -574,6 +613,9 @@ def _bench_config( sr_mode: str = "RN", rectangle_for_nowrite: bool = False, write_checkpoint: bool = True, + mode: str = "monolithic", + mix_samples_cpu=None, + mix_label: str = "", warmup_only: bool = False, ) -> None: """ @@ -831,18 +873,59 @@ def _parse_sweep(val): maxnreg_values = _parse_sweep(args.maxnreg) num_ctas_values = _parse_sweep(args.num_ctas) - # --- Replay kernel, one row per prev_k --- + # --- Replay kernel --- # Cache T-axis capacity (for prev_k validity check on the nowrite path). max_window = getattr(args, "max_window", 0) or mtp_len + + # Build the list of scenarios to time. A scenario is one cell in the + # output: pure-mode scenarios fill prev_tokens with one constant before + # the timing loop; mix-mode scenarios feed a pre-baked per-iter samples + # tensor, with the per-iter copy captured inside the CUDA graph. Pure + # and mix can coexist in one call so a single nsys trace covers both. + scenarios = [] for prev_k in prev_ks: - # On the nowrite path, new tokens append at [prev_k, prev_k+T) of the - # active buffer, so prev_k+T must fit within max_window. Skip - # silently for combinations that don't satisfy this — lets a single - # nsys run sweep both write modes against a shared prev_k list. - if not write_checkpoint and prev_k + mtp_len > max_window: + # On the nowrite path, new tokens append at [prev_k, prev_k+T) of + # the active buffer, so prev_k+T must fit within max_window. + # mode != monolithic dispatches per-slot from PNAT, so any + # prev_k <= max_window is valid for those modes. + if mode == "monolithic" and not write_checkpoint and prev_k + mtp_len > max_window: continue - prev_tokens.fill_(prev_k) - tag = f"incr_b{batch}_mtp{mtp_len}_k{prev_k}_s{state_dtype_name}_a{act_dtype_name}" + scenarios.append({ + "label": f"k{prev_k}", + "print_label": prev_k, + "fill": prev_k, + "pre_iter": None, + "iters": None, # use args.iters + }) + # Mix scenario: skip on monolithic (mono on mixed PNAT corrupts the + # wrong-mode slots). prev_tokens varies per iter; the per-iter copy + # is captured inside the CUDA graph from a pre-baked GPU samples + # tensor (warmup samples distinct from timed-iter samples so any + # nsys-included warmup leaks aren't biased). + if mix_samples_cpu is not None and mode != "monolithic": + device = state_work.device + samples_gpu = torch.from_numpy(mix_samples_cpu).to(device=device, dtype=torch.int32) + + def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): + _pt.copy_(_s[i]) + + # Mix iters override: if --mix-iters set, use it; else use args.iters. + mix_iters = getattr(args, "mix_iters", None) + scenarios.append({ + "label": f"mix{mix_label}", + "print_label": "mix", + "fill": None, + "pre_iter": _mix_pre_iter, + "iters": mix_iters, # None => use args.iters + }) + + for scn in scenarios: + if scn["fill"] is not None: + prev_tokens.fill_(scn["fill"]) + prev_k_for_print = scn["print_label"] + scenario_pre_iter = scn["pre_iter"] + scenario_iters = scn.get("iters") # None => use args.iters + tag = f"incr_b{batch}_mtp{mtp_len}_{scn['label']}_s{state_dtype_name}_a{act_dtype_name}" for ( block_size_m, @@ -865,7 +948,6 @@ def _parse_sweep(val): ): def _run_incr( - prev_k=prev_k, block_size_m=block_size_m, num_warps=num_warps, num_stages=num_stages, @@ -889,6 +971,7 @@ def _run_incr( if args.variant == "checkpointing": extra_kwargs["write_checkpoint"] = write_checkpoint extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite + extra_kwargs["mode"] = mode if state_scales_work is not None: extra_kwargs["state_scales"] = state_scales_work variant_fn( @@ -943,23 +1026,30 @@ def _run_incr( parts.append(f"SR={1 if use_philox else 0}") parts.append(f"RECT={1 if rectangle_for_nowrite else 0}") parts.append(f"WC={1 if write_checkpoint else 0}") + parts.append(f"MODE={mode}") sweep_suffix = (" " + ",".join(parts)) if parts else "" sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") reset_fn = _reset_conv1d_realistic if with_conv1d else _reset if warmup_only: reset_fn() + if scenario_pre_iter is not None: + scenario_pre_iter(0) _run_incr() torch.cuda.synchronize() else: - median_us, p95_us, p99_us = _time_kernel(args, _run_incr, reset_fn, sweep_tag) + median_us, p95_us, p99_us = _time_kernel( + args, _run_incr, reset_fn, sweep_tag, + pre_iter_fn=scenario_pre_iter, + iters_override=scenario_iters, + ) _print_row( show_kernel_col, args.variant, batch, mtp_len, - prev_k, + prev_k_for_print, state_dtype_name, act_dtype_name, median_us, @@ -1067,26 +1157,69 @@ def _run_benchmark(args) -> None: sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) + modes_list = getattr(args, "modes_list", ["monolithic"]) + + # Pre-load AL distribution for mix mode (if --mix-csv set). + mix_al = None + mix_label = "" + if args.mix_csv is not None: + from pathlib import Path as _Path + from checkpoint_mix_sim import load_al_distribution as _load_al + mix_label = _Path(args.mix_csv).stem + # T (= mtp_len) varies per cell; load once with the LARGEST mtp so + # we have enough columns; the loader normalizes the dist anyway. + mix_al = _load_al(_Path(args.mix_csv), T=max(mtp_lengths), column=args.mix_csv_column) for batch in batch_sizes: for mtp_len in mtp_lengths: # Resolve prev_k fractions → clamped integers in [0, mtp_len] prev_ks = _resolve_prev_ks(args, mtp_len) + + # Pre-generate mix samples once per (batch, mtp_len) cell so all + # tuning configs see the same per-iter prev_tokens vectors — + # tuning differences become signal, mix-noise is shared. + # Size the sample buffer for the LARGER of args.iters and + # args.mix_iters since mix scenarios use mix_iters. + mix_samples_cpu = None + if mix_al is not None: + from checkpoint_mix_sim import sample_steady_state_pnat as _sample_pnat + _max_window = getattr(args, "max_window", 0) or mtp_len + _max_iters = max(args.iters, getattr(args, "mix_iters", None) or args.iters) + mix_samples_cpu = _sample_pnat( + mix_al, T=mtp_len, window=_max_window, batch=batch, + K=args.warmup + _max_iters, seed=args.mix_seed, + ) + for state_dtype in state_dtypes: for act_dtype in act_dtypes: for sr_mode in sr_modes_list: - for write_ckpt in write_modes_list: - # Rectangle only meaningful for nowrite cells. - effective_rect_list = ( - [False] if write_ckpt else rect_list + for mode in modes_list: + # Non-monolithic modes ignore write_checkpoint + # (per-slot from PNAT) — collapse the sweep so we + # don't duplicate identical cells. + effective_write_modes = ( + write_modes_list if mode == "monolithic" else [True] ) - for rect in effective_rect_list: - _bench_config( - args, batch, mtp_len, prev_ks, state_dtype, act_dtype, - baseline_fn, sr_mode=sr_mode, - rectangle_for_nowrite=rect, - write_checkpoint=write_ckpt, - ) + for write_ckpt in effective_write_modes: + # Rectangle is meaningful for: nowrite cells in + # monolithic; always for dynamic / doublelaunch + # (constexpr knob). + if mode == "monolithic": + effective_rect_list = ( + [False] if write_ckpt else rect_list + ) + else: + effective_rect_list = rect_list + for rect in effective_rect_list: + _bench_config( + args, batch, mtp_len, prev_ks, state_dtype, act_dtype, + baseline_fn, sr_mode=sr_mode, + rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, + mode=mode, + mix_samples_cpu=mix_samples_cpu, + mix_label=mix_label, + ) if args.profile: torch.cuda.cudart().cudaProfilerStop() @@ -1323,6 +1456,58 @@ def _parse_args() -> argparse.Namespace: "write path always uses replay-style). Only applies to the " "checkpointing variant.", ) + parser.add_argument( + "--modes", + type=str, + default="monolithic", + help="Comma-separated dispatch modes to sweep, any of " + "{monolithic,dynamic,doublelaunch}. monolithic = today's behavior " + "(one kernel pair, write_checkpoint applied to whole batch); " + "dynamic = single kernel pair that dispatches per-slot at runtime " + "based on PNAT (rectangle_for_nowrite picks RECTANGLE constexpr); " + "doublelaunch = two kernel pairs launched in sequence with " + "EARLY_OUT=True, each handling slots whose mode matches it. " + "Only applies to the checkpointing variant; non-monolithic modes " + "ignore --write-modes (per-slot from PNAT).", + ) + parser.add_argument( + "--mix-csv", + type=str, + default=None, + help="Path to AL histogram CSV (cols: AL, count). When set, an " + "additional 'mix' cell is emitted per (batch, mtp, dtype, sr, " + "mode, RECT, M, W, ...) combo where prev_tokens varies per iter, " + "drawn from the steady-state PNAT distribution induced by the " + "AL histogram. Mix cells run only on dynamic and doublelaunch " + "modes (mono on a mixed batch corrupts wrong-mode slots). " + "Each iteration of the captured CUDA graph has a different " + "pre-baked prev_tokens vector; warmup iters use distinct samples " + "from the timed iters so nsys-included warmup leaks don't bias.", + ) + parser.add_argument( + "--mix-csv-column", + type=int, + default=1, + help="Column index (0-based) in the AL histogram CSV for the " + "count/probability column. Default 1 (second column).", + ) + parser.add_argument( + "--mix-seed", + type=int, + default=42, + help="RNG seed for the steady-state PNAT sampler. Same seed " + "across runs => same per-slot samples for reproducible " + "comparisons.", + ) + parser.add_argument( + "--mix-iters", + type=int, + default=None, + help="Iteration count override for mix scenarios (each iter is a " + "different per-slot prev_tokens draw). Default (None) uses " + "--iters. Mix scenarios benefit from more iters since each " + "iter samples a different mix; pure scenarios don't.", + ) parser.add_argument( "--philox-rounding", action="store_true", @@ -1393,6 +1578,15 @@ def _parse_args() -> argparse.Namespace: args.write_modes_list = write_list else: args.write_modes_list = [args.write_checkpoint] + + modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] + valid_modes = {"monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl"} + for m in modes_raw: + if m not in valid_modes: + parser.error( + f"--modes value must be one of {sorted(valid_modes)}, got {m!r}" + ) + args.modes_list = modes_raw or ["monolithic"] return args diff --git a/tests/unittest/_torch/modules/mamba/checkpoint_mix_sim.py b/tests/unittest/_torch/modules/mamba/checkpoint_mix_sim.py new file mode 100644 index 000000000000..104e5349a4a3 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/checkpoint_mix_sim.py @@ -0,0 +1,331 @@ +#!/usr/bin/env python3 +""" +Compute the steady-state checkpoint (write) fraction for a Mamba +replay-style cache, given an acceptance-length (AL) distribution. + +Model +----- +* T tokens per step (DL draft + 1 target). +* WINDOW = max_window: cache T-axis size. +* PNAT = "previous number of accepted tokens" already cached at start of + step. After accepting AL tokens this step: + - nowrite step (PNAT + T <= WINDOW): new tokens append at + [PNAT, PNAT+T) of the active buffer; PNAT_new = PNAT + AL. + - write step (PNAT + T > WINDOW): new tokens go to the staging + buffer at [0, T); cache_buf_idx flips; PNAT_new = AL. + +The write decision in the kernel is exactly `pnat + T > WINDOW`. + +Two methods +----------- +1. **Markov chain stationary** (exact): build the (WINDOW+1)x(WINDOW+1) + transition matrix from the AL distribution, solve for pi. + write_frac = sum_{p > WINDOW - T} pi(p) + +2. **Depletion sim** (approximate, matches the user's pen-and-paper + intuition): start with all mass at PNAT=0, propagate forward, count + mass that hits a write state at each step, remove that mass without + reinjecting. Compute E[N] = E[step at first checkpoint | start + PNAT=0]. Steady-state write_frac = 1 / (E[N] - 1). + The "-1" is because the first step from PNAT=0 is always nowrite + (free): subsequent cycles start from PNAT ~ AL_dist (post-write + distribution), one step shorter than the from-PNAT=0 cycle. + +Both methods should agree (verified on small examples). + +CSV format for the AL distribution +---------------------------------- +Two columns per row: AL (int 1..T), count or probability (float). +First row treated as a header if the first cell isn't numeric. +The histogram is auto-normalized to a probability distribution. +""" +from __future__ import annotations + +import argparse +import csv +import sys +from pathlib import Path + +import numpy as np + + +def load_al_distribution(path: Path, T: int, column: int = 1) -> np.ndarray: + """Return a length-(T+1) probability vector indexed by AL (0..T). + + Reads column 0 as AL, column ``column`` as count/probability. + Non-numeric rows (header, trailing summary rows like "total"/"mean") + are silently skipped. + """ + rows = list(csv.reader(open(path))) + if not rows: + sys.exit(f"empty CSV: {path}") + al_to_count: dict[int, float] = {} + for r in rows: + if not r or not r[0].strip(): + continue + try: + al = int(float(r[0])) + c = float(r[column]) + except (ValueError, IndexError): + continue # header / summary row + al_to_count[al] = al_to_count.get(al, 0.0) + c + if not al_to_count: + sys.exit(f"no numeric rows parsed from {path} (column index {column})") + al_max = max(al_to_count) + if al_max > T: + sys.exit( + f"CSV has AL={al_max} but --T={T} (a step processes only T tokens; " + f"AL > T is impossible). Mismatch likely indicates wrong --T or wrong CSV." + ) + if min(al_to_count) < 0: + sys.exit(f"CSV has negative AL values") + if al_to_count.get(0, 0) > 0: + print( + f"WARN: CSV has AL=0 mass ({al_to_count[0]:.4f} unnormalized). " + "AL=0 means no progress on a step; usually impossible in spec " + "decoding (target token always accepted). Treating as a real " + "value (will produce a chain that doesn't converge if AL=0 has " + "non-trivial mass).", + file=sys.stderr, + ) + dist = np.zeros(T + 1, dtype=np.float64) + for al, c in al_to_count.items(): + dist[al] = c + s = dist.sum() + if s == 0: + sys.exit(f"AL distribution sums to zero in {path}") + return dist / s + + +def markov_stationary(al_dist: np.ndarray, T: int, window: int) -> np.ndarray: + """ + Build the (window+1) x (window+1) transition matrix and return the + stationary distribution pi. + """ + n = window + 1 + P = np.zeros((n, n)) + for p in range(n): + is_write = (p + T > window) + for al in range(1, T + 1): + prob = al_dist[al] + if prob == 0: + continue + p_new = al if is_write else p + al + assert 0 <= p_new <= window, ( + f"unreachable transition: p={p} al={al} write={is_write} " + f"-> p_new={p_new} outside [0, {window}]" + ) + P[p, p_new] += prob + # Solve pi P = pi via the left eigenvector for eigenvalue 1. + # Equivalently: P^T pi = pi. + eigvals, eigvecs = np.linalg.eig(P.T) + idx = int(np.argmin(np.abs(eigvals - 1.0))) + if abs(eigvals[idx] - 1.0) > 1e-6: + sys.exit( + f"no stationary eigenvalue near 1 (closest is {eigvals[idx]}). " + "AL dist may not produce an ergodic chain." + ) + pi = np.real(eigvecs[:, idx]) + # Iterative refinement: start near the eigenvector, then iterate + # P^k to clean up any imaginary leakage from the eig solve. + pi = np.maximum(pi, 0.0) + if pi.sum() == 0: + # Fallback: power iteration from uniform. + pi = np.ones(n) / n + for _ in range(2000): + new_pi = pi @ P + if np.allclose(new_pi, pi, atol=1e-12, rtol=0): + pi = new_pi + break + pi = new_pi + pi = pi / pi.sum() + return pi + + +def sample_steady_state_pnat( + al_dist: np.ndarray, + T: int, + window: int, + batch: int, + K: int, + seed: int = 42, +) -> np.ndarray: + """ + Draw K independent (batch,)-shaped per-slot PNAT vectors from the + Markov chain's stationary distribution, given an AL distribution. + + Returns int64 array of shape (K, batch) with values in [0, window]. + PNAT=0 has weight ~0 in steady state so it is effectively never + sampled (it's purely a boot-up state). + """ + pi = markov_stationary(al_dist, T, window) + pi = np.maximum(pi, 0) + pi = pi / pi.sum() + states = np.arange(window + 1) + rng = np.random.default_rng(seed) + return rng.choice(states, size=(K, batch), p=pi).astype(np.int64) + + +def depletion_sim( + al_dist: np.ndarray, T: int, window: int, max_steps: int = 256 +) -> np.ndarray: + """ + Start with mass 1 at PNAT=0. At each step, mass at PNAT > WINDOW-T + is removed (= it checkpoints this step). Remaining mass advances + by AL. Returns array `removed[step-1]` = mass that checkpointed at + step. + """ + n = window + 1 + mass = np.zeros(n) + mass[0] = 1.0 + write_thresh = window - T # PNAT > this => write + removed = [] + for _ in range(max_steps): + write_mass = mass[write_thresh + 1:].sum() + removed.append(float(write_mass)) + # Strip the mass that checkpointed (it's been "removed"). + mass = mass.copy() + mass[write_thresh + 1:] = 0.0 + # Propagate the rest by AL. + new_mass = np.zeros(n) + for p in range(write_thresh + 1): + if mass[p] == 0: + continue + for al in range(1, T + 1): + prob = al_dist[al] + if prob == 0: + continue + p_new = p + al + assert p_new <= window, "unreachable" + new_mass[p_new] += mass[p] * prob + mass = new_mass + if mass.sum() < 1e-15: + break + return np.array(removed) + + +def _format_pi(pi: np.ndarray, T: int, window: int) -> str: + write_thresh = window - T + lines = [] + for p, prob in enumerate(pi): + if prob < 1e-9: + continue + marker = " <- write state" if p > write_thresh else "" + lines.append(f" pi(PNAT={p:2d}) = {prob:.4f}{marker}") + return "\n".join(lines) + + +def main() -> None: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("al_csv", type=Path, help="CSV with AL,count|prob columns") + ap.add_argument( + "--T", type=int, default=6, + help="Tokens per step (DL+1). Default 6 = production Nemotron config.", + ) + ap.add_argument( + "--window", type=int, default=16, + help="Cache T-axis size (max_window). Default 16 = production.", + ) + ap.add_argument( + "--method", choices=["markov", "depletion", "both"], default="both", + help="Which method(s) to compute.", + ) + ap.add_argument( + "--column", type=int, default=1, + help="Column index (0-based) of the count/prob in the CSV. " + "Use this to pick a specific variant when the CSV has multiple " + "count columns (e.g. replay_count, true_baseline_count, ...).", + ) + ap.add_argument( + "--quiet", action="store_true", + help="Print only the final write fraction (script-friendly).", + ) + args = ap.parse_args() + + if args.T >= args.window: + ap.error( + f"--T ({args.T}) must be < --window ({args.window}); otherwise " + "every step from PNAT>0 is a write step." + ) + + al = load_al_distribution(args.al_csv, args.T, column=args.column) + + if not args.quiet: + print(f"AL distribution (normalized):") + for i, p in enumerate(al): + if p > 0: + print(f" AL={i}: {p:.4f}") + e_al = float(sum(i * p for i, p in enumerate(al))) + print(f" E[AL] = {e_al:.3f} (per-step throughput)") + print() + print(f"T = {args.T}, window = {args.window}") + print(f"write threshold (exclusive): PNAT > {args.window - args.T}") + print(f" => write states: PNAT in {{{args.window - args.T + 1}, ..., {args.window}}}") + print() + + write_frac_mc = None + write_frac_dep = None + + if args.method in ("markov", "both"): + pi = markov_stationary(al, args.T, args.window) + write_frac_mc = float(pi[args.window - args.T + 1:].sum()) + if not args.quiet: + print("--- Markov chain stationary (exact) ---") + print(_format_pi(pi, args.T, args.window)) + print(f" write fraction = {write_frac_mc:.4f}") + if write_frac_mc > 0: + print(f" avg steps per checkpoint = {1/write_frac_mc:.2f}") + print() + + if args.method in ("depletion", "both"): + removed = depletion_sim(al, args.T, args.window) + total = float(removed.sum()) + if total < 0.99: + print( + f"WARN: depletion sim only depleted {total:.4f} of mass " + f"in {len(removed)} steps; may be missing tail.", + file=sys.stderr, + ) + e_step = float(sum((i + 1) * r for i, r in enumerate(removed)) / max(total, 1e-12)) + # Steady-state cycle length = E[N from PNAT=0] - 1 (the first + # nowrite step from PNAT=0 is the "free" one; subsequent cycles + # start from PNAT ~ AL_dist, one step shorter). + cycle = e_step - 1.0 + write_frac_dep = (1.0 / cycle) if cycle > 0 else float("inf") + if not args.quiet: + print("--- Depletion sim (start at PNAT=0) ---") + for i, r in enumerate(removed): + if r > 1e-6: + print(f" step {i+1:2d}: checkpoint mass = {r:.4f}") + print(f" total mass depleted = {total:.6f}") + print(f" E[step at first checkpoint] = {e_step:.3f}") + print(f" steady-state cycle (=E[N]-1) = {cycle:.3f}") + print(f" write fraction (1/cycle) = {write_frac_dep:.4f}") + print() + + if args.quiet: + wf = write_frac_mc if write_frac_mc is not None else write_frac_dep + print(f"{wf:.6f}") + return + + if write_frac_mc is not None and write_frac_dep is not None: + diff = abs(write_frac_mc - write_frac_dep) + rel = diff / max(write_frac_mc, 1e-12) + print( + f"--- cross-check ---\n" + f" markov: {write_frac_mc:.6f}\n" + f" depletion: {write_frac_dep:.6f}\n" + f" abs diff: {diff:.2e} (rel: {rel:.2%})" + ) + if rel > 0.01: + print( + " WARN: methods disagree by >1%. Possible: AL=0 mass or " + "non-ergodic chain — check the AL distribution.", + file=sys.stderr, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index 215e357a0888..9b1bbbf046d2 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -108,9 +108,14 @@ def _maybe_skip_dtype(state_dtype, use_sr): ], ids=["write", "no_write_replay", "no_write_rectangle"], ) +@pytest.mark.parametrize( + "mode", + ["monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl"], + ids=["monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl"], +) def test_checkpointing_state_update( nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T, - write_checkpoint, rectangle_for_nowrite, + write_checkpoint, rectangle_for_nowrite, mode, ): """ Verify that: @@ -326,6 +331,7 @@ def test_checkpointing_state_update( state_scales=test_scales, write_checkpoint=write_checkpoint, rectangle_for_nowrite=rectangle_for_nowrite, + mode=mode, ) # Tolerance rationale: the replay kernel uses bf16 tl.dot for four @@ -536,6 +542,193 @@ def test_checkpointing_state_update( ) +@pytest.mark.parametrize( + "mode,rectangle_for_nowrite", + [ + ("dynamic", False), + ("dynamic", True), + ("doublelaunch", False), + ("doublelaunch", True), + ("dlgrouped", False), + ("dlgrouped", True), + ("maindl", False), + ("maindl", True), + ], + ids=[ + "dynamic_replay", + "dynamic_rectangle", + "doublelaunch_replay", + "doublelaunch_rectangle", + "dlgrouped_replay", + "dlgrouped_rectangle", + "maindl_replay", + "maindl_rectangle", + ], +) +def test_checkpointing_state_update_mixed_mode(mode, rectangle_for_nowrite): + """ + Mixed-mode dispatch: a batch where some slots have PNAT triggering + write and others triggering nowrite, exercising the per-slot dispatch + of mode={dynamic, doublelaunch}. + + Setup: 4 slots, max_window=16, T=6. + pnat_per_slot = [3, 10, 12, 16] + slots 0, 1: nowrite (pnat + T <= max_window) + slots 2, 3: write (pnat + T > max_window) + + Reference: per-slot post-replay state computed from the captured + state evolution, then selective_state_update for the new step. + Output and post-replay state are verified per-slot. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + cache_size = batch + device = "cuda" + dtype = torch.bfloat16 + state_dtype = torch.bfloat16 + + # PNAT mix: write threshold is pnat + T > max_window → pnat >= 11. + pnat_per_slot = torch.tensor([3, 10, 12, 16], device=device, dtype=torch.int32) + # Per-slot dispatch destinations under each mode (for the postcondition checks). + pnat_means_write = (pnat_per_slot + T > max_window).tolist() # [F, F, T, T] + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + ref_input_state = state0.float() + + # Old inputs spanning the full window (so any pnat 0..max_window is exercised). + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + # Capture per-step intermediate SSM states across the window. + states_buffer_f32 = torch.zeros( + cache_size, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + # Build the cache tensors with old data on each slot's active buffer. + old_x = torch.zeros(cache_size, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + old_dA_cumsum = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(cache_size): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + # New-step inputs. + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + # Reference: per-slot post-replay state, then selective_state_update. + ref_state_f32 = ref_input_state.clone() + for i in range(cache_size): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, + x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + out=ref_out, + ) + + # Kernel call. + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_x_w = old_x.clone() + old_B_w = old_B.clone() + old_dt_w = old_dt.clone() + old_dA_cumsum_w = old_dA_cumsum.clone() + cache_buf_idx_w = cache_buf_idx.clone() + + checkpointing_state_update( + test_state, + old_x_w, old_B_w, old_dt_w, old_dA_cumsum_w, + cache_buf_idx_w, + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + rectangle_for_nowrite=rectangle_for_nowrite, + mode=mode, + # write_checkpoint is ignored in dynamic / doublelaunch. + ) + + # Output: every slot must match its reference (bf16 atol consistent + # with existing tests). + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (mode={mode}, rect={rectangle_for_nowrite})", + ) + + # State postconditions per slot: + # write slots (pnat + T > max_window): state in HBM is the post-replay + # fp32 reference (cast back to state_dtype with bf16 atol). + # nowrite slots: HBM state is unchanged (still state0 bitwise). + for i in range(cache_size): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), + ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (mode={mode})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], + rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM was modified (mode={mode})", + ) + + @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) @pytest.mark.parametrize( "state_dtype", From ea9ed6c00f75f37b96744cccb2c27c7e9fe37eec Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 6 May 2026 14:35:58 -0700 Subject: [PATCH 18/89] main kernels: vectorize HAS_Z output loop ((T, M) tile op) Replace per-T loop in both _checkpointing_main_kernel and _rectangle_main_kernel with a single (T, M) tile op: z_all = tl.load(z[T, M]) # (T, M) out_all_z = out_all * z_all * tl.sigmoid(z_all) tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) vs: for t in range(T): z_t = tl.load(z[t, M]) out_t = tl.sum(tl.where((offs_t == t)[:, None], out_all, 0.0), axis=0) out_t = out_t * z_t * tl.sigmoid(z_t) tl.store(out[t, M], out_t) Eliminates the awkward row-pick via where+sum and the T iterations. Note: HAS_Z is dead in TRT-LLM production (Mamba2 z-gate is folded into the downstream RMS norm; every call site in mamba2_mixer.py passes z=None). Kept as code-hygiene improvement; no measurable production impact. 542 tests pass. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 8594c13b9004..e38923a05278 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -1170,13 +1170,13 @@ def _checkpointing_main_kernel( out_all = out_all + x_all * D[None, :] if HAS_Z: - for t in range(T): - z_t = tl.load( - z_ptr + t * stride_z_T + offs_m * stride_z_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - out_t = tl.sum(tl.where((offs_t == t)[:, None], out_all, 0.0), axis=0) - out_t = out_t * z_t * tl.sigmoid(z_t) - tl.store(out_ptr + t * stride_out_T + offs_m * stride_out_dim, out_t, mask=m_mask) + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) else: out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) @@ -1425,13 +1425,13 @@ def _rectangle_main_kernel( out_all = out_all + x_all * D[None, :] if HAS_Z: - for t in range(T): - z_t = tl.load( - z_ptr + t * stride_z_T + offs_m * stride_z_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - out_t = tl.sum(tl.where((offs_t == t)[:, None], out_all, 0.0), axis=0) - out_t = out_t * z_t * tl.sigmoid(z_t) - tl.store(out_ptr + t * stride_out_T + offs_m * stride_out_dim, out_t, mask=m_mask) + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) else: out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) From cbebb654325444f42ca15c27f8e95eb9fc53297c Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 6 May 2026 20:45:23 -0700 Subject: [PATCH 19/89] Slot-perm dispatch: USE_PERM/REVERSE_PERM constexprs in dl-family kernels + bench flags + sorted-dispatch test Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 224 ++++++++++++++++-- ...benchmark_replay_selective_state_update.py | 155 ++++++++++-- .../mamba/test_checkpointing_state_update.py | 149 ++++++++++++ 3 files changed, 486 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 0f0134a6aa62..39adcad73547 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -108,6 +108,11 @@ def _replay_precompute_impl( # on no-checkpoint steps). prev_num_accepted_tokens_ptr, state_batch_indices_ptr, + # Slot permutation: maps grid program_id -> original slot index. + # When USE_PERM=False, pid_b = tl.program_id(0) (today's behavior) and + # this ptr is unused. When USE_PERM=True, pid_b = perm[pid_grid] (or + # perm[B-1-pid_grid] if REVERSE_PERM=True). + slot_perm_ptr, pad_slot_id, # Dimensions T: tl.constexpr, @@ -162,6 +167,13 @@ def _replay_precompute_impl( BLOCK_SIZE_T: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, HEADS_PER_BLOCK: tl.constexpr, + # Slot permutation flags. USE_PERM=True gates a slot_perm_ptr load that + # remaps grid program_id -> original slot index. REVERSE_PERM=True walks + # the perm from the tail (B-1-pid_grid). Used by sorted-dispatch + # variants of dl/dlgrouped/maindl to cluster early-outs at one end of + # the grid; ignored by monolithic / dynamic. + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, # Checkpointing flag — selects target buffer + offset for new-token # cache writes. See "Cache write semantics" block below. # Runtime (not constexpr): the only WRITE_CHECKPOINT-dependent code in @@ -171,7 +183,14 @@ def _replay_precompute_impl( # per-slot needs_write flag instead of inlining two specializations. write_checkpoint, ): - pid_b = tl.program_id(axis=0) + pid_grid = tl.program_id(axis=0) + if USE_PERM: + if REVERSE_PERM: + pid_b = tl.load(slot_perm_ptr + (tl.num_programs(axis=0) - 1 - pid_grid)) + else: + pid_b = tl.load(slot_perm_ptr + pid_grid) + else: + pid_b = pid_grid pid_hg = tl.program_id(axis=1) # head-group index first_head = pid_hg * HEADS_PER_BLOCK @@ -370,6 +389,7 @@ def _checkpointing_precompute_kernel( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, + slot_perm_ptr, pad_slot_id, # Dimensions T: tl.constexpr, @@ -428,6 +448,8 @@ def _checkpointing_precompute_kernel( HEADS_PER_BLOCK: tl.constexpr, WRITE_CHECKPOINT: tl.constexpr, EARLY_OUT: tl.constexpr, + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, ): # Hoisted PDL signal: fire as the first thing every program does, so # main can start its setup regardless of how this program ends (pad, @@ -441,7 +463,14 @@ def _checkpointing_precompute_kernel( # just an impl call. When True, this kernel only runs for slots whose # (PNAT + T > MAX) status matches WRITE_CHECKPOINT. if EARLY_OUT: - pid_b_eo = tl.program_id(axis=0) + pid_grid_eo = tl.program_id(axis=0) + if USE_PERM: + if REVERSE_PERM: + pid_b_eo = tl.load(slot_perm_ptr + (tl.num_programs(axis=0) - 1 - pid_grid_eo)) + else: + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo) + else: + pid_b_eo = pid_grid_eo if HAS_CACHE_BATCH_INDICES: cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) if cbi_eo == pad_slot_id: @@ -465,6 +494,7 @@ def _checkpointing_precompute_kernel( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, + slot_perm_ptr, pad_slot_id, T, dstate, @@ -509,6 +539,8 @@ def _checkpointing_precompute_kernel( BLOCK_SIZE_T, LAUNCH_WITH_PDL, HEADS_PER_BLOCK, + USE_PERM, + REVERSE_PERM, WRITE_CHECKPOINT, ) @@ -544,6 +576,8 @@ def _rectangle_precompute_impl( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, + # Slot permutation: see _replay_precompute_impl for semantics. + slot_perm_ptr, pad_slot_id, # Dimensions T: tl.constexpr, @@ -600,8 +634,18 @@ def _rectangle_precompute_impl( BLOCK_SIZE_K: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, HEADS_PER_BLOCK: tl.constexpr, + # Slot permutation flags — see _replay_precompute_impl. + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, ): - pid_b = tl.program_id(axis=0) + pid_grid = tl.program_id(axis=0) + if USE_PERM: + if REVERSE_PERM: + pid_b = tl.load(slot_perm_ptr + (tl.num_programs(axis=0) - 1 - pid_grid)) + else: + pid_b = tl.load(slot_perm_ptr + pid_grid) + else: + pid_b = pid_grid pid_hg = tl.program_id(axis=1) first_head = pid_hg * HEADS_PER_BLOCK @@ -939,6 +983,7 @@ def _rectangle_precompute_kernel( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, + slot_perm_ptr, pad_slot_id, # Dimensions T: tl.constexpr, @@ -997,6 +1042,8 @@ def _rectangle_precompute_kernel( LAUNCH_DEPENDENT_KERNELS: tl.constexpr, HEADS_PER_BLOCK: tl.constexpr, EARLY_OUT: tl.constexpr, + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, ): # Hoisted PDL signal: fire as the first thing every program does, so # main can start its setup regardless of how this program ends. @@ -1005,7 +1052,14 @@ def _rectangle_precompute_kernel( # Per-program early-out gate. Rectangle is nowrite-only, so EARLY_OUT # skips slots whose PNAT + T > MAX (slots that would need write). if EARLY_OUT: - pid_b_eo = tl.program_id(axis=0) + pid_grid_eo = tl.program_id(axis=0) + if USE_PERM: + if REVERSE_PERM: + pid_b_eo = tl.load(slot_perm_ptr + (tl.num_programs(axis=0) - 1 - pid_grid_eo)) + else: + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo) + else: + pid_b_eo = pid_grid_eo if HAS_CACHE_BATCH_INDICES: cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) if cbi_eo == pad_slot_id: @@ -1029,6 +1083,7 @@ def _rectangle_precompute_kernel( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, + slot_perm_ptr, pad_slot_id, T, MAX_REPLAY_BUFFER_LENGTH, @@ -1075,6 +1130,8 @@ def _rectangle_precompute_kernel( BLOCK_SIZE_K, LAUNCH_WITH_PDL, HEADS_PER_BLOCK, + USE_PERM, + REVERSE_PERM, ) @@ -1208,6 +1265,7 @@ def _dynamic_precompute_kernel( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, + state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) pad_slot_id, T, dstate, @@ -1252,6 +1310,8 @@ def _dynamic_precompute_kernel( BLOCK_SIZE_T, LAUNCH_WITH_PDL, HEADS_PER_BLOCK, + False, # USE_PERM + False, # REVERSE_PERM needs_write_runtime, ) else: @@ -1269,6 +1329,7 @@ def _dynamic_precompute_kernel( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, + state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) pad_slot_id, T, MAX_REPLAY_BUFFER_LENGTH, @@ -1315,6 +1376,8 @@ def _dynamic_precompute_kernel( BLOCK_SIZE_K, LAUNCH_WITH_PDL, HEADS_PER_BLOCK, + False, # USE_PERM + False, # REVERSE_PERM ) @@ -1347,6 +1410,8 @@ def _replay_main_impl( cb_scaled_ptr, decay_vec_ptr, state_batch_indices_ptr, + # Slot permutation: see _replay_precompute_impl for semantics. + slot_perm_ptr, # Stochastic rounding rand_seed_ptr, pad_slot_id, @@ -1447,6 +1512,9 @@ def _replay_main_impl( # start its setup while this main is still computing. Default False # for monolithic / dynamic / the LAST main in dl/maindl chains. LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + # Slot permutation flags — see _replay_precompute_impl. + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, ): # Hoisted PDL signal: fire as the first thing every program does, so # downstream kernels can start setup regardless of how this program @@ -1468,7 +1536,14 @@ def _replay_main_impl( ) pid_m = tl.program_id(axis=0) - pid_b = tl.program_id(axis=1) + pid_grid_b = tl.program_id(axis=1) + if USE_PERM: + if REVERSE_PERM: + pid_b = tl.load(slot_perm_ptr + (tl.num_programs(axis=1) - 1 - pid_grid_b)) + else: + pid_b = tl.load(slot_perm_ptr + pid_grid_b) + else: + pid_b = pid_grid_b pid_h = tl.program_id(axis=2) if HAS_CACHE_BATCH_INDICES: @@ -1817,6 +1892,7 @@ def _checkpointing_main_kernel( cb_scaled_ptr, decay_vec_ptr, state_batch_indices_ptr, + slot_perm_ptr, rand_seed_ptr, pad_slot_id, # Dimensions @@ -1902,6 +1978,8 @@ def _checkpointing_main_kernel( WRITE_CHECKPOINT: tl.constexpr, EARLY_OUT: tl.constexpr, LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, ): # Hoisted PDL signal: fire as the first thing every program does. if LAUNCH_DEPENDENT_KERNELS: @@ -1909,7 +1987,14 @@ def _checkpointing_main_kernel( # Per-program early-out gate. Signal-then-skip lets the next kernel # in dl/maindl chains start regardless of early-out outcome. if EARLY_OUT: - pid_b_eo = tl.program_id(axis=1) + pid_grid_eo = tl.program_id(axis=1) + if USE_PERM: + if REVERSE_PERM: + pid_b_eo = tl.load(slot_perm_ptr + (tl.num_programs(axis=1) - 1 - pid_grid_eo)) + else: + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo) + else: + pid_b_eo = pid_grid_eo if HAS_CACHE_BATCH_INDICES: cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) if cbi_eo == pad_slot_id: @@ -1936,6 +2021,7 @@ def _checkpointing_main_kernel( cb_scaled_ptr, decay_vec_ptr, state_batch_indices_ptr, + slot_perm_ptr, rand_seed_ptr, pad_slot_id, T, @@ -2005,6 +2091,8 @@ def _checkpointing_main_kernel( QUANT_MAX, WRITE_CHECKPOINT, LAUNCH_DEPENDENT_KERNELS, + USE_PERM, + REVERSE_PERM, ) @@ -2031,6 +2119,8 @@ def _rectangle_main_impl( cb_scaled_ptr, # rectangle (batch, nheads, T, K) decay_vec_ptr, # folded (batch, nheads, T) — total_decay * exp(cumAdt_new[t]) state_batch_indices_ptr, + # Slot permutation: see _replay_precompute_impl for semantics. + slot_perm_ptr, pad_slot_id, # Dimensions T: tl.constexpr, @@ -2095,12 +2185,22 @@ def _rectangle_main_impl( LAUNCH_WITH_PDL: tl.constexpr, QUANT_MAX: tl.constexpr, LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + # Slot permutation flags — see _replay_precompute_impl. + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, ): if LAUNCH_DEPENDENT_KERNELS: tl.extra.cuda.gdc_launch_dependents() pid_m = tl.program_id(axis=0) - pid_b = tl.program_id(axis=1) + pid_grid_b = tl.program_id(axis=1) + if USE_PERM: + if REVERSE_PERM: + pid_b = tl.load(slot_perm_ptr + (tl.num_programs(axis=1) - 1 - pid_grid_b)) + else: + pid_b = tl.load(slot_perm_ptr + pid_grid_b) + else: + pid_b = pid_grid_b pid_h = tl.program_id(axis=2) if HAS_CACHE_BATCH_INDICES: @@ -2286,6 +2386,7 @@ def _rectangle_main_kernel( cb_scaled_ptr, decay_vec_ptr, state_batch_indices_ptr, + slot_perm_ptr, pad_slot_id, # Dimensions T: tl.constexpr, @@ -2351,12 +2452,21 @@ def _rectangle_main_kernel( QUANT_MAX: tl.constexpr, EARLY_OUT: tl.constexpr, LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, ): if LAUNCH_DEPENDENT_KERNELS: tl.extra.cuda.gdc_launch_dependents() # Per-program early-out gate. Rectangle is nowrite-only. if EARLY_OUT: - pid_b_eo = tl.program_id(axis=1) + pid_grid_eo = tl.program_id(axis=1) + if USE_PERM: + if REVERSE_PERM: + pid_b_eo = tl.load(slot_perm_ptr + (tl.num_programs(axis=1) - 1 - pid_grid_eo)) + else: + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo) + else: + pid_b_eo = pid_grid_eo if HAS_CACHE_BATCH_INDICES: cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) if cbi_eo == pad_slot_id: @@ -2380,6 +2490,7 @@ def _rectangle_main_kernel( cb_scaled_ptr, decay_vec_ptr, state_batch_indices_ptr, + slot_perm_ptr, pad_slot_id, T, MAX_REPLAY_BUFFER_LENGTH, @@ -2432,6 +2543,8 @@ def _rectangle_main_kernel( LAUNCH_WITH_PDL, QUANT_MAX, LAUNCH_DEPENDENT_KERNELS, + USE_PERM, + REVERSE_PERM, ) @@ -2603,6 +2716,7 @@ def _dynamic_main_kernel( cb_scaled_ptr, decay_vec_ptr, state_batch_indices_ptr, + state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) rand_seed_ptr, pad_slot_id, T, @@ -2672,6 +2786,8 @@ def _dynamic_main_kernel( QUANT_MAX, True, # WRITE_CHECKPOINT (constexpr) False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top + False, # USE_PERM + False, # REVERSE_PERM ) else: if RECTANGLE: @@ -2689,6 +2805,7 @@ def _dynamic_main_kernel( cb_scaled_ptr, decay_vec_ptr, state_batch_indices_ptr, + state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) pad_slot_id, T, MAX_REPLAY_BUFFER_LENGTH, @@ -2741,6 +2858,8 @@ def _dynamic_main_kernel( LAUNCH_WITH_PDL, QUANT_MAX, False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top + False, # USE_PERM + False, # REVERSE_PERM ) else: # Replay-style nowrite (WRITE_CHECKPOINT=False constexpr). @@ -2761,6 +2880,7 @@ def _dynamic_main_kernel( cb_scaled_ptr, decay_vec_ptr, state_batch_indices_ptr, + state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) rand_seed_ptr, pad_slot_id, T, @@ -2830,6 +2950,8 @@ def _dynamic_main_kernel( QUANT_MAX, False, # WRITE_CHECKPOINT (constexpr) False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top + False, # USE_PERM + False, # REVERSE_PERM ) @@ -2871,6 +2993,17 @@ def checkpointing_state_update( write_checkpoint: bool = True, rectangle_for_nowrite: bool = False, mode: str = "monolithic", + # Slot permutation: int32 (batch,) tensor mapping grid program_id -> + # original slot index. When provided, dl-family kernels (doublelaunch / + # dlgrouped / maindl) read pid_b through this perm so callers can pre-sort + # slots (e.g. write-first) to cluster early-outs at one end of the grid. + # Ignored by monolithic / dynamic. None => identity (today's behavior). + slot_perm: torch.Tensor | None = None, + # When True and slot_perm is provided, the nowrite-side kernels in + # dlgrouped/doublelaunch traverse the perm in reverse (B-1-pid_grid). + # Combined with a write-first sort, this front-loads real work in BOTH + # halves of the dl chain (writes from the head, nowrites from the tail). + reverse_nowrite: bool = False, _block_size_m: int | None = None, _num_warps: int | None = None, _num_stages: int | None = None, @@ -3251,6 +3384,25 @@ def checkpointing_state_update( state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 state_scales_strides = (0, 0, 0) + # Slot permutation — pointer + USE_PERM gate. When the caller provides + # a perm tensor the dl-family launches read pid_b through it; otherwise + # we pass any valid pointer (state_batch_indices) and USE_PERM=False so + # the kernel falls back to pid_grid. Sort-driven dispatch (write-first + # clustering) is opt-in per call; monolithic / dynamic ignore the flag. + if slot_perm is not None: + assert slot_perm.dtype in (torch.int32, torch.int64), ( + f"slot_perm must be int32/int64, got {slot_perm.dtype}" + ) + assert slot_perm.numel() >= batch, ( + f"slot_perm has {slot_perm.numel()} entries; need >= batch ({batch})" + ) + slot_perm_arg = slot_perm + use_perm = True + else: + # Any valid ptr — gated by USE_PERM=False at compile time. + slot_perm_arg = state_batch_indices if state_batch_indices is not None else state + use_perm = False + # Grid for main kernels (M tiling × batch × nheads). def main_grid(META): return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) @@ -3263,13 +3415,14 @@ def main_grid(META): # full positional + kwarg argument list. Mode-dependent constexprs # (write_checkpoint, early_out, rectangle) are passed in. - def launch_replay_precompute(write_checkpoint: bool, early_out: bool): + def launch_replay_precompute(write_checkpoint: bool, early_out: bool, + reverse_perm: bool = False): _checkpointing_precompute_kernel[precomp_grid]( dt, dt_bias, A, B, C, cb_scaled, decay_vec, old_B, old_dt, old_dA_cumsum, cache_buf_idx, prev_num_accepted_tokens, - state_batch_indices, pad_slot_id, + state_batch_indices, slot_perm_arg, pad_slot_id, T, max_window, dstate, nheads // ngroups, dt.stride(0), dt.stride(1), dt.stride(2), dt_bias.stride(0) if dt_bias is not None else 0, @@ -3292,18 +3445,20 @@ def launch_replay_precompute(write_checkpoint: bool, early_out: bool): HEADS_PER_BLOCK=heads_per_block, WRITE_CHECKPOINT=write_checkpoint, EARLY_OUT=early_out, + USE_PERM=use_perm, + REVERSE_PERM=reverse_perm, num_warps=precompute_num_warps, **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), launch_pdl=launch_with_pdl, ) - def launch_rectangle_precompute(early_out: bool): + def launch_rectangle_precompute(early_out: bool, reverse_perm: bool = False): _rectangle_precompute_kernel[precomp_grid]( dt, dt_bias, A, B, C, cb_scaled, decay_vec, old_B, old_dt, old_dA_cumsum, cache_buf_idx, prev_num_accepted_tokens, - state_batch_indices, pad_slot_id, + state_batch_indices, slot_perm_arg, pad_slot_id, T, max_window, dstate, nheads // ngroups, dt.stride(0), dt.stride(1), dt.stride(2), dt_bias.stride(0) if dt_bias is not None else 0, @@ -3325,6 +3480,8 @@ def launch_rectangle_precompute(early_out: bool): LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, HEADS_PER_BLOCK=heads_per_block, EARLY_OUT=early_out, + USE_PERM=use_perm, + REVERSE_PERM=reverse_perm, num_warps=precompute_num_warps, **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), launch_pdl=launch_with_pdl, @@ -3364,14 +3521,15 @@ def launch_dynamic_precompute(rectangle: bool): ) def launch_replay_main(write_checkpoint: bool, early_out: bool, - launch_dependent_kernels: bool = False): + launch_dependent_kernels: bool = False, + reverse_perm: bool = False): _checkpointing_main_kernel[main_grid]( state, state_scales_arg, old_x, old_B, old_dt, old_dA_cumsum, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, cb_scaled, decay_vec, - state_batch_indices, rand_seed, pad_slot_id, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, T, max_window, dim, dstate, nheads // ngroups, state.stride(0), state.stride(1), state.stride(2), state.stride(3), state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], @@ -3397,6 +3555,8 @@ def launch_replay_main(write_checkpoint: bool, early_out: bool, WRITE_CHECKPOINT=write_checkpoint, EARLY_OUT=early_out, LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + REVERSE_PERM=reverse_perm, num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), @@ -3405,13 +3565,14 @@ def launch_replay_main(write_checkpoint: bool, early_out: bool, ) def launch_rectangle_main(early_out: bool, - launch_dependent_kernels: bool = False): + launch_dependent_kernels: bool = False, + reverse_perm: bool = False): _rectangle_main_kernel[main_grid]( state, state_scales_arg, old_x, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, cb_scaled, decay_vec, - state_batch_indices, pad_slot_id, + state_batch_indices, slot_perm_arg, pad_slot_id, T, max_window, dim, dstate, nheads // ngroups, state.stride(0), state.stride(1), state.stride(2), state.stride(3), state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], @@ -3429,6 +3590,8 @@ def launch_rectangle_main(early_out: bool, QUANT_MAX=quant_max, EARLY_OUT=early_out, LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + REVERSE_PERM=reverse_perm, num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), @@ -3495,30 +3658,37 @@ def launch_dynamic_main(rectangle: bool, # EARLY_OUT so each retains its constexpr-specialized reg # envelope. First main signals PDL dependents so the second # main can start its setup while the first is still computing. + # Dynamic precompute doesn't support sort (no early-out to + # cluster), so the perm only flows into the two EARLY_OUT mains. launch_dynamic_precompute(rectangle=rectangle_for_nowrite) launch_replay_main(write_checkpoint=True, early_out=True, launch_dependent_kernels=True) if rectangle_for_nowrite: - launch_rectangle_main(early_out=True) + launch_rectangle_main(early_out=True, reverse_perm=reverse_nowrite) else: - launch_replay_main(write_checkpoint=False, early_out=True) + launch_replay_main(write_checkpoint=False, early_out=True, + reverse_perm=reverse_nowrite) elif mode == "dlgrouped": # Same 4 kernels as doublelaunch but reordered: both precomputes # first, then both mains. Lets the GPU run precomp1 || precomp2 # in parallel (they're tiny grids) before the mains start, vs # doublelaunch's interleaved precomp1→main1→precomp2→main2. # First main signals PDL so the second main's setup overlaps. + # When slot_perm + reverse_nowrite are set, the nowrite-side + # walks the perm in reverse so both kernels front-load real work. launch_replay_precompute(write_checkpoint=True, early_out=True) if rectangle_for_nowrite: - launch_rectangle_precompute(early_out=True) + launch_rectangle_precompute(early_out=True, reverse_perm=reverse_nowrite) else: - launch_replay_precompute(write_checkpoint=False, early_out=True) + launch_replay_precompute(write_checkpoint=False, early_out=True, + reverse_perm=reverse_nowrite) launch_replay_main(write_checkpoint=True, early_out=True, launch_dependent_kernels=True) if rectangle_for_nowrite: - launch_rectangle_main(early_out=True) + launch_rectangle_main(early_out=True, reverse_perm=reverse_nowrite) else: - launch_replay_main(write_checkpoint=False, early_out=True) + launch_replay_main(write_checkpoint=False, early_out=True, + reverse_perm=reverse_nowrite) else: # mode == "doublelaunch" # Write half: always replay-style write. First main signals # PDL dependents so the second precompute can start its setup @@ -3528,8 +3698,10 @@ def launch_dynamic_main(rectangle: bool, launch_dependent_kernels=True) # Nowrite half: rectangle if asked, else replay-nowrite. if rectangle_for_nowrite: - launch_rectangle_precompute(early_out=True) - launch_rectangle_main(early_out=True) + launch_rectangle_precompute(early_out=True, reverse_perm=reverse_nowrite) + launch_rectangle_main(early_out=True, reverse_perm=reverse_nowrite) else: - launch_replay_precompute(write_checkpoint=False, early_out=True) - launch_replay_main(write_checkpoint=False, early_out=True) + launch_replay_precompute(write_checkpoint=False, early_out=True, + reverse_perm=reverse_nowrite) + launch_replay_main(write_checkpoint=False, early_out=True, + reverse_perm=reverse_nowrite) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 5ae190c6aebb..a17b42346ac5 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -56,6 +56,7 @@ from datetime import datetime from pathlib import Path +import numpy as np import torch from einops import repeat @@ -324,6 +325,11 @@ def _build_tensors( # prev_tokens placeholder — overwritten per-run prev_tokens = torch.zeros(batch, device=device, dtype=torch.int32) + # slot_perm placeholder — overwritten per-run by mix pre_iter_fn when + # sort_slots is enabled. Identity by default so cells that don't sort + # (or pure-batch cells) get a meaningful identity perm if the kernel + # ends up reading it (USE_PERM=False makes this path unused). + slot_perm_buf = torch.arange(batch, device=device, dtype=torch.int32) out_incr = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) out_base = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) @@ -373,6 +379,7 @@ def _build_tensors( dt_bias, D, prev_tokens, + slot_perm_buf, out_incr, out_base, intermediate_states_buffer, @@ -547,6 +554,8 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) modes_list = getattr(args, "modes_list", ["monolithic"]) + sort_list = getattr(args, "sort_slots_list", [False]) + rev_list = getattr(args, "reverse_nowrite_list", [False]) configs = [] for batch in batch_sizes: @@ -567,21 +576,39 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp else: effective_rect_list = rect_list for rect in effective_rect_list: - configs.append(( - batch, mtp_len, prev_ks, state_dtype, act_dtype, - sr_mode, rect, write_ckpt, mode, - )) + is_dl_family = mode in ( + "doublelaunch", "dlgrouped", "maindl" + ) + # Match the timed-run skip: sort=1 only + # makes sense when there's a mix scenario. + can_sort = is_dl_family and (args.mix_csv is not None) + effective_sort_list = ( + sort_list if can_sort else [False] + ) + for sort_slots in effective_sort_list: + effective_rev_list = ( + rev_list if sort_slots else [False] + ) + for reverse_nowrite in effective_rev_list: + configs.append(( + batch, mtp_len, prev_ks, + state_dtype, act_dtype, + sr_mode, rect, write_ckpt, mode, + sort_slots, reverse_nowrite, + )) print(f"[compile-warmup] {len(configs)} configs across {max_workers} threads") t0 = time.perf_counter() def _warm(cfg): (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, - rect, write_ckpt, mode) = cfg + rect, write_ckpt, mode, sort_slots, reverse_nowrite) = cfg _bench_config( args, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, sr_mode=sr_mode, rectangle_for_nowrite=rect, - write_checkpoint=write_ckpt, mode=mode, warmup_only=True, + write_checkpoint=write_ckpt, mode=mode, + sort_slots=sort_slots, reverse_nowrite=reverse_nowrite, + warmup_only=True, ) errors = [] @@ -616,6 +643,9 @@ def _bench_config( mode: str = "monolithic", mix_samples_cpu=None, mix_label: str = "", + sort_slots: bool = False, + reverse_nowrite: bool = False, + perm_samples_cpu=None, warmup_only: bool = False, ) -> None: """ @@ -649,6 +679,7 @@ def _bench_config( dt_bias, D, prev_tokens, + slot_perm_buf, out_incr, out_base, intermediate_states_buffer, @@ -906,8 +937,18 @@ def _parse_sweep(val): device = state_work.device samples_gpu = torch.from_numpy(mix_samples_cpu).to(device=device, dtype=torch.int32) - def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): - _pt.copy_(_s[i]) + if sort_slots and perm_samples_cpu is not None: + perm_samples_gpu = torch.from_numpy(perm_samples_cpu).to( + device=device, dtype=torch.int32 + ) + + def _mix_pre_iter(i, _s=samples_gpu, _ps=perm_samples_gpu, + _pt=prev_tokens, _pm=slot_perm_buf): + _pt.copy_(_s[i]) + _pm.copy_(_ps[i]) + else: + def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): + _pt.copy_(_s[i]) # Mix iters override: if --mix-iters set, use it; else use args.iters. mix_iters = getattr(args, "mix_iters", None) @@ -972,6 +1013,9 @@ def _run_incr( extra_kwargs["write_checkpoint"] = write_checkpoint extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite extra_kwargs["mode"] = mode + if sort_slots: + extra_kwargs["slot_perm"] = slot_perm_buf + extra_kwargs["reverse_nowrite"] = reverse_nowrite if state_scales_work is not None: extra_kwargs["state_scales"] = state_scales_work variant_fn( @@ -1027,6 +1071,8 @@ def _run_incr( parts.append(f"RECT={1 if rectangle_for_nowrite else 0}") parts.append(f"WC={1 if write_checkpoint else 0}") parts.append(f"MODE={mode}") + parts.append(f"SORT={1 if sort_slots else 0}") + parts.append(f"REVN={1 if reverse_nowrite else 0}") sweep_suffix = (" " + ",".join(parts)) if parts else "" sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") @@ -1158,6 +1204,8 @@ def _run_benchmark(args) -> None: rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) modes_list = getattr(args, "modes_list", ["monolithic"]) + sort_list = getattr(args, "sort_slots_list", [False]) + rev_list = getattr(args, "reverse_nowrite_list", [False]) # Pre-load AL distribution for mix mode (if --mix-csv set). mix_al = None @@ -1181,6 +1229,7 @@ def _run_benchmark(args) -> None: # Size the sample buffer for the LARGER of args.iters and # args.mix_iters since mix scenarios use mix_iters. mix_samples_cpu = None + perm_samples_cpu = None # per-iter slot perm sorted write-first if mix_al is not None: from checkpoint_mix_sim import sample_steady_state_pnat as _sample_pnat _max_window = getattr(args, "max_window", 0) or mtp_len @@ -1189,6 +1238,15 @@ def _run_benchmark(args) -> None: mix_al, T=mtp_len, window=_max_window, batch=batch, K=args.warmup + _max_iters, seed=args.mix_seed, ) + if any(sort_list): + # write-first stable argsort: kind='stable' preserves + # original-slot order within each mode group. + write_mask = ( + mix_samples_cpu + mtp_len > _max_window + ).astype(np.int8) # 1 = write, 0 = nowrite + perm_samples_cpu = np.argsort( + -write_mask, kind="stable", axis=-1 + ).astype(np.int32) for state_dtype in state_dtypes: for act_dtype in act_dtypes: @@ -1211,15 +1269,44 @@ def _run_benchmark(args) -> None: else: effective_rect_list = rect_list for rect in effective_rect_list: - _bench_config( - args, batch, mtp_len, prev_ks, state_dtype, act_dtype, - baseline_fn, sr_mode=sr_mode, - rectangle_for_nowrite=rect, - write_checkpoint=write_ckpt, - mode=mode, - mix_samples_cpu=mix_samples_cpu, - mix_label=mix_label, + # Sort/reverse only meaningful for the + # dl-family early-out kernels AND only + # against the mix scenario (the actual + # sort experiment). Pure k= scenarios + # under sort=1 would just run a + # USE_PERM=True kernel against an + # identity perm — same data point as + # sort=0 + extra compile. Skip sort=1 + # when no mix is configured; mono / + # dynamic also skip sort=1; reverse=1 + # with sort=0 is a no-op (skip). + is_dl_family = mode in ( + "doublelaunch", "dlgrouped", "maindl" ) + can_sort = ( + is_dl_family and mix_samples_cpu is not None + ) + effective_sort_list = ( + sort_list if can_sort else [False] + ) + for sort_slots in effective_sort_list: + effective_rev_list = ( + rev_list if sort_slots else [False] + ) + for reverse_nowrite in effective_rev_list: + _bench_config( + args, batch, mtp_len, prev_ks, + state_dtype, act_dtype, + baseline_fn, sr_mode=sr_mode, + rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, + mode=mode, + mix_samples_cpu=mix_samples_cpu, + mix_label=mix_label, + sort_slots=sort_slots, + reverse_nowrite=reverse_nowrite, + perm_samples_cpu=perm_samples_cpu, + ) if args.profile: torch.cuda.cudart().cudaProfilerStop() @@ -1499,6 +1586,26 @@ def _parse_args() -> argparse.Namespace: "across runs => same per-slot samples for reproducible " "comparisons.", ) + parser.add_argument( + "--sort-slots", + type=str, + default="0", + help="Comma-separated 0/1. When 1, mix scenarios pre-sort slots " + "write-first (write slots at the head of slot_perm, nowrite at the " + "tail) and the dl-family kernels read pid_b through that perm — " + "clusters early-outs at one end of the grid. Only meaningful for " + "doublelaunch/dlgrouped/maindl with mix scenarios; mono/dynamic " + "and pure-batch cells skip sort=1.", + ) + parser.add_argument( + "--reverse-nowrite", + type=str, + default="0", + help="Comma-separated 0/1. When 1 (and --sort-slots 1), the " + "nowrite-side kernels in dlgrouped/doublelaunch/maindl walk the " + "perm in reverse so both halves of the dl chain front-load real " + "work. reverse=1 with sort=0 is skipped (no perm to reverse).", + ) parser.add_argument( "--mix-iters", type=int, @@ -1568,6 +1675,22 @@ def _parse_args() -> argparse.Namespace: rect_list.append(v == "1") args.rectangle_for_nowrite_list = rect_list + sort_modes = [v.strip() for v in (args.sort_slots or "0").split(",") if v.strip()] + sort_list = [] + for v in sort_modes: + if v not in ("0", "1"): + parser.error(f"--sort-slots value must be 0 or 1, got {v!r}") + sort_list.append(v == "1") + args.sort_slots_list = sort_list + + rev_modes = [v.strip() for v in (args.reverse_nowrite or "0").split(",") if v.strip()] + rev_list = [] + for v in rev_modes: + if v not in ("0", "1"): + parser.error(f"--reverse-nowrite value must be 0 or 1, got {v!r}") + rev_list.append(v == "1") + args.reverse_nowrite_list = rev_list + if args.write_modes is not None: wm = [v.strip() for v in args.write_modes.split(",") if v.strip()] write_list = [] diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index 9b1bbbf046d2..b8bd29cfc56e 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -729,6 +729,155 @@ def test_checkpointing_state_update_mixed_mode(mode, rectangle_for_nowrite): ) +@pytest.mark.parametrize( + "mode,rectangle_for_nowrite,reverse_nowrite", + [ + ("doublelaunch", False, False), + ("doublelaunch", False, True), + ("doublelaunch", True, False), + ("doublelaunch", True, True), + ("dlgrouped", False, False), + ("dlgrouped", False, True), + ("dlgrouped", True, False), + ("dlgrouped", True, True), + ("maindl", False, False), + ("maindl", False, True), + ("maindl", True, False), + ("maindl", True, True), + ], + ids=lambda v: str(v), +) +def test_checkpointing_state_update_sorted_dispatch(mode, rectangle_for_nowrite, reverse_nowrite): + """ + Sort-driven dispatch: caller pre-sorts slots write-first via slot_perm. + Verifies the perm-aware kernels remap pid_b correctly so each slot's + work lands at the right grid program — i.e. slot S's output and HBM + state still match the reference under any permutation. + + Setup mirrors test_checkpointing_state_update_mixed_mode but with + slot_perm = [2, 3, 0, 1] (write slots 2, 3 first; nowrite 0, 1 after). + reverse_nowrite=True walks the perm tail-first on the nowrite-side, + so e.g. rectangle main reads perm[B-1-pid_grid] = [1, 0, 3, 2]. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + + pnat_per_slot = torch.tensor([3, 10, 12, 16], device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() # [F, F, T, T] + # Write-first perm: indices 2, 3 (write) then 0, 1 (nowrite). + slot_perm = torch.tensor([2, 3, 0, 1], device=device, dtype=torch.int32) + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=dtype) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + test_state, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + rectangle_for_nowrite=rectangle_for_nowrite, + mode=mode, + slot_perm=slot_perm, + reverse_nowrite=reverse_nowrite, + ) + + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (mode={mode}, rect={rectangle_for_nowrite}, rev={reverse_nowrite})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (mode={mode}, rev={reverse_nowrite})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM modified (mode={mode}, rev={reverse_nowrite})", + ) + + @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) @pytest.mark.parametrize( "state_dtype", From 96df45895c75af7caf22a213240e6ceced790b66 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 6 May 2026 17:36:29 -0700 Subject: [PATCH 20/89] =?UTF-8?q?rect:=20convention=20swap=20=E2=80=94=20n?= =?UTF-8?q?ew=20tokens=20at=20runtime=20[PNAT,=20PNAT+T)=20instead=20of=20?= =?UTF-8?q?K=5FNEW=5FSHIFT?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the compile-time-static K_NEW_SHIFT (= MAX - T) layout for new tokens in the rectangle K-axis with a runtime PNAT-based offset. Both rect_precompute and rect_main now place new tokens at K-axis positions [PNAT, PNAT+T), matching the cache write layout. Win: the cache write target [write_offset, write_offset+T) = [PNAT, PNAT+T) is now the same as the K-axis row positions for new tokens. rect_main loads x ONCE at the K-axis-shifted positions (a single (BLOCK_K, M) load masked to is_new_k = offs_k in [PNAT, PNAT+T)) and stores it directly to the cache at offs_k — no separate (T, M) load and no row reshuffling. The T-axis x_all needed for D feedthrough / Z-gating is extracted from x_K via a (T, K) selection tl.dot only when HAS_D or HAS_Z. Trade-off: precomp loses the compile-time K_NEW_SHIFT static-addressing optimization for new-side B/dt loads — they now use runtime PNAT-offset indexing. Same for the causal mask. Empirically the unification wins. Same-node sweep vs the prior commit (Z-loop vectorize, rect_main unchanged): 23 better, 6 worse, 11 same. Headlines: - b=64 across all dtypes: -7.9 to -10.1% (biggest wins) - b=32 across dtypes: -1.3 to -3.2% - b=16 fp16 / fp8: small wins Regressions: - b=1 fp8 RN/SR: +4.7-6.2% (autotune flipped) - b=4 fp8 RN/SR: +1.1-1.3% We also tried 3 alternative approaches (A: tile row-shift via tl.dot selection matrix; B: single K-axis load + tl.dot extraction for cache write; D: A + early D feedthrough) — all net regressions or wash. C wins decisively. Tests: 144 rectangle tests pass. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 54 +++++++++++-------- 1 file changed, 31 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index e38923a05278..5a6da1cea8db 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -467,10 +467,11 @@ def _rectangle_precompute_kernel( t_mask = offs_t < T n_mask = offs_n < dstate - # K-axis masks + # K-axis masks (approach C: runtime PNAT-offset instead of K_NEW_SHIFT) + # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. is_old_k = offs_k < prev_num_accepted_tokens safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - K_NEW_SHIFT + k_new_idx = offs_k - prev_num_accepted_tokens is_new_k = (k_new_idx >= 0) & (k_new_idx < T) safe_k_new = tl.where(is_new_k, k_new_idx, 0) @@ -682,10 +683,11 @@ def _rectangle_precompute_kernel( ) # Causal mask (BLOCK_SIZE_T × BLOCK_SIZE_K, shared across heads). + # Approach C: new tokens at runtime [PNAT, PNAT+T) instead of K_NEW_SHIFT. t_idx_2d = offs_t[:, None] k_idx_2d = offs_k[None, :] is_old_k_2d = k_idx_2d < prev_num_accepted_tokens - k_new_idx_2d = k_idx_2d - K_NEW_SHIFT + k_new_idx_2d = k_idx_2d - prev_num_accepted_tokens is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] @@ -1307,10 +1309,11 @@ def _rectangle_main_kernel( n_mask = offs_n < dstate t_mask = offs_t < T - # K-axis masks + # K-axis masks (approach C: PNAT-runtime offset, matches precompute). + # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. is_old_k = offs_k < prev_num_accepted_tokens safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - K_NEW_SHIFT + k_new_idx = offs_k - prev_num_accepted_tokens is_new_k = (k_new_idx >= 0) & (k_new_idx < T) safe_k_new = tl.where(is_new_k, k_new_idx, 0) @@ -1369,31 +1372,36 @@ def _rectangle_main_kernel( mask=t_mask[:, None] & n_mask[None, :], other=0.0, ) - x_all = tl.load( - x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=t_mask[:, None] & m_mask[None, :], + # Single (BLOCK_K, M) load at PNAT-offset positions (approach C): + # K-axis [PNAT, PNAT+T) gets new tokens directly from x[0:T, :] via + # safe_k_new = offs_k - PNAT. Cache layout matches K-axis layout, so + # one load serves both matmul and cache write. + x_K = tl.load( + x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=is_new_k[:, None] & m_mask[None, :], other=0.0, ) - # Append new x to cache at [PNAT, PNAT+T) of the active buffer. + # Cache write: store at offs_k directly (is_new_k mask makes offs_k land + # at [PNAT, PNAT+T) in the cache, which is exactly write_offset+0..T-1). tl.store( old_x_base - + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_k[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - x_all, - mask=t_mask[:, None] & m_mask[None, :], + x_K, + mask=is_new_k[:, None] & m_mask[None, :], ) - x_all = x_all.to(tl.float32) - # Build x_combined for the rectangle token_out matmul (STATIC layout). - # old_x_load is from cache (loaded above gdc_wait); new_x_shifted reads - # x (conv1d output) at compile-time-shifted [K_NEW_SHIFT, K_NEW_SHIFT+T). - # Disjoint masks → safe to add. - new_x_shifted = tl.load( - x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=is_new_k[:, None] & m_mask[None, :], - other=0.0, - ).to(tl.float32) - x_combined = old_x_load + new_x_shifted # (BLOCK_SIZE_K, BLOCK_SIZE_M) + x_K_f32 = x_K.to(tl.float32) + # Matmul side: K-axis aligned; sum with old_x_load. + x_combined = old_x_load + x_K_f32 + + # T-axis view for D feedthrough / Z-gating: extract via (T, K) selection. + # Only materialized when needed. + if HAS_D or HAS_Z: + sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) + x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) + else: + x_all = x_K_f32 # placeholder; unused # Load precomputed rectangle CB and folded decay_vec. cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head From 3d1e4c17e487f3a9380aa2126950297ebbc30fca Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 6 May 2026 22:09:01 -0700 Subject: [PATCH 21/89] hardcode-sort: bench-side per-iter PNAT pre-sort (kernel unchanged) + REVERSE_PERM applies without USE_PERM Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 72 +++++----- ...benchmark_replay_selective_state_update.py | 127 ++++++++++++++---- 2 files changed, 133 insertions(+), 66 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 39adcad73547..e9a45fcf55a5 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -184,13 +184,15 @@ def _replay_precompute_impl( write_checkpoint, ): pid_grid = tl.program_id(axis=0) + # REVERSE_PERM walks the grid tail-first regardless of USE_PERM — + # combined with hardcode-sorted prev_tokens (kernel-side USE_PERM=False + # but PNAT pre-sorted write-first), reverse traversal makes the nowrite + # half front-load real work. + pid_grid_eff = (tl.num_programs(axis=0) - 1 - pid_grid) if REVERSE_PERM else pid_grid if USE_PERM: - if REVERSE_PERM: - pid_b = tl.load(slot_perm_ptr + (tl.num_programs(axis=0) - 1 - pid_grid)) - else: - pid_b = tl.load(slot_perm_ptr + pid_grid) + pid_b = tl.load(slot_perm_ptr + pid_grid_eff) else: - pid_b = pid_grid + pid_b = pid_grid_eff pid_hg = tl.program_id(axis=1) # head-group index first_head = pid_hg * HEADS_PER_BLOCK @@ -464,13 +466,11 @@ def _checkpointing_precompute_kernel( # (PNAT + T > MAX) status matches WRITE_CHECKPOINT. if EARLY_OUT: pid_grid_eo = tl.program_id(axis=0) + pid_grid_eo_eff = (tl.num_programs(axis=0) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo if USE_PERM: - if REVERSE_PERM: - pid_b_eo = tl.load(slot_perm_ptr + (tl.num_programs(axis=0) - 1 - pid_grid_eo)) - else: - pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo) + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) else: - pid_b_eo = pid_grid_eo + pid_b_eo = pid_grid_eo_eff if HAS_CACHE_BATCH_INDICES: cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) if cbi_eo == pad_slot_id: @@ -639,13 +639,15 @@ def _rectangle_precompute_impl( REVERSE_PERM: tl.constexpr, ): pid_grid = tl.program_id(axis=0) + # REVERSE_PERM walks the grid tail-first regardless of USE_PERM — + # combined with hardcode-sorted prev_tokens (kernel-side USE_PERM=False + # but PNAT pre-sorted write-first), reverse traversal makes the nowrite + # half front-load real work. + pid_grid_eff = (tl.num_programs(axis=0) - 1 - pid_grid) if REVERSE_PERM else pid_grid if USE_PERM: - if REVERSE_PERM: - pid_b = tl.load(slot_perm_ptr + (tl.num_programs(axis=0) - 1 - pid_grid)) - else: - pid_b = tl.load(slot_perm_ptr + pid_grid) + pid_b = tl.load(slot_perm_ptr + pid_grid_eff) else: - pid_b = pid_grid + pid_b = pid_grid_eff pid_hg = tl.program_id(axis=1) first_head = pid_hg * HEADS_PER_BLOCK @@ -1053,13 +1055,11 @@ def _rectangle_precompute_kernel( # skips slots whose PNAT + T > MAX (slots that would need write). if EARLY_OUT: pid_grid_eo = tl.program_id(axis=0) + pid_grid_eo_eff = (tl.num_programs(axis=0) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo if USE_PERM: - if REVERSE_PERM: - pid_b_eo = tl.load(slot_perm_ptr + (tl.num_programs(axis=0) - 1 - pid_grid_eo)) - else: - pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo) + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) else: - pid_b_eo = pid_grid_eo + pid_b_eo = pid_grid_eo_eff if HAS_CACHE_BATCH_INDICES: cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) if cbi_eo == pad_slot_id: @@ -1537,13 +1537,11 @@ def _replay_main_impl( pid_m = tl.program_id(axis=0) pid_grid_b = tl.program_id(axis=1) + pid_grid_b_eff = (tl.num_programs(axis=1) - 1 - pid_grid_b) if REVERSE_PERM else pid_grid_b if USE_PERM: - if REVERSE_PERM: - pid_b = tl.load(slot_perm_ptr + (tl.num_programs(axis=1) - 1 - pid_grid_b)) - else: - pid_b = tl.load(slot_perm_ptr + pid_grid_b) + pid_b = tl.load(slot_perm_ptr + pid_grid_b_eff) else: - pid_b = pid_grid_b + pid_b = pid_grid_b_eff pid_h = tl.program_id(axis=2) if HAS_CACHE_BATCH_INDICES: @@ -1988,13 +1986,11 @@ def _checkpointing_main_kernel( # in dl/maindl chains start regardless of early-out outcome. if EARLY_OUT: pid_grid_eo = tl.program_id(axis=1) + pid_grid_eo_eff = (tl.num_programs(axis=1) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo if USE_PERM: - if REVERSE_PERM: - pid_b_eo = tl.load(slot_perm_ptr + (tl.num_programs(axis=1) - 1 - pid_grid_eo)) - else: - pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo) + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) else: - pid_b_eo = pid_grid_eo + pid_b_eo = pid_grid_eo_eff if HAS_CACHE_BATCH_INDICES: cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) if cbi_eo == pad_slot_id: @@ -2194,13 +2190,11 @@ def _rectangle_main_impl( pid_m = tl.program_id(axis=0) pid_grid_b = tl.program_id(axis=1) + pid_grid_b_eff = (tl.num_programs(axis=1) - 1 - pid_grid_b) if REVERSE_PERM else pid_grid_b if USE_PERM: - if REVERSE_PERM: - pid_b = tl.load(slot_perm_ptr + (tl.num_programs(axis=1) - 1 - pid_grid_b)) - else: - pid_b = tl.load(slot_perm_ptr + pid_grid_b) + pid_b = tl.load(slot_perm_ptr + pid_grid_b_eff) else: - pid_b = pid_grid_b + pid_b = pid_grid_b_eff pid_h = tl.program_id(axis=2) if HAS_CACHE_BATCH_INDICES: @@ -2460,13 +2454,11 @@ def _rectangle_main_kernel( # Per-program early-out gate. Rectangle is nowrite-only. if EARLY_OUT: pid_grid_eo = tl.program_id(axis=1) + pid_grid_eo_eff = (tl.num_programs(axis=1) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo if USE_PERM: - if REVERSE_PERM: - pid_b_eo = tl.load(slot_perm_ptr + (tl.num_programs(axis=1) - 1 - pid_grid_eo)) - else: - pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo) + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) else: - pid_b_eo = pid_grid_eo + pid_b_eo = pid_grid_eo_eff if HAS_CACHE_BATCH_INDICES: cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) if cbi_eo == pad_slot_id: diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index a17b42346ac5..02b6ab58616b 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -556,6 +556,7 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp modes_list = getattr(args, "modes_list", ["monolithic"]) sort_list = getattr(args, "sort_slots_list", [False]) rev_list = getattr(args, "reverse_nowrite_list", [False]) + hsort_list = getattr(args, "hardcode_sort_list", [False]) configs = [] for batch in batch_sizes: @@ -585,29 +586,37 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp effective_sort_list = ( sort_list if can_sort else [False] ) + effective_hsort_list = ( + hsort_list if can_sort else [False] + ) for sort_slots in effective_sort_list: effective_rev_list = ( rev_list if sort_slots else [False] ) for reverse_nowrite in effective_rev_list: - configs.append(( - batch, mtp_len, prev_ks, - state_dtype, act_dtype, - sr_mode, rect, write_ckpt, mode, - sort_slots, reverse_nowrite, - )) + for hardcode_sort in effective_hsort_list: + if sort_slots and hardcode_sort: + continue + configs.append(( + batch, mtp_len, prev_ks, + state_dtype, act_dtype, + sr_mode, rect, write_ckpt, mode, + sort_slots, reverse_nowrite, + hardcode_sort, + )) print(f"[compile-warmup] {len(configs)} configs across {max_workers} threads") t0 = time.perf_counter() def _warm(cfg): (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, - rect, write_ckpt, mode, sort_slots, reverse_nowrite) = cfg + rect, write_ckpt, mode, sort_slots, reverse_nowrite, hardcode_sort) = cfg _bench_config( args, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, sr_mode=sr_mode, rectangle_for_nowrite=rect, write_checkpoint=write_ckpt, mode=mode, sort_slots=sort_slots, reverse_nowrite=reverse_nowrite, + hardcode_sort=hardcode_sort, warmup_only=True, ) @@ -646,6 +655,8 @@ def _bench_config( sort_slots: bool = False, reverse_nowrite: bool = False, perm_samples_cpu=None, + hardcode_sort: bool = False, + mix_samples_sorted_cpu=None, warmup_only: bool = False, ) -> None: """ @@ -935,7 +946,13 @@ def _parse_sweep(val): # nsys-included warmup leaks aren't biased). if mix_samples_cpu is not None and mode != "monolithic": device = state_work.device - samples_gpu = torch.from_numpy(mix_samples_cpu).to(device=device, dtype=torch.int32) + # Hardcode-sort: per-iter prev_tokens are CPU-sorted write-first. + # Kernel runs USE_PERM=False but the EO gate sees clustered modes. + # Output is scrambled (we don't permute x/B/C/dt to match) but + # timing is meaningful — isolates clustering benefit from the + # per-program perm-load overhead in --sort-slots. + src = mix_samples_sorted_cpu if (hardcode_sort and mix_samples_sorted_cpu is not None) else mix_samples_cpu + samples_gpu = torch.from_numpy(src).to(device=device, dtype=torch.int32) if sort_slots and perm_samples_cpu is not None: perm_samples_gpu = torch.from_numpy(perm_samples_cpu).to( @@ -1015,6 +1032,11 @@ def _run_incr( extra_kwargs["mode"] = mode if sort_slots: extra_kwargs["slot_perm"] = slot_perm_buf + # reverse_nowrite is meaningful in two ways: + # - with slot_perm: walk the perm tail-first + # - without slot_perm (hardcode-sort): walk pid_b + # itself tail-first via the REVERSE_PERM constexpr + if sort_slots or (hardcode_sort and reverse_nowrite): extra_kwargs["reverse_nowrite"] = reverse_nowrite if state_scales_work is not None: extra_kwargs["state_scales"] = state_scales_work @@ -1073,6 +1095,7 @@ def _run_incr( parts.append(f"MODE={mode}") parts.append(f"SORT={1 if sort_slots else 0}") parts.append(f"REVN={1 if reverse_nowrite else 0}") + parts.append(f"HSORT={1 if hardcode_sort else 0}") sweep_suffix = (" " + ",".join(parts)) if parts else "" sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") @@ -1206,6 +1229,7 @@ def _run_benchmark(args) -> None: modes_list = getattr(args, "modes_list", ["monolithic"]) sort_list = getattr(args, "sort_slots_list", [False]) rev_list = getattr(args, "reverse_nowrite_list", [False]) + hsort_list = getattr(args, "hardcode_sort_list", [False]) # Pre-load AL distribution for mix mode (if --mix-csv set). mix_al = None @@ -1230,6 +1254,7 @@ def _run_benchmark(args) -> None: # args.mix_iters since mix scenarios use mix_iters. mix_samples_cpu = None perm_samples_cpu = None # per-iter slot perm sorted write-first + mix_samples_sorted_cpu = None # per-iter prev_tokens, write-first if mix_al is not None: from checkpoint_mix_sim import sample_steady_state_pnat as _sample_pnat _max_window = getattr(args, "max_window", 0) or mtp_len @@ -1238,15 +1263,24 @@ def _run_benchmark(args) -> None: mix_al, T=mtp_len, window=_max_window, batch=batch, K=args.warmup + _max_iters, seed=args.mix_seed, ) - if any(sort_list): + if any(sort_list) or any(hsort_list): # write-first stable argsort: kind='stable' preserves # original-slot order within each mode group. write_mask = ( mix_samples_cpu + mtp_len > _max_window ).astype(np.int8) # 1 = write, 0 = nowrite - perm_samples_cpu = np.argsort( + perm_idx = np.argsort( -write_mask, kind="stable", axis=-1 ).astype(np.int32) + if any(sort_list): + perm_samples_cpu = perm_idx + if any(hsort_list): + # Apply the perm to the prev_tokens samples themselves. + # Result row i = mix_samples_cpu[i] reordered such + # that write-mode entries come first. + mix_samples_sorted_cpu = np.take_along_axis( + mix_samples_cpu, perm_idx, axis=-1 + ).astype(mix_samples_cpu.dtype) for state_dtype in state_dtypes: for act_dtype in act_dtypes: @@ -1289,24 +1323,44 @@ def _run_benchmark(args) -> None: effective_sort_list = ( sort_list if can_sort else [False] ) + effective_hsort_list = ( + hsort_list if can_sort else [False] + ) for sort_slots in effective_sort_list: - effective_rev_list = ( - rev_list if sort_slots else [False] - ) - for reverse_nowrite in effective_rev_list: - _bench_config( - args, batch, mtp_len, prev_ks, - state_dtype, act_dtype, - baseline_fn, sr_mode=sr_mode, - rectangle_for_nowrite=rect, - write_checkpoint=write_ckpt, - mode=mode, - mix_samples_cpu=mix_samples_cpu, - mix_label=mix_label, - sort_slots=sort_slots, - reverse_nowrite=reverse_nowrite, - perm_samples_cpu=perm_samples_cpu, + for hardcode_sort in effective_hsort_list: + # sort_slots and hardcode_sort + # are alternative experiments + # for the same idea — skip the + # combined cell to avoid double + # interpretation. + if sort_slots and hardcode_sort: + continue + # rev=1 is meaningful with EITHER + # sort_slots=1 (perm-based) or + # hardcode_sort=1 (raw pid_b + # subtraction in unsorted-perm + # path). rev=1 with both 0 is + # a no-op. + effective_rev_list = ( + rev_list if (sort_slots or hardcode_sort) else [False] ) + for reverse_nowrite in effective_rev_list: + _bench_config( + args, batch, mtp_len, + prev_ks, state_dtype, + act_dtype, baseline_fn, + sr_mode=sr_mode, + rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, + mode=mode, + mix_samples_cpu=mix_samples_cpu, + mix_label=mix_label, + sort_slots=sort_slots, + reverse_nowrite=reverse_nowrite, + perm_samples_cpu=perm_samples_cpu, + hardcode_sort=hardcode_sort, + mix_samples_sorted_cpu=mix_samples_sorted_cpu, + ) if args.profile: torch.cuda.cudart().cudaProfilerStop() @@ -1606,6 +1660,19 @@ def _parse_args() -> argparse.Namespace: "perm in reverse so both halves of the dl chain front-load real " "work. reverse=1 with sort=0 is skipped (no perm to reverse).", ) + parser.add_argument( + "--hardcode-sort", + type=str, + default="0", + help="Comma-separated 0/1. When 1, the per-iter prev_tokens " + "samples are pre-sorted write-first OFFLINE (CPU-side) before " + "the timed region — kernel runs unchanged (USE_PERM=False) but " + "the EO gate sees sorted PNAT so early-outs cluster naturally. " + "Zero per-program load cost vs --sort-slots; output is " + "scrambled (we don't permute x/B/C/dt) but timing is meaningful. " + "Used to isolate whether clustering helps independent of the " + "perm-load overhead in the sort-slots path.", + ) parser.add_argument( "--mix-iters", type=int, @@ -1691,6 +1758,14 @@ def _parse_args() -> argparse.Namespace: rev_list.append(v == "1") args.reverse_nowrite_list = rev_list + hsort_modes = [v.strip() for v in (args.hardcode_sort or "0").split(",") if v.strip()] + hsort_list = [] + for v in hsort_modes: + if v not in ("0", "1"): + parser.error(f"--hardcode-sort value must be 0 or 1, got {v!r}") + hsort_list.append(v == "1") + args.hardcode_sort_list = hsort_list + if args.write_modes is not None: wm = [v.strip() for v in args.write_modes.split(",") if v.strip()] write_list = [] From a5311d00fb5a9c65696681cb49569188b8d173fa Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 6 May 2026 22:30:23 -0700 Subject: [PATCH 22/89] test/bench: --max-window default 16 matches Nemotron-Super production MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The benchmark default was 0 → mtp_len (= 6 for typical MTP=6 sweeps), a placeholder for the degenerate every-step-checkpoint case. Real production uses max_window=16, and the WC=0 nowrite path silently skips configs with prev_k+T>max_window — meaning a default-flagged sweep at MTP=6 with prev-tokens-fracs=1.0 produced ZERO WC=0 measurements. Match Nemotron-Super-120B production by defaulting to 16; users can still pass 0 explicitly for the degenerate case. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/benchmark_replay_selective_state_update.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 62ce155086ef..4e30974f42ab 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -1239,11 +1239,10 @@ def _parse_args() -> argparse.Namespace: parser.add_argument( "--max-window", type=int, - default=0, - help="Cache T-axis capacity (max replay buffer length). 0 (default) " - "= use mtp_len, the placeholder/degenerate every-step-checkpoint " - "case. Set to e.g. 16 for real replay-style checkpointing on " - "Nemotron-3-Super-120B.", + default=16, + help="Cache T-axis capacity (max replay buffer length). Default 16 " + "matches Nemotron-3-Super-120B production. Pass 0 to fall back to " + "mtp_len (degenerate every-step-checkpoint case, mostly unused).", ) parser.add_argument( "--prev-tokens-int", From 33a5d437a2eea7f49133058ab90969becd1ad7f9 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 6 May 2026 23:02:43 -0700 Subject: [PATCH 23/89] dl_write_only debug mode + EO regression isolation runs Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../_torch/modules/mamba/checkpointing_state_update.py | 10 +++++++++- .../mamba/benchmark_replay_selective_state_update.py | 8 +++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index e9a45fcf55a5..1cac35348298 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -3096,7 +3096,7 @@ def checkpointing_state_update( # and nowrite halves. Strictly fewer kernel launches than # doublelaunch (3 vs 4) at the cost of dispatch precompute's wider # reg envelope. write_checkpoint ignored. - assert mode in ("monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl"), ( + assert mode in ("monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", "dl_write_only"), ( f"unknown mode {mode!r}; expected one of " "'monolithic', 'dynamic', 'doublelaunch', 'dlgrouped', or 'maindl'" ) @@ -3681,6 +3681,14 @@ def launch_dynamic_main(rectangle: bool, else: launch_replay_main(write_checkpoint=False, early_out=True, reverse_perm=reverse_nowrite) + elif mode == "dl_write_only": + # Debug-only: just the write half of doublelaunch. EARLY_OUT=True + # means nowrite slots still pay the EO-gate tax (PNAT load + branch), + # but no nowrite-side kernels run. Used to isolate "is the sort + # regression in the write-side kernels?". + launch_replay_precompute(write_checkpoint=True, early_out=True) + launch_replay_main(write_checkpoint=True, early_out=True, + launch_dependent_kernels=False) else: # mode == "doublelaunch" # Write half: always replay-style write. First main signals # PDL dependents so the second precompute can start its setup diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 02b6ab58616b..0cadb0a406d8 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -578,7 +578,8 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp effective_rect_list = rect_list for rect in effective_rect_list: is_dl_family = mode in ( - "doublelaunch", "dlgrouped", "maindl" + "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", ) # Match the timed-run skip: sort=1 only # makes sense when there's a mix scenario. @@ -1315,7 +1316,8 @@ def _run_benchmark(args) -> None: # dynamic also skip sort=1; reverse=1 # with sort=0 is a no-op (skip). is_dl_family = mode in ( - "doublelaunch", "dlgrouped", "maindl" + "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", ) can_sort = ( is_dl_family and mix_samples_cpu is not None @@ -1778,7 +1780,7 @@ def _parse_args() -> argparse.Namespace: args.write_modes_list = [args.write_checkpoint] modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] - valid_modes = {"monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl"} + valid_modes = {"monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", "dl_write_only"} for m in modes_raw: if m not in valid_modes: parser.error( From bbcc22e7518ac56e674a785eafb257684c9b1567 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 6 May 2026 23:53:41 -0700 Subject: [PATCH 24/89] TMA state load/store toggles for rect/replay main kernels (backlog #17) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three independent constexpr toggles, all default OFF, exposed via wrapper kwargs: _use_tma_state — rect main state load _use_tma_state_load_replay — replay main state load (WC=0/1) _use_tma_state_store_replay — replay main state store (WC=1 only) Implementation: host-side `TensorDescriptor.from_tensor` over a flat 2D view (state.view(-1, dstate)). In-kernel offset compute via existing strides — no new kernel args except a state_desc_ptr second arg on replay main (set to a dummy when both flags off). triton.set_allocator configured lazily once. Diagnostic sweep findings (with conv1d pipeline, full autotune, batches 1/16/64/512, int8+fp16, RN+SR): - rect main TMA load: not a reproducible win. fp16 hurts +1-12%; int8 wash to slight worse on this node (a prior int8-only sweep on a different node showed -1 to -4% wins at b=1, within noise). - replay main load (nowrite): int8 b>=64 wins -8 to -12% ⭐; fp16 b=64 -4%; small batch hurts. - replay main load (write): mixed; int8 b>=64 small wins. - replay main store (write): fp16 b=512 -5 to -11% ⭐; int8 SR wins. - replay main load+store: combines load and store wins. Why replay benefits more than rect: replay's `state += dot(coeff, dB)` work hides TMA's mbarrier sync latency; rect is too tight to hide it. Prior "+5x branches" anomaly resolved (host-side 2D descriptor adds +1024 SASS branches = 0.16% of total). Source likely: in-kernel descriptor + 4D leading-1 block_shape in prior attempt; we avoid both. No production runtime dispatch wired — would need per-(dtype, batch) selection. Infrastructure kept for persistent-kernels integration (other agent's work) where compute-to-load ratio is different and TMA may apply more broadly. Tests: 144 rect + 398 replay tests pass with all flags default-off. TMA paths smoke-tested correct (out + stored state byte-equal to non-TMA at int8 nowrite + write). Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 183 +++++++++++++++--- ...benchmark_replay_selective_state_update.py | 26 +++ 2 files changed, 182 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 5a6da1cea8db..1eeb75d1683e 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -29,6 +29,27 @@ from .softplus import softplus +# Lazy global allocator for Triton TMA tensor descriptors. Required by any +# host- or device-built tensor_descriptor; without it Triton raises at first +# launch. See TMA backlog item #17 / scratch experiment notes. +_TMA_ALLOCATOR_SET = False + + +def _ensure_tma_allocator() -> None: + global _TMA_ALLOCATOR_SET + if _TMA_ALLOCATOR_SET: + return + + def _alloc_fn(size, alignment, stream): + # Triton expects an int8 buffer of `size` bytes; alignment is enforced + # by the allocator returning a buffer satisfying it (PyTorch's + # cudaMalloc-backed tensors are 256B-aligned, so we're fine). + return torch.empty(size, device="cuda", dtype=torch.int8) + + triton.set_allocator(_alloc_fn) + _TMA_ALLOCATOR_SET = True + + @triton.jit def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: """Stochastic rounding: fp32 pair → fp16x2 using Philox random bits. @@ -739,6 +760,10 @@ def _rectangle_precompute_kernel( def _checkpointing_main_kernel( # Pointers state_ptr, + # state_desc_ptr: TMA tensor_descriptor over state's flat 2D view, or + # the same `state_ptr` tensor when neither USE_TMA_STATE_LOAD nor + # USE_TMA_STATE_STORE is enabled (kernel ignores it via constexpr). + state_desc_ptr, # Per-(cache, head, dim) decode scale, fp32, only consulted when QUANT_MAX>0. # Layout (cache, nheads, dim) — broadcast over dstate at load/store. state_scales_ptr, @@ -854,6 +879,12 @@ def _checkpointing_main_kernel( # The rectangle non-checkpoint path is implemented in # _rectangle_main_kernel (separate kernel pair, picked by # the wrapper via rectangle_for_nowrite=True). + # TMA toggles (independent). When True, state_ptr is a host-built + # tensor_descriptor over a flat (cache_size*nheads*dim, dstate) view of + # state; otherwise it's a regular state pointer. The two flags share + # state_ptr — caller must pass a descriptor if either is True. + USE_TMA_STATE_LOAD: tl.constexpr = False, + USE_TMA_STATE_STORE: tl.constexpr = False, ): # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized # state dtype (int8 / int16 / float8e4nv) and only those. Cheap @@ -902,13 +933,29 @@ def _checkpointing_main_kernel( n_mask = offs_n < dstate t_mask = offs_t < T - # Load state - state_ptr += cache_batch_idx * stride_state_batch + pid_h * stride_state_head + # Load state. + # When TMA load is enabled, state_ptr is a host-built tensor_descriptor + # over a flat (cache*nheads*dim, dstate) view of state; otherwise it's + # a raw pointer. Same descriptor is reused for the WC=True store path. + state_mask = m_mask[:, None] & n_mask[None, :] + if USE_TMA_STATE_LOAD or USE_TMA_STATE_STORE: + # offs_y = (c * stride_state_batch + h * stride_state_head + m * stride_state_dim) / stride_state_dim + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + # state_ptrs is the raw-pointer view (used for !TMA load and !TMA store). + # We compute it unconditionally; it costs only int math and Triton may DCE + # if neither path consumes it. + state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head state_ptrs = ( - state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate ) - state_mask = m_mask[:, None] & n_mask[None, :] - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + if USE_TMA_STATE_LOAD: + state = state_desc_ptr.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) # Dequantize on load (per-(head, dim) decode scale, broadcast over dstate). # Only consulted when QUANT_MAX>0 — non-quantized paths skip entirely. if QUANT_MAX > 0.0: @@ -1051,11 +1098,11 @@ def _checkpointing_main_kernel( if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): # fp8_e4m3fn + SR — PTX cvt.rs.satfinite.e4m3x4.f32. Output # is final fp8 (saturate included); store directly. - tl.store( - state_ptrs, - _stochastic_round_fp8x4_e4m3(state_q, rand), - mask=state_mask, - ) + _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) + if USE_TMA_STATE_STORE: + state_desc_ptr.store([offs_y, 0], _state_q_fp8sr) + else: + tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) else: if USE_RS_ROUNDING: # int8 / int16 + SR — uniform-noise + floor. @@ -1083,11 +1130,11 @@ def _checkpointing_main_kernel( # fp8 RN reaches here without prior round() — .to(float8e4nv) # does native RN at the fp8 grid. state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) - tl.store( - state_ptrs, - state_q.to(state_ptrs.dtype.element_ty), - mask=state_mask, - ) + _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) + if USE_TMA_STATE_STORE: + state_desc_ptr.store([offs_y, 0], _state_q_cast) + else: + tl.store(state_ptrs, _state_q_cast, mask=state_mask) elif USE_RS_ROUNDING: # Non-quantized + SR: only fp16 (bf16 has no PTX SR cast; fp32 # doesn't need rounding). @@ -1095,10 +1142,18 @@ def _checkpointing_main_kernel( state_ptrs.dtype.element_ty == tl.float16, "Non-quantized SR only supports fp16 state.", ) - tl.store(state_ptrs, _stochastic_round_fp16x2(state, rand), mask=state_mask) + _state_sr = _stochastic_round_fp16x2(state, rand) + if USE_TMA_STATE_STORE: + state_desc_ptr.store([offs_y, 0], _state_sr) + else: + tl.store(state_ptrs, _state_sr, mask=state_mask) else: # Non-quantized + RN: fp16 / bf16 / fp32 native cast. - tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=state_mask) + _state_cast = state.to(state_ptrs.dtype.element_ty) + if USE_TMA_STATE_STORE: + state_desc_ptr.store([offs_y, 0], _state_cast) + else: + tl.store(state_ptrs, _state_cast, mask=state_mask) # Phase 2: Output using precomputed CB_scaled and decay_vec x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head @@ -1281,6 +1336,7 @@ def _rectangle_main_kernel( BLOCK_SIZE_K: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, QUANT_MAX: tl.constexpr, + USE_TMA_STATE: tl.constexpr = False, ): pid_m = tl.program_id(axis=0) pid_b = tl.program_id(axis=1) @@ -1317,14 +1373,40 @@ def _rectangle_main_kernel( is_new_k = (k_new_idx >= 0) & (k_new_idx < T) safe_k_new = tl.where(is_new_k, k_new_idx, 0) - # Load state (with dequant on load if quantized). Read-only — no HBM - # write on the nowrite path. - state_ptr += cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - ) - state_mask = m_mask[:, None] & n_mask[None, :] - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + # Load state. Read-only — no HBM write on the nowrite path. + # Quant scale hoist (backlog #16): for QUANT_MAX > 0 paths, defer the + # `* decode_scale` to AFTER the C @ state dot — applied to the (T, M) + # dot output instead of broadcast-multiplied into the (M, dstate) state + # tile. Algebra-equivalent (decode_scale is per-M, commutes with the + # matmul over dstate). Saves M·dstate fp32 muls (replaced by T·M), + # but the bigger potential win is shorter register lifetime for state + # (kept as native int8/int16/fp8 until just before the dot, where Triton + # casts to bf16 — vs current fp32 tile across the whole kernel). Only + # applies in rectangle main (no `state += dot` here). + if USE_TMA_STATE: + # TMA descriptor (host-built, passed via state_ptr as a + # tensor_descriptor) over the flat 2D view of state: + # shape=[cache_size * nheads * dim, dstate], strides=[dstate, 1]. + # Convert (cache, head, m) → flat row index using existing strides: + # rows-per-cache-slot = stride_state_batch / stride_state_dim + # rows-per-head = stride_state_head / stride_state_dim = dim (constexpr) + # rows-per-m = 1 + # Diagnostic: prior in-kernel descriptor attempts emitted + # ttng.tensormap_create setup (divergent shared-mem write) which + # blows up branch count; host-built descriptors avoid that. + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state = state_ptr.load([offs_y, 0]) + else: + state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0) if QUANT_MAX > 0.0: state_scales_base = ( state_scales_ptr @@ -1335,7 +1417,9 @@ def _rectangle_main_kernel( state_scales_base + offs_m * stride_state_scales_dim, mask=m_mask, other=1.0, ).to(tl.float32) - state = state * decode_scale[:, None] + # state stays in native quant dtype — cast happens inside the dot below. + else: + state = state.to(tl.float32) # Group / pointer offset setup group_idx = pid_h // nheads_ngroups_ratio @@ -1419,10 +1503,16 @@ def _rectangle_main_kernel( # state_out: state_prev contribution to output, with decay folded post-matmul. # No state_prev_decayed (M, dstate) materialization — state is consumed # directly by the matmul, then decay_vec_full multiplies the (T, M) result. + # For QUANT_MAX > 0 (#16 hoist): decode_scale also applies post-matmul + # at (T, M) granularity instead of pre-multiplied into the (M, dstate) + # state tile. Triton's tl.dot(a.to(bf16), b.to(bf16)) handles the + # int8/fp8 → bf16 cast inside the dot's input prep. state_out = ( tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec_full[:, None] ) + if QUANT_MAX > 0.0: + state_out = state_out * decode_scale[None, :] # token_out: combined old + new tokens contribution via the rectangle. token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) @@ -1490,6 +1580,9 @@ def checkpointing_state_update( _heads_per_block: int | None = None, _maxnreg: int | None = None, _num_ctas: int | None = None, + _use_tma_state: bool = False, + _use_tma_state_load_replay: bool = False, + _use_tma_state_store_replay: bool = False, ): """ Replay SSM state update with precomputed CB and tl.dot fast-forward. @@ -1844,6 +1937,38 @@ def checkpointing_state_update( state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 state_scales_strides = (0, 0, 0) + # TMA descriptor for state (rectangle path only, currently). Built + # host-side over a flat 2D view of state — shape (cache*nheads*dim, + # dstate), inner stride 1. Triton's set_allocator() must be called + # before any descriptor-using kernel launches. + if _use_tma_state and use_rectangle: + from triton.tools.tensor_descriptor import TensorDescriptor + _ensure_tma_allocator() + assert state.is_contiguous(), "TMA state load requires contiguous state" + assert state.stride(-1) == 1, "TMA state load requires inner stride 1" + state_for_kernel = TensorDescriptor.from_tensor( + state.view(-1, state.shape[-1]), + block_shape=[BLOCK_SIZE_M, triton.next_power_of_2(dstate)], + ) + else: + state_for_kernel = state + + # Replay-main TMA descriptor: built once if either load or store TMA + # is enabled. Same flat 2D view as rectangle. Passed as the + # state_desc_ptr arg; raw `state` tensor is still passed as state_ptr + # for the !TMA paths (kernel branches via constexpr). + if (_use_tma_state_load_replay or _use_tma_state_store_replay) and not use_rectangle: + from triton.tools.tensor_descriptor import TensorDescriptor + _ensure_tma_allocator() + assert state.is_contiguous(), "TMA state load/store requires contiguous state" + assert state.stride(-1) == 1, "TMA state load/store requires inner stride 1" + state_desc_replay = TensorDescriptor.from_tensor( + state.view(-1, state.shape[-1]), + block_shape=[BLOCK_SIZE_M, triton.next_power_of_2(dstate)], + ) + else: + state_desc_replay = state # dummy; kernel ignores via constexpr + if use_rectangle: _rectangle_precompute_kernel[(batch, nheads // heads_per_block)]( dt, @@ -1919,7 +2044,7 @@ def grid(META): return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) _rectangle_main_kernel[grid]( - state, + state_for_kernel, state_scales_arg, old_x, prev_num_accepted_tokens, @@ -1986,6 +2111,7 @@ def grid(META): BLOCK_SIZE_M, LAUNCH_WITH_PDL=use_internal_pdl, QUANT_MAX=quant_max, + USE_TMA_STATE=bool(_use_tma_state), num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), @@ -2070,6 +2196,7 @@ def grid(META): _checkpointing_main_kernel[grid]( state, + state_desc_replay, state_scales_arg, old_x, old_B, @@ -2158,6 +2285,8 @@ def grid(META): PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, QUANT_MAX=quant_max, WRITE_CHECKPOINT=write_checkpoint, + USE_TMA_STATE_LOAD=bool(_use_tma_state_load_replay), + USE_TMA_STATE_STORE=bool(_use_tma_state_store_replay and write_checkpoint), num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 4e30974f42ab..8c6eee963c2a 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -891,6 +891,12 @@ def _run_incr( extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite if state_scales_work is not None: extra_kwargs["state_scales"] = state_scales_work + if getattr(args, "use_tma_state", False): + extra_kwargs["_use_tma_state"] = True + if getattr(args, "use_tma_state_load_replay", False): + extra_kwargs["_use_tma_state_load_replay"] = True + if getattr(args, "use_tma_state_store_replay", False): + extra_kwargs["_use_tma_state_store_replay"] = True variant_fn( state_work, old_x_work, @@ -1322,6 +1328,26 @@ def _parse_args() -> argparse.Namespace: "write path always uses replay-style). Only applies to the " "checkpointing variant.", ) + parser.add_argument( + "--use-tma-state", + action="store_true", + help="Use TMA (host-built tensor descriptor) for the state load in " + "_rectangle_main_kernel. Only applies to the rectangle nowrite path " + "of the checkpointing variant; ignored otherwise.", + ) + parser.add_argument( + "--use-tma-state-load-replay", + action="store_true", + help="Use TMA for state LOAD in _checkpointing_main_kernel (replay " + "main, both WC=0 and WC=1 paths). Independent from rect TMA.", + ) + parser.add_argument( + "--use-tma-state-store-replay", + action="store_true", + help="Use TMA for state STORE in _checkpointing_main_kernel (replay " + "main, WC=1 path only — no-op for WC=0). Independent from rect TMA " + "and from --use-tma-state-load-replay.", + ) parser.add_argument( "--philox-rounding", action="store_true", From 4bf337a8be52e60087868a50efa7b2af4fe77fa6 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 6 May 2026 23:33:39 -0700 Subject: [PATCH 25/89] persistent_main: maindl pattern, n_writes/batch_total kernel args, cta_per_sm + flatten + warp_specialize knobs, bench plumbing, test (3 scenarios pass) Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 761 +++++++++++++++++- ...benchmark_replay_selective_state_update.py | 134 ++- .../mamba/test_checkpointing_state_update.py | 171 ++++ 3 files changed, 1056 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 1cac35348298..1b6dbbba662a 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -2957,6 +2957,635 @@ def _dynamic_main_kernel( } +# ============================================================================ +# Persistent main kernel — 1D grid, persistent CTA loop with tl.range +# ============================================================================ +# +# Design (see ~/dev/scripts/mamba_replay/kernel_microbenchmarks/PERSISTENT_KERNELS.md +# for the full strawman): +# +# * Outer 1D grid of `NUM_PERSISTENT` CTAs (start at NUM_SMS, sweep upward). +# * Inside the kernel, a `tl.range(pid, total_work, NUM_PERSISTENT, flatten=True, +# num_stages=NUM_STAGES)` loop iterates over (slot, M_tile, head) work units. +# * Hard-sort PNAT host-side and pass `n_writes` as a runtime int32 scalar: +# the launcher invokes the kernel twice — once with slot_offset=0, +# n_slots=n_writes, WRITE_CHECKPOINT=True, and once with +# slot_offset=n_writes, n_slots=B-n_writes, WRITE_CHECKPOINT=False. +# * `_persistent_main_impl` is a copy of `_replay_main_impl`'s body with the +# program_id reads replaced by parameters and the slot_perm logic moved into +# the persistent loop wrapper. No code shared with the existing kernels; +# easy to delete if the experiment is abandoned. +# +# Notes: +# * `flatten=True` is canonical for Triton 3.6 persistent kernels (matches the +# upstream `_p_matmul_ogs.py` and tutorial 09). Combined with `num_stages=2` +# it pipelines the loop body — but watch open issue triton-lang/triton#8259 +# which reports this combo can corrupt stores in non-dot loops. First run +# correctness check is critical. +# * Warp specialization (`warp_specialize=True`) is NOT enabled — Triton 3.6 +# only supports it for simple matmul loops and our scan won't pattern-match. +# * No 2CTA cluster mode — that's dot-only per the kernel-tileir-optimization +# skill classification. + + +@triton.jit() +def _persistent_main_impl( + # Per-work-unit indices (computed by the persistent wrapper). + # `pid_b` is the post-perm slot index (caller has already applied any + # slot permutation and slot_offset). + pid_m, + pid_b, + pid_h, + # Pointers + state_ptr, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, +): + # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized + # state dtype (int8 / int16 / float8e4nv) and only those. + tl.static_assert( + (QUANT_MAX > 0.0) + == ( + (state_ptr.dtype.element_ty == tl.int8) + or (state_ptr.dtype.element_ty == tl.int16) + or (state_ptr.dtype.element_ty == tl.float8e4nv) + ), + "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", + ) + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if WRITE_CHECKPOINT: + write_buf = 1 - active_buf # noqa: F841 + write_offset = 0 + else: + write_buf = active_buf # noqa: F841 + write_offset = prev_num_accepted_tokens + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # Load state + state_ptr += cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, + other=1.0, + ).to(tl.float32) + state = state * decode_scale[:, None] + + # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) + group_idx = pid_h // nheads_ngroups_ratio + + old_window_mask = offs_window < prev_num_accepted_tokens + + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + active_buf * stride_old_dt_dbuf + + pid_h * stride_old_dt_head + ) + old_dt_all = tl.load( + old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 + ).to(tl.float32) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + active_buf * stride_old_dA_cumsum_dbuf + + pid_h * stride_old_dA_cumsum_head + ) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, + mask=old_window_mask, other=0.0, + ).to(tl.float32) + + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( + tl.float32 + ) + + coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all + + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + old_x_all = tl.load( + old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + mask=old_window_mask[:, None] & m_mask[None, :], + other=0.0, + ) + + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + active_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_all = tl.load( + old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + mask=old_window_mask[:, None] & n_mask[None, :], + other=0.0, + ).to(tl.float32) + + dB_scaled = coeff[:, None] * old_B_all + + total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) + state *= total_decay + + state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) + + if WRITE_CHECKPOINT: + if USE_RS_ROUNDING: + rand_seed = tl.load(rand_seed_ptr) + base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head + offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // 4) + rand_offsets_q = ( + base_rand + + offs_m[:, None] * stride_state_dim + + offs_n_q[None, :] * (stride_state_dstate * 4) + ) + if PHILOX_ROUNDS > 0: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + else: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) + r01 = tl.join(r0, r1) + r23 = tl.join(r2, r3) + r0123 = tl.join(r01, r23) + rand = tl.reshape(r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + + if QUANT_MAX > 0.0: + amax = tl.max(tl.abs(state), axis=1) + encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) + decode_scale = 1.0 / encode_scale + state_scales_ptrs = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + + offs_m * stride_state_scales_dim + ) + tl.store(state_scales_ptrs, decode_scale, mask=m_mask) + state_q = state * encode_scale[:, None] + if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): + tl.store( + state_ptrs, + _stochastic_round_fp8x4_e4m3(state_q, rand), + mask=state_mask, + ) + else: + if USE_RS_ROUNDING: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized SR fall-through expects int8 or int16; " + "fp8 SR is handled by the prior branch.", + ) + rand01 = (rand & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) + state_q = tl.extra.cuda.libdevice.floor(state_q + rand01) + elif state_ptrs.dtype.element_ty != tl.float8e4nv: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized RN with explicit round() expects int8 or int16.", + ) + state_q = tl.extra.cuda.libdevice.round(state_q) + state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) + tl.store( + state_ptrs, + state_q.to(state_ptrs.dtype.element_ty), + mask=state_mask, + ) + elif USE_RS_ROUNDING: + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "Non-quantized SR only supports fp16 state.", + ) + tl.store(state_ptrs, _stochastic_round_fp16x2(state, rand), mask=state_mask) + else: + tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=state_mask) + + # Phase 2: Output + x_ptr_local = x_ptr + pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr_local = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr_local = z_ptr + pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr_local = out_ptr + pid_b * stride_out_batch + pid_h * stride_out_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + C_all = tl.load( + C_ptr_local + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + x_all = tl.load( + x_ptr_local + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_base + + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_all, + mask=t_mask[:, None] & m_mask[None, :], + ) + x_all = x_all.to(tl.float32) + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + tl.float32 + ) + + init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) + out_all = init_out + cb_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + for t in range(T): + z_t = tl.load( + z_ptr_local + t * stride_z_T + offs_m * stride_z_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + out_t = tl.sum(tl.where((offs_t == t)[:, None], out_all, 0.0), axis=0) + out_t = out_t * z_t * tl.sigmoid(z_t) + tl.store(out_ptr_local + t * stride_out_T + offs_m * stride_out_dim, out_t, mask=m_mask) + else: + out_all_ptrs = ( + out_ptr_local + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + ) + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# Persistent main kernel: 1D grid, persistent CTA loop. +# Heuristics mirror those of `_checkpointing_main_kernel`. +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"NUM_PID_M_BLOCKS": lambda args: triton.cdiv(args["dim"], args["BLOCK_SIZE_M"])} +) +@triton.jit() +def _persistent_main_kernel( + # Pointers + state_ptr, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + rand_seed_ptr, + pad_slot_id, + # Persistent-loop work-distribution scalars. Caller pre-sorts the batch + # write-first; the kernel uses (n_writes, batch_total, WRITE_CHECKPOINT) + # to derive its own slot range. Write half processes [0, n_writes), + # nowrite half processes [n_writes, batch_total). + n_writes, # int32: count of write-mode slots in the pre-sorted batch + batch_total, # int32: total slot count + nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + USE_PERM: tl.constexpr, + NUM_PERSISTENT: tl.constexpr, + NUM_LOOP_STAGES: tl.constexpr, + NUM_PID_M_BLOCKS: tl.constexpr, + FLATTEN: tl.constexpr, + WARP_SPECIALIZE: tl.constexpr, +): + # PDL signal: fire once at kernel entry (not per work unit). + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + # Derive this kernel's slot range from (n_writes, batch_total, WRITE_CHECKPOINT). + # Hard-sort contract: caller has pre-sorted so [0, n_writes) are write-mode + # and [n_writes, batch_total) are nowrite-mode. Each launch knows which + # half it owns from its WRITE_CHECKPOINT constexpr. + if WRITE_CHECKPOINT: + slot_lo = 0 + slot_hi = n_writes + else: + slot_lo = n_writes + slot_hi = batch_total + n_slots_local = slot_hi - slot_lo + + pid = tl.program_id(axis=0) + total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads + + # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) + # with pid_m varying fastest (M-tile cache locality on state load), then + # slot, then head — mirrors the existing 3D grid's axis ordering + # (axis=0 fastest = pid_m). + for tile_id in tl.range( + pid, total_work, NUM_PERSISTENT, + flatten=FLATTEN, num_stages=NUM_LOOP_STAGES, warp_specialize=WARP_SPECIALIZE, + ): + pid_m = tile_id % NUM_PID_M_BLOCKS + pid_b_local = (tile_id // NUM_PID_M_BLOCKS) % n_slots_local + pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) + # Translate local slot index → global slot index. When USE_PERM is + # set, the caller-provided slot_perm gives the original slot index + # for the post-sort position. + pid_b_grid = pid_b_local + slot_lo + if USE_PERM: + pid_b = tl.load(slot_perm_ptr + pid_b_grid) + else: + pid_b = pid_b_grid + + _persistent_main_impl( + pid_m, pid_b, pid_h, + state_ptr, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + WRITE_CHECKPOINT, + ) + + +# ============================================================================ +# Python wrapper +# ============================================================================ + + def checkpointing_state_update( state: torch.Tensor, old_x: torch.Tensor, @@ -3004,6 +3633,26 @@ def checkpointing_state_update( _heads_per_block: int | None = None, _maxnreg: int | None = None, _num_ctas: int | None = None, + # Persistent-mode bench kwargs (only consulted when mode == "persistent_main"): + # _n_writes : int — count of write-mode slots in the (pre-sorted) batch. + # Required when mode == "persistent_main"; the persistent kernel uses + # it as a runtime int32 to compute total_work for write/nowrite halves. + # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally + # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. Default = 1. + # _num_loop_stages : int — `num_stages` arg on the inner `tl.range(...)` + # persistent loop. Note: this is loop-level, NOT the kernel-arg + # `num_stages` (which only pipelines dot-feeding loads). Default 2. + # _flatten : bool — `flatten` arg on `tl.range(...)`. Default True + # (the canonical Triton 3.6 persistent idiom). + # _warp_specialize : bool — `warp_specialize` arg on `tl.range(...)`. + # Default False. Triton 3.6 only supports it on simple matmul loops; + # our scan loop probably won't pattern-match — but exposed as a knob + # for sweep experiments. Requires num_warps >= 4 if True. + _n_writes: int | None = None, + _cta_per_sm: int | None = None, + _num_loop_stages: int | None = None, + _flatten: bool | None = None, + _warp_specialize: bool | None = None, ): """ Replay SSM state update with precomputed CB and tl.dot fast-forward. @@ -3096,9 +3745,13 @@ def checkpointing_state_update( # and nowrite halves. Strictly fewer kernel launches than # doublelaunch (3 vs 4) at the cost of dispatch precompute's wider # reg envelope. write_checkpoint ignored. - assert mode in ("monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", "dl_write_only"), ( + assert mode in ( + "monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", "persistent_main", + ), ( f"unknown mode {mode!r}; expected one of " - "'monolithic', 'dynamic', 'doublelaunch', 'dlgrouped', or 'maindl'" + "'monolithic', 'dynamic', 'doublelaunch', 'dlgrouped', 'maindl', " + "'dl_write_only', or 'persistent_main'" ) use_rectangle = rectangle_for_nowrite and not write_checkpoint @@ -3631,6 +4284,79 @@ def launch_dynamic_main(rectangle: bool, launch_pdl=use_internal_pdl, ) + # ---- launch_persistent_main ------------------------------------------ + # Persistent-CTA main kernel. Single launch covers `n_slots` slots + # starting at `slot_offset`. Caller invokes twice: once for the write + # half (slot_offset=0, n_slots=n_writes, write_checkpoint=True) and + # once for the nowrite half (slot_offset=n_writes, + # n_slots=batch-n_writes, write_checkpoint=False). Hard-sort + # contract: caller has pre-sorted slots so [0, n_writes) are writes + # and [n_writes, batch) are nowrites. + + # Resolve persistent-mode bench knobs. Defaults: cta_per_sm = 1 + # (one CTA per SM, matches upstream `_p_matmul_ogs.py`); num_loop_stages + # = 2 (matches in-tree `swiglu` precedent for non-dot persistent loops); + # flatten = True (canonical Triton 3.6 idiom); warp_specialize = False. + _num_sms = torch.cuda.get_device_properties(device).multi_processor_count + cta_per_sm_arg = _cta_per_sm if _cta_per_sm else 1 + num_persistent_arg = cta_per_sm_arg * _num_sms + num_loop_stages_arg = _num_loop_stages if _num_loop_stages else 2 + flatten_arg = True if _flatten is None else bool(_flatten) + warp_specialize_arg = False if _warp_specialize is None else bool(_warp_specialize) + + def launch_persistent_main(write_checkpoint: bool, n_writes: int, + launch_dependent_kernels: bool = False): + # This launch's effective slot count: n_writes for write half, + # batch - n_writes for nowrite half. Skip the launch entirely if + # it has no work — empty halves are normal at the boundaries. + n_slots_for_kernel = n_writes if write_checkpoint else (batch - n_writes) + if n_slots_for_kernel <= 0: + return + grid = (num_persistent_arg,) + _persistent_main_kernel[grid]( + state, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + n_writes, batch, nheads, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=write_checkpoint, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + NUM_PERSISTENT=num_persistent_arg, + NUM_LOOP_STAGES=num_loop_stages_arg, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + # ---- Mode dispatch ---------------------------------------------------- with torch.cuda.device(device.index): if mode == "monolithic": @@ -3689,6 +4415,37 @@ def launch_dynamic_main(rectangle: bool, launch_replay_precompute(write_checkpoint=True, early_out=True) launch_replay_main(write_checkpoint=True, early_out=True, launch_dependent_kernels=False) + elif mode == "persistent_main": + # Experimental: persistent-CTA main kernel. Reuses maindl's + # precompute structure (one shared dynamic_precompute that + # dispatches per-slot at runtime based on PNAT) followed by two + # persistent_main launches (write half + nowrite half). + # + # Hard-sort contract: caller has pre-sorted slots host-side so + # PNAT is monotone (writes first). Pass the perm via + # slot_perm + USE_PERM and the count via _n_writes. + # + # The persistent main kernel takes (n_writes, batch_total, + # WRITE_CHECKPOINT) and computes its own slot range. Each + # launch is skipped at the Python level if its half is empty + # (n_writes=0 or n_writes=batch), so degenerate batches + # incur zero kernel launch overhead on the empty side. + assert _n_writes is not None, ( + "mode='persistent_main' requires _n_writes (count of write " + "slots in the pre-sorted batch). Provide via the bench's " + "--hardcode-sort path." + ) + assert 0 <= _n_writes <= batch, ( + f"_n_writes={_n_writes} must be in [0, batch={batch}]" + ) + n_writes_v = _n_writes + launch_dynamic_precompute(rectangle=False) + launch_persistent_main(write_checkpoint=True, + n_writes=n_writes_v, + launch_dependent_kernels=True) + launch_persistent_main(write_checkpoint=False, + n_writes=n_writes_v, + launch_dependent_kernels=False) else: # mode == "doublelaunch" # Write half: always replay-style write. First main signals # PDL dependents so the second precompute can start its setup diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 0cadb0a406d8..216e0a3a3908 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -577,9 +577,18 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp else: effective_rect_list = rect_list for rect in effective_rect_list: + # Note: "persistent_main" is included in + # the dl-family for sort/hsort sweep + # eligibility — it consumes the same + # slot_perm and benefits from the same + # write-first clustering. It additionally + # requires _n_writes (count of write + # slots) which the bench computes from + # the pure-scenario PNAT (mix scenarios + # not yet supported for persistent_main). is_dl_family = mode in ( "doublelaunch", "dlgrouped", "maindl", - "dl_write_only", + "dl_write_only", "persistent_main", ) # Match the timed-run skip: sort=1 only # makes sense when there's a mix scenario. @@ -915,6 +924,11 @@ def _parse_sweep(val): heads_per_block_values = _parse_sweep(args.heads_per_block) maxnreg_values = _parse_sweep(args.maxnreg) num_ctas_values = _parse_sweep(args.num_ctas) + # Persistent-only sweep dims; ignored when the cell's mode != persistent_main. + cta_per_sm_values = _parse_sweep(args.cta_per_sm) + num_loop_stages_values = _parse_sweep(args.num_loop_stages) + flatten_values = _parse_sweep(args.flatten) + warp_specialize_values = _parse_sweep(args.warp_specialize) # --- Replay kernel --- # Cache T-axis capacity (for prev_k validity check on the nowrite path). @@ -941,11 +955,15 @@ def _parse_sweep(val): "iters": None, # use args.iters }) # Mix scenario: skip on monolithic (mono on mixed PNAT corrupts the - # wrong-mode slots). prev_tokens varies per iter; the per-iter copy - # is captured inside the CUDA graph from a pre-baked GPU samples - # tensor (warmup samples distinct from timed-iter samples so any - # nsys-included warmup leaks aren't biased). - if mix_samples_cpu is not None and mode != "monolithic": + # wrong-mode slots). Also skip for persistent_main — the kernel takes + # `n_writes` as a runtime int32 scalar, but in mix mode `n_writes` + # varies per iter. Plumbing per-iter n_writes through the CUDA graph + # would require a 1-elem int32 GPU tensor scanned host-side or via a + # tiny scan kernel; out of scope for the skeleton. prev_tokens varies + # per iter; the per-iter copy is captured inside the CUDA graph from a + # pre-baked GPU samples tensor (warmup samples distinct from + # timed-iter samples so any nsys-included warmup leaks aren't biased). + if mix_samples_cpu is not None and mode not in ("monolithic", "persistent_main"): device = state_work.device # Hardcode-sort: per-iter prev_tokens are CPU-sorted write-first. # Kernel runs USE_PERM=False but the EO gate sees clustered modes. @@ -995,6 +1013,10 @@ def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): heads_per_block, maxnreg, num_ctas, + cta_per_sm, + num_loop_stages, + flatten, + warp_specialize, ) in itertools.product( block_size_m_values, num_warps_values, @@ -1004,6 +1026,10 @@ def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): heads_per_block_values, maxnreg_values, num_ctas_values, + cta_per_sm_values, + num_loop_stages_values, + flatten_values, + warp_specialize_values, ): def _run_incr( @@ -1015,6 +1041,10 @@ def _run_incr( heads_per_block=heads_per_block, maxnreg=maxnreg, num_ctas=num_ctas, + cta_per_sm=cta_per_sm, + num_loop_stages=num_loop_stages, + flatten=flatten, + warp_specialize=warp_specialize, ): if with_conv1d: x_call, B_call, C_call = _conv1d_split( @@ -1041,6 +1071,33 @@ def _run_incr( extra_kwargs["reverse_nowrite"] = reverse_nowrite if state_scales_work is not None: extra_kwargs["state_scales"] = state_scales_work + # persistent_main needs n_writes (count of write-mode + # slots in the pre-sorted batch) as a host-side int. + # Pure scenarios: every slot has the same PNAT, so + # n_writes is either 0 (all nowrite) or batch (all + # write) depending on whether PNAT+T overflows the + # window. Mix scenarios are skipped earlier. + if mode == "persistent_main": + scn_fill = scn["fill"] + assert scn_fill is not None, ( + "persistent_main does not yet support mix " + "scenarios (per-iter n_writes plumbing not " + "implemented)." + ) + is_write_scenario = (scn_fill + mtp_len) > max_window + extra_kwargs["_n_writes"] = batch if is_write_scenario else 0 + # Per-cell sweep values for persistent-only knobs. + # _parse_sweep returns [None] when the user didn't + # pass the flag, in which case we leave the wrapper's + # defaults in place. + if cta_per_sm is not None: + extra_kwargs["_cta_per_sm"] = cta_per_sm + if num_loop_stages is not None: + extra_kwargs["_num_loop_stages"] = num_loop_stages + if flatten is not None: + extra_kwargs["_flatten"] = bool(flatten) + if warp_specialize is not None: + extra_kwargs["_warp_specialize"] = bool(warp_specialize) variant_fn( state_work, old_x_work, @@ -1090,6 +1147,17 @@ def _run_incr( parts.append(f"R={maxnreg}") if num_ctas is not None: parts.append(f"CT={num_ctas}") + # Persistent-only knobs (only meaningful when MODE=persistent_main; + # printed unconditionally so output rows are uniformly comparable + # across modes when the user passed these sweeps). + if cta_per_sm is not None: + parts.append(f"CPS={cta_per_sm}") + if num_loop_stages is not None: + parts.append(f"LS={num_loop_stages}") + if flatten is not None: + parts.append(f"FL={flatten}") + if warp_specialize is not None: + parts.append(f"WS={warp_specialize}") parts.append(f"SR={1 if use_philox else 0}") parts.append(f"RECT={1 if rectangle_for_nowrite else 0}") parts.append(f"WC={1 if write_checkpoint else 0}") @@ -1315,9 +1383,18 @@ def _run_benchmark(args) -> None: # when no mix is configured; mono / # dynamic also skip sort=1; reverse=1 # with sort=0 is a no-op (skip). + # Note: "persistent_main" is included in + # the dl-family for sort/hsort sweep + # eligibility — it consumes the same + # slot_perm and benefits from the same + # write-first clustering. It additionally + # requires _n_writes (count of write + # slots) which the bench computes from + # the pure-scenario PNAT (mix scenarios + # not yet supported for persistent_main). is_dl_family = mode in ( "doublelaunch", "dlgrouped", "maindl", - "dl_write_only", + "dl_write_only", "persistent_main", ) can_sort = ( is_dl_family and mix_samples_cpu is not None @@ -1580,6 +1657,44 @@ def _parse_args() -> argparse.Namespace: default=None, help="Override num_ctas for the main kernel (comma-separated sweep).", ) + parser.add_argument( + "--cta-per-sm", + type=str, + default=None, + help="CTAs per SM in the 1D persistent grid for mode=persistent_main " + "(comma-separated sweep). num_persistent = cta_per_sm × NUM_SMS. " + "Default = 1 (one CTA per SM). Replaces the old --num-persistent. " + "Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--num-loop-stages", + type=str, + default=None, + help="num_stages on the inner tl.range(...) persistent loop for " + "mode=persistent_main (comma-separated sweep). Default = 2. Note: " + "this is loop-level, NOT the kernel-arg num_stages (which only " + "pipelines dot-feeding loads). Watch Triton issue #8259 — " + "num_stages>1 + flatten=True can corrupt stores in non-dot kernels. " + "Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--flatten", + type=str, + default=None, + help="`flatten` arg on tl.range(...) for mode=persistent_main " + "(comma-separated 0/1 sweep). Default = 1. Ignored for " + "non-persistent_main modes.", + ) + parser.add_argument( + "--warp-specialize", + type=str, + default=None, + help="`warp_specialize` arg on tl.range(...) for mode=persistent_main " + "(comma-separated 0/1 sweep). Default = 0. Triton 3.6 only " + "supports it on simple matmul loops; our scan loop probably won't " + "pattern-match — exposed as a knob for sweep experiments. Requires " + "num_warps >= 4 if 1. Ignored for non-persistent_main modes.", + ) parser.add_argument( "--sr-modes", type=str, @@ -1780,7 +1895,10 @@ def _parse_args() -> argparse.Namespace: args.write_modes_list = [args.write_checkpoint] modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] - valid_modes = {"monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", "dl_write_only"} + valid_modes = { + "monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", "persistent_main", + } for m in modes_raw: if m not in valid_modes: parser.error( diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index b8bd29cfc56e..54aebe659d55 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -878,6 +878,177 @@ def test_checkpointing_state_update_sorted_dispatch(mode, rectangle_for_nowrite, ) +@pytest.mark.parametrize( + "scenario,pnat_per_slot_list,n_writes_expected,slot_perm_list", + [ + # All-write: every slot has PNAT triggering write + # (PNAT + T > max_window). No permutation needed. + ("all_write", [12, 13, 14, 15], 4, [0, 1, 2, 3]), + # All-nowrite: every slot fits in the window. n_writes = 0. + ("all_nowrite", [3, 4, 5, 6], 0, [0, 1, 2, 3]), + # Mixed (write-first sorted via slot_perm): physical slots 2, 3 + # are writes; physical slots 0, 1 are nowrites. slot_perm + # remaps grid pid_b 0..3 to physical slots 2, 3, 0, 1 — so the + # first n_writes=2 grid programs hit write slots and the rest + # hit nowrite slots. + ("mixed_sorted", [3, 10, 12, 16], 2, [2, 3, 0, 1]), + ], + ids=["all_write", "all_nowrite", "mixed_sorted"], +) +def test_checkpointing_state_update_persistent_main( + scenario, pnat_per_slot_list, n_writes_expected, slot_perm_list, +): + """ + Persistent-CTA main kernel: 1D-grid kernel that loops over + (slot, M-tile, head) work units via tl.range. Caller pre-sorts + slots write-first and passes _n_writes (count of write slots) so + the kernel can split the persistent loop into write and nowrite + halves with the right WRITE_CHECKPOINT constexpr each time. + + Setup mirrors test_checkpointing_state_update_sorted_dispatch + (same fixed seeds, same input shapes) so the reference state + evolution is identical and we can compare per-slot output and + HBM-state postconditions to the same reference. + + Cases: + - all_write (n_writes=B): every slot exercises the + WRITE_CHECKPOINT=True branch of the persistent loop. + - all_nowrite (n_writes=0): every slot exercises the + WRITE_CHECKPOINT=False branch. Verifies the kernel handles + the "write half is empty" launch (n_slots=0 → early return). + - mixed_sorted: slots [2, 3] are writes, slots [0, 1] are + nowrites. slot_perm = [2, 3, 0, 1]. Persistent kernel + should call its impl with pid_b ∈ {2, 3} for the write half + and pid_b ∈ {0, 1} for the nowrite half, even though the + grid pid_b_grid is 0..n_slots-1 in each. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + + pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() + slot_perm = torch.tensor(slot_perm_list, device=device, dtype=torch.int32) + # Sanity: caller-supplied n_writes must match the actual count of + # write slots in the post-perm order. + write_count = sum(pnat_means_write) + assert write_count == n_writes_expected, ( + f"test setup error: expected {n_writes_expected} writes, got {write_count}" + ) + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=dtype + ) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + batch, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + test_state, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + mode="persistent_main", + slot_perm=slot_perm, + _n_writes=n_writes_expected, + ) + + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (scenario={scenario})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (scenario={scenario})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", + ) + + @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) @pytest.mark.parametrize( "state_dtype", From f078fc7e858b690a066f0cb948cb9cbb23bf9374 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 8 May 2026 10:44:17 -0700 Subject: [PATCH 26/89] SR randomness over-generation fix (backlog #19) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cvt.rs.f16x2 / cvt.rs.satfinite.e4m3x4 PTX instructions consume one b32 random and split the bits internally for 2 (fp16) or 4 (fp8) outputs. But Triton's tl.inline_asm_elementwise has a uniform pack value across all args, so it provides 2 (fp16) or 4 (fp8) rand inputs per asm call — only the first is read by the cvt instruction; the others are dead inputs occupying register slots. Previously: generate (M, dstate) randoms via M*dstate/4 randint4x calls (4 randoms per Philox round). Per dtype actual consumption: fp16 SR: dstate/2 randoms used per row (50% waste) fp8 SR: dstate/4 randoms used per row (75% waste) int8/int16 SR: full dstate (uniform-noise + floor uses full bits) Fix: generate only what's consumed (gated by RAND_DIVISOR constexpr, selected from QUANT_MAX + state dtype), then broadcast to (M, dstate) to fill Triton's dead slots. Same-node SR write sweep (sr_baseline vs sr_postfix, fp16+fp8, batches 1/16/64/512, full autotune, --with-conv1d): b=512 fp8 SR: -11.9% better b=512 fp16 SR: -9.3% better b=64 fp8 SR: -8.8% better b=16 fp8 SR: -10.6% better b=64 fp16 SR: -3.7% better b=16 fp16 SR: -3.3% better b=1 fp8 SR: -0.9% same b=1 fp16 SR: +4.9% worse (autotune drift, +0.32us absolute) Summary: 6 better, 1 worse, 1 same. Headline wins on b>=16 fp8 (-8 to -12%). Note: baseline was on prior compute node, but win magnitudes clearly survive cross-node noise. Register count up at one cell (80 -> 96 at fp16 M=32 W=4) due to rand_compact's lifetime across the asm call; tested zero-fill variant (tl.join with zeros) gave 95 — essentially equivalent. ptxas doesn't use RZ for the dead asm input slots in either case. Occupancy still healthy (96 regs / 128 threads / SM = ~5 CTAs per SM at this M). int8/int16 SR path unchanged (uniform-noise + floor uses 24 bits per element; cannot share randomness across elements). Tests: 542 unit tests pass; 48 dedicated philox SR tests pass. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .bash_profile | 0 .bashrc | 0 .claude/commands | 0 .claude/settings.json | 0 .gitconfig | 0 .idea | 0 .mcp.json | 0 .profile | 0 .ripgreprc | 0 .zprofile | 0 .zshrc | 0 .../mamba/checkpointing_state_update.py | 64 +++++++++++++++---- 12 files changed, 52 insertions(+), 12 deletions(-) create mode 100644 .bash_profile create mode 100644 .bashrc create mode 100644 .claude/commands create mode 100644 .claude/settings.json create mode 100644 .gitconfig create mode 100644 .idea create mode 100644 .mcp.json create mode 100644 .profile create mode 100644 .ripgreprc create mode 100644 .zprofile create mode 100644 .zshrc diff --git a/.bash_profile b/.bash_profile new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.bashrc b/.bashrc new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.claude/commands b/.claude/commands new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.claude/settings.json b/.claude/settings.json new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.gitconfig b/.gitconfig new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.idea b/.idea new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.mcp.json b/.mcp.json new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.profile b/.profile new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.ripgreprc b/.ripgreprc new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.zprofile b/.zprofile new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.zshrc b/.zshrc new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 1eeb75d1683e..2e48b9299e66 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -1052,27 +1052,67 @@ def _checkpointing_main_kernel( # win of replay-style checkpointing on the common (non-checkpoint) step. if WRITE_CHECKPOINT: if USE_RS_ROUNDING: - # Generate (M, dstate) random tensor for stochastic rounding. - # Used by fp16 / int8 / int16 / fp8 SR paths below. Quarter-sized - # randint4x calls produce 4 u32s each; we interleave them into the - # full (M, dstate) tensor — 4x fewer PRNG rounds vs per-element. + # Generate random tensor for stochastic rounding. The amount of + # randomness needed depends on the SR codegen path: + # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) + # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) + # int8/int16 SR (uniform-noise + floor): 1 b32 per output + # The PTX cvt.rs.* instructions consume a single 32-bit random + # and split the bits internally for 2 or 4 conversions. The + # tl.inline_asm_elementwise wrapper has uniform `pack` across all + # args, so it provides 2 (fp16) or 4 (fp8) rand inputs per asm + # call but only the first is read; the others are dead. Generate + # only what's actually consumed and broadcast to fill the unused + # slots — saves Philox rounds proportionally. + if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: + RAND_DIVISOR: tl.constexpr = 4 # fp8 SR + elif QUANT_MAX == 0.0: + RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) + else: + RAND_DIVISOR: tl.constexpr = 1 # int8/int16 SR (full per-element) + rand_seed = tl.load(rand_seed_ptr) base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head - offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // 4) + # Number of unique randoms per row = dstate / RAND_DIVISOR. + # randint4x emits 4 randoms per offset, so use that / 4 offsets. + offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) rand_offsets_q = ( base_rand + offs_m[:, None] * stride_state_dim - + offs_n_q[None, :] * (stride_state_dstate * 4) - ) # (M, dstate//4) + + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) + ) # (M, dstate / (4*RAND_DIVISOR)) if PHILOX_ROUNDS > 0: r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) else: r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) - # Interleave 4 quarter-sized tensors → full (M, dstate) random tensor - r01 = tl.join(r0, r1) # (M, dstate//4, 2) - r23 = tl.join(r2, r3) # (M, dstate//4, 2) - r0123 = tl.join(r01, r23) # (M, dstate//4, 2, 2) - rand = tl.reshape(r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + r01 = tl.join(r0, r1) + r23 = tl.join(r2, r3) + r0123 = tl.join(r01, r23) + rand_compact = tl.reshape( + r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) + ) + # Broadcast each unique rand to RAND_DIVISOR adjacent positions + # in the dstate axis. Pack-group (pack=2 fp16 / pack=4 fp8) + # consumes adjacent positions; the unique rand lands at the + # asm's read slot ($3 fp16 / $5 fp8); duplicates feed the dead + # slots ($4 fp16; $6/$7/$8 fp8). Triton's broadcast_to is + # stride-0 in IR. + # + # Tested zero-fill alternative (tl.join with zeros): essentially + # equivalent register count (95 vs 96 at one config) and same + # timing. ptxas does not use RZ for the dead asm input slots + # in either case; the extra ~15 regs vs pre-fix come from + # rand_compact's lifetime across the asm call, not from the + # fill pattern. Broadcast wins on simplicity. + if RAND_DIVISOR > 1: + rand_3d = rand_compact[:, :, None] + rand_3d = tl.broadcast_to( + rand_3d, + (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), + ) + rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + else: + rand = rand_compact if QUANT_MAX > 0.0: # Quantized state path: int8 / int16 / fp8_e4m3fn (RN or SR). From 6d5a6633fb20b79fe7300b11c0fd41652fbb59c0 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 8 May 2026 10:11:11 -0700 Subject: [PATCH 27/89] persistent_dynamic: single-launch persistent kernel with runtime per-slot WRITE_CHECKPOINT branch via IS_DYNAMIC constexpr; tests pass; +5%/+23% slower than persistent_main on WRITE/NOWRITE at b=128 Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 110 +++++++++++-- ...benchmark_replay_selective_state_update.py | 32 ++-- .../mamba/test_checkpointing_state_update.py | 144 ++++++++++++++++++ 3 files changed, 261 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 1b6dbbba662a..98f7c0f0ce9d 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -3096,7 +3096,14 @@ def _persistent_main_impl( PHILOX_ROUNDS: tl.constexpr, QUANT_MAX: tl.constexpr, WRITE_CHECKPOINT: tl.constexpr, + IS_DYNAMIC: tl.constexpr, ): + # IS_DYNAMIC: when False, WRITE_CHECKPOINT is the constexpr write/nowrite + # selector (caller pre-sorts and splits halves). When True, WRITE_CHECKPOINT + # is ignored; the impl computes is_write at runtime per work-item from + # the loaded PNAT. Used by mode="persistent_dynamic" — single kernel, + # one launch, no half-split, runtime per-slot dispatch. + # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized # state dtype (int8 / int16 / float8e4nv) and only those. tl.static_assert( @@ -3118,7 +3125,15 @@ def _persistent_main_impl( active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - if WRITE_CHECKPOINT: + # Resolve is_write: constexpr from caller (persistent_main) or runtime + # from PNAT (persistent_dynamic). When IS_DYNAMIC=False, is_write + # collapses to a constexpr 0/1 and the downstream `if is_write:` + # blocks DCE the dead path at compile time. + if IS_DYNAMIC: + is_write = (prev_num_accepted_tokens + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_write = WRITE_CHECKPOINT + if is_write: write_buf = 1 - active_buf # noqa: F841 write_offset = 0 else: @@ -3214,7 +3229,7 @@ def _persistent_main_impl( state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) - if WRITE_CHECKPOINT: + if is_write: if USE_RS_ROUNDING: rand_seed = tl.load(rand_seed_ptr) base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head @@ -3486,21 +3501,27 @@ def _persistent_main_kernel( NUM_PID_M_BLOCKS: tl.constexpr, FLATTEN: tl.constexpr, WARP_SPECIALIZE: tl.constexpr, + IS_DYNAMIC: tl.constexpr, ): # PDL signal: fire once at kernel entry (not per work unit). if LAUNCH_DEPENDENT_KERNELS: tl.extra.cuda.gdc_launch_dependents() - # Derive this kernel's slot range from (n_writes, batch_total, WRITE_CHECKPOINT). - # Hard-sort contract: caller has pre-sorted so [0, n_writes) are write-mode - # and [n_writes, batch_total) are nowrite-mode. Each launch knows which - # half it owns from its WRITE_CHECKPOINT constexpr. - if WRITE_CHECKPOINT: + # Derive this kernel's slot range. Two modes: + # IS_DYNAMIC=False (persistent_main): caller pre-sorts and splits halves; + # slot range is [0, n_writes) when WRITE_CHECKPOINT else [n_writes, batch_total) + # IS_DYNAMIC=True (persistent_dynamic): single launch covers full batch; + # each work-item dispatches via runtime PNAT check inside the impl. + if IS_DYNAMIC: slot_lo = 0 - slot_hi = n_writes - else: - slot_lo = n_writes slot_hi = batch_total + else: + if WRITE_CHECKPOINT: + slot_lo = 0 + slot_hi = n_writes + else: + slot_lo = n_writes + slot_hi = batch_total n_slots_local = slot_hi - slot_lo pid = tl.program_id(axis=0) @@ -3578,6 +3599,7 @@ def _persistent_main_kernel( PHILOX_ROUNDS, QUANT_MAX, WRITE_CHECKPOINT, + IS_DYNAMIC, ) @@ -3747,11 +3769,11 @@ def checkpointing_state_update( # reg envelope. write_checkpoint ignored. assert mode in ( "monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", - "dl_write_only", "persistent_main", + "dl_write_only", "persistent_main", "persistent_dynamic", ), ( f"unknown mode {mode!r}; expected one of " "'monolithic', 'dynamic', 'doublelaunch', 'dlgrouped', 'maindl', " - "'dl_write_only', or 'persistent_main'" + "'dl_write_only', 'persistent_main', or 'persistent_dynamic'" ) use_rectangle = rectangle_for_nowrite and not write_checkpoint @@ -4350,6 +4372,61 @@ def launch_persistent_main(write_checkpoint: bool, n_writes: int, NUM_LOOP_STAGES=num_loop_stages_arg, FLATTEN=flatten_arg, WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=False, + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + def launch_persistent_dynamic_main(launch_dependent_kernels: bool = False): + # Single-launch persistent kernel covering the whole batch with + # runtime per-slot WRITE_CHECKPOINT branch. No half-split, no + # n_writes needed. Reuses _persistent_main_kernel with + # IS_DYNAMIC=True. + grid = (num_persistent_arg,) + # n_writes is unused in the kernel when IS_DYNAMIC=True; pass 0. + # WRITE_CHECKPOINT is also unused; pass False. is_write is computed + # at runtime per work-item from the loaded PNAT. + _persistent_main_kernel[grid]( + state, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + 0, batch, nheads, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=False, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + NUM_PERSISTENT=num_persistent_arg, + NUM_LOOP_STAGES=num_loop_stages_arg, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=True, num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), @@ -4415,6 +4492,15 @@ def launch_persistent_main(write_checkpoint: bool, n_writes: int, launch_replay_precompute(write_checkpoint=True, early_out=True) launch_replay_main(write_checkpoint=True, early_out=True, launch_dependent_kernels=False) + elif mode == "persistent_dynamic": + # Single-launch persistent kernel covering the full batch. + # Each work-item dispatches via runtime PNAT check (is_write = + # (pnat + T) > MAX). No n_writes/half-split, no PDL hop + # between mains — only one main launch. Precompute is + # dynamic_precompute (per-slot dispatch), so the whole pipeline + # is dynamic_precompute → persistent_dynamic_main. + launch_dynamic_precompute(rectangle=False) + launch_persistent_dynamic_main(launch_dependent_kernels=False) elif mode == "persistent_main": # Experimental: persistent-CTA main kernel. Reuses maindl's # precompute structure (one shared dynamic_precompute that diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 216e0a3a3908..7efaf33eb9f0 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -589,6 +589,7 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp is_dl_family = mode in ( "doublelaunch", "dlgrouped", "maindl", "dl_write_only", "persistent_main", + "persistent_dynamic", ) # Match the timed-run skip: sort=1 only # makes sense when there's a mix scenario. @@ -1077,19 +1078,11 @@ def _run_incr( # n_writes is either 0 (all nowrite) or batch (all # write) depending on whether PNAT+T overflows the # window. Mix scenarios are skipped earlier. - if mode == "persistent_main": - scn_fill = scn["fill"] - assert scn_fill is not None, ( - "persistent_main does not yet support mix " - "scenarios (per-iter n_writes plumbing not " - "implemented)." - ) - is_write_scenario = (scn_fill + mtp_len) > max_window - extra_kwargs["_n_writes"] = batch if is_write_scenario else 0 + if mode in ("persistent_main", "persistent_dynamic"): # Per-cell sweep values for persistent-only knobs. - # _parse_sweep returns [None] when the user didn't - # pass the flag, in which case we leave the wrapper's - # defaults in place. + # Apply to both persistent variants. _parse_sweep + # returns [None] when the user didn't pass the flag, + # in which case we leave the wrapper's defaults. if cta_per_sm is not None: extra_kwargs["_cta_per_sm"] = cta_per_sm if num_loop_stages is not None: @@ -1098,6 +1091,18 @@ def _run_incr( extra_kwargs["_flatten"] = bool(flatten) if warp_specialize is not None: extra_kwargs["_warp_specialize"] = bool(warp_specialize) + if mode == "persistent_main": + # n_writes is needed only by persistent_main's + # half-split. persistent_dynamic dispatches per-slot + # at runtime and doesn't need it. + scn_fill = scn["fill"] + assert scn_fill is not None, ( + "persistent_main does not yet support mix " + "scenarios (per-iter n_writes plumbing not " + "implemented)." + ) + is_write_scenario = (scn_fill + mtp_len) > max_window + extra_kwargs["_n_writes"] = batch if is_write_scenario else 0 variant_fn( state_work, old_x_work, @@ -1395,6 +1400,7 @@ def _run_benchmark(args) -> None: is_dl_family = mode in ( "doublelaunch", "dlgrouped", "maindl", "dl_write_only", "persistent_main", + "persistent_dynamic", ) can_sort = ( is_dl_family and mix_samples_cpu is not None @@ -1897,7 +1903,7 @@ def _parse_args() -> argparse.Namespace: modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] valid_modes = { "monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", - "dl_write_only", "persistent_main", + "dl_write_only", "persistent_main", "persistent_dynamic", } for m in modes_raw: if m not in valid_modes: diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index 54aebe659d55..5fa87feca225 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -1049,6 +1049,150 @@ def test_checkpointing_state_update_persistent_main( ) +@pytest.mark.parametrize( + "scenario,pnat_per_slot_list", + [ + # All-write: every slot has PNAT triggering write (PNAT + T > max_window). + ("all_write", [12, 13, 14, 15]), + # All-nowrite: every slot fits in the window. + ("all_nowrite", [3, 4, 5, 6]), + # Mixed: some slots write, some nowrite. No pre-sort needed; the + # dynamic kernel dispatches per-slot at runtime via PNAT load. + ("mixed_unsorted", [3, 12, 10, 15]), + ], + ids=["all_write", "all_nowrite", "mixed_unsorted"], +) +def test_checkpointing_state_update_persistent_dynamic( + scenario, pnat_per_slot_list, +): + """ + Persistent-dynamic kernel: 1D persistent-CTA grid covering the full + batch, with runtime per-slot WRITE_CHECKPOINT branch derived from + each slot's PNAT. Single launch, no half-split, no n_writes needed, + no slot_perm needed (handles unsorted batches natively). + + Same setup as test_checkpointing_state_update_persistent_main; we + verify all three scenarios — including a mixed-unsorted batch the + persistent_main kernel can't handle without pre-sorting. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + + pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=dtype + ) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + batch, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + test_state, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + mode="persistent_dynamic", + ) + + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (scenario={scenario})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (scenario={scenario})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", + ) + + @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) @pytest.mark.parametrize( "state_dtype", From 3d7ee6f806597a7511cf9d99c13ea30b6708423a Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 8 May 2026 12:51:59 -0700 Subject: [PATCH 28/89] cupti: in-process kernel timing for benchmark_replay_selective_state_update Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 539 +++++++++++++++--- 1 file changed, 471 insertions(+), 68 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 8c6eee963c2a..8046308bc51e 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -27,22 +27,93 @@ Calls selective_state_update with T=mtp_len tokens and disable_state_update=True, matching the MTP scoring pass in mamba2_mixer.py exactly. +Timing methodology +================== + +All in-bench timing comes from CUPTI's Activity API (1 ns kernel +timestamps from the GPU profiling fabric). cudaEvent.elapsed_time() was +removed — its ~0.5 us resolution overshoots CUPTI by ~50% on short kernels +in graphs, and we have no other use for it here. See the CUPTI block +lower in this file for the timer source. + +Three modes: + + --cupti --cuda-graph (default) + Capture one CUDA graph per cell (warmup + timed iters inlined), + replay once, read kernel start/end from CUPTI. ~20× faster than + nsys-wrapped capture and matches it to within ~1% / noise floor. + + --cupti --no-cuda-graph + Eager loop with CUPTI. Per-kernel timestamps are still accurate, but + the per-iter SPAN (max(end) - min(start)) now includes the Python + launch latency BETWEEN consecutive kernels in run_fn (~100 µs on + Hopper/Blackwell). Graph capture and PDL hide that latency; eager + mode honestly reports it. For per-kernel timing in eager mode, look + at per_kernel.start_us/end_us in --json-detailed output rather than + the span percentiles. Useful when graph capture is undesirable. + + --no-cupti (with or without --cuda-graph) + No in-bench timing — just runs the kernels for an external profiler + (nsys / ncu) to time. In-process CUPTI conflicts with nsys's own + subscriber, so disable ours when wrapping in nsys. Bench output + reports zeros for median/p95/p99; trust the external trace. + +JSON output schema (--json-output PATH) +======================================= + +Designed to be parsed by collect.py / report.py without touching sqlite or +NVTX traces. Future agents: prefer reading this JSON over re-running nsys. + + { + "metadata": {timestamp, cmd, tp_size, warmup, iters, variant, cupti}, + "results": { + "": {median, p95, p99, n, [iters_us], [per_kernel]} + } + } + +Key format mirrors collect.py's kernel_data.json convention: + incremental/{batch}/{mtp}/{sd}/k{prev_k}/{sweep_parts}/tp{tp} + triton/{batch}/{mtp}/{sd}/tp{tp} + flashinfer/{batch}/{mtp}/{sd}/tp{tp} + + - is normalized: bf16 / fp16 / fp32 / int8 / int16 / fp8. + - is e.g. "M16_W1_S3_SR0_RECT0_WC1" — flags concatenated by + underscore in canonical (M, W, S, pW, pS, H, R, CT, SR, RECT, WC) order. + - All numeric values in microseconds (us). + +Per-record fields: + - median, p95, p99: span statistics (us). Span = max(kernel_end_ns) - + min(kernel_start_ns) across the iter's kernels — same convention as + nsys-derived collect.py used to use. + - n: number of timed iters that contributed. + - iters_us: list of length n, raw per-iter spans (only with --json-detailed). + - per_kernel: {: {start_us: [...], end_us: [...]}} where + timestamps are RELATIVE to that iter's first kernel start, in us. Lets + you see PDL overlap directly without an external profiler. Only with + --json-detailed. + Example usage: - # Basic sweep + # Basic sweep (default = --cupti, just summary stats) python benchmark_replay_selective_state_update.py \\ --batch-sizes 1,2,4 --mtp-lengths 1,4,8 --warmup 5 --iters 20 - # With CUDA graph (default) and Triton baseline: - python benchmark_replay_selective_state_update.py --baseline \\ - --batch-sizes 1,2,4 --mtp-lengths 5,10,20 + # JSON output, summary stats only (compact) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 16 --mtp-lengths 6 --json-output /tmp/out.json + + # JSON output, full per-iter / per-kernel data (for PDL analysis etc.) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 16 --mtp-lengths 6 \\ + --json-output /tmp/out.json --json-detailed - # nsys capture (NVTX ranges visible in timeline) + # nsys capture (--no-cupti so our subscriber doesn't conflict) nsys profile --capture-range=cudaProfilerApi \\ - python benchmark_replay_selective_state_update.py --profile + python benchmark_replay_selective_state_update.py --profile --no-cupti - # ncu capture + # ncu capture (--no-cupti --no-cuda-graph: each kernel replayable solo) ncu --target-processes all \\ python benchmark_replay_selective_state_update.py --profile \\ + --no-cupti --no-cuda-graph \\ --batch-sizes 1 --mtp-lengths 4 --warmup 5 --iters 5 """ @@ -385,16 +456,201 @@ def _build_tensors( ) +# ============================================================================= +# CUPTI in-process kernel timing +# +# Self-contained module-in-a-file. Reads kernel start/end timestamps directly +# from the GPU profiling fabric via NVIDIA's cupti-python bindings (1 ns +# resolution), avoiding two pitfalls of the cuda-events path: +# +# 1. cudaEvent.elapsed_time() resolution (~0.5 us) is too coarse for the +# short kernels we care about, especially with PDL + cuda graphs at +# small batch — events recorded inside a graph have proven noisy. +# 2. nsys is the only known accurate alternative, but the +# profile-export-sqlite-parse pipeline is heavy and out-of-process. +# +# This is functionally equivalent to wrapping each cell in nsys, except it +# runs in the same Python process with no serialization. When this proves +# out, lift `CuptiKernelTimer` and `_time_kernel_cuda_graph_cupti` into a +# proper TRT-LLM utility module — there is no benchmark-specific code below. +# ============================================================================= + + +# Substring match: kernels run_fn launches that we want to time. Mirrors +# the parser in scripts/.../collect.py so cupti and nsys-based outputs agree. +_CUPTI_KEEP_KERNEL_SUBSTRINGS = ( + "_replay_precompute", + "_checkpointing_precompute", + "_rectangle_precompute", + "_dynamic_precompute", + "_replay_state_update", + "_checkpointing_main", + "_rectangle_main", + "_dynamic_main", + "selective_scan_update", + "selective_state_update", + "causal_conv1d_update", +) + + +class CuptiKernelTimer: + """Process-singleton wrapper around CUPTI's CONCURRENT_KERNEL activity. + + CUPTI's callbacks are global (one subscriber per process), so the timer + is constructed lazily once via `CuptiKernelTimer.get()`. cupti-python + parses the activity buffer for us — `buffer_completed` receives a Python + list of typed activity objects, not a raw byte buffer — so no FFI is + needed. + + Usage: + timer = CuptiKernelTimer.get() + timer.start() # arms; drops any stale records + + records = timer.stop() # flush; list of tuples per kernel + # (name, start_ns, end_ns, corr, + # graph_id, graph_node_id, stream) + + The callback fires from a CUPTI worker thread, so a lock guards the + record buffer. Records are kept tiny (tuple of ints + str) to minimize + Python overhead in the hot path of the callback. + """ + + _instance = None + _import_error = None + + @classmethod + def get(cls) -> "CuptiKernelTimer": + if cls._instance is not None: + return cls._instance + if cls._import_error is not None: + raise cls._import_error + try: + from cupti import cupti as _c + except ImportError as e: # pragma: no cover — env-dependent + cls._import_error = e + raise + cls._instance = cls._init(_c) + return cls._instance + + @classmethod + def _init(cls, _c) -> "CuptiKernelTimer": + import threading + + self = object.__new__(cls) + self._c = _c + self._records: list[tuple] = [] + self._lock = threading.Lock() + + # CUPTI callback contract (from cupti-python-samples/cupti_common.py): + # buffer_requested() -> (buffer_size, max_num_records) + # buffer_completed(activities: list) + # Setting max_num_records=0 (unbounded) avoids spurious buffer + # requests. 8 MiB matches the sample defaults. + def _buf_req(): + return (8 * 1024 * 1024, 0) + + kernel_kinds = (_c.ActivityKind.CONCURRENT_KERNEL, _c.ActivityKind.KERNEL) + + def _buf_done(activities): + recs = [] + for a in activities: + if a.kind not in kernel_kinds: + continue + # start/end == 0 means CUPTI couldn't time this kernel. + if a.start == 0 or a.end == 0: + continue + recs.append(( + a.name, + int(a.start), + int(a.end), + int(a.correlation_id), + int(a.graph_id), + int(a.graph_node_id), + int(a.stream_id), + )) + if recs: + with self._lock: + self._records.extend(recs) + + # Hold strong refs so the C side never sees GC'd Python callables. + self._buf_req = _buf_req + self._buf_done = _buf_done + + _c.activity_register_callbacks(_buf_req, _buf_done) + _c.activity_enable(_c.ActivityKind.CONCURRENT_KERNEL) + return self + + def start(self) -> None: + """Arm capture: flush any stale records, then clear the buffer.""" + self._c.activity_flush_all(1) + with self._lock: + self._records.clear() + + def stop(self) -> list[tuple]: + """Flush and return all kernel records since the last start().""" + self._c.activity_flush_all(1) + with self._lock: + return list(self._records) + + +# ============================================================================= # Timing helpers +# ============================================================================= -def _compute_stats(latencies_us: list[float]) -> tuple[float, float, float]: - """Return (median_us, p95_us, p99_us) from a list of latencies.""" - median_us = statistics.median(latencies_us) - s = sorted(latencies_us) - p95_us = s[int(0.95 * len(s))] - p99_us = s[int(0.99 * len(s))] - return median_us, p95_us, p99_us +def _stats_from_spans(spans_us: list[float]) -> dict: + """Compute median / p95 / p99 / n from a per-iter span list.""" + s = sorted(spans_us) + return { + "median": statistics.median(s), + "p95": s[int(0.95 * len(s))], + "p99": s[int(0.99 * len(s))], + "n": len(s), + } + + +def _stats_from_cupti_records(records, warmup, iters, tag): + """Bin a flat CUPTI kernel record stream into per-iter spans + per-kernel + relative timestamps. Used by both graph and eager CUPTI paths. + + `records` are tuples (name, start_ns, end_ns, ...) — see CuptiKernelTimer. + The first warmup*K records are dropped; the rest are chunked into K-tuples. + """ + records = [ + r for r in records + if any(s in r[0] for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS) + ] + records.sort(key=lambda r: r[1]) # by start_ns + + total = len(records) + expected_iters = warmup + iters + if total == 0 or total % expected_iters != 0: + names_seen = sorted({r[0] for r in records}) + raise RuntimeError( + f"CUPTI capture mismatch for {tag!r}: got {total} records, " + f"expected a multiple of {expected_iters} (warmup+iters={expected_iters}). " + f"Kernel names captured: {names_seen}" + ) + K = total // expected_iters + timed = records[warmup * K:] + + spans_us: list[float] = [] + per_kernel: dict[str, dict[str, list[float]]] = {} + for i in range(iters): + chunk = timed[i * K:(i + 1) * K] + iter_start_ns = min(r[1] for r in chunk) + iter_end_ns = max(r[2] for r in chunk) + spans_us.append((iter_end_ns - iter_start_ns) / 1000.0) + for r in chunk: + name = r[0] + slot = per_kernel.setdefault(name, {"start_us": [], "end_us": []}) + slot["start_us"].append((r[1] - iter_start_ns) / 1000.0) + slot["end_us"].append((r[2] - iter_start_ns) / 1000.0) + + out = _stats_from_spans(spans_us) + out["iters_us"] = spans_us + out["per_kernel"] = per_kernel + return out def _time_kernel_cuda_graph( @@ -402,21 +658,17 @@ def _time_kernel_cuda_graph( run_fn, reset_fn, tag: str, -) -> tuple[float, float, float]: - """ - All-in-one CUDA graph timing. +) -> dict: + """CUDA-graph CUPTI timer. - Captures a single graph containing warmup iterations followed by timed - iterations with per-iteration event pairs recorded inside the graph. - One replay, one sync, then all timings are read. + Captures one CUDA graph (warmup + timed iters inlined), replays once, + reads kernel start/end timestamps from CUPTI (1 ns resolution). """ + timer = CuptiKernelTimer.get() warmup = args.warmup iters = args.iters - start_events = [torch.cuda.Event(enable_timing=True, external=True) for _ in range(iters)] - end_events = [torch.cuda.Event(enable_timing=True, external=True) for _ in range(iters)] - - # Eager warmup before graph capture (triggers Triton autotune if active) + # Eager warmup before graph capture (triggers Triton autotune if active). reset_fn() run_fn() torch.cuda.synchronize() @@ -426,32 +678,22 @@ def _time_kernel_cuda_graph( g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): - # Warmup iterations (unrolled into the graph) - for _ in range(warmup): - reset_fn() - if args.l2_flush: - _l2_flush.fill_(0.0) - run_fn() - - # Timed iterations with events inside the graph - for i in range(iters): + for _ in range(warmup + iters): reset_fn() if args.l2_flush: _l2_flush.fill_(0.0) - start_events[i].record() run_fn() - end_events[i].record() torch.cuda.synchronize() - # Single replay + timer.start() torch.cuda.nvtx.range_push(tag) g.replay() torch.cuda.synchronize() torch.cuda.nvtx.range_pop() + records = timer.stop() - latencies_us = [start_events[i].elapsed_time(end_events[i]) * 1000.0 for i in range(iters)] - return _compute_stats(latencies_us) + return _stats_from_cupti_records(records, warmup, iters, tag) def _time_kernel_eager( @@ -459,35 +701,81 @@ def _time_kernel_eager( run_fn, reset_fn, tag: str, -) -> tuple[float, float, float]: - """Non-CUDA-graph timing path (for debugging, ncu, etc.).""" - # Warmup - for _ in range(args.warmup): - reset_fn() - run_fn() - torch.cuda.synchronize() +) -> dict: + """Non-graph CUPTI timer (for ncu wrapping, debugging, etc.). - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) + Each iter runs serially with sync between, but kernel start/end still + come from CUPTI — same accuracy as the graph path, just slower per-iter + (extra Python + sync overhead). + """ + timer = CuptiKernelTimer.get() + warmup = args.warmup + iters = args.iters - latencies_us: list[float] = [] + timer.start() torch.cuda.nvtx.range_push(tag) - for _ in range(args.iters): + for _ in range(warmup + iters): reset_fn() if args.l2_flush: _flush_l2() # includes synchronize - start_event.record() run_fn() - end_event.record() - torch.cuda.synchronize() - latencies_us.append(start_event.elapsed_time(end_event) * 1000.0) + torch.cuda.synchronize() torch.cuda.nvtx.range_pop() + records = timer.stop() + + return _stats_from_cupti_records(records, warmup, iters, tag) + + +def _run_kernel_untimed(args, run_fn, reset_fn, tag: str) -> dict: + """No in-bench timing: just run the kernels for an external profiler + (nsys / ncu) to time externally. Returns a stats dict full of zeros so + downstream code (table, JSON) doesn't break. + """ + warmup = args.warmup + iters = args.iters + + if args.cuda_graph: + # Eager warmup before capture (Triton autotune) + reset_fn(); run_fn(); torch.cuda.synchronize() + reset_fn(); torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + for _ in range(warmup + iters): + reset_fn() + if args.l2_flush: + _l2_flush.fill_(0.0) + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_push(tag) + g.replay() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + else: + torch.cuda.nvtx.range_push(tag) + for _ in range(warmup + iters): + reset_fn() + if args.l2_flush: + _flush_l2() + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + + spans_us = [0.0] * iters + out = _stats_from_spans(spans_us) + out["iters_us"] = spans_us + out["per_kernel"] = {} + return out - return _compute_stats(latencies_us) +def _time_kernel(args, run_fn, reset_fn, tag: str) -> dict: + """Dispatch to graph-CUPTI / eager-CUPTI / no-timer path. -def _time_kernel(args, run_fn, reset_fn, tag: str) -> tuple[float, float, float]: - """Dispatch to CUDA-graph or eager timing path.""" + --cupti: in-process CUPTI Activity API timing (default). Use --no-cupti + when running under nsys (in-process CUPTI conflicts with nsys's own + subscriber); the bench then runs the kernels for nsys to time externally. + """ + if not getattr(args, "cupti", True): + return _run_kernel_untimed(args, run_fn, reset_fn, tag) if args.cuda_graph: return _time_kernel_cuda_graph(args, run_fn, reset_fn, tag) return _time_kernel_eager(args, run_fn, reset_fn, tag) @@ -801,7 +1089,7 @@ def _run_baseline(): _run_baseline() torch.cuda.synchronize() else: - median_us, p95_us, p99_us = _time_kernel(args, _run_baseline, reset_fn, tag) + stats = _time_kernel(args, _run_baseline, reset_fn, tag) _print_row( show_kernel_col, @@ -811,9 +1099,10 @@ def _run_baseline(): "N/A", state_dtype_name, act_dtype_name, - median_us, - p95_us, - p99_us, + stats, + json_results=getattr(args, "_json_results", None), + tp_size=args.tp_size, + json_detailed=getattr(args, "json_detailed", False), ) # --- Sweep parameter parsing (invariant across prev_k) --- @@ -958,7 +1247,7 @@ def _run_incr( _run_incr() torch.cuda.synchronize() else: - median_us, p95_us, p99_us = _time_kernel(args, _run_incr, reset_fn, sweep_tag) + stats = _time_kernel(args, _run_incr, reset_fn, sweep_tag) _print_row( show_kernel_col, @@ -968,13 +1257,55 @@ def _run_incr( prev_k, state_dtype_name, act_dtype_name, - median_us, - p95_us, - p99_us, + stats, sweep_suffix, + json_results=getattr(args, "_json_results", None), + tp_size=args.tp_size, + json_detailed=getattr(args, "json_detailed", False), ) +# Map full torch dtype name → short tag used in JSON keys (matches collect.py). +_DTYPE_SHORT = { + "float32": "fp32", "bfloat16": "bf16", "float16": "fp16", + "int8": "int8", "int16": "int16", "float8_e4m3fn": "fp8", +} + + +def _build_json_key( + kernel_name, batch, mtp_len, prev_k, state_dtype_name, sweep_suffix, tp_size +): + """Build a key matching collect.py's kernel_data.json convention: + + incremental/{batch}/{mtp}/{sd}/k{k}/{sweep_parts}/tp{tp} + triton/{batch}/{mtp}/{sd}/tp{tp} + flashinfer/{batch}/{mtp}/{sd}/tp{tp} + + `kernel_name` is what _print_row receives: variant name for the timed + kernel (replay/checkpointing) or baseline name for the baseline row. + Variant rows collapse to "incremental" — the variant choice is captured + by the sweep flags collect.py would otherwise apply via --variant. + """ + if kernel_name in ("replay", "checkpointing"): + kind = "incremental" + else: + kind = kernel_name # "triton" / "flashinfer" + + sd = _DTYPE_SHORT.get(state_dtype_name, state_dtype_name) + parts = [kind, str(batch), str(mtp_len), sd] + if prev_k != "N/A": + parts.append(f"k{prev_k}") + if sweep_suffix: + # sweep_suffix format: " M=4,W=1,S=1,SR=0,RECT=0,WC=1" + # collect.py format: "M4_W1_S1_SR0_RECT0_WC0" + # Strip leading/trailing whitespace, drop '=', commas → underscores. + parts.append( + sweep_suffix.strip().replace("=", "").replace(",", "_") + ) + parts.append(f"tp{tp_size}") + return "/".join(parts) + + def _print_row( show_kernel_col, kernel_name, @@ -983,24 +1314,48 @@ def _print_row( prev_k, state_dtype_name, act_dtype_name, - median_us, - p95_us, - p99_us, + stats, sweep_suffix="", + json_results=None, + tp_size=None, + json_detailed=False, ): + """Print one summary row and optionally accumulate stats for JSON output. + + `stats` is a dict from _time_kernel: {median, p95, p99, n, iters_us, + [per_kernel]}. The summary table only shows the headline percentiles. + JSON output captures median/p95/p99/n by default; with json_detailed=True + it also captures the per-iter and per-kernel data. + """ kernel_col = f"{kernel_name:>11} | " if show_kernel_col else "" print( f"| {kernel_col}{batch:>5} | {mtp_len:>7} | {str(prev_k):>6} | " f"{state_dtype_name:>11} | {act_dtype_name:>9} | " - f"{median_us:>9.2f} | {p95_us:>7.2f} | {p99_us:>7.2f} |" + f"{stats['median']:>9.2f} | {stats['p95']:>7.2f} | {stats['p99']:>7.2f} |" f"{sweep_suffix}" ) + if json_results is not None: + key = _build_json_key( + kernel_name, batch, mtp_len, prev_k, state_dtype_name, + sweep_suffix, tp_size, + ) + if json_detailed: + json_results[key] = stats + else: + json_results[key] = { + k: stats[k] for k in ("median", "p95", "p99", "n") + if k in stats + } # Main benchmark loop def _run_benchmark(args) -> None: + # JSON accumulator — populated by _print_row when --json-output is set. + # Stash on args so we don't need to thread a dict through every helper. + args._json_results = {} if getattr(args, "json_output", None) else None + assert args.nheads % args.tp_size == 0, ( f"nheads ({args.nheads}) must be divisible by tp_size ({args.tp_size})" ) @@ -1097,6 +1452,27 @@ def _run_benchmark(args) -> None: if args.profile: torch.cuda.cudart().cudaProfilerStop() + if args.json_output and args._json_results is not None: + import json + payload = { + "metadata": { + "timestamp": datetime.now().isoformat(), + "cmd": " ".join(sys.argv), + "tp_size": args.tp_size, + "warmup": args.warmup, + "iters": args.iters, + "variant": args.variant, + "cupti": getattr(args, "cupti", False), + }, + "results": args._json_results, + } + tmp = args.json_output + ".tmp" + with open(tmp, "w") as f: + json.dump(payload, f, indent=2) + os.replace(tmp, args.json_output) + print(f"\nJSON results written to: {args.json_output} " + f"({len(args._json_results)} entries)") + # CLI @@ -1181,6 +1557,33 @@ def _parse_args() -> argparse.Namespace: "single CUDA graph with per-iteration events " "inside the graph, eliminating all host overhead.", ) + parser.add_argument( + "--cupti", + action=argparse.BooleanOptionalAction, + default=True, + help="Time kernels via CUPTI Activity API (1 ns from the GPU " + "profiling fabric); per-iter span = max(kernel_end) - " + "min(kernel_start). Default ON. --no-cupti disables in-bench " + "timing entirely (kernels still run, but median/p95/p99 are zero) " + "— use when wrapping the bench in nsys/ncu, where the external " + "profiler provides timings and our CUPTI subscriber would conflict.", + ) + parser.add_argument( + "--json-output", + default=None, + help="If set, write per-cell results to this JSON file in the " + "shape consumed by collect.py / report.py. See the 'JSON output " + "schema' section at the top of this file.", + ) + parser.add_argument( + "--json-detailed", + action=argparse.BooleanOptionalAction, + default=False, + help="When --json-output is set, also include iters_us (raw per-iter " + "spans) and per_kernel (per-iter relative start/end timestamps for " + "each kernel) — useful for PDL overlap analysis but adds ~4 KB/cell. " + "Default off keeps records to ~40 bytes (median/p95/p99/n only).", + ) parser.add_argument( "--prev-tokens-fracs", default="0,0.5,1.0", From a1460dee3603bce45b436f9c794fa3e4ba2d1333 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 11 May 2026 15:01:12 -0700 Subject: [PATCH 29/89] persistent: cap launch grid at min(NUM_PERSISTENT, total_work) Wraps the launch grid for persistent_main + persistent_dynamic kernels with min(NUM_PERSISTENT, total_work) where total_work = n_slots * num_pid_m * nheads. At small batch the full persistent grid (CPS * num_sms, up to ~1184 CTAs at CPS=8) far exceeds actual tile count, so most CTAs would launch empty. The kernel's tl.range(pid, total_work, NUM_PERSISTENT) loop already handles grid < NUM_PERSISTENT correctly: each live pid covers exactly one tile (loop step >= total_work exits immediately). NUM_PERSISTENT stays a constexpr = CPS * num_sms so kernel compilation is unaffected. For pure scenarios where host knows host_n_writes, n_slots is tightened to the half-specific count (host_n_writes for the write half, batch - host_n_writes for the nowrite half). Mix scenarios use the upper bound 'batch' (host can't read device n_writes without a sync). Empirical at bf16 b=16 mtp=6 max_window=16 prev_k=full WC=1, persistent_main default knobs: CPS=8 17.66us -> 10.22us (-42%) by avoiding 672 empty CTAs. At b=1, CPS=1 vs CPS=8 timings converge (both now use grid=16, no wasted CTAs). At b=128, work >> grid so timing unchanged. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 29 +++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 91533582e3a5..6fdb72d8b6f0 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -4952,6 +4952,15 @@ def launch_dynamic_main(rectangle: bool, num_loop_stages_arg = _num_loop_stages if _num_loop_stages else 2 flatten_arg = True if _flatten is None else bool(_flatten) warp_specialize_arg = False if _warp_specialize is None else bool(_warp_specialize) + # Per-launch work-item count. At small batch, total_work may be < the + # full persistent grid; capping `grid` at `min(NUM_PERSISTENT, total_work)` + # avoids launching empty CTAs that pay setup cost for no work. Correctness: + # the kernel's `tl.range(pid, total_work, NUM_PERSISTENT)` ensures each + # tile_id is covered exactly once across all live pids in [0, grid) when + # grid <= NUM_PERSISTENT (each CTA does 1 tile; loop step >= total_work + # exits immediately) AND when grid == NUM_PERSISTENT (each CTA loops over + # multiple tiles). NUM_PERSISTENT stays a constexpr = cta_per_sm * num_sms. + _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M def launch_persistent_main(write_checkpoint: bool, n_writes_dev: torch.Tensor, @@ -4972,7 +4981,19 @@ def launch_persistent_main(write_checkpoint: bool, n_slots_for_kernel = host_n_writes if write_checkpoint else (batch - host_n_writes) if n_slots_for_kernel <= 0: return - grid = (num_persistent_arg,) + # Grid sizing: cap at min(full persistent grid, actual total_work). + # `n_slots` for this launch is `host_n_writes` (write half) / `batch - + # host_n_writes` (nowrite half) when host knows it (pure); else upper + # bound `batch` for mix scenarios where host can't read n_writes_dev + # without a sync. Upper-bound is fine — the kernel's runtime check + # only iterates actual work; the only cost of overcounting is a few + # extra CTAs. + if host_n_writes is not None: + _n_slots_for_launch = host_n_writes if write_checkpoint else (batch - host_n_writes) + else: + _n_slots_for_launch = batch + _total_work_launch = max(1, _n_slots_for_launch * _num_pid_m * nheads) + grid = (min(num_persistent_arg, _total_work_launch),) _persistent_main_kernel[grid]( state, state_tma_descriptor, state_scales_arg, old_x, old_B, old_dt, old_dA_cumsum, @@ -5041,7 +5062,11 @@ def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, # We still pass `n_writes_dev` (the same tensor the persistent_main # path uses) so the kernel signature is uniform; the value is # immaterial. - grid = (num_persistent_arg,) + # Grid sizing: cap at total_work (= batch * num_pid_m * nheads) for + # the dynamic case (full-batch coverage); see launch_persistent_main + # comment for correctness rationale. + _total_work_launch = max(1, batch * _num_pid_m * nheads) + grid = (min(num_persistent_arg, _total_work_launch),) _persistent_main_kernel[grid]( state, state_tma_descriptor, state_scales_arg, old_x, old_B, old_dt, old_dA_cumsum, From f0e805bd4fdb5135bd702f7bdcaa554d6ed9f01c Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 11 May 2026 15:04:54 -0700 Subject: [PATCH 30/89] bench + persistent kernel: scale compile-warmup parallelism MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three related changes that together let the persistent-kernel tuning sweep run efficiently: 1) ProcessPoolExecutor for compile-warmup (replaces ThreadPoolExecutor) The prior ThreadPoolExecutor hit GIL serialization in Triton's Python AST/codegen phase: with 28 threads in the pool, only ~4 were actually running at any moment ("4 R / 24 S" in /proc//task), limiting the parent process to ~1.2 cores worth of compile work and producing ~70 configs/min. Empirical test (2026-05-11) with external compile-only probes sharing the same TRITON_CACHE_DIR showed 6.6x speedup (134/min -> 890/min combined across two GPUs), confirming the GIL hypothesis. Replaces with ProcessPoolExecutor + spawn start method (avoids CUDA fork-corruption from active parent CUDA context). Hoists the nested _warm helper to module level (_warm_one_config) so it's picklable. Workers share the on-disk Triton cache via TRITON_CACHE_DIR; first to write any (kernel source x constexpr set) hash wins; concurrent same-hash writes are wasteful but not corrupting. baseline_fn is NOT passed to workers: avoids pickling C-extension function refs (flashinfer.mamba.selective_state_update). When --baseline is set, the baseline kernel compiles lazily in the parent during timing — one extra compile, negligible. 2) NUM_PERSISTENT runtime (was tl.constexpr) — drops cta_per_sm from the compile-time constexpr space NUM_PERSISTENT in _persistent_main_kernel is ONLY used as the stride in 'tl.range(pid, total_work, NUM_PERSISTENT, ...)' — i.e., the persistent-loop step. Work-item decomposition (pid_m, pid_b_local, pid_h) uses constexpr NUM_PID_M_BLOCKS and runtime n_slots_local, not NUM_PERSISTENT. Triton's loop-pipelining flags on tl.range (flatten=, num_stages=, warp_specialize=) need their OWN args constexpr, but the step itself can be runtime. Verified by empirical sweep: persistent kernel still works correctly with NUM_PERSISTENT as an int32 runtime arg. Effect: cta_per_sm (CPS) was multiplying the compile space by 8 (dense sweep) or 4 (sparse). Now CPS is purely a runtime/launch-grid knob; all CPS values share one compiled kernel. Persistent tuning sweeps go from O(128K) constexpr-combos per dtype to O(16K). 3) Task-level parallelism for compile-warmup (was outer-config-level) Prior _compile_warmup_phase enumerated only the OUTER axes (batch, mode, dtype, sr, rect, write_ckpt, sort_*, ...). Each worker received one outer config and serially iterated the INNER knob sweep (block_size_m x num_warps x num_stages x precompute_num_warps x heads_per_block x ...) within itself. For typical stage-A configs with 10 outer x 1728 inner knob combos, only 10 workers were ever concurrent (one per outer config), regardless of --compile-threads N value. --compile-threads 50 gave the same wall time as --compile-threads 10. Fix: enumerate (outer x inner) cartesian and submit each as a separate ProcessPoolExecutor task. Each worker handles ONE inner-knob-combo by cloning args with single-value sweep lists for that combo, then calling _bench_config(..., warmup_only=True) which iterates a 1x1x...x1 inner cartesian (= 1 compile). N workers now achieve true N-way concurrency. Dedupe non-compile axes: * batch: runtime int, doesn't affect kernel hash -> use only first batch in compile-warmup enumeration. 5x reduction at typical --batch-sizes 1,16,64,128,512. * cta_per_sm (CPS): runtime after change (2) above -> collapse to first value. 4-8x reduction. * prev_k: already a list passed to _bench_config (not enumerated). Shuffle task list before submitting: reduces the chance two workers pick adjacent (same-mode, neighboring-knob) tasks that hash to the same kernel. With shuffled order, workers' kernel-hash sets diverge quickly and same-hash races become rare. Together: stage-A compile-warmup task count drops from ~691K to ~34K (20x reduction); each remaining task now actually parallelizes 50-way instead of 10-way. Net 100x potential speedup on cold-cache full sweeps. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 11 +- ...benchmark_replay_selective_state_update.py | 177 +++++++++++++++--- 2 files changed, 158 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 6fdb72d8b6f0..f664166718e2 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -3990,7 +3990,16 @@ def _persistent_main_kernel( WRITE_CHECKPOINT: tl.constexpr, LAUNCH_DEPENDENT_KERNELS: tl.constexpr, USE_PERM: tl.constexpr, - NUM_PERSISTENT: tl.constexpr, + # NUM_PERSISTENT: runtime int (not constexpr). Used ONLY as the loop + # stride in `tl.range(pid, total_work, NUM_PERSISTENT, ...)`. Making it + # runtime collapses the cta_per_sm tuning dim from the kernel's compile + # signature: 8 CPS values used to mean 8x recompiles; now they share one + # compiled kernel. Work decomposition (pid_m, pid_b_local, pid_h) does + # NOT depend on NUM_PERSISTENT — it uses constexpr NUM_PID_M_BLOCKS and + # runtime n_slots_local — so loop unrolling and flatten=/num_stages=/ + # warp_specialize= optimizations on `tl.range` operate independently of + # the stride value. + NUM_PERSISTENT, NUM_LOOP_STAGES: tl.constexpr, NUM_PID_M_BLOCKS: tl.constexpr, FLATTEN: tl.constexpr, diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 0e6e7380c984..c9d4f8d0ecb4 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -932,20 +932,76 @@ def _time_kernel( # Per-config benchmark (consolidated baseline + replay) -def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, - baseline_fn, max_workers: int) -> None: - """Run each config once in parallel to compile + cache Triton kernels. +def _warm_one_config(args, cfg, baseline_fn) -> None: + """Module-level worker for the compile-warmup process pool. + + Module-level so ProcessPoolExecutor can pickle it (nested functions + aren't picklable). Each worker process holds its own GIL → no + serialization between concurrent compiles. + + ``cfg`` is a tuple of (outer_cfg, inner_overrides): + * outer_cfg = (batch, mtp_len, prev_ks, state_dtype, act_dtype, + sr_mode, rect, write_ckpt, mode, + sort_slots, reverse_nowrite, hardcode_sort) + * inner_overrides = dict of args attribute name -> single-value string + to clamp the per-cell inner-knob sweep to ONE + combination. Triggers exactly one Triton compile + per worker invocation, so N workers achieve + N-way concurrency regardless of outer config + count. (Prior design fanned out only at outer + granularity, capping concurrency at ~10 even + with --compile-threads 50.) + + ``baseline_fn`` is optional — when ``None``, only the checkpointing + kernel is warmed (the baseline-selection kernel can be warmed once in + the parent if needed). This lets us avoid pickling C-extension + function references across processes. + """ + outer_cfg, inner_overrides = cfg + (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, + rect, write_ckpt, mode, sort_slots, reverse_nowrite, hardcode_sort) = outer_cfg + # Clone args and override inner-knob sweep lists to single values. + # _bench_config then iterates a 1×1×...×1 cartesian inside. + import argparse as _ap + args_copy = _ap.Namespace(**vars(args)) + for k, v in inner_overrides.items(): + setattr(args_copy, k, v) + _bench_config( + args_copy, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, + sr_mode=sr_mode, rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, mode=mode, + sort_slots=sort_slots, reverse_nowrite=reverse_nowrite, + hardcode_sort=hardcode_sort, + warmup_only=True, + ) - Triton's ``compile()`` releases the GIL, so a ThreadPoolExecutor - fans out shape compilations across CPU cores in one shared CUDA - context. Compiled binaries land in Triton's on-disk cache (default - ``~/.triton/cache``) and the subsequent sequential measurement - phase loads them with no compile cost. - Parallel measurement would race for GPU time and skew numbers, so - only the warmup is parallelized; timing stays serial. +def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, + baseline_fn, max_workers: int) -> None: + """Parallel compile-warmup using a ProcessPoolExecutor with `spawn` + start method. + + Each worker process holds its own GIL and its own CUDA context, so + Triton compiles (Python AST/codegen + LLVM/ptxas) run truly in + parallel. Previous ThreadPoolExecutor design hit GIL contention + in the Python codegen phase, capping throughput at ~1-2 cores even + with 28 threads (observed: 4 R threads vs 28 in pool). + + Compiled binaries land in Triton's on-disk cache (TRITON_CACHE_DIR + or default ~/.triton/cache). Workers share the cache via filesystem + — first to write any given (kernel_source × constexpr_set) hash + wins; concurrent writes to the SAME hash are wasteful but not + corrupting. + + spawn start method avoids inheriting parent CUDA state (which is + unsafe after fork on Linux with active CUDA contexts). Per-worker + import + CUDA init costs ~10s, amortized over each worker's many + compiles. baseline_fn is intentionally NOT passed to workers to + avoid pickling complications; the parent compiles the baseline + kernel itself before launching the pool when applicable. """ - from concurrent.futures import ThreadPoolExecutor + from concurrent.futures import ProcessPoolExecutor + import multiprocessing sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) @@ -956,8 +1012,17 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp rev_list = getattr(args, "reverse_nowrite_list", [False]) hsort_list = getattr(args, "hardcode_sort_list", [False]) + # Compile-warmup task enumeration: outer × inner cartesian. + # CRITICAL: only enumerate axes that change the kernel's COMPILE signature. + # Drop runtime axes (batch, prev_k) that produce identical kernel hashes — + # otherwise we'd pay ~50-100ms of bench setup per redundant cache-hit task. + # + # Batches collapsed to first only: batch is a runtime int passed to the + # kernel, not a constexpr; all batches share the same compiled kernel. + # prev_k is already a list passed into _bench_config (not enumerated here). configs = [] - for batch in batch_sizes: + _compile_batches = batch_sizes[:1] # collapse runtime axis + for batch in _compile_batches: for mtp_len in mtp_lengths: prev_ks = _resolve_prev_ks(args, mtp_len) for state_dtype in state_dtypes: @@ -1014,29 +1079,83 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp hardcode_sort, )) - print(f"[compile-warmup] {len(configs)} configs across {max_workers} threads") + # Enumerate inner-knob cartesian — same axes _bench_config iterates + # internally. Each (outer × inner) tuple becomes one task; workers + # then trigger exactly one Triton compile per task, giving true + # N-way concurrency with --compile-threads N. + def _ps(val): + if val is None or (isinstance(val, str) and not val): + return [None] + if isinstance(val, str): + return [v.strip() for v in val.split(",") if v.strip()] + return [val] + + # CPS (cta_per_sm) collapsed to first value: NUM_PERSISTENT is runtime now, + # so different CPS values share the same compiled kernel. Collapsing here + # avoids enumerating 4-8x redundant tasks that would each pay ~50-100ms + # bench setup overhead for a cache hit on the same kernel hash. + _cps_for_compile = _ps(args.cta_per_sm)[:1] + knob_axes = [ + ("block_size_m", _ps(args.block_size_m)), + ("num_warps", _ps(args.num_warps)), + ("num_stages", _ps(args.num_stages)), + ("precompute_num_warps", _ps(args.precompute_num_warps)), + ("precompute_num_stages", _ps(args.precompute_num_stages)), + ("heads_per_block", _ps(args.heads_per_block)), + ("maxnreg", _ps(args.maxnreg)), + ("num_ctas", _ps(args.num_ctas)), + ("cta_per_sm", _cps_for_compile), # collapsed (runtime) + ("num_loop_stages", _ps(args.num_loop_stages)), + ("flatten", _ps(args.flatten)), + ("warp_specialize", _ps(args.warp_specialize)), + ("use_tma_rect_load", _ps(args.use_tma_rect_load)), + ("use_tma_replay_write_load", _ps(args.use_tma_replay_write_load)), + ("use_tma_replay_nowrite_load", _ps(args.use_tma_replay_nowrite_load)), + ("use_tma_replay_write_store", _ps(args.use_tma_replay_write_store)), + ] + import itertools as _it + inner_combos = list(_it.product(*(values for _, values in knob_axes))) + + # Cross product outer × inner. Override only knobs that have an + # explicit value (skip None — those leave args. at its CLI default, + # which _bench_config handles via its own _parse_sweep). + tasks = [] + for outer in configs: + for inner_tuple in inner_combos: + inner_overrides = { + name: str(val) + for (name, _), val in zip(knob_axes, inner_tuple) + if val is not None + } + tasks.append((outer, inner_overrides)) + + # Shuffle to reduce cross-worker race on the same kernel hash. Two + # workers picking adjacent tasks (same mode, neighboring knob value) + # could both miss + compile the same kernel hash; shuffling spreads + # workloads across different kernel hash families. + import random as _r + _r.shuffle(tasks) + + print(f"[compile-warmup] {len(tasks)} compile tasks " + f"({len(configs)} outer × {len(inner_combos)} inner combos) " + f"across {max_workers} processes (ProcessPoolExecutor, spawn start)") t0 = time.perf_counter() - def _warm(cfg): - (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, - rect, write_ckpt, mode, sort_slots, reverse_nowrite, hardcode_sort) = cfg - _bench_config( - args, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, - sr_mode=sr_mode, rectangle_for_nowrite=rect, - write_checkpoint=write_ckpt, mode=mode, - sort_slots=sort_slots, reverse_nowrite=reverse_nowrite, - hardcode_sort=hardcode_sort, - warmup_only=True, - ) - + ctx = multiprocessing.get_context("spawn") errors = [] - with ThreadPoolExecutor(max_workers=max_workers) as ex: - futures = [ex.submit(_warm, cfg) for cfg in configs] - for cfg, fut in zip(configs, futures): + with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as ex: + # baseline_fn=None: workers compile only the checkpointing kernel. + # Baseline kernels (if any) get compiled lazily in the parent during + # the timing phase — usually just one extra compile, negligible. + futures = { + ex.submit(_warm_one_config, args, task, None): task + for task in tasks + } + for fut in futures: try: fut.result() except Exception as e: - errors.append((cfg, e)) + errors.append((futures[fut], e)) if errors: for cfg, e in errors: From bef49814198bb170fbea7e9ebf7aa56715edd870 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 11 May 2026 17:02:27 -0700 Subject: [PATCH 31/89] bench: cache + grow-in-place batch-scaled tensors Per-cell tensor allocation was the headline host-side cost in the timing phase (~10-30ms per cell on b=512). Bench's _bench_config called _build_tensors which freshly allocates ~22 tensors (state, old_x, old_B, old_dt, old_dA_cumsum, x, dt, B, C, prev_tokens, out_incr, out_base, intermediate_states_buffer, xbc_input, conv_state, etc) at the cell's batch size on every call. For a 5-batch sweep (1, 16, 64, 128, 512), each cell triggered a full re-alloc. Add a module-level _TENSOR_CACHE keyed by all fixed dimensions (state_dtype, act_dtype, max_window, mtp_len, nheads, head_dim, d_state, ngroups). Within a key, the batch dim grows in place: - request_batch <= cached max_batch: return slices [:request_batch] of cached tensors (zero alloc cost). - request_batch > cached max_batch: re-alloc at new batch, store as new cache entry, return slices. Old buffers released to PyTorch's caching allocator (typically reused for the new alloc on same device). Pre-warm the cache at the largest requested batch up front in _run_benchmark before the timing loop starts. Without pre-warm, the loop would trigger multiple growth allocations as it encountered progressively-larger batches. Pre-warm at max(batch_sizes) makes every subsequent cell a pure view-slice. Tensors that don't scale with batch (A, dt_bias, D, conv_weight, conv_bias, d_inner, conv_dim) are stored as-is and returned by reference. Reset state (state_work = state0.copy_) still runs per cell in the caller; cached state0 is purely a read-only reference whose contents stay fixed once allocated. This is correct for reproducibility and removes any contention. Expected impact at 5-batch / 6-mode / 1728-inner-combo sweep: - Before: ~50-70ms per cell (tensor alloc dominant) - After (pre-warm + cache): per-cell timing should drop closer to 20-30ms, with the remainder being CUPTI flush + Python kernel dispatch + CUDA graph capture overhead. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 142 +++++++++++++++--- 1 file changed, 124 insertions(+), 18 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index c9d4f8d0ecb4..f713335e6b1a 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -300,6 +300,23 @@ def _resolve_prev_ks(args, mtp_len: int) -> list[int]: # Tensor construction helpers +# Module-level cache for tensor buffers shared across cells. Keyed by all +# the "fixed" dimensions (state_dtype, act_dtype, max_window, mtp_len, +# nheads, head_dim, d_state, ngroups). Within a key, the batch dim grows +# in place: if a new cell requests a batch <= cached max_batch, we return +# views (slices) of the existing tensors; if batch > cached max_batch, we +# realloc at the new batch (which becomes the new max). Tensors never shrink. +# +# Rationale: torch.randn/zeros for these tensor shapes at b=512 takes +# ~10-30ms per call. At ~895 cells/min with 5 different batch sizes, +# we were re-allocating every cell. Caching saves the bulk of that per-cell +# overhead, raising GPU util in the timing phase. +# +# Reset state lives in caller (state_work = state0.copy_), so cached state0 +# is purely a reference whose contents stay fixed once allocated. This is +# fine: it's only read by the reset path. +_TENSOR_CACHE: dict = {} + def _build_tensors( batch: int, @@ -328,6 +345,49 @@ def _build_tensors( """ device = "cuda" + # Cache lookup — grow batch in place if needed; else return views. + cache_key = (state_dtype, act_dtype, max_window, mtp_len, + nheads, head_dim, d_state, ngroups) + cached = _TENSOR_CACHE.get(cache_key) + if cached is not None and cached["max_batch"] >= batch: + # Hit — return slices for current batch. + b = batch + return ( + cached["state0"][:b], + cached["state_scales0"][:b] if cached["state_scales0"] is not None else None, + cached["old_x"][:b], + cached["old_B"][:b], + cached["old_dt"][:b], + cached["old_dA_cumsum"][:b], + cached["cache_buf_idx"][:b], + cached["x"][:b], + cached["dt"][:b], + cached["B"][:b], + cached["C"][:b], + cached["A"], + cached["dt_bias"], + cached["D"], + cached["prev_tokens"][:b], + cached["slot_perm_buf"][:b], + cached["out_incr"][:b], + cached["out_base"][:b], + cached["intermediate_states_buffer"][:b], + cached["xbc_input"][:b], + cached["conv_state"][:b], + cached["conv_weight"], + cached["conv_bias"], + cached["d_inner"], + cached["conv_dim"], + ) + + # Miss or grow. Allocate at new max_batch (existing data, if any, is + # released — caller code re-fills via reset paths anyway). Rebind + # `batch` locally to alloc_batch so the existing allocation code below + # uses the larger size; keep request_batch for the final slice. + request_batch = batch + alloc_batch = batch if cached is None else max(batch, cached["max_batch"]) + batch = alloc_batch + torch.manual_seed(42) # --- SSM parameters (float32, tie_hdim strides) --- @@ -434,28 +494,58 @@ def _build_tensors( # conv_bias: (conv_dim,) — parameter conv_bias = torch.randn(conv_dim, device=device, dtype=act_dtype) + # Store full-batch buffers in cache and return slices at request_batch. + _TENSOR_CACHE[cache_key] = { + "max_batch": alloc_batch, + "state0": state0, + "state_scales0": state_scales0, + "old_x": old_x, + "old_B": old_B, + "old_dt": old_dt, + "old_dA_cumsum": old_dA_cumsum, + "cache_buf_idx": cache_buf_idx, + "x": x, + "dt": dt, + "B": B, + "C": C, + "A": A, + "dt_bias": dt_bias, + "D": D, + "prev_tokens": prev_tokens, + "slot_perm_buf": slot_perm_buf, + "out_incr": out_incr, + "out_base": out_base, + "intermediate_states_buffer": intermediate_states_buffer, + "xbc_input": xbc_input, + "conv_state": conv_state, + "conv_weight": conv_weight, + "conv_bias": conv_bias, + "d_inner": d_inner, + "conv_dim": conv_dim, + } + rb = request_batch return ( - state0, - state_scales0, - old_x, - old_B, - old_dt, - old_dA_cumsum, - cache_buf_idx, - x, - dt, - B, - C, + state0[:rb], + state_scales0[:rb] if state_scales0 is not None else None, + old_x[:rb], + old_B[:rb], + old_dt[:rb], + old_dA_cumsum[:rb], + cache_buf_idx[:rb], + x[:rb], + dt[:rb], + B[:rb], + C[:rb], A, dt_bias, D, - prev_tokens, - slot_perm_buf, - out_incr, - out_base, - intermediate_states_buffer, - xbc_input, - conv_state, + prev_tokens[:rb], + slot_perm_buf[:rb], + out_incr[:rb], + out_base[:rb], + intermediate_states_buffer[:rb], + xbc_input[:rb], + conv_state[:rb], conv_weight, conv_bias, d_inner, @@ -2037,6 +2127,22 @@ def _run_benchmark(args) -> None: baseline_fn, max_workers=args.compile_threads, ) + # Pre-warm the per-(state_dtype, act_dtype, mtp_len, ...) tensor cache at + # the largest requested batch size. Without this, the timing loop would + # progressively grow the cache as it encounters larger batches (e.g., + # iterate 1 -> 16 -> 64 -> 128 -> 512 = 5 separate growth allocations, + # each freeing the previous buffers). Pre-warming at max-batch up front + # makes every subsequent timing cell a view-slice (zero alloc cost). + _max_batch = max(batch_sizes) + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + for mtp_len in mtp_lengths: + _build_tensors( + _max_batch, mtp_len, state_dtype, act_dtype, + args.tp_nheads, args.head_dim, args.d_state, args.tp_ngroups, + max_window=getattr(args, "max_window", None) or None, + ) + if args.profile: torch.cuda.cudart().cudaProfilerStart() From 658485b7394bd504a143a1aa2bde804af887acad Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 11 May 2026 20:19:21 -0700 Subject: [PATCH 32/89] bench: crash-safe JSONL incremental persistence + resume + per-main knob split MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit bench (benchmark_replay_selective_state_update.py): - --json-output now also writes .jsonl in append-only JSONL. Each cell appends one record on completion (open-per-write to avoid pickling file handles to ProcessPoolExecutor workers). flush() per line for crash-safe persistence. - Resume: on startup, read .jsonl if present, populate _json_results + _done_keys skip set, skip cells already there. Tolerates partial last line from a crash mid-write via json.JSONDecodeError pass. - Cross-host guard: each JSONL record carries a "host" field; resume only uses records from the current hostname. Mismatches warn + ignore. - --cupti-retry N: inline retry on CUPTI capture mismatch. Cell skipped to sidecar after exhausting retries. Default 1. - --skipped-output / --retry-cells: sidecar JSON of cells that still failed + filter to re-time only those tags on a follow-up invocation. - Per-main knob split: --block-size-m-{write,nowrite}, --num-warps-*, --num-stages-*, --cta-per-sm-*, --num-loop-stages-* let the two main kernel launches (write half vs nowrite half) take independent values. Default: tied to the shared knob (backward compat). --skip-diagonal drops the tied subset for incremental sweeps that extend earlier tied-knob results. wrapper (checkpointing_state_update.py): - _block_size_m_{write,nowrite} + _num_warps_*/_num_stages_*/_cta_per_sm_*/ _num_loop_stages_* args. None = tied to shared (backward compat). - launch_replay_main / launch_rectangle_main / launch_persistent_main use the WC-appropriate knob via local closure over BLOCK_SIZE_M_WRITE / BLOCK_SIZE_M_NOWRITE etc. Precompute knobs (heads_per_block, precompute_num_warps) intentionally not split — shared precompute wins. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 97 +- ...benchmark_replay_selective_state_update.py | 639 +++- ...selective_state_update.py.tunings_proposal | 3253 +++++++++++++++++ 3 files changed, 3878 insertions(+), 111 deletions(-) create mode 100644 tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py.tunings_proposal diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index f664166718e2..b1de20f5f103 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -4240,6 +4240,23 @@ def checkpointing_state_update( _heads_per_block: int | None = None, _maxnreg: int | None = None, _num_ctas: int | None = None, + # Per-main knobs (override shared values for one half of the dl-family / + # persistent_main launches). Default None = tied to the shared value + # (backward compat). The two main kernels (write vs nowrite) have + # different per-slot work — write does a state shift + store, nowrite + # just appends — so the optimum (M, W, S, H) can differ. Precompute + # knobs are intentionally NOT split: shared precompute wins (cheaper + # launch, hotter precompute outputs in L2). Persistent CPS / LS knobs + # are also split per-main since the two persistent_main launches have + # different grid sizes. + _block_size_m_write: int | None = None, + _block_size_m_nowrite: int | None = None, + _num_warps_write: int | None = None, + _num_warps_nowrite: int | None = None, + _num_stages_write: int | None = None, + _num_stages_nowrite: int | None = None, + # Note: heads_per_block / precompute_num_warps are NOT split — they only + # affect the precompute kernel, which is shared across write/nowrite. # TMA state-tensor toggles — 4 independent paths (see CHECKPOINTING_DESIGN.md # item #17 for measured perf profiles). Each is False=raw load/store, True= # use a host-built TMA tensor_descriptor for that path. @@ -4279,6 +4296,14 @@ def checkpointing_state_update( _num_loop_stages: int | None = None, _flatten: bool | None = None, _warp_specialize: bool | None = None, + # Per-main persistent-specific knobs. Same rationale as the BLOCK_SIZE_M + # split above: the two persistent_main launches (write half vs nowrite + # half) have different grid sizes and per-work-item costs, so they may + # want different cta_per_sm / num_loop_stages. + _cta_per_sm_write: int | None = None, + _cta_per_sm_nowrite: int | None = None, + _num_loop_stages_write: int | None = None, + _num_loop_stages_nowrite: int | None = None, ): """ Replay SSM state update with precomputed CB and tl.dot fast-forward. @@ -4633,6 +4658,21 @@ def checkpointing_state_update( if _precompute_num_warps is not None: precompute_num_warps = _precompute_num_warps + # Per-main knob resolution: each _*_{write,nowrite} arg, if not None, + # overrides the corresponding shared value for ONE main launch only. + # Default (None) = tied to shared value (current behavior). + BLOCK_SIZE_M_WRITE = _block_size_m_write if _block_size_m_write is not None else BLOCK_SIZE_M + BLOCK_SIZE_M_NOWRITE = _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M + NUM_WARPS_WRITE = _num_warps_write if _num_warps_write is not None else num_warps + NUM_WARPS_NOWRITE = _num_warps_nowrite if _num_warps_nowrite is not None else num_warps + NUM_STAGES_WRITE = _num_stages_write if _num_stages_write is not None else _num_stages + NUM_STAGES_NOWRITE = _num_stages_nowrite if _num_stages_nowrite is not None else _num_stages + # Persistent-only per-main: + CTA_PER_SM_WRITE = _cta_per_sm_write if _cta_per_sm_write is not None else _cta_per_sm + CTA_PER_SM_NOWRITE = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _cta_per_sm + NUM_LOOP_STAGES_WRITE = _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages + NUM_LOOP_STAGES_NOWRITE = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages + HAS_CACHE_BATCH_INDICES = state_batch_indices is not None assert nheads % heads_per_block == 0, ( @@ -4815,7 +4855,15 @@ def launch_dynamic_precompute(rectangle: bool): def launch_replay_main(write_checkpoint: bool, early_out: bool, launch_dependent_kernels: bool = False, reverse_perm: bool = False): - _checkpointing_main_kernel[main_grid]( + # Per-main knob selection: write vs nowrite branches use independent + # M / num_warps / num_stages / heads_per_block values. Grid is + # M-dependent so it must be a closure over the selected M. + _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE + _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE + _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE + def _main_grid_local(META, _bsm=_bsm): + return (triton.cdiv(dim, _bsm), batch, nheads) + _checkpointing_main_kernel[_main_grid_local]( state, state_tma_descriptor, state_scales_arg, old_x, old_B, old_dt, old_dA_cumsum, prev_num_accepted_tokens, cache_buf_idx, @@ -4840,7 +4888,7 @@ def launch_replay_main(write_checkpoint: bool, early_out: bool, cb_scaled.stride(0), cb_scaled.stride(1), cb_scaled.stride(2), cb_scaled.stride(3), decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - BLOCK_SIZE_M, + _bsm, LAUNCH_WITH_PDL=use_internal_pdl, PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, QUANT_MAX=quant_max, @@ -4854,8 +4902,8 @@ def launch_replay_main(write_checkpoint: bool, early_out: bool, USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), USE_TMA_LOAD_NOWRITE=bool(_use_tma_replay_nowrite_load and not write_checkpoint), USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), - num_warps=num_warps, - **({"num_stages": _num_stages} if _num_stages else {}), + num_warps=_nw, + **({"num_stages": _ns} if _ns else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), **({"maxnreg": _maxnreg} if _maxnreg else {}), launch_pdl=use_internal_pdl, @@ -4864,7 +4912,13 @@ def launch_replay_main(write_checkpoint: bool, early_out: bool, def launch_rectangle_main(early_out: bool, launch_dependent_kernels: bool = False, reverse_perm: bool = False): - _rectangle_main_kernel[main_grid]( + # Rectangle is the nowrite-side path; use the nowrite-main knobs. + _bsm = BLOCK_SIZE_M_NOWRITE + _nw = NUM_WARPS_NOWRITE + _ns = NUM_STAGES_NOWRITE + def _main_grid_local(META, _bsm=_bsm): + return (triton.cdiv(dim, _bsm), batch, nheads) + _rectangle_main_kernel[_main_grid_local]( state, state_tma_descriptor, state_scales_arg, old_x, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, @@ -4882,7 +4936,7 @@ def launch_rectangle_main(early_out: bool, cb_scaled.stride(0), cb_scaled.stride(1), cb_scaled.stride(2), cb_scaled.stride(3), decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - BLOCK_SIZE_M, + _bsm, LAUNCH_WITH_PDL=use_internal_pdl, QUANT_MAX=quant_max, EARLY_OUT=early_out, @@ -4890,8 +4944,8 @@ def launch_rectangle_main(early_out: bool, USE_PERM=use_perm, REVERSE_PERM=reverse_perm, USE_TMA_LOAD=bool(_use_tma_rect_load), - num_warps=num_warps, - **({"num_stages": _num_stages} if _num_stages else {}), + num_warps=_nw, + **({"num_stages": _ns} if _ns else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), **({"maxnreg": _maxnreg} if _maxnreg else {}), launch_pdl=use_internal_pdl, @@ -4990,6 +5044,19 @@ def launch_persistent_main(write_checkpoint: bool, n_slots_for_kernel = host_n_writes if write_checkpoint else (batch - host_n_writes) if n_slots_for_kernel <= 0: return + # Per-main knob selection. The two persistent_main launches (write + # half vs nowrite half) get independent BLOCK_SIZE_M / num_warps / + # num_stages / cta_per_sm / num_loop_stages. See the per-main args + # block in the wrapper signature. + _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE + _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE + _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE + _cps = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE + _cps = _cps if _cps else 1 + _nls = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE + _nls = _nls if _nls else 2 + _num_persistent = _cps * _num_sms + _num_pid_m_local = (dim + _bsm - 1) // _bsm # Grid sizing: cap at min(full persistent grid, actual total_work). # `n_slots` for this launch is `host_n_writes` (write half) / `batch - # host_n_writes` (nowrite half) when host knows it (pure); else upper @@ -5001,8 +5068,8 @@ def launch_persistent_main(write_checkpoint: bool, _n_slots_for_launch = host_n_writes if write_checkpoint else (batch - host_n_writes) else: _n_slots_for_launch = batch - _total_work_launch = max(1, _n_slots_for_launch * _num_pid_m * nheads) - grid = (min(num_persistent_arg, _total_work_launch),) + _total_work_launch = max(1, _n_slots_for_launch * _num_pid_m_local * nheads) + grid = (min(_num_persistent, _total_work_launch),) _persistent_main_kernel[grid]( state, state_tma_descriptor, state_scales_arg, old_x, old_B, old_dt, old_dA_cumsum, @@ -5029,15 +5096,15 @@ def launch_persistent_main(write_checkpoint: bool, cb_scaled.stride(0), cb_scaled.stride(1), cb_scaled.stride(2), cb_scaled.stride(3), decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - BLOCK_SIZE_M, + _bsm, LAUNCH_WITH_PDL=use_internal_pdl, PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, QUANT_MAX=quant_max, WRITE_CHECKPOINT=write_checkpoint, LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, USE_PERM=use_perm, - NUM_PERSISTENT=num_persistent_arg, - NUM_LOOP_STAGES=num_loop_stages_arg, + NUM_PERSISTENT=_num_persistent, + NUM_LOOP_STAGES=_nls, FLATTEN=flatten_arg, WARP_SPECIALIZE=warp_specialize_arg, IS_DYNAMIC=False, @@ -5053,8 +5120,8 @@ def launch_persistent_main(write_checkpoint: bool, and not write_checkpoint ), USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), - num_warps=num_warps, - **({"num_stages": _num_stages} if _num_stages else {}), + num_warps=_nw, + **({"num_stages": _ns} if _ns else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), **({"maxnreg": _maxnreg} if _maxnreg else {}), launch_pdl=use_internal_pdl, diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index f713335e6b1a..3baeb1247ae2 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -120,6 +120,7 @@ import argparse import importlib import itertools +import json import os import statistics import sys @@ -655,9 +656,14 @@ class CuptiKernelTimer: timer = CuptiKernelTimer.get() timer.start() # arms; drops any stale records - records = timer.stop() # flush; list of tuples per kernel - # (name, start_ns, end_ns, corr, - # graph_id, graph_node_id, stream) + records, zero_ts_count, zero_ts_names = timer.stop() + # records: list of tuples per kernel + # (name, start_ns, end_ns, corr, + # graph_id, graph_node_id, stream) + # zero_ts_count: kernel records CUPTI + # delivered with start=0 or end=0 + # (couldn't timestamp); zero_ts_names: + # name → count breakdown. The callback fires from a CUPTI worker thread, so a lock guards the record buffer. Records are kept tiny (tuple of ints + str) to minimize @@ -688,6 +694,8 @@ def _init(cls, _c) -> "CuptiKernelTimer": self = object.__new__(cls) self._c = _c self._records: list[tuple] = [] + self._zero_ts_count = 0 # how many kernel records were dropped due to start/end==0 + self._zero_ts_names: dict = {} # name -> count of zero-ts drops (for diagnostics) self._lock = threading.Lock() # CUPTI callback contract (from cupti-python-samples/cupti_common.py): @@ -702,11 +710,18 @@ def _buf_req(): def _buf_done(activities): recs = [] + zero_drops_local = 0 + zero_names_local: dict = {} for a in activities: if a.kind not in kernel_kinds: continue # start/end == 0 means CUPTI couldn't time this kernel. + # Track these instead of silently dropping — they're a sign + # CUPTI is failing to record kernels we DID launch. if a.start == 0 or a.end == 0: + zero_drops_local += 1 + name = getattr(a, "name", "?") + zero_names_local[name] = zero_names_local.get(name, 0) + 1 continue recs.append(( a.name, @@ -717,9 +732,14 @@ def _buf_done(activities): int(a.graph_node_id), int(a.stream_id), )) - if recs: + if recs or zero_drops_local: with self._lock: - self._records.extend(recs) + if recs: + self._records.extend(recs) + if zero_drops_local: + self._zero_ts_count += zero_drops_local + for k, v in zero_names_local.items(): + self._zero_ts_names[k] = self._zero_ts_names.get(k, 0) + v # Hold strong refs so the C side never sees GC'd Python callables. self._buf_req = _buf_req @@ -734,12 +754,19 @@ def start(self) -> None: self._c.activity_flush_all(1) with self._lock: self._records.clear() + self._zero_ts_count = 0 + self._zero_ts_names = {} - def stop(self) -> list[tuple]: - """Flush and return all kernel records since the last start().""" + def stop(self) -> tuple[list[tuple], int, dict]: + """Flush and return all kernel records + count of zero-timestamp drops. + + Returns (records, zero_ts_count, zero_ts_names_dict). The latter two + are diagnostic: nonzero values mean CUPTI delivered records with + start=0 or end=0, indicating it failed to timestamp the kernel. + """ self._c.activity_flush_all(1) with self._lock: - return list(self._records) + return (list(self._records), self._zero_ts_count, dict(self._zero_ts_names)) # ============================================================================= @@ -758,7 +785,9 @@ def _stats_from_spans(spans_us: list[float]) -> dict: } -def _stats_from_cupti_records(records, warmup, iters, tag, expected_K): +def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, + zero_ts_count: int = 0, + zero_ts_names: dict | None = None): """Bin a flat CUPTI kernel record stream into per-iter spans + per-kernel relative timestamps. Used by both graph and eager CUPTI paths. @@ -782,14 +811,43 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K): if total != expected_total: from collections import Counter name_counts = dict(Counter(r[0] for r in records)) - raise RuntimeError( - f"CUPTI capture mismatch for {tag!r}: expected {expected_K} " - f"kernels/iter × {expected_iters} iters (warmup+iters) = " - f"{expected_total} records, got {total}. Kernel record counts: " - f"{name_counts}. If a kernel name is missing, add a substring " - f"to _CUPTI_KEEP_KERNEL_SUBSTRINGS; if present but the count is " - f"wrong, adjust _kernels_per_iter_* for this mode." + # Non-fatal: skip this cell instead of killing the whole sweep. + # Mismatch may be a CUPTI dropped-records issue (rare configs), + # not necessarily a K-table bug. Log so the user can investigate + # the specific cell post-hoc; return None so the caller can skip + # writing a JSON row. + zero_msg = "" + if zero_ts_count: + zero_msg = ( + f" + {zero_ts_count} records with start/end=0 " + f"(dropped by callback, breakdown {zero_ts_names}). " + f"Total observed kernel records (timed + zero-ts) = " + f"{total + zero_ts_count} / {expected_total}." + ) + print( + f"[WARN] CUPTI capture mismatch for {tag!r}: expected " + f"{expected_K} kernels/iter × {expected_iters} iters " + f"(warmup+iters) = {expected_total} records, got {total}. " + f"Kernel record counts: {name_counts}.{zero_msg} SKIPPING cell.", + file=sys.stderr, ) + # Per-record dump: (name, start_ns_rel, end_ns_rel, corr_id, graph_id, stream_id). + # Times relative to first record so absolute ns isn't drowning output. + # Limit dump to first 30 records to avoid flooding logs at high K. + if records: + t0_ns = records[0][1] + for i, r in enumerate(records[:30]): + # r = (name, start_ns, end_ns, corr_id, graph_id, graph_node_id, stream_id) + rel_start = (r[1] - t0_ns) / 1000.0 # us + rel_end = (r[2] - t0_ns) / 1000.0 + print( + f" rec[{i:3d}] name={r[0]!r} start={rel_start:.2f}us " + f"end={rel_end:.2f}us corr={r[3]} graph={r[4]} stream={r[6]}", + file=sys.stderr, + ) + if len(records) > 30: + print(f" ... ({len(records) - 30} more records elided)", file=sys.stderr) + return None K = expected_K timed = records[warmup * K:] @@ -893,9 +951,11 @@ def _time_kernel_cuda_graph( g.replay() torch.cuda.synchronize() torch.cuda.nvtx.range_pop() - records = timer.stop() + records, zero_ts_count, zero_ts_names = timer.stop() - return _stats_from_cupti_records(records, warmup, iters, tag, expected_K) + return _stats_from_cupti_records(records, warmup, iters, tag, expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names) def _time_kernel_eager( @@ -930,9 +990,11 @@ def _time_kernel_eager( run_fn() torch.cuda.synchronize() torch.cuda.nvtx.range_pop() - records = timer.stop() + records, zero_ts_count, zero_ts_names = timer.stop() - return _stats_from_cupti_records(records, warmup, iters, tag, expected_K) + return _stats_from_cupti_records(records, warmup, iters, tag, expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names) def _run_kernel_untimed(args, run_fn, reset_fn, tag: str) -> dict: @@ -1508,19 +1570,22 @@ def _run_baseline(): expected_K=_kernels_per_iter_baseline(with_conv1d), ) - _print_row( - show_kernel_col, - args.baseline, - batch, - mtp_len, - "N/A", - state_dtype_name, - act_dtype_name, - stats, - json_results=getattr(args, "_json_results", None), - tp_size=args.tp_size, - json_detailed=getattr(args, "json_detailed", False), - ) + if stats is not None: + _print_row( + show_kernel_col, + args.baseline, + batch, + mtp_len, + "N/A", + state_dtype_name, + act_dtype_name, + stats, + json_results=getattr(args, "_json_results", None), + tp_size=args.tp_size, + json_detailed=getattr(args, "json_detailed", False), + jsonl_path=getattr(args, "_jsonl_path", None), + jsonl_host=getattr(args, "_jsonl_host", None), + ) # --- Sweep parameter parsing (invariant across prev_k) --- def _parse_sweep(val): @@ -1541,6 +1606,32 @@ def _parse_sweep(val): num_loop_stages_values = _parse_sweep(args.num_loop_stages) flatten_values = _parse_sweep(args.flatten) warp_specialize_values = _parse_sweep(args.warp_specialize) + # Per-main split-knob sweeps. Default = same as the shared sweep (so each + # combo is tied). When set independently, the inner loop sweeps the + # cross-product (write × nowrite); --skip-diagonal drops the tied subset. + def _split_or_share(split_csv, shared_values): + return _parse_sweep(split_csv) if split_csv else shared_values + block_size_m_write_values = _split_or_share(args.block_size_m_write, block_size_m_values) + block_size_m_nowrite_values = _split_or_share(args.block_size_m_nowrite, block_size_m_values) + num_warps_write_values = _split_or_share(args.num_warps_write, num_warps_values) + num_warps_nowrite_values = _split_or_share(args.num_warps_nowrite, num_warps_values) + num_stages_write_values = _split_or_share(args.num_stages_write, num_stages_values) + num_stages_nowrite_values = _split_or_share(args.num_stages_nowrite, num_stages_values) + cta_per_sm_write_values = _split_or_share(args.cta_per_sm_write, cta_per_sm_values) + cta_per_sm_nowrite_values = _split_or_share(args.cta_per_sm_nowrite, cta_per_sm_values) + num_loop_stages_write_values = _split_or_share(args.num_loop_stages_write, num_loop_stages_values) + num_loop_stages_nowrite_values = _split_or_share(args.num_loop_stages_nowrite, num_loop_stages_values) + # Whether any *_write / *_nowrite knob was independently set — used by + # --skip-diagonal to know if the cross-product is non-trivial. Without + # any split, the per-main values == shared values and skip-diagonal is + # a no-op (which is correct). + _any_split = any(getattr(args, name) for name in ( + "block_size_m_write", "block_size_m_nowrite", + "num_warps_write", "num_warps_nowrite", + "num_stages_write", "num_stages_nowrite", + "cta_per_sm_write", "cta_per_sm_nowrite", + "num_loop_stages_write", "num_loop_stages_nowrite", + )) # TMA toggles — independent 0/1 sweep per path. The skip-dupe at the # top of the inner loop body collapses cells where a flag's path is # unreachable, so e.g. monolithic + WC=True only runs the value=0 @@ -1602,10 +1693,14 @@ def _parse_sweep(val): # allocate a sentinel scratch so the wrapper API is uniform. n_writes_samples_gpu = None n_writes_dev_mix = None + # n_writes per iter = number of slots that overflow the window. + # Computed for ALL mix scenarios so the JSON output (--json-detailed) + # can pair each iter's span_us with its mix composition for downstream + # analysis (group iters by # writes → per-bucket median → analytic + # expectation under the steady-state PNAT distribution). + n_writes_per_iter_all = ((src + mtp_len) > max_window).sum(axis=1).astype(np.int32) if mode in ("persistent_main", "persistent_dynamic"): - # n_writes per iter = number of slots that overflow the window. - n_writes_per_iter = ((src + mtp_len) > max_window).sum(axis=1).astype(np.int32) - n_writes_samples_gpu = torch.from_numpy(n_writes_per_iter).to( + n_writes_samples_gpu = torch.from_numpy(n_writes_per_iter_all).to( device=device, dtype=torch.int32 ) n_writes_dev_mix = torch.zeros(1, dtype=torch.int32, device=device) @@ -1650,6 +1745,10 @@ def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): # Pass through to _run_incr so the wrapper receives _n_writes_dev # (mix scenarios) instead of _n_writes (pure scenarios). "n_writes_dev": n_writes_dev_mix, + # Full per-iter n_writes array (size = warmup + iters). Used by + # the JSON-detailed output to pair each iter's span with its + # mix composition for post-hoc bucketing analysis. + "n_writes_per_iter": n_writes_per_iter_all, }) # Pure scenarios don't pre-allocate n_writes_dev; mix scenarios do. @@ -1666,41 +1765,99 @@ def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): scenario_iters = scn.get("iters") # None => use args.iters tag = f"incr_b{batch}_mtp{mtp_len}_{scn['label']}_s{state_dtype_name}_a{act_dtype_name}" + # Iteration over per-cell knob combos. + # When NO per-main split is requested (_any_split=False), each row in + # the cross-product gives the same value to both write_main and + # nowrite_main (current behavior — backward-compat). When ANY split + # IS requested, we iterate the write and nowrite axes independently + # (cross-product blowup is the user's responsibility — they typically + # pair this with --skip-diagonal to drop the tied subset). + if _any_split: + _iter_axes = ( + block_size_m_write_values, block_size_m_nowrite_values, + num_warps_write_values, num_warps_nowrite_values, + num_stages_write_values, num_stages_nowrite_values, + precompute_num_warps_values, + precompute_num_stages_values, + heads_per_block_values, + maxnreg_values, num_ctas_values, + cta_per_sm_write_values, cta_per_sm_nowrite_values, + num_loop_stages_write_values, num_loop_stages_nowrite_values, + flatten_values, warp_specialize_values, + use_tma_rect_load_values, + use_tma_replay_write_load_values, + use_tma_replay_nowrite_load_values, + use_tma_replay_write_store_values, + ) + else: + # Tied: one value per shared knob. Wrap in single-element list for + # uniform iteration; the body sets w/nw both to the shared value. + _iter_axes = ( + block_size_m_values, [None], + num_warps_values, [None], + num_stages_values, [None], + precompute_num_warps_values, + precompute_num_stages_values, + heads_per_block_values, + maxnreg_values, num_ctas_values, + cta_per_sm_values, [None], + num_loop_stages_values, [None], + flatten_values, warp_specialize_values, + use_tma_rect_load_values, + use_tma_replay_write_load_values, + use_tma_replay_nowrite_load_values, + use_tma_replay_write_store_values, + ) for ( - block_size_m, - num_warps, - num_stages, + block_size_m_w, + block_size_m_nw, + num_warps_w, + num_warps_nw, + num_stages_w, + num_stages_nw, precompute_num_warps, precompute_num_stages, heads_per_block, maxnreg, num_ctas, - cta_per_sm, - num_loop_stages, + cta_per_sm_w, + cta_per_sm_nw, + num_loop_stages_w, + num_loop_stages_nw, flatten, warp_specialize, use_tma_rect_load, use_tma_replay_write_load, use_tma_replay_nowrite_load, use_tma_replay_write_store, - ) in itertools.product( - block_size_m_values, - num_warps_values, - num_stages_values, - precompute_num_warps_values, - precompute_num_stages_values, - heads_per_block_values, - maxnreg_values, - num_ctas_values, - cta_per_sm_values, - num_loop_stages_values, - flatten_values, - warp_specialize_values, - use_tma_rect_load_values, - use_tma_replay_write_load_values, - use_tma_replay_nowrite_load_values, - use_tma_replay_write_store_values, - ): + ) in itertools.product(*_iter_axes): + # When tied, _nw values were placeholder None; fill from _w (the + # shared value). When split, _w and _nw came from independent lists. + if not _any_split: + block_size_m_nw = block_size_m_w + num_warps_nw = num_warps_w + num_stages_nw = num_stages_w + cta_per_sm_nw = cta_per_sm_w + num_loop_stages_nw = num_loop_stages_w + # Skip-diagonal: when split is on, drop the tied subset (same as a + # prior shared-knob sweep would cover). + if _any_split and args.skip_diagonal and ( + block_size_m_w == block_size_m_nw and + num_warps_w == num_warps_nw and + num_stages_w == num_stages_nw and + cta_per_sm_w == cta_per_sm_nw and + num_loop_stages_w == num_loop_stages_nw + ): + continue + # Backward-compat aliases used by the existing body below. When + # tied, these are simply the shared value. When split, the + # _write copy is used for sweep_tag and grouping (a stable choice + # so the tag is unique per (write, nowrite) combo). + block_size_m = block_size_m_w + num_warps = num_warps_w + num_stages = num_stages_w + cta_per_sm = cta_per_sm_w + num_loop_stages = num_loop_stages_w # Skip-dupe for TMA flag sweeps: a flag whose code path isn't # reachable in this cell produces identical timing for value=0 # and value=1. We canonicalize by skipping value=1 cells when @@ -1899,16 +2056,37 @@ def _run_incr( _heads_per_block=heads_per_block, _maxnreg=maxnreg, _num_ctas=num_ctas, + # Per-main overrides (None = tied to shared above; explicit + # only when the inner loop is iterating split axes). + _block_size_m_write=block_size_m_w if _any_split else None, + _block_size_m_nowrite=block_size_m_nw if _any_split else None, + _num_warps_write=num_warps_w if _any_split else None, + _num_warps_nowrite=num_warps_nw if _any_split else None, + _num_stages_write=num_stages_w if _any_split else None, + _num_stages_nowrite=num_stages_nw if _any_split else None, + _cta_per_sm_write=cta_per_sm_w if _any_split else None, + _cta_per_sm_nowrite=cta_per_sm_nw if _any_split else None, + _num_loop_stages_write=num_loop_stages_w if _any_split else None, + _num_loop_stages_nowrite=num_loop_stages_nw if _any_split else None, **extra_kwargs, ) parts = [] - if block_size_m is not None: - parts.append(f"M={block_size_m}") - if num_warps is not None: - parts.append(f"W={num_warps}") - if num_stages is not None: - parts.append(f"S={num_stages}") + # When tied (not _any_split), emit the shared single-value tag + # (M=8 etc). When split, emit explicit Mw / Mnw tags so cells + # with the same shared value but different per-main values get + # unique JSON keys. + def _emit_split(name_w, name_nw, val_w, val_nw): + if val_w is None and val_nw is None: + return + if not _any_split or val_w == val_nw: + parts.append(f"{name_w[:-1]}={val_w}") # strip the 'w' suffix + else: + parts.append(f"{name_w}={val_w}") + parts.append(f"{name_nw}={val_nw}") + _emit_split("Mw", "Mnw", block_size_m_w, block_size_m_nw) + _emit_split("Ww", "Wnw", num_warps_w, num_warps_nw) + _emit_split("Sw", "Snw", num_stages_w, num_stages_nw) if precompute_num_warps is not None: parts.append(f"pW={precompute_num_warps}") if precompute_num_stages is not None: @@ -1922,10 +2100,8 @@ def _run_incr( # Persistent-only knobs (only meaningful when MODE=persistent_main; # printed unconditionally so output rows are uniformly comparable # across modes when the user passed these sweeps). - if cta_per_sm is not None: - parts.append(f"CPS={cta_per_sm}") - if num_loop_stages is not None: - parts.append(f"LS={num_loop_stages}") + _emit_split("CPSw", "CPSnw", cta_per_sm_w, cta_per_sm_nw) + _emit_split("LSw", "LSnw", num_loop_stages_w, num_loop_stages_nw) if flatten is not None: parts.append(f"FL={flatten}") if warp_specialize is not None: @@ -1958,6 +2134,26 @@ def _run_incr( sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") reset_fn = _reset_conv1d_realistic if with_conv1d else _reset + # --retry-cells filter: only time cells whose tag is in the + # retry set. Cheaper than re-enumerating cells in the orchestrator. + retry_set = getattr(args, "_retry_cells_set", None) + if retry_set and sweep_tag not in retry_set: + continue + # Resume from JSONL: skip cells already recorded. Built the same + # way _print_row builds JSON keys; must stay in sync. + done_keys = getattr(args, "_done_keys", None) + if done_keys: + # One key per scenario (k=6, k=11, mix) — skip the whole cell + # only if ALL of its scenarios are already done. We don't + # know which scenarios will be emitted here without + # re-evaluating the inner scenario loop; conservatively skip + # only when the prev_k_for_print's specific key is done. + _resume_key = _build_json_key( + args.variant, batch, mtp_len, prev_k_for_print, + state_dtype_name, sweep_suffix, args.tp_size, + ) + if _resume_key in done_keys: + continue if warmup_only: reset_fn() if scenario_pre_iter is not None: @@ -1965,30 +2161,71 @@ def _run_incr( _run_incr() torch.cuda.synchronize() else: - stats = _time_kernel( - args, _run_incr, reset_fn, sweep_tag, - expected_K=_kernels_per_iter_incremental( - mode, with_conv1d=with_conv1d, - persistent_skip_empty=scenario_skip_empty, - ), - pre_iter_fn=scenario_pre_iter, - iters_override=scenario_iters, - ) - - _print_row( - show_kernel_col, - args.variant, - batch, - mtp_len, - prev_k_for_print, - state_dtype_name, - act_dtype_name, - stats, - sweep_suffix, - json_results=getattr(args, "_json_results", None), - tp_size=args.tp_size, - json_detailed=getattr(args, "json_detailed", False), - ) + # Inline retry: CUPTI sometimes loses records under PDL + + # high cell count; retrying the SAME cell often catches it + # because the failure is transient at the kernel-launch level. + # Per --cupti-retry budget. On final failure, append tag to + # the skipped list for an external rerun in a fresh process. + retry_budget = max(0, getattr(args, "cupti_retry", 1)) + stats = None + for attempt in range(retry_budget + 1): + stats = _time_kernel( + args, _run_incr, reset_fn, sweep_tag, + expected_K=_kernels_per_iter_incremental( + mode, with_conv1d=with_conv1d, + persistent_skip_empty=scenario_skip_empty, + ), + pre_iter_fn=scenario_pre_iter, + iters_override=scenario_iters, + ) + if stats is not None: + break + if attempt < retry_budget: + print( + f"[retry] CUPTI mismatch on {sweep_tag!r}; " + f"retrying ({attempt + 1}/{retry_budget})", + file=sys.stderr, + ) + if stats is None: + args._skipped_cells.append(sweep_tag) + + # Attach n_writes_per_iter for --json-detailed bucketing. + # For pure scenarios, n_writes is constant: 0 (nowrite) or + # batch (write), determined by scn["fill"] + mtp_len > max_window. + # For mix, scn carries the precomputed per-iter array. + if stats is not None and getattr(args, "json_detailed", False): + eff_iters = scenario_iters if scenario_iters is not None else args.iters + if scn["fill"] is not None: + # Pure scenario: constant n_writes for every iter. + is_write = (scn["fill"] + mtp_len > max_window) + per_iter_nw = [batch if is_write else 0] * eff_iters + else: + # Mix scenario: slice off warmup, keep timed iters. + nw_full = scn.get("n_writes_per_iter") + if nw_full is not None: + per_iter_nw = nw_full[args.warmup:args.warmup + eff_iters].tolist() + else: + per_iter_nw = None + if per_iter_nw is not None: + stats["n_writes_per_iter"] = per_iter_nw + + if stats is not None: + _print_row( + show_kernel_col, + args.variant, + batch, + mtp_len, + prev_k_for_print, + state_dtype_name, + act_dtype_name, + stats, + sweep_suffix, + json_results=getattr(args, "_json_results", None), + tp_size=args.tp_size, + json_detailed=getattr(args, "json_detailed", False), + jsonl_path=getattr(args, "_jsonl_path", None), + jsonl_host=getattr(args, "_jsonl_host", None), + ) # Map full torch dtype name → short tag used in JSON keys (matches collect.py). @@ -2045,6 +2282,8 @@ def _print_row( json_results=None, tp_size=None, json_detailed=False, + jsonl_path=None, + jsonl_host=None, ): """Print one summary row and optionally accumulate stats for JSON output. @@ -2052,6 +2291,12 @@ def _print_row( [per_kernel]}. The summary table only shows the headline percentiles. JSON output captures median/p95/p99/n by default; with json_detailed=True it also captures the per-iter and per-kernel data. + + When `jsonl_path` is provided, also appends one JSON line per row to the + JSONL sidecar (crash-safe incremental persistence; lets a killed sweep + resume from the last completed cell on rerun). Open-per-write because + `args` is pickled to ProcessPoolExecutor workers and file handles aren't + picklable. """ kernel_col = f"{kernel_name:>11} | " if show_kernel_col else "" print( @@ -2066,12 +2311,25 @@ def _print_row( sweep_suffix, tp_size, ) if json_detailed: - json_results[key] = stats + row_stats = stats else: - json_results[key] = { + row_stats = { k: stats[k] for k in ("median", "p95", "p99", "n") if k in stats } + json_results[key] = row_stats + # Append to JSONL sidecar if a path is set (incremental persistence). + # Open per-write because args is pickled to ProcessPoolExecutor + # workers, and file handles aren't picklable. A clean SIGTERM or + # Python exception will leave the file consistent up to the last + # newline; catastrophic kills can leave a partial last line, which + # the resume reader tolerates via json.JSONDecodeError pass. + if jsonl_path is not None: + rec = {"key": key, "stats": row_stats} + if jsonl_host is not None: + rec["host"] = jsonl_host + with open(jsonl_path, "a") as f: + f.write(json.dumps(rec) + "\n") # Main benchmark loop @@ -2082,6 +2340,94 @@ def _run_benchmark(args) -> None: # Stash on args so we don't need to thread a dict through every helper. args._json_results = {} if getattr(args, "json_output", None) else None + # JSONL incremental sidecar. Path = `.jsonl`. Each completed + # cell appends one line `{"key": , "stats": {...}}` to this file + # as it finishes timing. On startup we read this sidecar (if present) and + # populate _json_results + _done_keys so a killed bench can resume without + # redoing already-timed cells. Crash-safe by construction: append-only + # writes survive SIGTERM/SIGKILL/reboot mid-sweep. + # Note: we store only paths/strings on `args` because args is pickled to + # ProcessPoolExecutor workers during compile-warmup, and file handles + # (TextIOWrapper) aren't picklable. _print_row open-appends per cell. + args._jsonl_path = None + args._done_keys: set[str] = set() + args._jsonl_host = None # hostname stamp for the current run + if getattr(args, "json_output", None): + import socket + args._jsonl_host = socket.gethostname() + args._jsonl_path = args.json_output + ".jsonl" + # Read existing JSONL if present; build skip set IFF hostname matches. + # Cross-node timings aren't directly comparable (CPU/GPU clock/topology + # differ), so resuming on a new host would mix incompatible numbers. + # If the JSONL was recorded on a different host, log + skip resume so + # the user can decide (rename / archive the old file). + if os.path.exists(args._jsonl_path): + n_loaded = 0 + n_skipped_host = 0 + recorded_hosts = set() + with open(args._jsonl_path) as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except json.JSONDecodeError: + # Tolerate partial last line from a crash mid-write. + continue + rec_host = rec.get("host") + if rec_host is not None: + recorded_hosts.add(rec_host) + if rec_host is not None and rec_host != args._jsonl_host: + n_skipped_host += 1 + continue + k = rec.get("key") + if k is None: + continue + args._json_results[k] = rec.get("stats", {}) + args._done_keys.add(k) + n_loaded += 1 + if n_loaded: + print( + f"[resume] {args._jsonl_path}: loaded {n_loaded} prior " + f"cell results from host={args._jsonl_host}; sweep will " + f"skip them.", + file=sys.stderr, + ) + if n_skipped_host: + print( + f"[resume] WARNING: {args._jsonl_path} also contains " + f"{n_skipped_host} records from other hosts " + f"{sorted(recorded_hosts - {args._jsonl_host})!r}; ignoring " + f"them. If you want a clean run, remove or archive the " + f"jsonl file first.", + file=sys.stderr, + ) + + # Skipped cells accumulator — populated by _bench_config when CUPTI capture + # mismatch causes a cell to be skipped. Written to args.skipped_output + # (or derived from json_output) at end of run. Pair with --retry-cells + # to re-time only the skipped cells in a fresh process. + args._skipped_cells = [] + + # Retry-cells filter set. When non-empty, only cells whose tag is in this + # set will be timed; all others are silently skipped. Cells are matched + # against the sweep_tag string built in _bench_config. + args._retry_cells_set: set[str] = set() + if getattr(args, "retry_cells", None): + with open(args.retry_cells) as f: + data = json.load(f) + # Accept either a list of tag strings, or the same shape we write + # (dict with "skipped" key). Tolerant of both for hand-edited files. + if isinstance(data, dict) and "skipped" in data: + data = data["skipped"] + args._retry_cells_set = set(data) + print( + f"[retry] --retry-cells loaded {len(args._retry_cells_set)} tags " + f"from {args.retry_cells}; sweep will skip all other cells.", + file=sys.stderr, + ) + assert args.nheads % args.tp_size == 0, ( f"nheads ({args.nheads}) must be divisible by tp_size ({args.tp_size})" ) @@ -2321,7 +2667,6 @@ def _run_benchmark(args) -> None: torch.cuda.cudart().cudaProfilerStop() if args.json_output and args._json_results is not None: - import json payload = { "metadata": { "timestamp": datetime.now().isoformat(), @@ -2341,6 +2686,32 @@ def _run_benchmark(args) -> None: print(f"\nJSON results written to: {args.json_output} " f"({len(args._json_results)} entries)") + # Write the skipped-cells sidecar. Pair with --retry-cells in a separate + # invocation to re-time the failed cells in a fresh process. + skipped_path = getattr(args, "skipped_output", None) + if skipped_path is None and args.json_output: + # Derive default: foo.json -> foo.skipped.json + skipped_path = args.json_output.rsplit(".", 1)[0] + ".skipped.json" + if skipped_path is not None and args._skipped_cells: + payload = { + "metadata": { + "timestamp": datetime.now().isoformat(), + "cmd": " ".join(sys.argv), + "skipped_count": len(args._skipped_cells), + }, + "skipped": args._skipped_cells, + } + tmp = skipped_path + ".tmp" + with open(tmp, "w") as f: + json.dump(payload, f, indent=2) + os.replace(tmp, skipped_path) + print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " + f"tags written to: {skipped_path}", file=sys.stderr) + elif args._skipped_cells: + # No output path but there are skipped cells — emit a stderr summary. + print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " + f"first 5: {args._skipped_cells[:5]}", file=sys.stderr) + # CLI @@ -2452,6 +2823,33 @@ def _parse_args() -> argparse.Namespace: "each kernel) — useful for PDL overlap analysis but adds ~4 KB/cell. " "Default off keeps records to ~40 bytes (median/p95/p99/n only).", ) + parser.add_argument( + "--cupti-retry", + type=int, + default=1, + help="On CUPTI capture mismatch (kernel record count != expected), " + "retry the cell this many times in-process before giving up. CUPTI " + "gets racy after thousands of cells in one process (PDL + small " + "kernels occasionally lose records); a single retry usually catches " + "transient cases. Set 0 to disable and skip on first mismatch.", + ) + parser.add_argument( + "--skipped-output", + default=None, + help="Path to write the list of cells that failed CUPTI capture even " + "after --cupti-retry retries (JSON list of sweep_tag strings). " + "Default: derived from --json-output by replacing .json with " + ".skipped.json. Pair with --retry-cells in a separate invocation " + "(fresh process = fresh CUPTI subscriber) to re-time these cells.", + ) + parser.add_argument( + "--retry-cells", + default=None, + help="Path to a JSON list of cell tags (as written by --skipped-output " + "in a prior invocation). When set, the sweep iterates as normal but " + "skips any cell whose tag is NOT in the listed set. Lets collect.py " + "drive a retry pass over only the cells that failed the first time.", + ) parser.add_argument( "--prev-tokens-fracs", default="0,0.5,1.0", @@ -2501,6 +2899,55 @@ def _parse_args() -> argparse.Namespace: default=None, help="Override num_stages for the main kernel (comma-separated sweep).", ) + parser.add_argument( + "--block-size-m-write", type=str, default=None, + help="Sweep BLOCK_SIZE_M for the WRITE main only (overrides --block-size-m " + "for the write half). Tied to --block-size-m if unset.", + ) + parser.add_argument( + "--block-size-m-nowrite", type=str, default=None, + help="Sweep BLOCK_SIZE_M for the NOWRITE main only. Tied to --block-size-m if unset.", + ) + parser.add_argument( + "--num-warps-write", type=str, default=None, + help="Sweep num_warps for the WRITE main only. Tied to --num-warps if unset.", + ) + parser.add_argument( + "--num-warps-nowrite", type=str, default=None, + help="Sweep num_warps for the NOWRITE main only. Tied to --num-warps if unset.", + ) + parser.add_argument( + "--num-stages-write", type=str, default=None, + help="Sweep num_stages for the WRITE main only. Tied to --num-stages if unset.", + ) + parser.add_argument( + "--num-stages-nowrite", type=str, default=None, + help="Sweep num_stages for the NOWRITE main only. Tied to --num-stages if unset.", + ) + parser.add_argument( + "--cta-per-sm-write", type=str, default=None, + help="Sweep cta_per_sm for the WRITE persistent_main only. Tied to --cta-per-sm if unset.", + ) + parser.add_argument( + "--cta-per-sm-nowrite", type=str, default=None, + help="Sweep cta_per_sm for the NOWRITE persistent_main only. Tied to --cta-per-sm if unset.", + ) + parser.add_argument( + "--num-loop-stages-write", type=str, default=None, + help="Sweep num_loop_stages for the WRITE persistent_main only. Tied to --num-loop-stages if unset.", + ) + parser.add_argument( + "--num-loop-stages-nowrite", type=str, default=None, + help="Sweep num_loop_stages for the NOWRITE persistent_main only. Tied to --num-loop-stages if unset.", + ) + parser.add_argument( + "--skip-diagonal", action=argparse.BooleanOptionalAction, default=False, + help="When sweeping any per-main *_write / *_nowrite knobs, skip cells " + "where ALL splittable knobs satisfy write_value == nowrite_value (i.e. " + "the 'diagonal' that's already covered by a prior shared-knob sweep). " + "Useful for incremental sweeps that extend earlier results without redoing " + "the tied-knob cells.", + ) parser.add_argument( "--precompute-num-warps", type=str, diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py.tunings_proposal b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py.tunings_proposal new file mode 100644 index 000000000000..b59b0af822e4 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py.tunings_proposal @@ -0,0 +1,3253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Standalone benchmark for replay_selective_state_update (Triton kernel). + +Suitable for nsight-compute (ncu) and nsight-systems (nsys) capture. + +Fixed model config: NVIDIA-Nemotron-3-Super-120B-A12B at TP=8 + nheads=16, head_dim=64, d_state=128, ngroups=1 + +mtp_len is the per-request sequence length processed by replay: in MTP it +equals num_draft_tokens + 1 target token, so --mtp-lengths 6 models 5 drafts ++ 1 target. + +Baseline kernel (--baseline [triton|flashinfer]): + Calls selective_state_update with T=mtp_len tokens and disable_state_update=True, + matching the MTP scoring pass in mamba2_mixer.py exactly. + +Timing methodology +================== + +All in-bench timing comes from CUPTI's Activity API (1 ns kernel +timestamps from the GPU profiling fabric). cudaEvent.elapsed_time() was +removed — its ~0.5 us resolution overshoots CUPTI by ~50% on short kernels +in graphs, and we have no other use for it here. See the CUPTI block +lower in this file for the timer source. + +Three modes: + + --cupti --cuda-graph (default) + Capture one CUDA graph per cell (warmup + timed iters inlined), + replay once, read kernel start/end from CUPTI. ~20× faster than + nsys-wrapped capture and matches it to within ~1% / noise floor. + + --cupti --no-cuda-graph + Eager loop with CUPTI. Per-kernel timestamps are still accurate, but + the per-iter SPAN (max(end) - min(start)) now includes the Python + launch latency BETWEEN consecutive kernels in run_fn (~100 µs on + Hopper/Blackwell). Graph capture and PDL hide that latency; eager + mode honestly reports it. For per-kernel timing in eager mode, look + at per_kernel.start_us/end_us in --json-detailed output rather than + the span percentiles. Useful when graph capture is undesirable. + + --no-cupti (with or without --cuda-graph) + No in-bench timing — just runs the kernels for an external profiler + (nsys / ncu) to time. In-process CUPTI conflicts with nsys's own + subscriber, so disable ours when wrapping in nsys. Bench output + reports zeros for median/p95/p99; trust the external trace. + +JSON output schema (--json-output PATH) +======================================= + +Designed to be parsed by collect.py / report.py without touching sqlite or +NVTX traces. Future agents: prefer reading this JSON over re-running nsys. + + { + "metadata": {timestamp, cmd, tp_size, warmup, iters, variant, cupti}, + "results": { + "": {median, p95, p99, n, [iters_us], [per_kernel]} + } + } + +Key format mirrors collect.py's kernel_data.json convention: + incremental/{batch}/{mtp}/{sd}/k{prev_k}/{sweep_parts}/tp{tp} + triton/{batch}/{mtp}/{sd}/tp{tp} + flashinfer/{batch}/{mtp}/{sd}/tp{tp} + + - is normalized: bf16 / fp16 / fp32 / int8 / int16 / fp8. + - is e.g. "M16_W1_S3_SR0_RECT0_WC1" — flags concatenated by + underscore in canonical (M, W, S, pW, pS, H, R, CT, SR, RECT, WC) order. + - All numeric values in microseconds (us). + +Per-record fields: + - median, p95, p99: span statistics (us). Span = max(kernel_end_ns) - + min(kernel_start_ns) across the iter's kernels — same convention as + nsys-derived collect.py used to use. + - n: number of timed iters that contributed. + - iters_us: list of length n, raw per-iter spans (only with --json-detailed). + - per_kernel: {: {start_us: [...], end_us: [...]}} where + timestamps are RELATIVE to that iter's first kernel start, in us. Lets + you see PDL overlap directly without an external profiler. Only with + --json-detailed. + +Example usage: + # Basic sweep (default = --cupti, just summary stats) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 1,2,4 --mtp-lengths 1,4,8 --warmup 5 --iters 20 + + # JSON output, summary stats only (compact) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 16 --mtp-lengths 6 --json-output /tmp/out.json + + # JSON output, full per-iter / per-kernel data (for PDL analysis etc.) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 16 --mtp-lengths 6 \\ + --json-output /tmp/out.json --json-detailed + + # nsys capture (--no-cupti so our subscriber doesn't conflict) + nsys profile --capture-range=cudaProfilerApi \\ + python benchmark_replay_selective_state_update.py --profile --no-cupti + + # ncu capture (--no-cupti --no-cuda-graph: each kernel replayable solo) + ncu --target-processes all \\ + python benchmark_replay_selective_state_update.py --profile \\ + --no-cupti --no-cuda-graph \\ + --batch-sizes 1 --mtp-lengths 4 --warmup 5 --iters 5 +""" + +import argparse +import importlib +import itertools +import json +import os +import statistics +import sys +import time +from datetime import datetime +from pathlib import Path + +import numpy as np +import torch +from einops import repeat + + +def _import_mamba_kernels_fast(): + """Load kernel modules directly (~40s faster than a full tensorrt_llm init). + Use --full-import as the fallback if module dependencies change. + + Stub the parent packages (tensorrt_llm, tensorrt_llm._torch, + tensorrt_llm._torch.modules) in sys.modules with __path__ set, but do NOT + execute their __init__.py. Then load the leaf kernel modules. When a + kernel body imports e.g. tensorrt_llm._utils.get_sm_version, Python's + machinery resolves it against our stub's __path__ and loads only _utils.py + — skipping the heavy tensorrt_llm package init. + """ + import types + + repo_root = Path(__file__).resolve().parents[5] + trtllm_dir = repo_root / "tensorrt_llm" + mamba_pkg = "tensorrt_llm._torch.modules.mamba" + mamba_dir = trtllm_dir / "_torch" / "modules" / "mamba" + + def _stub_pkg(fqn: str, pkg_dir: Path): + if fqn in sys.modules: + return + stub = types.ModuleType(fqn) + stub.__path__ = [str(pkg_dir)] + sys.modules[fqn] = stub + + _stub_pkg("tensorrt_llm", trtllm_dir) + _stub_pkg("tensorrt_llm._torch", trtllm_dir / "_torch") + _stub_pkg("tensorrt_llm._torch.modules", trtllm_dir / "_torch" / "modules") + + def _load(mod_name: str, file_name: str): + fqn = f"{mamba_pkg}.{mod_name}" if mod_name else mamba_pkg + if fqn in sys.modules: + return sys.modules[fqn] + spec = importlib.util.spec_from_file_location( + fqn, + mamba_dir / file_name, + submodule_search_locations=[str(mamba_dir)] if file_name == "__init__.py" else [], + ) + mod = importlib.util.module_from_spec(spec) + sys.modules[fqn] = mod + spec.loader.exec_module(mod) + return mod + + # 1. Package __init__ (defines PAD_SLOT_ID = -1) + _load("", "__init__.py") + # 2. softplus helper (used by both kernel modules) + _load("softplus", "softplus.py") + # 3. The actual kernels + replay_mod = _load("replay_selective_state_update", "replay_selective_state_update.py") + checkpoint_mod = _load("checkpointing_state_update", "checkpointing_state_update.py") + base_mod = _load("selective_state_update", "selective_state_update.py") + conv1d_mod = _load("causal_conv1d_triton", "causal_conv1d_triton.py") + + return ( + replay_mod.replay_selective_state_update, + checkpoint_mod.checkpointing_state_update, + base_mod.selective_state_update, + conv1d_mod.causal_conv1d_update, + ) + + +def _import_mamba_kernels_full(): + """Import via the standard tensorrt_llm package (slow but safe).""" + from tensorrt_llm._torch.modules.mamba.causal_conv1d_triton import causal_conv1d_update + from tensorrt_llm._torch.modules.mamba.checkpointing_state_update import ( + checkpointing_state_update, + ) + from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( + replay_selective_state_update, + ) + from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update + + return ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) + + +# Use fast import by default; --full-import parsed later but we need the +# functions at module level. Check sys.argv early. +if "--full-import" in sys.argv: + ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) = _import_mamba_kernels_full() +else: + try: + ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) = _import_mamba_kernels_fast() + except Exception as e: # noqa: BLE001 - exit loudly; don't hide a fast-import regression + print( + f"ERROR: fast import failed ({type(e).__name__}: {e})\n" + "Re-run with --full-import for the slow but stable path, " + "then file a bug or fix _import_mamba_kernels_fast.", + file=sys.stderr, + ) + sys.exit(1) + + +_VARIANT_FNS = { + "replay": lambda: replay_selective_state_update, + "checkpointing": lambda: checkpointing_state_update, +} + +# Model config defaults (Nemotron-3-Super-120B full model). +# --tp-size divides nheads and ngroups to get the per-GPU slice. +# TP=1: nheads=128, ngroups=8 +# TP=4: nheads=32, ngroups=2 +# TP=8: nheads=16, ngroups=1 (default) +NHEADS = 128 +HEAD_DIM = 64 +D_STATE = 128 +NGROUPS = 8 +TP_SIZE = 8 # default; overridden by --tp-size + +# L2 flush buffer: ~128 MB — larger than L2 on A100/H100/B200 +_L2_FLUSH_SIZE = 32 * 1024 * 1024 # float32 elements → 128 MB +_l2_flush: torch.Tensor | None = None + + +def _init_l2_flush() -> None: + global _l2_flush + _l2_flush = torch.empty(_L2_FLUSH_SIZE, dtype=torch.float32, device="cuda") + + +def _flush_l2() -> None: + """Evict L2 by writing to a large buffer then synchronising.""" + assert _l2_flush is not None + _l2_flush.fill_(0.0) + torch.cuda.synchronize() + + +def _resolve_prev_ks(args, mtp_len: int) -> list[int]: + """Resolve prev_k values for one mtp_len cell. + + Two input modes (mutually exclusive in spirit; absolute wins if both given): + --prev-tokens-int "0,10,11,16" → use literal integers, clamped to + [0, max_window] (where max_window is the cache T-axis capacity). + --prev-tokens-fracs "0,0.5,1.0" → fractions of mtp_len, clamped to + [0, mtp_len] (current behavior). + + For replay-style checkpointing the cache holds up to max_window old + tokens, so absolute integers are the right knob. Fractions are kept + for back-compat with prior placeholder runs. + """ + upper = getattr(args, "max_window", 0) or mtp_len + if getattr(args, "prev_tokens_int", None): + return sorted(set(max(0, min(upper, int(v))) for v in args.prev_tokens_int)) + return sorted( + set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs) + ) + + +# Tensor construction helpers + +# Module-level cache for tensor buffers shared across cells. Keyed by all +# the "fixed" dimensions (state_dtype, act_dtype, max_window, mtp_len, +# nheads, head_dim, d_state, ngroups). Within a key, the batch dim grows +# in place: if a new cell requests a batch <= cached max_batch, we return +# views (slices) of the existing tensors; if batch > cached max_batch, we +# realloc at the new batch (which becomes the new max). Tensors never shrink. +# +# Rationale: torch.randn/zeros for these tensor shapes at b=512 takes +# ~10-30ms per call. At ~895 cells/min with 5 different batch sizes, +# we were re-allocating every cell. Caching saves the bulk of that per-cell +# overhead, raising GPU util in the timing phase. +# +# Reset state lives in caller (state_work = state0.copy_), so cached state0 +# is purely a reference whose contents stay fixed once allocated. This is +# fine: it's only read by the reset path. +_TENSOR_CACHE: dict = {} + + +def _build_tensors( + batch: int, + mtp_len: int, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + nheads: int, + head_dim: int, + d_state: int, + ngroups: int, + max_window: int | None = None, +): + """ + Build all tensors for one benchmark configuration. + + nheads/ngroups are already TP-split (i.e. full_nheads // tp_size). + + Returns: + state0 : (batch, nheads, head_dim, d_state) – initial SSM state + x, dt, B, C : (batch, mtp_len, ...) – token inputs for both kernels + A, dt_bias, D : SSM parameters (float32, tie_hdim strides) + prev_tokens : (batch,) + out_incr : pre-allocated output for replay kernel (batch, mtp_len, nheads, head_dim) + out_base : pre-allocated output for baseline kernel (batch, mtp_len, nheads, head_dim) + intermediate_states_buffer: for baseline kernel (batch, mtp_len, nheads, head_dim, d_state) + """ + device = "cuda" + + # Cache lookup — grow batch in place if needed; else return views. + cache_key = (state_dtype, act_dtype, max_window, mtp_len, + nheads, head_dim, d_state, ngroups) + cached = _TENSOR_CACHE.get(cache_key) + if cached is not None and cached["max_batch"] >= batch: + # Hit — return slices for current batch. + b = batch + return ( + cached["state0"][:b], + cached["state_scales0"][:b] if cached["state_scales0"] is not None else None, + cached["old_x"][:b], + cached["old_B"][:b], + cached["old_dt"][:b], + cached["old_dA_cumsum"][:b], + cached["cache_buf_idx"][:b], + cached["x"][:b], + cached["dt"][:b], + cached["B"][:b], + cached["C"][:b], + cached["A"], + cached["dt_bias"], + cached["D"], + cached["prev_tokens"][:b], + cached["slot_perm_buf"][:b], + cached["out_incr"][:b], + cached["out_base"][:b], + cached["intermediate_states_buffer"][:b], + cached["xbc_input"][:b], + cached["conv_state"][:b], + cached["conv_weight"], + cached["conv_bias"], + cached["d_inner"], + cached["conv_dim"], + ) + + # Miss or grow. Allocate at new max_batch (existing data, if any, is + # released — caller code re-fills via reset paths anyway). Rebind + # `batch` locally to alloc_batch so the existing allocation code below + # uses the larger size; keep request_batch for the final slice. + request_batch = batch + alloc_batch = batch if cached is None else max(batch, cached["max_batch"]) + batch = alloc_batch + + torch.manual_seed(42) + + # --- SSM parameters (float32, tie_hdim strides) --- + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) # stride(-1)=0, stride(-2)=0 + + dt_bias_base = torch.randn(nheads, device=device, dtype=torch.float32) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) # stride(-1)=0 + + D_base = torch.randn(nheads, device=device, dtype=torch.float32) + D = repeat(D_base, "h -> h p", p=head_dim) + + # --- SSM state --- + # Quantized dtypes need their own initializer (torch.randn doesn't accept + # int) and a parallel fp32 scales tensor (per-(head, dim) channel decode + # scale, broadcast over dstate). Quant state is filled with realistic- + # range values via fp32 → quant; scales are derived consistently so the + # initial state isn't garbage on dequant. + _QUANT_BENCH = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, + } + if state_dtype in _QUANT_BENCH: + quant_max = _QUANT_BENCH[state_dtype] + state_fp32 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + amax = state_fp32.abs().amax(dim=-1) # (batch, nheads, head_dim) + encode_scale = quant_max / amax.clamp(min=1e-30) + state_scales0 = (1.0 / encode_scale).to(torch.float32) # decode scale + scaled = state_fp32 * encode_scale.unsqueeze(-1) + if state_dtype == torch.float8_e4m3fn: + state0 = scaled.clamp(-quant_max, quant_max).to(state_dtype) + else: + state0 = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) + else: + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state_scales0 = None + + # --- Cache tensors for replay kernel --- + # max_window is the cache T-axis capacity; defaults to mtp_len (the + # placeholder/degenerate case where every step is a checkpoint step). + # For real replay-style checkpointing, max_window > mtp_len. + cache_T = max_window if max_window is not None else mtp_len + # old_x: single-buffered (cache, max_window, nheads, dim) + old_x = torch.randn(batch, cache_T, nheads, head_dim, device=device, dtype=act_dtype) + # old_B: double-buffered (cache, 2, max_window, ngroups, dstate) + old_B = torch.randn(batch, 2, cache_T, ngroups, d_state, device=device, dtype=act_dtype) + # old_dt: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous + old_dt = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) + # old_dA_cumsum: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous + old_dA_cumsum = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) + # cache_buf_idx: which buffer to read (0 or 1) + cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) + + # --- Token inputs (used by both replay and baseline kernels) --- + x = torch.randn(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) + # dt must match D's dtype (fp32) for flashinfer — force it for all paths. + dt_base = torch.randn(batch, mtp_len, nheads, device=device, dtype=torch.float32) + dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) # tie_hdim + B = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) + C = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) + + # prev_tokens placeholder — overwritten per-run + prev_tokens = torch.zeros(batch, device=device, dtype=torch.int32) + # slot_perm placeholder — overwritten per-run by mix pre_iter_fn when + # sort_slots is enabled. Identity by default so cells that don't sort + # (or pure-batch cells) get a meaningful identity perm if the kernel + # ends up reading it (USE_PERM=False makes this path unused). + slot_perm_buf = torch.arange(batch, device=device, dtype=torch.int32) + + out_incr = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) + out_base = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) + + # intermediate_states_buffer is only consumed by the fp/baseline path; + # for quantized state dtypes we'll skip baselines entirely, so the buffer + # dtype falls back to fp32 to keep selective_state_update happy. + int_buffer_dtype = state_dtype if state_dtype not in _QUANT_BENCH else torch.float32 + intermediate_states_buffer = torch.zeros( + batch, mtp_len, nheads, head_dim, d_state, device=device, dtype=int_buffer_dtype + ) + + # --- Conv1d tensors (for --with-conv1d mode) --- + d_inner = nheads * head_dim + conv_dim = d_inner + 2 * ngroups * d_state + d_conv = 4 # conv kernel width for Nemotron/Mamba2 + + # xbc_input: (batch, conv_dim, mtp_len) — "hot" input from in_proj. + # Match production layout: in_proj output is (batch*mtp_len, conv_dim) + # contiguous, then .view(batch, mtp_len, conv_dim).transpose(1, 2) + # gives strides (mtp_len*conv_dim, 1, conv_dim) — NOT the standard + # (conv_dim*mtp_len, mtp_len, 1) of a freshly allocated 3D tensor. + # Conv1d preserves input strides in its output, so downstream split + # + view inherits the correct layout without needing .contiguous(). + xbc_input_flat = torch.randn(batch * mtp_len, conv_dim, device=device, dtype=act_dtype) + xbc_input = xbc_input_flat.view(batch, mtp_len, conv_dim).transpose(1, 2) + # conv_state: (batch, conv_dim, d_conv) — "cold" cache + conv_state = torch.randn(batch, conv_dim, d_conv, device=device, dtype=act_dtype) + # conv_weight: (conv_dim, d_conv) — parameter + conv_weight = torch.randn(conv_dim, d_conv, device=device, dtype=act_dtype) + # conv_bias: (conv_dim,) — parameter + conv_bias = torch.randn(conv_dim, device=device, dtype=act_dtype) + + # Store full-batch buffers in cache and return slices at request_batch. + _TENSOR_CACHE[cache_key] = { + "max_batch": alloc_batch, + "state0": state0, + "state_scales0": state_scales0, + "old_x": old_x, + "old_B": old_B, + "old_dt": old_dt, + "old_dA_cumsum": old_dA_cumsum, + "cache_buf_idx": cache_buf_idx, + "x": x, + "dt": dt, + "B": B, + "C": C, + "A": A, + "dt_bias": dt_bias, + "D": D, + "prev_tokens": prev_tokens, + "slot_perm_buf": slot_perm_buf, + "out_incr": out_incr, + "out_base": out_base, + "intermediate_states_buffer": intermediate_states_buffer, + "xbc_input": xbc_input, + "conv_state": conv_state, + "conv_weight": conv_weight, + "conv_bias": conv_bias, + "d_inner": d_inner, + "conv_dim": conv_dim, + } + rb = request_batch + return ( + state0[:rb], + state_scales0[:rb] if state_scales0 is not None else None, + old_x[:rb], + old_B[:rb], + old_dt[:rb], + old_dA_cumsum[:rb], + cache_buf_idx[:rb], + x[:rb], + dt[:rb], + B[:rb], + C[:rb], + A, + dt_bias, + D, + prev_tokens[:rb], + slot_perm_buf[:rb], + out_incr[:rb], + out_base[:rb], + intermediate_states_buffer[:rb], + xbc_input[:rb], + conv_state[:rb], + conv_weight, + conv_bias, + d_inner, + conv_dim, + ) + + +# ============================================================================= +# CUPTI in-process kernel timing +# +# Self-contained module-in-a-file. Reads kernel start/end timestamps directly +# from the GPU profiling fabric via NVIDIA's cupti-python bindings (1 ns +# resolution), avoiding two pitfalls of the cuda-events path: +# +# 1. cudaEvent.elapsed_time() resolution (~0.5 us) is too coarse for the +# short kernels we care about, especially with PDL + cuda graphs at +# small batch — events recorded inside a graph have proven noisy. +# 2. nsys is the only known accurate alternative, but the +# profile-export-sqlite-parse pipeline is heavy and out-of-process. +# +# This is functionally equivalent to wrapping each cell in nsys, except it +# runs in the same Python process with no serialization. When this proves +# out, lift `CuptiKernelTimer` and `_time_kernel_cuda_graph_cupti` into a +# proper TRT-LLM utility module — there is no benchmark-specific code below. +# ============================================================================= + + +# Substring match: kernels run_fn launches that we want to time. Mirrors +# the parser in scripts/.../collect.py so cupti and nsys-based outputs agree. +_CUPTI_KEEP_KERNEL_SUBSTRINGS = ( + "_replay_precompute", + "_checkpointing_precompute", + "_rectangle_precompute", + "_dynamic_precompute", + "_replay_state_update", + "_checkpointing_main", + "_rectangle_main", + "_dynamic_main", + "_persistent_main", + "selective_scan_update", + "selective_state_update", + "causal_conv1d_update", +) + + +def _kernels_per_iter_incremental( + mode: str, + with_conv1d: bool, + *, + persistent_skip_empty: bool = True, +) -> int: + """Expected number of CUPTI-tracked kernels per iter for the incremental + kernel chain, given the dispatch mode and the conv1d flag. + + Used to validate CUPTI record counts (no auto-inference — silent + mis-timing is the failure mode we're guarding against). + + `persistent_skip_empty=True` (today's behavior): the + `mode='persistent_main'` launch helper host-early-outs when its half + is empty (n_writes=0 or n_writes=batch in pure scenarios), so only + one of the two persistent_main_kernel launches actually fires per + iter. With `persistent_skip_empty=False` (future no-eo mode), both + halves always launch and K bumps by 1. + + `persistent_dynamic` always launches 1 main; not affected by the flag. + """ + if mode == "monolithic": + k = 2 # precomp + main + elif mode == "dynamic": + k = 2 # dynamic_precomp + dynamic_main + elif mode == "maindl": + k = 3 # 1 dynamic_precomp + 2 mains (write + nowrite) + elif mode in ("doublelaunch", "dlgrouped"): + k = 4 # 2 precomp + 2 main + elif mode == "dl_write_only": + k = 2 # 1 precomp + 1 main (write only) + elif mode == "persistent_dynamic": + k = 2 # 1 dynamic_precomp + 1 persistent_main + elif mode == "persistent_main": + k = 2 if persistent_skip_empty else 3 # see docstring + else: + raise ValueError(f"_kernels_per_iter_incremental: unknown mode {mode!r}") + if with_conv1d: + k += 1 + return k + + +def _kernels_per_iter_baseline(with_conv1d: bool) -> int: + """Expected kernels per iter for triton / flashinfer baselines. + + Both baselines run a single state-update kernel; `--with-conv1d` + prepends one conv1d kernel. + """ + return 2 if with_conv1d else 1 + + +class CuptiKernelTimer: + """Process-singleton wrapper around CUPTI's CONCURRENT_KERNEL activity. + + CUPTI's callbacks are global (one subscriber per process), so the timer + is constructed lazily once via `CuptiKernelTimer.get()`. cupti-python + parses the activity buffer for us — `buffer_completed` receives a Python + list of typed activity objects, not a raw byte buffer — so no FFI is + needed. + + Usage: + timer = CuptiKernelTimer.get() + timer.start() # arms; drops any stale records + + records, zero_ts_count, zero_ts_names = timer.stop() + # records: list of tuples per kernel + # (name, start_ns, end_ns, corr, + # graph_id, graph_node_id, stream) + # zero_ts_count: kernel records CUPTI + # delivered with start=0 or end=0 + # (couldn't timestamp); zero_ts_names: + # name → count breakdown. + + The callback fires from a CUPTI worker thread, so a lock guards the + record buffer. Records are kept tiny (tuple of ints + str) to minimize + Python overhead in the hot path of the callback. + """ + + _instance = None + _import_error = None + + @classmethod + def get(cls) -> "CuptiKernelTimer": + if cls._instance is not None: + return cls._instance + if cls._import_error is not None: + raise cls._import_error + try: + from cupti import cupti as _c + except ImportError as e: # pragma: no cover — env-dependent + cls._import_error = e + raise + cls._instance = cls._init(_c) + return cls._instance + + @classmethod + def _init(cls, _c) -> "CuptiKernelTimer": + import threading + + self = object.__new__(cls) + self._c = _c + self._records: list[tuple] = [] + self._zero_ts_count = 0 # how many kernel records were dropped due to start/end==0 + self._zero_ts_names: dict = {} # name -> count of zero-ts drops (for diagnostics) + self._lock = threading.Lock() + + # CUPTI callback contract (from cupti-python-samples/cupti_common.py): + # buffer_requested() -> (buffer_size, max_num_records) + # buffer_completed(activities: list) + # Setting max_num_records=0 (unbounded) avoids spurious buffer + # requests. 8 MiB matches the sample defaults. + def _buf_req(): + return (8 * 1024 * 1024, 0) + + kernel_kinds = (_c.ActivityKind.CONCURRENT_KERNEL, _c.ActivityKind.KERNEL) + + def _buf_done(activities): + recs = [] + zero_drops_local = 0 + zero_names_local: dict = {} + for a in activities: + if a.kind not in kernel_kinds: + continue + # start/end == 0 means CUPTI couldn't time this kernel. + # Track these instead of silently dropping — they're a sign + # CUPTI is failing to record kernels we DID launch. + if a.start == 0 or a.end == 0: + zero_drops_local += 1 + name = getattr(a, "name", "?") + zero_names_local[name] = zero_names_local.get(name, 0) + 1 + continue + recs.append(( + a.name, + int(a.start), + int(a.end), + int(a.correlation_id), + int(a.graph_id), + int(a.graph_node_id), + int(a.stream_id), + )) + if recs or zero_drops_local: + with self._lock: + if recs: + self._records.extend(recs) + if zero_drops_local: + self._zero_ts_count += zero_drops_local + for k, v in zero_names_local.items(): + self._zero_ts_names[k] = self._zero_ts_names.get(k, 0) + v + + # Hold strong refs so the C side never sees GC'd Python callables. + self._buf_req = _buf_req + self._buf_done = _buf_done + + _c.activity_register_callbacks(_buf_req, _buf_done) + _c.activity_enable(_c.ActivityKind.CONCURRENT_KERNEL) + return self + + def start(self) -> None: + """Arm capture: flush any stale records, then clear the buffer.""" + self._c.activity_flush_all(1) + with self._lock: + self._records.clear() + self._zero_ts_count = 0 + self._zero_ts_names = {} + + def stop(self) -> tuple[list[tuple], int, dict]: + """Flush and return all kernel records + count of zero-timestamp drops. + + Returns (records, zero_ts_count, zero_ts_names_dict). The latter two + are diagnostic: nonzero values mean CUPTI delivered records with + start=0 or end=0, indicating it failed to timestamp the kernel. + """ + self._c.activity_flush_all(1) + with self._lock: + return (list(self._records), self._zero_ts_count, dict(self._zero_ts_names)) + + +# ============================================================================= +# Timing helpers +# ============================================================================= + + +def _stats_from_spans(spans_us: list[float]) -> dict: + """Compute median / p95 / p99 / n from a per-iter span list.""" + s = sorted(spans_us) + return { + "median": statistics.median(s), + "p95": s[int(0.95 * len(s))], + "p99": s[int(0.99 * len(s))], + "n": len(s), + } + + +def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, + zero_ts_count: int = 0, + zero_ts_names: dict | None = None): + """Bin a flat CUPTI kernel record stream into per-iter spans + per-kernel + relative timestamps. Used by both graph and eager CUPTI paths. + + `records` are tuples (name, start_ns, end_ns, ...) — see CuptiKernelTimer. + `expected_K` is the kernels-per-iter count the caller declares; we + validate the CUPTI total matches `expected_K * (warmup + iters)` exactly. + On mismatch we dump per-name record counts so missing or extra kernels + are obvious (most common cause: a new dispatch mode whose kernels lack + a matching entry in `_CUPTI_KEEP_KERNEL_SUBSTRINGS`, silently filtering + them out). + """ + records = [ + r for r in records + if any(s in r[0] for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS) + ] + records.sort(key=lambda r: r[1]) # by start_ns + + total = len(records) + expected_iters = warmup + iters + expected_total = expected_K * expected_iters + if total != expected_total: + from collections import Counter + name_counts = dict(Counter(r[0] for r in records)) + # Non-fatal: skip this cell instead of killing the whole sweep. + # Mismatch may be a CUPTI dropped-records issue (rare configs), + # not necessarily a K-table bug. Log so the user can investigate + # the specific cell post-hoc; return None so the caller can skip + # writing a JSON row. + zero_msg = "" + if zero_ts_count: + zero_msg = ( + f" + {zero_ts_count} records with start/end=0 " + f"(dropped by callback, breakdown {zero_ts_names}). " + f"Total observed kernel records (timed + zero-ts) = " + f"{total + zero_ts_count} / {expected_total}." + ) + print( + f"[WARN] CUPTI capture mismatch for {tag!r}: expected " + f"{expected_K} kernels/iter × {expected_iters} iters " + f"(warmup+iters) = {expected_total} records, got {total}. " + f"Kernel record counts: {name_counts}.{zero_msg} SKIPPING cell.", + file=sys.stderr, + ) + # Per-record dump: (name, start_ns_rel, end_ns_rel, corr_id, graph_id, stream_id). + # Times relative to first record so absolute ns isn't drowning output. + # Limit dump to first 30 records to avoid flooding logs at high K. + if records: + t0_ns = records[0][1] + for i, r in enumerate(records[:30]): + # r = (name, start_ns, end_ns, corr_id, graph_id, graph_node_id, stream_id) + rel_start = (r[1] - t0_ns) / 1000.0 # us + rel_end = (r[2] - t0_ns) / 1000.0 + print( + f" rec[{i:3d}] name={r[0]!r} start={rel_start:.2f}us " + f"end={rel_end:.2f}us corr={r[3]} graph={r[4]} stream={r[6]}", + file=sys.stderr, + ) + if len(records) > 30: + print(f" ... ({len(records) - 30} more records elided)", file=sys.stderr) + return None + K = expected_K + timed = records[warmup * K:] + + spans_us: list[float] = [] + per_kernel: dict[str, dict[str, list[float]]] = {} + for i in range(iters): + chunk = timed[i * K:(i + 1) * K] + iter_start_ns = min(r[1] for r in chunk) + iter_end_ns = max(r[2] for r in chunk) + spans_us.append((iter_end_ns - iter_start_ns) / 1000.0) + for r in chunk: + name = r[0] + slot = per_kernel.setdefault(name, {"start_us": [], "end_us": []}) + slot["start_us"].append((r[1] - iter_start_ns) / 1000.0) + slot["end_us"].append((r[2] - iter_start_ns) / 1000.0) + + out = _stats_from_spans(spans_us) + out["iters_us"] = spans_us + out["per_kernel"] = per_kernel + return out + + +_PRE_GRAPH_WARMUP_ITERS = 3 # standard practice; see commit msg / design doc + + +def _time_kernel_cuda_graph( + args, + run_fn, + reset_fn, + tag: str, + *, + expected_K: int, + pre_iter_fn=None, + iters_override: int | None = None, +) -> dict: + """CUDA-graph CUPTI timer (graph-per-iter design). + + Captures one CUDA graph holding a single iter's worth of work (reset + + l2_flush + run_fn) and replays it `warmup + iters` times. Per-iter + setup (`pre_iter_fn`, e.g. mix-mode PNAT/n_writes copies) runs OUTSIDE + the graph, on the same CUDA stream so the order + `pre_iter_fn → l2_flush → run_fn` is preserved on every replay. + + Why graph-per-iter (vs the older "one giant graph holding all iters" + design): instantiating a CUDA graph is expensive — proportional to + graph size — so a single small graph instantiated once is much + cheaper than one big graph instantiated for each cell of a sweep. + Replays are cheap regardless. + + Why pre_iter_fn outside the graph: it depends on the iter index `i` + (different sample per iter), but graph capture would bake in the + capture-time `i`. Putting the per-iter copy outside the graph also + has a useful side effect: PNAT (the per-iter input) is loaded into + L2 by the copy, then evicted by the in-graph L2 flush, so the kernel + reads PNAT cold — closer to production behavior than the old design. + + Pre-graph eager warmup (3 iters): forces PyTorch's caching allocator + + Triton's autotune cache to settle before capture so the graph + doesn't bake in init-only allocations. + + ``iters_override`` (if not None) overrides ``args.iters`` for this + call. Used to give mix scenarios a higher iter count than pure + (more iters = more independent mix draws averaged in). + """ + timer = CuptiKernelTimer.get() + warmup = args.warmup + iters = iters_override if iters_override is not None else args.iters + + # Pre-graph eager warmup: full per-iter chain × N. Settles allocator + # + autotune; pre_iter_fn included so any side effects it has are + # exercised before capture. + for _ in range(_PRE_GRAPH_WARMUP_ITERS): + reset_fn() + if pre_iter_fn is not None: + pre_iter_fn(0) + run_fn() + torch.cuda.synchronize() + + # Reset just before capture so warmup state changes don't bleed in. + reset_fn() + torch.cuda.synchronize() + + # Capture ONE iter. pre_iter_fn deliberately not in here. + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + reset_fn() + if args.l2_flush: + _l2_flush.fill_(0.0) + run_fn() + torch.cuda.synchronize() + + # Time: replay the per-iter graph `warmup + iters` times, with + # pre_iter_fn called between replays on the same stream. CUPTI + # records every kernel launch; _stats_from_cupti_records validates + # against expected_K and slices warmup off the front. + timer.start() + torch.cuda.nvtx.range_push(tag) + for i in range(warmup + iters): + if pre_iter_fn is not None: + pre_iter_fn(i) + g.replay() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + records, zero_ts_count, zero_ts_names = timer.stop() + + return _stats_from_cupti_records(records, warmup, iters, tag, expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names) + + +def _time_kernel_eager( + args, + run_fn, + reset_fn, + tag: str, + *, + expected_K: int, + pre_iter_fn=None, + iters_override: int | None = None, +) -> dict: + """Non-graph CUPTI timer (for ncu wrapping, debugging, etc.). + + Each iter runs serially with sync between, but kernel start/end still + come from CUPTI — same accuracy as the graph path, just slower per-iter + (extra Python + sync overhead). + """ + timer = CuptiKernelTimer.get() + warmup = args.warmup + iters = iters_override if iters_override is not None else args.iters + + timer.start() + torch.cuda.nvtx.range_push(tag) + # Unified warmup+iters loop; CUPTI filters by warmup count internally. + for i in range(warmup + iters): + reset_fn() + if args.l2_flush: + _flush_l2() # includes synchronize + if pre_iter_fn is not None: + pre_iter_fn(i) + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + records, zero_ts_count, zero_ts_names = timer.stop() + + return _stats_from_cupti_records(records, warmup, iters, tag, expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names) + + +def _run_kernel_untimed(args, run_fn, reset_fn, tag: str) -> dict: + """No in-bench timing: just run the kernels for an external profiler + (nsys / ncu) to time externally. Returns a stats dict full of zeros so + downstream code (table, JSON) doesn't break. + + Note: pre_iter_fn / iters_override aren't plumbed here yet — mix-mode + benchmarking relies on CUPTI. Add when a use-case lands. + """ + warmup = args.warmup + iters = args.iters + + if args.cuda_graph: + # Eager warmup before capture (Triton autotune) + reset_fn(); run_fn(); torch.cuda.synchronize() + reset_fn(); torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + for _ in range(warmup + iters): + reset_fn() + if args.l2_flush: + _l2_flush.fill_(0.0) + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_push(tag) + g.replay() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + else: + torch.cuda.nvtx.range_push(tag) + for _ in range(warmup + iters): + reset_fn() + if args.l2_flush: + _flush_l2() + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + + spans_us = [0.0] * iters + out = _stats_from_spans(spans_us) + out["iters_us"] = spans_us + out["per_kernel"] = {} + return out + + +def _time_kernel( + args, run_fn, reset_fn, tag: str, + *, + expected_K: int, + pre_iter_fn=None, + iters_override: int | None = None, +) -> dict: + """Dispatch to graph-CUPTI / eager-CUPTI / no-timer path. + + --cupti: in-process CUPTI Activity API timing (default). Use --no-cupti + when running under nsys (in-process CUPTI conflicts with nsys's own + subscriber); the bench then runs the kernels for nsys to time externally. + + `expected_K` is the kernels-per-iter count the caller declares + (computed via _kernels_per_iter_*). CUPTI paths validate against it + explicitly; the no-timer fallback ignores it (no records to validate). + """ + if not getattr(args, "cupti", True): + if pre_iter_fn is not None: + raise RuntimeError( + "_time_kernel: pre_iter_fn requires CUPTI (mix-mode); " + "got --no-cupti. Re-run with CUPTI on or plumb pre_iter_fn " + "through _run_kernel_untimed." + ) + return _run_kernel_untimed(args, run_fn, reset_fn, tag) + if args.cuda_graph: + return _time_kernel_cuda_graph( + args, run_fn, reset_fn, tag, + expected_K=expected_K, + pre_iter_fn=pre_iter_fn, + iters_override=iters_override, + ) + return _time_kernel_eager( + args, run_fn, reset_fn, tag, + expected_K=expected_K, + pre_iter_fn=pre_iter_fn, + iters_override=iters_override, + ) + + +# Per-config benchmark (consolidated baseline + replay) + + +def _warm_one_config(args, cfg, baseline_fn) -> None: + """Module-level worker for the compile-warmup process pool. + + Module-level so ProcessPoolExecutor can pickle it (nested functions + aren't picklable). Each worker process holds its own GIL → no + serialization between concurrent compiles. + + ``cfg`` is a tuple of (outer_cfg, inner_overrides): + * outer_cfg = (batch, mtp_len, prev_ks, state_dtype, act_dtype, + sr_mode, rect, write_ckpt, mode, + sort_slots, reverse_nowrite, hardcode_sort) + * inner_overrides = dict of args attribute name -> single-value string + to clamp the per-cell inner-knob sweep to ONE + combination. Triggers exactly one Triton compile + per worker invocation, so N workers achieve + N-way concurrency regardless of outer config + count. (Prior design fanned out only at outer + granularity, capping concurrency at ~10 even + with --compile-threads 50.) + + ``baseline_fn`` is optional — when ``None``, only the checkpointing + kernel is warmed (the baseline-selection kernel can be warmed once in + the parent if needed). This lets us avoid pickling C-extension + function references across processes. + """ + outer_cfg, inner_overrides = cfg + (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, + rect, write_ckpt, mode, sort_slots, reverse_nowrite, hardcode_sort) = outer_cfg + # Clone args and override inner-knob sweep lists to single values. + # _bench_config then iterates a 1×1×...×1 cartesian inside. + import argparse as _ap + args_copy = _ap.Namespace(**vars(args)) + for k, v in inner_overrides.items(): + setattr(args_copy, k, v) + _bench_config( + args_copy, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, + sr_mode=sr_mode, rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, mode=mode, + sort_slots=sort_slots, reverse_nowrite=reverse_nowrite, + hardcode_sort=hardcode_sort, + warmup_only=True, + ) + + +def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, + baseline_fn, max_workers: int) -> None: + """Parallel compile-warmup using a ProcessPoolExecutor with `spawn` + start method. + + Each worker process holds its own GIL and its own CUDA context, so + Triton compiles (Python AST/codegen + LLVM/ptxas) run truly in + parallel. Previous ThreadPoolExecutor design hit GIL contention + in the Python codegen phase, capping throughput at ~1-2 cores even + with 28 threads (observed: 4 R threads vs 28 in pool). + + Compiled binaries land in Triton's on-disk cache (TRITON_CACHE_DIR + or default ~/.triton/cache). Workers share the cache via filesystem + — first to write any given (kernel_source × constexpr_set) hash + wins; concurrent writes to the SAME hash are wasteful but not + corrupting. + + spawn start method avoids inheriting parent CUDA state (which is + unsafe after fork on Linux with active CUDA contexts). Per-worker + import + CUDA init costs ~10s, amortized over each worker's many + compiles. baseline_fn is intentionally NOT passed to workers to + avoid pickling complications; the parent compiles the baseline + kernel itself before launching the pool when applicable. + """ + from concurrent.futures import ProcessPoolExecutor + import multiprocessing + + sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) + + rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) + write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) + modes_list = getattr(args, "modes_list", ["monolithic"]) + sort_list = getattr(args, "sort_slots_list", [False]) + rev_list = getattr(args, "reverse_nowrite_list", [False]) + hsort_list = getattr(args, "hardcode_sort_list", [False]) + + # Compile-warmup task enumeration: outer × inner cartesian. + # CRITICAL: only enumerate axes that change the kernel's COMPILE signature. + # Drop runtime axes (batch, prev_k) that produce identical kernel hashes — + # otherwise we'd pay ~50-100ms of bench setup per redundant cache-hit task. + # + # Batches collapsed to first only: batch is a runtime int passed to the + # kernel, not a constexpr; all batches share the same compiled kernel. + # prev_k is already a list passed into _bench_config (not enumerated here). + configs = [] + _compile_batches = batch_sizes[:1] # collapse runtime axis + for batch in _compile_batches: + for mtp_len in mtp_lengths: + prev_ks = _resolve_prev_ks(args, mtp_len) + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + for sr_mode in sr_modes_list: + for mode in modes_list: + effective_write_modes = ( + write_modes_list if mode == "monolithic" else [True] + ) + for write_ckpt in effective_write_modes: + if mode == "monolithic": + effective_rect_list = ( + [False] if write_ckpt else rect_list + ) + else: + effective_rect_list = rect_list + for rect in effective_rect_list: + # Note: "persistent_main" is included in + # the dl-family for sort/hsort sweep + # eligibility — it consumes the same + # slot_perm and benefits from the same + # write-first clustering. It additionally + # requires _n_writes (count of write + # slots) which the bench computes from + # the pure-scenario PNAT (mix scenarios + # not yet supported for persistent_main). + is_dl_family = mode in ( + "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", "persistent_main", + "persistent_dynamic", + ) + # Match the timed-run skip: sort=1 only + # makes sense when there's a mix scenario. + can_sort = is_dl_family and (args.mix_csv is not None) + effective_sort_list = ( + sort_list if can_sort else [False] + ) + effective_hsort_list = ( + hsort_list if can_sort else [False] + ) + for sort_slots in effective_sort_list: + effective_rev_list = ( + rev_list if sort_slots else [False] + ) + for reverse_nowrite in effective_rev_list: + for hardcode_sort in effective_hsort_list: + if sort_slots and hardcode_sort: + continue + configs.append(( + batch, mtp_len, prev_ks, + state_dtype, act_dtype, + sr_mode, rect, write_ckpt, mode, + sort_slots, reverse_nowrite, + hardcode_sort, + )) + + # Enumerate inner-knob cartesian — same axes _bench_config iterates + # internally. Each (outer × inner) tuple becomes one task; workers + # then trigger exactly one Triton compile per task, giving true + # N-way concurrency with --compile-threads N. + def _ps(val): + if val is None or (isinstance(val, str) and not val): + return [None] + if isinstance(val, str): + return [v.strip() for v in val.split(",") if v.strip()] + return [val] + + # CPS (cta_per_sm) collapsed to first value: NUM_PERSISTENT is runtime now, + # so different CPS values share the same compiled kernel. Collapsing here + # avoids enumerating 4-8x redundant tasks that would each pay ~50-100ms + # bench setup overhead for a cache hit on the same kernel hash. + _cps_for_compile = _ps(args.cta_per_sm)[:1] + knob_axes = [ + ("block_size_m", _ps(args.block_size_m)), + ("num_warps", _ps(args.num_warps)), + ("num_stages", _ps(args.num_stages)), + ("precompute_num_warps", _ps(args.precompute_num_warps)), + ("precompute_num_stages", _ps(args.precompute_num_stages)), + ("heads_per_block", _ps(args.heads_per_block)), + ("maxnreg", _ps(args.maxnreg)), + ("num_ctas", _ps(args.num_ctas)), + ("cta_per_sm", _cps_for_compile), # collapsed (runtime) + ("num_loop_stages", _ps(args.num_loop_stages)), + ("flatten", _ps(args.flatten)), + ("warp_specialize", _ps(args.warp_specialize)), + ("use_tma_rect_load", _ps(args.use_tma_rect_load)), + ("use_tma_replay_write_load", _ps(args.use_tma_replay_write_load)), + ("use_tma_replay_nowrite_load", _ps(args.use_tma_replay_nowrite_load)), + ("use_tma_replay_write_store", _ps(args.use_tma_replay_write_store)), + ] + import itertools as _it + inner_combos = list(_it.product(*(values for _, values in knob_axes))) + + # Cross product outer × inner. Override only knobs that have an + # explicit value (skip None — those leave args. at its CLI default, + # which _bench_config handles via its own _parse_sweep). + tasks = [] + for outer in configs: + for inner_tuple in inner_combos: + inner_overrides = { + name: str(val) + for (name, _), val in zip(knob_axes, inner_tuple) + if val is not None + } + tasks.append((outer, inner_overrides)) + + # Shuffle to reduce cross-worker race on the same kernel hash. Two + # workers picking adjacent tasks (same mode, neighboring knob value) + # could both miss + compile the same kernel hash; shuffling spreads + # workloads across different kernel hash families. + import random as _r + _r.shuffle(tasks) + + print(f"[compile-warmup] {len(tasks)} compile tasks " + f"({len(configs)} outer × {len(inner_combos)} inner combos) " + f"across {max_workers} processes (ProcessPoolExecutor, spawn start)") + t0 = time.perf_counter() + + ctx = multiprocessing.get_context("spawn") + errors = [] + with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as ex: + # baseline_fn=None: workers compile only the checkpointing kernel. + # Baseline kernels (if any) get compiled lazily in the parent during + # the timing phase — usually just one extra compile, negligible. + futures = { + ex.submit(_warm_one_config, args, task, None): task + for task in tasks + } + for fut in futures: + try: + fut.result() + except Exception as e: + errors.append((futures[fut], e)) + + if errors: + for cfg, e in errors: + print(f"[compile-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", + file=sys.stderr) + raise errors[0][1] + + print(f"[compile-warmup] done in {time.perf_counter() - t0:.1f}s") + + +def _bench_config( + args, + batch: int, + mtp_len: int, + prev_ks: list[int], + state_dtype: torch.dtype, + act_dtype: torch.dtype, + baseline_fn, + sr_mode: str = "RN", + rectangle_for_nowrite: bool = False, + write_checkpoint: bool = True, + mode: str = "monolithic", + mix_samples_cpu=None, + mix_label: str = "", + sort_slots: bool = False, + reverse_nowrite: bool = False, + perm_samples_cpu=None, + hardcode_sort: bool = False, + mix_samples_sorted_cpu=None, + warmup_only: bool = False, +) -> None: + """ + Benchmark one (batch, mtp_len, dtype) configuration. + + Runs the baseline kernel (if baseline_fn is not None) followed by the + replay kernel for each prev_k value. Tensors are built once and + shared across all runs in this config. + + When ``warmup_only`` is True, calls each kernel exactly once instead of + timing it. Used by the parallel-warmup phase to populate Triton's + persistent compile cache across all configs concurrently. No timing + output is produced. + """ + state_dtype_name = str(state_dtype).split(".")[-1] + act_dtype_name = str(act_dtype).split(".")[-1] + + ( + state0, + state_scales0, + old_x0, + old_B0, + old_dt0, + old_dA_cumsum0, + cache_buf_idx0, + x, + dt, + B, + C, + A, + dt_bias, + D, + prev_tokens, + slot_perm_buf, + out_incr, + out_base, + intermediate_states_buffer, + xbc_input0, + conv_state0, + conv_weight, + conv_bias, + d_inner, + conv_dim, + ) = _build_tensors( + batch, + mtp_len, + state_dtype, + act_dtype, + args.tp_nheads, + args.head_dim, + args.d_state, + args.tp_ngroups, + max_window=getattr(args, "max_window", None) or None, + ) + + nheads = args.tp_nheads + ngroups = args.tp_ngroups + head_dim = args.head_dim + d_state = args.d_state + with_conv1d = getattr(args, "with_conv1d", False) + use_philox = (sr_mode == "SR") + variant_fn = _VARIANT_FNS[args.variant]() + + # SR rounding: allow fp16 and the quantized dtypes (int8/int16/fp8). + # bf16/fp32 SR is not supported (no PTX path for bf16; fp32 doesn't need + # rounding). When sweeping --sr-modes RN,SR over a mixed dtype set, + # silently skip the SR cell for unsupported dtypes — the RN cell still + # prints, and other dtypes still get their SR row. + rand_seed = None + _SR_SUPPORTED = ( + torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn, + ) + if use_philox: + if state_dtype not in _SR_SUPPORTED: + return + rand_seed = torch.randint(0, 2**62, (1,), device="cuda", dtype=torch.int64) + + is_quantized = state_dtype in (torch.int8, torch.int16, torch.float8_e4m3fn) + + state_work = state0.clone() + state_scales_work = state_scales0.clone() if state_scales0 is not None else None + old_x_work = old_x0.clone() + old_B_work = old_B0.clone() + old_dt_work = old_dt0.clone() + old_dA_cumsum_work = old_dA_cumsum0.clone() + cache_buf_idx_work = cache_buf_idx0.clone() + xbc_input_work = xbc_input0.clone() + conv_state_work = conv_state0.clone() + + def _reset(): + state_work.copy_(state0) + if state_scales_work is not None: + state_scales_work.copy_(state_scales0) + old_x_work.copy_(old_x0) + old_B_work.copy_(old_B0) + old_dt_work.copy_(old_dt0) + old_dA_cumsum_work.copy_(old_dA_cumsum0) + cache_buf_idx_work.copy_(cache_buf_idx0) + if with_conv1d: + conv_state_work.copy_(conv_state0) + + def _reset_conv1d_realistic(): + """Realistic reset: cold cache, L2 flush, then hot in_proj output.""" + # 1. Reset cold state (cache tensors, SSM state) + state_work.copy_(state0) + if state_scales_work is not None: + state_scales_work.copy_(state_scales0) + old_x_work.copy_(old_x0) + old_B_work.copy_(old_B0) + old_dt_work.copy_(old_dt0) + old_dA_cumsum_work.copy_(old_dA_cumsum0) + cache_buf_idx_work.copy_(cache_buf_idx0) + conv_state_work.copy_(conv_state0) + # 2. L2 flush (evicts cold state from cache) + if _l2_flush is not None: + _l2_flush.fill_(0.0) + # 3. Write hot tensors (simulates in_proj output landing in L2) + xbc_input_work.copy_(xbc_input0) + + # Silently skip the baseline row for any (baseline, state_dtype, SR) + # combo it can't run. Better than erroring on a partial sweep — our + # kernel rows still print. Compatibility: + # * Quantized states (int8 / int16 / fp8): no baseline supports them. + # * Triton baseline (selective_state_update): no rand_seed kwarg. + # * flashinfer baseline: rand_seed only on fp16 state. + def _baseline_supports() -> bool: + if baseline_fn is None: + return False + if is_quantized: + return False + if use_philox: + if args.baseline == "triton": + return False + if args.baseline == "flashinfer" and state_dtype != torch.float16: + return False + return True + + if baseline_fn is not None and not _baseline_supports(): + if not warmup_only: + sr_tag = " + SR" if use_philox else "" + print( + f"# Skipping {args.baseline} baseline for " + f"state_dtype={state_dtype_name}{sr_tag} (unsupported)." + ) + baseline_fn = None + + show_kernel_col = baseline_fn is not None + + def _conv1d_split(xbc_in, conv_st, launch_dependent_kernels=False): + """Run conv1d update and split output into (x, B, C) views. + + The input tensor's strides are preserved through conv1d and the + transpose+view chain. With the production-matching layout + (contiguous (batch*T, conv_dim) viewed as (batch, conv_dim, T)), + the output after transpose+view has stride(-1)==1 and + stride(1)==dim, satisfying both our kernel and flashinfer. + """ + xbc_result = causal_conv1d_update( + xbc_in, + conv_st, + conv_weight, + conv_bias, + activation="silu", + launch_dependent_kernels=launch_dependent_kernels, + ) + xbc_flat = xbc_result.transpose(1, 2).view(batch * mtp_len, conv_dim) + x_flat, B_flat, C_flat = torch.split( + xbc_flat, [d_inner, ngroups * d_state, ngroups * d_state], dim=-1 + ) + x_conv = x_flat.view(batch, mtp_len, nheads, head_dim) + B_conv = B_flat.view(batch, mtp_len, ngroups, d_state) + C_conv = C_flat.view(batch, mtp_len, ngroups, d_state) + return x_conv, B_conv, C_conv + + # --- Baseline --- + if baseline_fn is not None: + tag = f"base_b{batch}_mtp{mtp_len}_s{state_dtype_name}_a{act_dtype_name}" + + philox_kwargs = {} + if rand_seed is not None and args.baseline == "flashinfer": + philox_kwargs = {"rand_seed": rand_seed, "philox_rounds": args.philox_rounds} + + if with_conv1d: + + def _run_baseline(): + x_conv, B_conv, C_conv = _conv1d_split(xbc_input_work, conv_state_work) + baseline_fn( + state_work, + x=x_conv, + dt=dt, + A=A, + B=B_conv, + C=C_conv, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + out=out_base, + disable_state_update=True, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=mtp_len, + **philox_kwargs, + ) + else: + + def _run_baseline(): + baseline_fn( + state_work, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + out=out_base, + disable_state_update=True, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=mtp_len, + **philox_kwargs, + ) + + reset_fn = _reset_conv1d_realistic if with_conv1d else _reset + if warmup_only: + reset_fn() + _run_baseline() + torch.cuda.synchronize() + else: + stats = _time_kernel( + args, _run_baseline, reset_fn, tag, + expected_K=_kernels_per_iter_baseline(with_conv1d), + ) + + if stats is not None: + _print_row( + show_kernel_col, + args.baseline, + batch, + mtp_len, + "N/A", + state_dtype_name, + act_dtype_name, + stats, + json_results=getattr(args, "_json_results", None), + tp_size=args.tp_size, + json_detailed=getattr(args, "json_detailed", False), + ) + + # --- Sweep parameter parsing (invariant across prev_k) --- + def _parse_sweep(val): + if val is None: + return [None] + return [int(v) for v in val.split(",")] + + block_size_m_values = _parse_sweep(args.block_size_m) + num_warps_values = _parse_sweep(args.num_warps) + num_stages_values = _parse_sweep(args.num_stages) + precompute_num_warps_values = _parse_sweep(args.precompute_num_warps) + precompute_num_stages_values = _parse_sweep(args.precompute_num_stages) + heads_per_block_values = _parse_sweep(args.heads_per_block) + maxnreg_values = _parse_sweep(args.maxnreg) + num_ctas_values = _parse_sweep(args.num_ctas) + # Persistent-only sweep dims; ignored when the cell's mode != persistent_main. + cta_per_sm_values = _parse_sweep(args.cta_per_sm) + num_loop_stages_values = _parse_sweep(args.num_loop_stages) + flatten_values = _parse_sweep(args.flatten) + warp_specialize_values = _parse_sweep(args.warp_specialize) + # Per-main split-knob sweeps. Default = same as the shared sweep (so each + # combo is tied). When set independently, the inner loop sweeps the + # cross-product (write × nowrite); --skip-diagonal drops the tied subset. + def _split_or_share(split_csv, shared_values): + return _parse_sweep(split_csv) if split_csv else shared_values + block_size_m_write_values = _split_or_share(args.block_size_m_write, block_size_m_values) + block_size_m_nowrite_values = _split_or_share(args.block_size_m_nowrite, block_size_m_values) + num_warps_write_values = _split_or_share(args.num_warps_write, num_warps_values) + num_warps_nowrite_values = _split_or_share(args.num_warps_nowrite, num_warps_values) + num_stages_write_values = _split_or_share(args.num_stages_write, num_stages_values) + num_stages_nowrite_values = _split_or_share(args.num_stages_nowrite, num_stages_values) + cta_per_sm_write_values = _split_or_share(args.cta_per_sm_write, cta_per_sm_values) + cta_per_sm_nowrite_values = _split_or_share(args.cta_per_sm_nowrite, cta_per_sm_values) + num_loop_stages_write_values = _split_or_share(args.num_loop_stages_write, num_loop_stages_values) + num_loop_stages_nowrite_values = _split_or_share(args.num_loop_stages_nowrite, num_loop_stages_values) + # Whether any *_write / *_nowrite knob was independently set — used by + # --skip-diagonal to know if the cross-product is non-trivial. Without + # any split, the per-main values == shared values and skip-diagonal is + # a no-op (which is correct). + _any_split = any(getattr(args, name) for name in ( + "block_size_m_write", "block_size_m_nowrite", + "num_warps_write", "num_warps_nowrite", + "num_stages_write", "num_stages_nowrite", + "cta_per_sm_write", "cta_per_sm_nowrite", + "num_loop_stages_write", "num_loop_stages_nowrite", + )) + # TMA toggles — independent 0/1 sweep per path. The skip-dupe at the + # top of the inner loop body collapses cells where a flag's path is + # unreachable, so e.g. monolithic + WC=True only runs the value=0 + # cells for nowrite-load and rect-load. + use_tma_rect_load_values = _parse_sweep(args.use_tma_rect_load) + use_tma_replay_write_load_values = _parse_sweep(args.use_tma_replay_write_load) + use_tma_replay_nowrite_load_values = _parse_sweep(args.use_tma_replay_nowrite_load) + use_tma_replay_write_store_values = _parse_sweep(args.use_tma_replay_write_store) + + # --- Replay kernel --- + # Cache T-axis capacity (for prev_k validity check on the nowrite path). + max_window = getattr(args, "max_window", 0) or mtp_len + + # Build the list of scenarios to time. A scenario is one cell in the + # output: pure-mode scenarios fill prev_tokens with one constant before + # the timing loop; mix-mode scenarios feed a pre-baked per-iter samples + # tensor, with the per-iter copy captured inside the CUDA graph. Pure + # and mix can coexist in one call so a single nsys trace covers both. + scenarios = [] + for prev_k in prev_ks: + # On the nowrite path, new tokens append at [prev_k, prev_k+T) of + # the active buffer, so prev_k+T must fit within max_window. + # mode != monolithic dispatches per-slot from PNAT, so any + # prev_k <= max_window is valid for those modes. + if mode == "monolithic" and not write_checkpoint and prev_k + mtp_len > max_window: + continue + scenarios.append({ + "label": f"k{prev_k}", + "print_label": prev_k, + "fill": prev_k, + "pre_iter": None, + "iters": None, # use args.iters + }) + # Mix scenario: skip on monolithic (mono on mixed PNAT corrupts the + # wrong-mode slots). Persistent_main + mix is now supported: bench + # pre-bakes both a per-iter PNAT samples tensor and a per-iter + # n_writes samples tensor; pre_iter_fn copies row i of each into the + # kernel-input tensors (PNAT and n_writes_dev) on the same stream as + # the captured CUDA graph, so they're cold w.r.t. the in-graph L2 + # flush. + if mix_samples_cpu is not None and mode != "monolithic": + device = state_work.device + # Hardcode-sort: per-iter prev_tokens are CPU-sorted write-first. + # Kernel runs USE_PERM=False but the EO gate sees clustered modes. + # Output is scrambled (we don't permute x/B/C/dt to match) but + # timing is meaningful — isolates clustering benefit from the + # per-program perm-load overhead in --sort-slots. + src = mix_samples_sorted_cpu if (hardcode_sort and mix_samples_sorted_cpu is not None) else mix_samples_cpu + samples_gpu = torch.from_numpy(src).to(device=device, dtype=torch.int32) + + # For persistent_main + mix: pre-compute the per-iter n_writes + # (count of slots needing the write path = PNAT+T > max_window) + # and the (1,) scratch the kernel reads from. Both halves of + # persistent_main always launch in mix scenarios (host can't + # cheaply read n_writes per iter without a sync), so the kernel's + # slot-range derivation must be correct from device n_writes. + # persistent_dynamic doesn't need n_writes (the kernel ignores + # n_writes_dev when IS_DYNAMIC=True via Triton DCE), but we still + # allocate a sentinel scratch so the wrapper API is uniform. + n_writes_samples_gpu = None + n_writes_dev_mix = None + # n_writes per iter = number of slots that overflow the window. + # Computed for ALL mix scenarios so the JSON output (--json-detailed) + # can pair each iter's span_us with its mix composition for downstream + # analysis (group iters by # writes → per-bucket median → analytic + # expectation under the steady-state PNAT distribution). + n_writes_per_iter_all = ((src + mtp_len) > max_window).sum(axis=1).astype(np.int32) + if mode in ("persistent_main", "persistent_dynamic"): + n_writes_samples_gpu = torch.from_numpy(n_writes_per_iter_all).to( + device=device, dtype=torch.int32 + ) + n_writes_dev_mix = torch.zeros(1, dtype=torch.int32, device=device) + + # Build _mix_pre_iter — the closure that runs OUTSIDE the captured + # graph between replays. Updates: prev_tokens (always), + # slot_perm_buf (when sort_slots), n_writes_dev_mix (persistent). + if sort_slots and perm_samples_cpu is not None: + perm_samples_gpu = torch.from_numpy(perm_samples_cpu).to( + device=device, dtype=torch.int32 + ) + if n_writes_samples_gpu is not None: + def _mix_pre_iter(i, _s=samples_gpu, _ps=perm_samples_gpu, + _ns=n_writes_samples_gpu, _pt=prev_tokens, + _pm=slot_perm_buf, _nw=n_writes_dev_mix): + _pt.copy_(_s[i]) + _pm.copy_(_ps[i]) + _nw.copy_(_ns[i:i+1]) + else: + def _mix_pre_iter(i, _s=samples_gpu, _ps=perm_samples_gpu, + _pt=prev_tokens, _pm=slot_perm_buf): + _pt.copy_(_s[i]) + _pm.copy_(_ps[i]) + else: + if n_writes_samples_gpu is not None: + def _mix_pre_iter(i, _s=samples_gpu, _ns=n_writes_samples_gpu, + _pt=prev_tokens, _nw=n_writes_dev_mix): + _pt.copy_(_s[i]) + _nw.copy_(_ns[i:i+1]) + else: + def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): + _pt.copy_(_s[i]) + + # Mix iters override: if --mix-iters set, use it; else use args.iters. + mix_iters = getattr(args, "mix_iters", None) + scenarios.append({ + "label": f"mix{mix_label}", + "print_label": "mix", + "fill": None, + "pre_iter": _mix_pre_iter, + "iters": mix_iters, # None => use args.iters + # Pass through to _run_incr so the wrapper receives _n_writes_dev + # (mix scenarios) instead of _n_writes (pure scenarios). + "n_writes_dev": n_writes_dev_mix, + # Full per-iter n_writes array (size = warmup + iters). Used by + # the JSON-detailed output to pair each iter's span with its + # mix composition for post-hoc bucketing analysis. + "n_writes_per_iter": n_writes_per_iter_all, + }) + + # Pure scenarios don't pre-allocate n_writes_dev; mix scenarios do. + # Default empty-halves skip: True for pure (host knows n_writes, + # production-equivalent host-skip), False for mix (host can't read + # device n_writes per iter without sync, must always launch both). + for scn in scenarios: + scenario_n_writes_dev = scn.get("n_writes_dev") # None for pure + scenario_skip_empty = scenario_n_writes_dev is None + if scn["fill"] is not None: + prev_tokens.fill_(scn["fill"]) + prev_k_for_print = scn["print_label"] + scenario_pre_iter = scn["pre_iter"] + scenario_iters = scn.get("iters") # None => use args.iters + tag = f"incr_b{batch}_mtp{mtp_len}_{scn['label']}_s{state_dtype_name}_a{act_dtype_name}" + + # Iteration over per-cell knob combos. + # When NO per-main split is requested (_any_split=False), each row in + # the cross-product gives the same value to both write_main and + # nowrite_main (current behavior — backward-compat). When ANY split + # IS requested, we iterate the write and nowrite axes independently + # (cross-product blowup is the user's responsibility — they typically + # pair this with --skip-diagonal to drop the tied subset). + if _any_split: + _iter_axes = ( + block_size_m_write_values, block_size_m_nowrite_values, + num_warps_write_values, num_warps_nowrite_values, + num_stages_write_values, num_stages_nowrite_values, + precompute_num_warps_values, + precompute_num_stages_values, + heads_per_block_values, + maxnreg_values, num_ctas_values, + cta_per_sm_write_values, cta_per_sm_nowrite_values, + num_loop_stages_write_values, num_loop_stages_nowrite_values, + flatten_values, warp_specialize_values, + use_tma_rect_load_values, + use_tma_replay_write_load_values, + use_tma_replay_nowrite_load_values, + use_tma_replay_write_store_values, + ) + else: + # Tied: one value per shared knob. Wrap in single-element list for + # uniform iteration; the body sets w/nw both to the shared value. + _iter_axes = ( + block_size_m_values, [None], + num_warps_values, [None], + num_stages_values, [None], + precompute_num_warps_values, + precompute_num_stages_values, + heads_per_block_values, + maxnreg_values, num_ctas_values, + cta_per_sm_values, [None], + num_loop_stages_values, [None], + flatten_values, warp_specialize_values, + use_tma_rect_load_values, + use_tma_replay_write_load_values, + use_tma_replay_nowrite_load_values, + use_tma_replay_write_store_values, + ) + for ( + block_size_m_w, + block_size_m_nw, + num_warps_w, + num_warps_nw, + num_stages_w, + num_stages_nw, + precompute_num_warps, + precompute_num_stages, + heads_per_block, + maxnreg, + num_ctas, + cta_per_sm_w, + cta_per_sm_nw, + num_loop_stages_w, + num_loop_stages_nw, + flatten, + warp_specialize, + use_tma_rect_load, + use_tma_replay_write_load, + use_tma_replay_nowrite_load, + use_tma_replay_write_store, + ) in itertools.product(*_iter_axes): + # When tied, _nw values were placeholder None; fill from _w (the + # shared value). When split, _w and _nw came from independent lists. + if not _any_split: + block_size_m_nw = block_size_m_w + num_warps_nw = num_warps_w + num_stages_nw = num_stages_w + cta_per_sm_nw = cta_per_sm_w + num_loop_stages_nw = num_loop_stages_w + # Skip-diagonal: when split is on, drop the tied subset (same as a + # prior shared-knob sweep would cover). + if _any_split and args.skip_diagonal and ( + block_size_m_w == block_size_m_nw and + num_warps_w == num_warps_nw and + num_stages_w == num_stages_nw and + cta_per_sm_w == cta_per_sm_nw and + num_loop_stages_w == num_loop_stages_nw + ): + continue + # Backward-compat aliases used by the existing body below. When + # tied, these are simply the shared value. When split, the + # _write copy is used for sweep_tag and grouping (a stable choice + # so the tag is unique per (write, nowrite) combo). + block_size_m = block_size_m_w + num_warps = num_warps_w + num_stages = num_stages_w + cta_per_sm = cta_per_sm_w + num_loop_stages = num_loop_stages_w + # Skip-dupe for TMA flag sweeps: a flag whose code path isn't + # reachable in this cell produces identical timing for value=0 + # and value=1. We canonicalize by skipping value=1 cells when + # the flag's path is unreachable. Path reachability rules: + # * write path (replay write-load + write-store): mono+WC=True, + # OR any non-monolithic mode. + # * rect path (rect-load): rectangle_for_nowrite=True AND a + # nowrite path exists in this mode (mono+WC=False, OR any + # non-monolithic mode). + # * replay-nowrite path (nowrite-load): nowrite path exists + # AND rect isn't taking it: mono+WC=False+rect=False, OR + # any non-monolithic mode with rect=False. + _is_mono = (mode == "monolithic") + _write_path = (_is_mono and write_checkpoint) or (not _is_mono) + _rect_path = rectangle_for_nowrite and ( + (not _is_mono) or (_is_mono and not write_checkpoint) + ) + _replay_nowrite_path = ( + (_is_mono and not write_checkpoint and not rectangle_for_nowrite) + or ((not _is_mono) and not rectangle_for_nowrite) + ) + def _set(v): # flag set to a non-zero sweep value + return v is not None and v != 0 + if (_set(use_tma_rect_load) and not _rect_path + or _set(use_tma_replay_write_load) and not _write_path + or _set(use_tma_replay_nowrite_load) and not _replay_nowrite_path + or _set(use_tma_replay_write_store) and not _write_path): + continue + + # Pre-allocate n_writes_dev tensor OUTSIDE the captured graph for + # persistent modes in pure scenarios. Mix scenarios already have + # `scenario_n_writes_dev` pre-allocated. The wrapper's fallback + # `torch.tensor([...], device=...)` allocation would invalidate + # the CUDA-graph capture stream — must allocate here, before the + # `_run_incr` lambda (which is what gets captured) is defined. + # For persistent_dynamic the kernel ignores the value (IS_DYNAMIC + # DCE's the load); we still need a valid pointer. For + # persistent_main pure, the value is constant per cell so we set + # it once here. + _n_writes_dev_pure: torch.Tensor | None = None + _host_n_writes_pure: int | None = None + if mode in ("persistent_main", "persistent_dynamic") and scenario_n_writes_dev is None: + _n_writes_dev_pure = torch.zeros(1, dtype=torch.int32, device=state_work.device) + if mode == "persistent_main": + scn_fill = scn["fill"] + is_write_scenario_local = (scn_fill + mtp_len) > max_window + _host_n_writes_pure = batch if is_write_scenario_local else 0 + _n_writes_dev_pure.fill_(_host_n_writes_pure) + + def _run_incr( + block_size_m=block_size_m, + num_warps=num_warps, + num_stages=num_stages, + precompute_num_warps=precompute_num_warps, + precompute_num_stages=precompute_num_stages, + heads_per_block=heads_per_block, + maxnreg=maxnreg, + num_ctas=num_ctas, + cta_per_sm=cta_per_sm, + num_loop_stages=num_loop_stages, + flatten=flatten, + warp_specialize=warp_specialize, + use_tma_rect_load=use_tma_rect_load, + use_tma_replay_write_load=use_tma_replay_write_load, + use_tma_replay_nowrite_load=use_tma_replay_nowrite_load, + use_tma_replay_write_store=use_tma_replay_write_store, + ): + if with_conv1d: + x_call, B_call, C_call = _conv1d_split( + xbc_input_work, conv_state_work, launch_dependent_kernels=args.external_pdl + ) + extra_kwargs = {"launch_with_pdl": args.external_pdl} + else: + x_call, B_call, C_call = x, B, C + extra_kwargs = {} + # write_checkpoint is only meaningful for the checkpointing + # variant; replay variant ignores the kwarg. state_scales + # is also checkpointing-only (replay kernel doesn't quantize). + if args.variant == "checkpointing": + extra_kwargs["write_checkpoint"] = write_checkpoint + extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite + extra_kwargs["mode"] = mode + if sort_slots: + extra_kwargs["slot_perm"] = slot_perm_buf + # reverse_nowrite is meaningful in two ways: + # - with slot_perm: walk the perm tail-first + # - without slot_perm (hardcode-sort): walk pid_b + # itself tail-first via the REVERSE_PERM constexpr + if sort_slots or (hardcode_sort and reverse_nowrite): + extra_kwargs["reverse_nowrite"] = reverse_nowrite + if state_scales_work is not None: + extra_kwargs["state_scales"] = state_scales_work + if use_tma_rect_load: # 1 → True, 0/None → False + extra_kwargs["_use_tma_rect_load"] = True + if use_tma_replay_write_load: + extra_kwargs["_use_tma_replay_write_load"] = True + if use_tma_replay_nowrite_load: + extra_kwargs["_use_tma_replay_nowrite_load"] = True + if use_tma_replay_write_store: + extra_kwargs["_use_tma_replay_write_store"] = True + # persistent_main needs n_writes (count of write-mode + # slots in the pre-sorted batch) as a host-side int. + # Pure scenarios: every slot has the same PNAT, so + # n_writes is either 0 (all nowrite) or batch (all + # write) depending on whether PNAT+T overflows the + # window. Mix scenarios are skipped earlier. + if mode in ("persistent_main", "persistent_dynamic"): + # Per-cell sweep values for persistent-only knobs. + # Apply to both persistent variants. _parse_sweep + # returns [None] when the user didn't pass the flag, + # in which case we leave the wrapper's defaults. + if cta_per_sm is not None: + extra_kwargs["_cta_per_sm"] = cta_per_sm + if num_loop_stages is not None: + extra_kwargs["_num_loop_stages"] = num_loop_stages + if flatten is not None: + extra_kwargs["_flatten"] = bool(flatten) + if warp_specialize is not None: + extra_kwargs["_warp_specialize"] = bool(warp_specialize) + if mode in ("persistent_main", "persistent_dynamic"): + # persistent_main + mix REQUIRES sort: the kernel + # partitions slots [0, n_writes) = write half, + # [n_writes, batch) = nowrite half. This only holds + # if PNAT is monotone (writes first), which sort + # provides via either: + # sort_slots=1 → USE_PERM reads slot_perm to remap + # hardcode_sort=1 → PNAT itself is CPU-pre-sorted + # persistent_dynamic doesn't need sort (per-slot + # runtime dispatch); persistent_main pure scenarios + # are trivially sorted (homogeneous PNAT). + if (mode == "persistent_main" + and scenario_n_writes_dev is not None + and not (sort_slots or hardcode_sort)): + raise AssertionError( + "persistent_main + mix requires sort_slots=1 " + "or hardcode_sort=1 — kernel partitions slots " + "by index, which is only valid when PNAT is " + "monotone (writes first). Without sort, the " + "partition silently mismatches actual slot " + "modes. Re-run with --sort-slots 1 or " + "--hardcode-sort 1." + ) + # n_writes plumbing: pure scenarios pass an int + # (host knows the value, can host-skip empty halves); + # mix scenarios pass a (1,) device tensor that the + # bench's pre_iter_fn updates per replay. + # _persistent_skip_empty_halves=False on mix so both + # halves always launch (kernel uses device n_writes + # to derive its slot range). + if scenario_n_writes_dev is not None: + # Mix path: caller-allocated tensor, updated per + # iter by scenario_pre_iter outside capture. + extra_kwargs["_n_writes_dev"] = scenario_n_writes_dev + extra_kwargs["_persistent_skip_empty_halves"] = False + elif mode == "persistent_main": + # Pure: caller pre-allocated `_n_writes_dev_pure` + # outside this lambda (so the alloc doesn't land + # inside the captured graph). Pass both the + # tensor and the host int so the wrapper can use + # host-skip when `_persistent_skip_empty_halves`. + extra_kwargs["_n_writes"] = _host_n_writes_pure + extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure + extra_kwargs["_persistent_skip_empty_halves"] = scenario_skip_empty + elif mode == "persistent_dynamic": + # persistent_dynamic pure: kernel ignores n_writes + # via IS_DYNAMIC DCE, but the wrapper needs a + # valid (1,) tensor pointer. Pass the pre-allocated + # zero tensor to avoid any in-capture alloc. + extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure + variant_fn( + state_work, + old_x_work, + old_B_work, + old_dt_work, + old_dA_cumsum_work, + cache_buf_idx_work, + prev_tokens, + x=x_call, + dt=dt, + A=A, + B=B_call, + C=C_call, + out=out_incr, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + rand_seed=rand_seed, + philox_rounds=args.philox_rounds, + use_internal_pdl=args.internal_pdl, + _block_size_m=block_size_m, + _num_warps=num_warps, + _num_stages=num_stages, + _precompute_num_warps=precompute_num_warps, + _precompute_num_stages=precompute_num_stages, + _heads_per_block=heads_per_block, + _maxnreg=maxnreg, + _num_ctas=num_ctas, + # Per-main overrides (None = tied to shared above; explicit + # only when the inner loop is iterating split axes). + _block_size_m_write=block_size_m_w if _any_split else None, + _block_size_m_nowrite=block_size_m_nw if _any_split else None, + _num_warps_write=num_warps_w if _any_split else None, + _num_warps_nowrite=num_warps_nw if _any_split else None, + _num_stages_write=num_stages_w if _any_split else None, + _num_stages_nowrite=num_stages_nw if _any_split else None, + _cta_per_sm_write=cta_per_sm_w if _any_split else None, + _cta_per_sm_nowrite=cta_per_sm_nw if _any_split else None, + _num_loop_stages_write=num_loop_stages_w if _any_split else None, + _num_loop_stages_nowrite=num_loop_stages_nw if _any_split else None, + **extra_kwargs, + ) + + parts = [] + # When tied (not _any_split), emit the shared single-value tag + # (M=8 etc). When split, emit explicit Mw / Mnw tags so cells + # with the same shared value but different per-main values get + # unique JSON keys. + def _emit_split(name_w, name_nw, val_w, val_nw): + if val_w is None and val_nw is None: + return + if not _any_split or val_w == val_nw: + parts.append(f"{name_w[:-1]}={val_w}") # strip the 'w' suffix + else: + parts.append(f"{name_w}={val_w}") + parts.append(f"{name_nw}={val_nw}") + _emit_split("Mw", "Mnw", block_size_m_w, block_size_m_nw) + _emit_split("Ww", "Wnw", num_warps_w, num_warps_nw) + _emit_split("Sw", "Snw", num_stages_w, num_stages_nw) + if precompute_num_warps is not None: + parts.append(f"pW={precompute_num_warps}") + if precompute_num_stages is not None: + parts.append(f"pS={precompute_num_stages}") + if heads_per_block is not None: + parts.append(f"H={heads_per_block}") + if maxnreg is not None: + parts.append(f"R={maxnreg}") + if num_ctas is not None: + parts.append(f"CT={num_ctas}") + # Persistent-only knobs (only meaningful when MODE=persistent_main; + # printed unconditionally so output rows are uniformly comparable + # across modes when the user passed these sweeps). + _emit_split("CPSw", "CPSnw", cta_per_sm_w, cta_per_sm_nw) + _emit_split("LSw", "LSnw", num_loop_stages_w, num_loop_stages_nw) + if flatten is not None: + parts.append(f"FL={flatten}") + if warp_specialize is not None: + parts.append(f"WS={warp_specialize}") + # TMA sweep tags. Four wrapper-level flags map to three + # kernel-level constexprs (rect-load and replay-nowrite-load + # share `USE_TMA_LOAD_NOWRITE`, picked by the wrapper based on + # RECTANGLE). TMARL specifically gates the rectangle path's + # state load; TMANL specifically gates the replay-style + # nowrite path's state load. Distinct because their measured + # perf profiles differ (see CHECKPOINTING_DESIGN.md item #17: + # rect TMA is "not a win" while replay-nowrite TMA is the + # biggest measured win at int8 b>=64). + if use_tma_rect_load is not None: + parts.append(f"TMARL={use_tma_rect_load}") # rect path load + if use_tma_replay_write_load is not None: + parts.append(f"TMAWL={use_tma_replay_write_load}") # replay-write load + if use_tma_replay_nowrite_load is not None: + parts.append(f"TMANL={use_tma_replay_nowrite_load}") # replay-NOWRITE load (NOT rect) + if use_tma_replay_write_store is not None: + parts.append(f"TMAWS={use_tma_replay_write_store}") # replay-write store + parts.append(f"SR={1 if use_philox else 0}") + parts.append(f"RECT={1 if rectangle_for_nowrite else 0}") + parts.append(f"WC={1 if write_checkpoint else 0}") + parts.append(f"MODE={mode}") + parts.append(f"SORT={1 if sort_slots else 0}") + parts.append(f"REVN={1 if reverse_nowrite else 0}") + parts.append(f"HSORT={1 if hardcode_sort else 0}") + sweep_suffix = (" " + ",".join(parts)) if parts else "" + sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") + + reset_fn = _reset_conv1d_realistic if with_conv1d else _reset + # --retry-cells filter: only time cells whose tag is in the + # retry set. Cheaper than re-enumerating cells in the orchestrator. + retry_set = getattr(args, "_retry_cells_set", None) + if retry_set and sweep_tag not in retry_set: + continue + if warmup_only: + reset_fn() + if scenario_pre_iter is not None: + scenario_pre_iter(0) + _run_incr() + torch.cuda.synchronize() + else: + # Inline retry: CUPTI sometimes loses records under PDL + + # high cell count; retrying the SAME cell often catches it + # because the failure is transient at the kernel-launch level. + # Per --cupti-retry budget. On final failure, append tag to + # the skipped list for an external rerun in a fresh process. + retry_budget = max(0, getattr(args, "cupti_retry", 1)) + stats = None + for attempt in range(retry_budget + 1): + stats = _time_kernel( + args, _run_incr, reset_fn, sweep_tag, + expected_K=_kernels_per_iter_incremental( + mode, with_conv1d=with_conv1d, + persistent_skip_empty=scenario_skip_empty, + ), + pre_iter_fn=scenario_pre_iter, + iters_override=scenario_iters, + ) + if stats is not None: + break + if attempt < retry_budget: + print( + f"[retry] CUPTI mismatch on {sweep_tag!r}; " + f"retrying ({attempt + 1}/{retry_budget})", + file=sys.stderr, + ) + if stats is None: + args._skipped_cells.append(sweep_tag) + + # Attach n_writes_per_iter for --json-detailed bucketing. + # For pure scenarios, n_writes is constant: 0 (nowrite) or + # batch (write), determined by scn["fill"] + mtp_len > max_window. + # For mix, scn carries the precomputed per-iter array. + if stats is not None and getattr(args, "json_detailed", False): + eff_iters = scenario_iters if scenario_iters is not None else args.iters + if scn["fill"] is not None: + # Pure scenario: constant n_writes for every iter. + is_write = (scn["fill"] + mtp_len > max_window) + per_iter_nw = [batch if is_write else 0] * eff_iters + else: + # Mix scenario: slice off warmup, keep timed iters. + nw_full = scn.get("n_writes_per_iter") + if nw_full is not None: + per_iter_nw = nw_full[args.warmup:args.warmup + eff_iters].tolist() + else: + per_iter_nw = None + if per_iter_nw is not None: + stats["n_writes_per_iter"] = per_iter_nw + + if stats is not None: + _print_row( + show_kernel_col, + args.variant, + batch, + mtp_len, + prev_k_for_print, + state_dtype_name, + act_dtype_name, + stats, + sweep_suffix, + json_results=getattr(args, "_json_results", None), + tp_size=args.tp_size, + json_detailed=getattr(args, "json_detailed", False), + ) + + +# Map full torch dtype name → short tag used in JSON keys (matches collect.py). +_DTYPE_SHORT = { + "float32": "fp32", "bfloat16": "bf16", "float16": "fp16", + "int8": "int8", "int16": "int16", "float8_e4m3fn": "fp8", +} + + +def _build_json_key( + kernel_name, batch, mtp_len, prev_k, state_dtype_name, sweep_suffix, tp_size +): + """Build a key matching collect.py's kernel_data.json convention: + + incremental/{batch}/{mtp}/{sd}/k{k}/{sweep_parts}/tp{tp} + triton/{batch}/{mtp}/{sd}/tp{tp} + flashinfer/{batch}/{mtp}/{sd}/tp{tp} + + `kernel_name` is what _print_row receives: variant name for the timed + kernel (replay/checkpointing) or baseline name for the baseline row. + Variant rows collapse to "incremental" — the variant choice is captured + by the sweep flags collect.py would otherwise apply via --variant. + """ + if kernel_name in ("replay", "checkpointing"): + kind = "incremental" + else: + kind = kernel_name # "triton" / "flashinfer" + + sd = _DTYPE_SHORT.get(state_dtype_name, state_dtype_name) + parts = [kind, str(batch), str(mtp_len), sd] + if prev_k != "N/A": + parts.append(f"k{prev_k}") + if sweep_suffix: + # sweep_suffix format: " M=4,W=1,S=1,SR=0,RECT=0,WC=1" + # collect.py format: "M4_W1_S1_SR0_RECT0_WC0" + # Strip leading/trailing whitespace, drop '=', commas → underscores. + parts.append( + sweep_suffix.strip().replace("=", "").replace(",", "_") + ) + parts.append(f"tp{tp_size}") + return "/".join(parts) + + +def _print_row( + show_kernel_col, + kernel_name, + batch, + mtp_len, + prev_k, + state_dtype_name, + act_dtype_name, + stats, + sweep_suffix="", + json_results=None, + tp_size=None, + json_detailed=False, +): + """Print one summary row and optionally accumulate stats for JSON output. + + `stats` is a dict from _time_kernel: {median, p95, p99, n, iters_us, + [per_kernel]}. The summary table only shows the headline percentiles. + JSON output captures median/p95/p99/n by default; with json_detailed=True + it also captures the per-iter and per-kernel data. + """ + kernel_col = f"{kernel_name:>11} | " if show_kernel_col else "" + print( + f"| {kernel_col}{batch:>5} | {mtp_len:>7} | {str(prev_k):>6} | " + f"{state_dtype_name:>11} | {act_dtype_name:>9} | " + f"{stats['median']:>9.2f} | {stats['p95']:>7.2f} | {stats['p99']:>7.2f} |" + f"{sweep_suffix}" + ) + if json_results is not None: + key = _build_json_key( + kernel_name, batch, mtp_len, prev_k, state_dtype_name, + sweep_suffix, tp_size, + ) + if json_detailed: + json_results[key] = stats + else: + json_results[key] = { + k: stats[k] for k in ("median", "p95", "p99", "n") + if k in stats + } + + +# Main benchmark loop + + +def _run_benchmark(args) -> None: + # JSON accumulator — populated by _print_row when --json-output is set. + # Stash on args so we don't need to thread a dict through every helper. + args._json_results = {} if getattr(args, "json_output", None) else None + + # Skipped cells accumulator — populated by _bench_config when CUPTI capture + # mismatch causes a cell to be skipped. Written to args.skipped_output + # (or derived from json_output) at end of run. Pair with --retry-cells + # to re-time only the skipped cells in a fresh process. + args._skipped_cells = [] + + # Retry-cells filter set. When non-empty, only cells whose tag is in this + # set will be timed; all others are silently skipped. Cells are matched + # against the sweep_tag string built in _bench_config. + args._retry_cells_set: set[str] = set() + if getattr(args, "retry_cells", None): + with open(args.retry_cells) as f: + data = json.load(f) + # Accept either a list of tag strings, or the same shape we write + # (dict with "skipped" key). Tolerant of both for hand-edited files. + if isinstance(data, dict) and "skipped" in data: + data = data["skipped"] + args._retry_cells_set = set(data) + print( + f"[retry] --retry-cells loaded {len(args._retry_cells_set)} tags " + f"from {args.retry_cells}; sweep will skip all other cells.", + file=sys.stderr, + ) + + assert args.nheads % args.tp_size == 0, ( + f"nheads ({args.nheads}) must be divisible by tp_size ({args.tp_size})" + ) + assert args.ngroups % args.tp_size == 0, ( + f"ngroups ({args.ngroups}) must be divisible by tp_size ({args.tp_size})" + ) + args.tp_nheads = args.nheads // args.tp_size + args.tp_ngroups = args.ngroups // args.tp_size + + batch_sizes = [int(x) for x in args.batch_sizes.split(",")] + mtp_lengths = [int(x) for x in args.mtp_lengths.split(",")] + + dtype_map = { + "bf16": torch.bfloat16, + "fp32": torch.float32, + "fp16": torch.float16, + "int8": torch.int8, + "int16": torch.int16, + "fp8": torch.float8_e4m3fn, + } + state_dtypes = [dtype_map[s] for s in args.state_dtypes.split(",")] + act_dtypes = [dtype_map[s] for s in args.act_dtypes.split(",")] + + # Resolve baseline function + if args.baseline == "flashinfer": + from flashinfer.mamba import selective_state_update as baseline_fn + elif args.baseline == "triton": + baseline_fn = selective_state_update + else: + baseline_fn = None + + # --with-conv1d uses its own realistic L2 flush (cold cache flush then + # hot in_proj write). Override the generic l2_flush to avoid double-flushing. + if args.with_conv1d: + args.l2_flush = False + _init_l2_flush() # still needed for the realistic reset's flush step + elif args.l2_flush: + _init_l2_flush() + + if args.compile_threads > 0: + _compile_warmup_phase( + args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, + baseline_fn, max_workers=args.compile_threads, + ) + + # Pre-warm the per-(state_dtype, act_dtype, mtp_len, ...) tensor cache at + # the largest requested batch size. Without this, the timing loop would + # progressively grow the cache as it encounters larger batches (e.g., + # iterate 1 -> 16 -> 64 -> 128 -> 512 = 5 separate growth allocations, + # each freeing the previous buffers). Pre-warming at max-batch up front + # makes every subsequent timing cell a view-slice (zero alloc cost). + _max_batch = max(batch_sizes) + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + for mtp_len in mtp_lengths: + _build_tensors( + _max_batch, mtp_len, state_dtype, act_dtype, + args.tp_nheads, args.head_dim, args.d_state, args.tp_ngroups, + max_window=getattr(args, "max_window", None) or None, + ) + + if args.profile: + torch.cuda.cudart().cudaProfilerStart() + + # Print header + if baseline_fn is not None: + print( + f"| {'kernel':>11} | {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " + f"{'state_dtype':>11} | {'act_dtype':>9} | " + f"{'median_us':>9} | {'p95_us':>7} | {'p99_us':>7} |" + ) + print( + f"|{'-' * 13}|{'-' * 7}|{'-' * 9}|{'-' * 8}|" + f"{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" + ) + else: + print( + f"| {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " + f"{'state_dtype':>11} | {'act_dtype':>9} | " + f"{'median_us':>9} | {'p95_us':>7} | {'p99_us':>7} |" + ) + print( + f"|{'-' * 7}|{'-' * 9}|{'-' * 8}|{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" + ) + + sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) + rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) + write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) + modes_list = getattr(args, "modes_list", ["monolithic"]) + sort_list = getattr(args, "sort_slots_list", [False]) + rev_list = getattr(args, "reverse_nowrite_list", [False]) + hsort_list = getattr(args, "hardcode_sort_list", [False]) + + # Pre-load AL distribution for mix mode (if --mix-csv set). + mix_al = None + mix_label = "" + if args.mix_csv is not None: + from pathlib import Path as _Path + from checkpoint_mix_sim import load_al_distribution as _load_al + mix_label = _Path(args.mix_csv).stem + # T (= mtp_len) varies per cell; load once with the LARGEST mtp so + # we have enough columns; the loader normalizes the dist anyway. + mix_al = _load_al(_Path(args.mix_csv), T=max(mtp_lengths), column=args.mix_csv_column) + + for batch in batch_sizes: + for mtp_len in mtp_lengths: + # Resolve prev_k fractions → clamped integers in [0, mtp_len] + prev_ks = _resolve_prev_ks(args, mtp_len) + + # Pre-generate mix samples once per (batch, mtp_len) cell so all + # tuning configs see the same per-iter prev_tokens vectors — + # tuning differences become signal, mix-noise is shared. + # Size the sample buffer for the LARGER of args.iters and + # args.mix_iters since mix scenarios use mix_iters. + mix_samples_cpu = None + perm_samples_cpu = None # per-iter slot perm sorted write-first + mix_samples_sorted_cpu = None # per-iter prev_tokens, write-first + if mix_al is not None: + from checkpoint_mix_sim import sample_steady_state_pnat as _sample_pnat + _max_window = getattr(args, "max_window", 0) or mtp_len + _max_iters = max(args.iters, getattr(args, "mix_iters", None) or args.iters) + mix_samples_cpu = _sample_pnat( + mix_al, T=mtp_len, window=_max_window, batch=batch, + K=args.warmup + _max_iters, seed=args.mix_seed, + ) + if any(sort_list) or any(hsort_list): + # write-first stable argsort: kind='stable' preserves + # original-slot order within each mode group. + write_mask = ( + mix_samples_cpu + mtp_len > _max_window + ).astype(np.int8) # 1 = write, 0 = nowrite + perm_idx = np.argsort( + -write_mask, kind="stable", axis=-1 + ).astype(np.int32) + if any(sort_list): + perm_samples_cpu = perm_idx + if any(hsort_list): + # Apply the perm to the prev_tokens samples themselves. + # Result row i = mix_samples_cpu[i] reordered such + # that write-mode entries come first. + mix_samples_sorted_cpu = np.take_along_axis( + mix_samples_cpu, perm_idx, axis=-1 + ).astype(mix_samples_cpu.dtype) + + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + for sr_mode in sr_modes_list: + for mode in modes_list: + # Non-monolithic modes ignore write_checkpoint + # (per-slot from PNAT) — collapse the sweep so we + # don't duplicate identical cells. + effective_write_modes = ( + write_modes_list if mode == "monolithic" else [True] + ) + for write_ckpt in effective_write_modes: + # Rectangle is meaningful for: nowrite cells in + # monolithic; always for dynamic / doublelaunch + # (constexpr knob). + if mode == "monolithic": + effective_rect_list = ( + [False] if write_ckpt else rect_list + ) + else: + effective_rect_list = rect_list + for rect in effective_rect_list: + # Sort/reverse only meaningful for the + # dl-family early-out kernels AND only + # against the mix scenario (the actual + # sort experiment). Pure k= scenarios + # under sort=1 would just run a + # USE_PERM=True kernel against an + # identity perm — same data point as + # sort=0 + extra compile. Skip sort=1 + # when no mix is configured; mono / + # dynamic also skip sort=1; reverse=1 + # with sort=0 is a no-op (skip). + # Note: "persistent_main" is included in + # the dl-family for sort/hsort sweep + # eligibility — it consumes the same + # slot_perm and benefits from the same + # write-first clustering. It additionally + # requires _n_writes (count of write + # slots) which the bench computes from + # the pure-scenario PNAT (mix scenarios + # not yet supported for persistent_main). + is_dl_family = mode in ( + "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", "persistent_main", + "persistent_dynamic", + ) + can_sort = ( + is_dl_family and mix_samples_cpu is not None + ) + effective_sort_list = ( + sort_list if can_sort else [False] + ) + effective_hsort_list = ( + hsort_list if can_sort else [False] + ) + for sort_slots in effective_sort_list: + for hardcode_sort in effective_hsort_list: + # sort_slots and hardcode_sort + # are alternative experiments + # for the same idea — skip the + # combined cell to avoid double + # interpretation. + if sort_slots and hardcode_sort: + continue + # rev=1 is meaningful with EITHER + # sort_slots=1 (perm-based) or + # hardcode_sort=1 (raw pid_b + # subtraction in unsorted-perm + # path). rev=1 with both 0 is + # a no-op. + effective_rev_list = ( + rev_list if (sort_slots or hardcode_sort) else [False] + ) + for reverse_nowrite in effective_rev_list: + _bench_config( + args, batch, mtp_len, + prev_ks, state_dtype, + act_dtype, baseline_fn, + sr_mode=sr_mode, + rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, + mode=mode, + mix_samples_cpu=mix_samples_cpu, + mix_label=mix_label, + sort_slots=sort_slots, + reverse_nowrite=reverse_nowrite, + perm_samples_cpu=perm_samples_cpu, + hardcode_sort=hardcode_sort, + mix_samples_sorted_cpu=mix_samples_sorted_cpu, + ) + + if args.profile: + torch.cuda.cudart().cudaProfilerStop() + + if args.json_output and args._json_results is not None: + payload = { + "metadata": { + "timestamp": datetime.now().isoformat(), + "cmd": " ".join(sys.argv), + "tp_size": args.tp_size, + "warmup": args.warmup, + "iters": args.iters, + "variant": args.variant, + "cupti": getattr(args, "cupti", False), + }, + "results": args._json_results, + } + tmp = args.json_output + ".tmp" + with open(tmp, "w") as f: + json.dump(payload, f, indent=2) + os.replace(tmp, args.json_output) + print(f"\nJSON results written to: {args.json_output} " + f"({len(args._json_results)} entries)") + + # Write the skipped-cells sidecar. Pair with --retry-cells in a separate + # invocation to re-time the failed cells in a fresh process. + skipped_path = getattr(args, "skipped_output", None) + if skipped_path is None and args.json_output: + # Derive default: foo.json -> foo.skipped.json + skipped_path = args.json_output.rsplit(".", 1)[0] + ".skipped.json" + if skipped_path is not None and args._skipped_cells: + payload = { + "metadata": { + "timestamp": datetime.now().isoformat(), + "cmd": " ".join(sys.argv), + "skipped_count": len(args._skipped_cells), + }, + "skipped": args._skipped_cells, + } + tmp = skipped_path + ".tmp" + with open(tmp, "w") as f: + json.dump(payload, f, indent=2) + os.replace(tmp, skipped_path) + print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " + f"tags written to: {skipped_path}", file=sys.stderr) + elif args._skipped_cells: + # No output path but there are skipped cells — emit a stderr summary. + print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " + f"first 5: {args._skipped_cells[:5]}", file=sys.stderr) + + +# CLI + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Benchmark replay_selective_state_update Triton kernel", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--nheads", + type=int, + default=NHEADS, + help="Full-model nheads (divided by --tp-size for per-GPU slice)", + ) + parser.add_argument( + "--ngroups", + type=int, + default=NGROUPS, + help="Full-model ngroups (divided by --tp-size for per-GPU slice)", + ) + parser.add_argument( + "--head-dim", type=int, default=HEAD_DIM, help="Head dimension (not TP-split)" + ) + parser.add_argument( + "--d-state", type=int, default=D_STATE, help="SSM state dimension (not TP-split)" + ) + parser.add_argument( + "--tp-size", + type=int, + default=TP_SIZE, + help="Tensor parallel size; divides nheads and ngroups", + ) + parser.add_argument( + "--batch-sizes", default="1,2,4,8", help="Comma-separated decode batch sizes" + ) + parser.add_argument( + "--mtp-lengths", + default="1,2,4,8", + help="Comma-separated per-request sequence lengths (num_draft_tokens + 1 target)", + ) + parser.add_argument( + "--state-dtypes", + default="fp32", + help="Comma-separated state dtypes: fp16,bf16,fp32,int8,int16,fp8. " + "Quantized dtypes (int8/int16/fp8) require the checkpointing variant " + "and skip baselines (selective_state_update doesn't accept them).", + ) + parser.add_argument( + "--act-dtypes", + default="bf16", + help="Comma-separated activation dtypes for x/B/C/dt: fp32,bf16", + ) + parser.add_argument("--warmup", type=int, default=20, help="Number of warmup iterations") + parser.add_argument("--iters", type=int, default=100, help="Number of timed iterations") + parser.add_argument( + "--compile-threads", + type=int, + default=64, + help="Number of THREADS used in the compile-warmup phase (one call " + "per (batch, mtp_len, prev_k, dtype, sweep) cell, parallelized over " + "N threads). Triton compile releases the GIL, so threads compile " + "in parallel and populate the persistent cache for free hits during " + "the sequential timed phase. 0 disables the phase. Default 64.", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Wrap timed region in cudaProfilerStart/Stop (for ncu --target-processes all)", + ) + parser.add_argument( + "--l2-flush", + action=argparse.BooleanOptionalAction, + default=True, + help="L2 eviction between iterations", + ) + parser.add_argument( + "--cuda-graph", + action=argparse.BooleanOptionalAction, + default=True, + help="Capture all warmup + timed iterations in a " + "single CUDA graph with per-iteration events " + "inside the graph, eliminating all host overhead.", + ) + parser.add_argument( + "--cupti", + action=argparse.BooleanOptionalAction, + default=True, + help="Time kernels via CUPTI Activity API (1 ns from the GPU " + "profiling fabric); per-iter span = max(kernel_end) - " + "min(kernel_start). Default ON. --no-cupti disables in-bench " + "timing entirely (kernels still run, but median/p95/p99 are zero) " + "— use when wrapping the bench in nsys/ncu, where the external " + "profiler provides timings and our CUPTI subscriber would conflict.", + ) + parser.add_argument( + "--json-output", + default=None, + help="If set, write per-cell results to this JSON file in the " + "shape consumed by collect.py / report.py. See the 'JSON output " + "schema' section at the top of this file.", + ) + parser.add_argument( + "--json-detailed", + action=argparse.BooleanOptionalAction, + default=False, + help="When --json-output is set, also include iters_us (raw per-iter " + "spans) and per_kernel (per-iter relative start/end timestamps for " + "each kernel) — useful for PDL overlap analysis but adds ~4 KB/cell. " + "Default off keeps records to ~40 bytes (median/p95/p99/n only).", + ) + parser.add_argument( + "--cupti-retry", + type=int, + default=1, + help="On CUPTI capture mismatch (kernel record count != expected), " + "retry the cell this many times in-process before giving up. CUPTI " + "gets racy after thousands of cells in one process (PDL + small " + "kernels occasionally lose records); a single retry usually catches " + "transient cases. Set 0 to disable and skip on first mismatch.", + ) + parser.add_argument( + "--skipped-output", + default=None, + help="Path to write the list of cells that failed CUPTI capture even " + "after --cupti-retry retries (JSON list of sweep_tag strings). " + "Default: derived from --json-output by replacing .json with " + ".skipped.json. Pair with --retry-cells in a separate invocation " + "(fresh process = fresh CUPTI subscriber) to re-time these cells.", + ) + parser.add_argument( + "--retry-cells", + default=None, + help="Path to a JSON list of cell tags (as written by --skipped-output " + "in a prior invocation). When set, the sweep iterates as normal but " + "skips any cell whose tag is NOT in the listed set. Lets collect.py " + "drive a retry pass over only the cells that failed the first time.", + ) + parser.add_argument( + "--prev-tokens-fracs", + default="0,0.5,1.0", + type=lambda s: [float(x) for x in s.split(",")], + help="Fractions of mtp_len to use as prev_num_accepted_tokens " + "for the replay kernel sweep. Values are rounded " + "and clamped to [0, mtp_len].", + ) + parser.add_argument( + "--baseline", + default=None, + nargs="?", + const="triton", + choices=[None, "triton", "flashinfer"], + help="Baseline to benchmark alongside the replay kernel. " + "'triton': native Triton selective_state_update. " + "'flashinfer': flashinfer selective_state_update (same signature). " + "Pass --baseline alone for 'triton'. Default: no baseline.", + ) + parser.add_argument( + "--output", + default=None, + help="Path to save results (file or directory). " + "If a directory, writes benchmark_replay_.txt inside it.", + ) + parser.add_argument( + "--block-size-m", + type=str, + default=None, + help="Override BLOCK_SIZE_M: single value or comma-separated sweep (e.g. '4,8,16,32').", + ) + parser.add_argument( + "--num-warps", + type=str, + default=None, + help="Override num_warps: single value or comma-separated sweep (e.g. '1,2,4').", + ) + parser.add_argument( + "--internal-pdl", + action=argparse.BooleanOptionalAction, + default=True, + help="Internal PDL between precompute and main kernels (default: on).", + ) + parser.add_argument( + "--num-stages", + type=str, + default=None, + help="Override num_stages for the main kernel (comma-separated sweep).", + ) + parser.add_argument( + "--block-size-m-write", type=str, default=None, + help="Sweep BLOCK_SIZE_M for the WRITE main only (overrides --block-size-m " + "for the write half). Tied to --block-size-m if unset.", + ) + parser.add_argument( + "--block-size-m-nowrite", type=str, default=None, + help="Sweep BLOCK_SIZE_M for the NOWRITE main only. Tied to --block-size-m if unset.", + ) + parser.add_argument( + "--num-warps-write", type=str, default=None, + help="Sweep num_warps for the WRITE main only. Tied to --num-warps if unset.", + ) + parser.add_argument( + "--num-warps-nowrite", type=str, default=None, + help="Sweep num_warps for the NOWRITE main only. Tied to --num-warps if unset.", + ) + parser.add_argument( + "--num-stages-write", type=str, default=None, + help="Sweep num_stages for the WRITE main only. Tied to --num-stages if unset.", + ) + parser.add_argument( + "--num-stages-nowrite", type=str, default=None, + help="Sweep num_stages for the NOWRITE main only. Tied to --num-stages if unset.", + ) + parser.add_argument( + "--cta-per-sm-write", type=str, default=None, + help="Sweep cta_per_sm for the WRITE persistent_main only. Tied to --cta-per-sm if unset.", + ) + parser.add_argument( + "--cta-per-sm-nowrite", type=str, default=None, + help="Sweep cta_per_sm for the NOWRITE persistent_main only. Tied to --cta-per-sm if unset.", + ) + parser.add_argument( + "--num-loop-stages-write", type=str, default=None, + help="Sweep num_loop_stages for the WRITE persistent_main only. Tied to --num-loop-stages if unset.", + ) + parser.add_argument( + "--num-loop-stages-nowrite", type=str, default=None, + help="Sweep num_loop_stages for the NOWRITE persistent_main only. Tied to --num-loop-stages if unset.", + ) + parser.add_argument( + "--skip-diagonal", action=argparse.BooleanOptionalAction, default=False, + help="When sweeping any per-main *_write / *_nowrite knobs, skip cells " + "where ALL splittable knobs satisfy write_value == nowrite_value (i.e. " + "the 'diagonal' that's already covered by a prior shared-knob sweep). " + "Useful for incremental sweeps that extend earlier results without redoing " + "the tied-knob cells.", + ) + parser.add_argument( + "--precompute-num-warps", + type=str, + default=None, + help="Override num_warps for precompute kernel (comma-separated sweep).", + ) + parser.add_argument( + "--precompute-num-stages", + type=str, + default=None, + help="Override num_stages for precompute kernel (comma-separated sweep).", + ) + parser.add_argument( + "--max-window", + type=int, + default=16, + help="Cache T-axis capacity (max replay buffer length). Default 16 " + "matches Nemotron-3-Super-120B production. Pass 0 to fall back to " + "mtp_len (degenerate every-step-checkpoint case, mostly unused).", + ) + parser.add_argument( + "--prev-tokens-int", + type=lambda s: [int(x) for x in s.split(",")] if s else None, + default=None, + help="Absolute prev_num_accepted_tokens values to test, comma-separated " + "(e.g. '0,10,11,16'). Clamped to [0, max_window]. When set, " + "overrides --prev-tokens-fracs.", + ) + parser.add_argument( + "--write-checkpoint", + action=argparse.BooleanOptionalAction, + default=True, + help="Whether the checkpointing kernel should write the post-replay " + "state to HBM. True = checkpoint step (default). False = " + "non-checkpoint step (skip state HBM write + Philox). No effect on " + "the replay variant. Ignored if --write-modes is set.", + ) + parser.add_argument( + "--write-modes", + type=str, + default=None, + help="Comma-separated 0/1 values to sweep both write modes in a " + "single nsys process — for apples-to-apples comparison of write " + "vs nowrite (replay) vs nowrite (rectangle) within one timeline. " + "Skips silently for (write=False, prev_k+T>max_window) combos. " + "When set, overrides --write-checkpoint.", + ) + parser.add_argument( + "--with-conv1d", + action="store_true", + help="Include conv1d kernel before replay SSM. " + "Uses realistic L2 flush: cold caches flushed, hot in_proj output " + "kept warm. Measures conv1d → precompute → main span.", + ) + parser.add_argument( + "--external-pdl", + action=argparse.BooleanOptionalAction, + default=True, + help="External PDL: conv1d launches dependents, precompute waits. " + "Only relevant with --with-conv1d. --no-external-pdl disables.", + ) + parser.add_argument( + "--heads-per-block", + type=str, + default=None, + help="Override HEADS_PER_BLOCK for precompute kernel (comma-separated sweep).", + ) + parser.add_argument( + "--maxnreg", + type=str, + default=None, + help="Override maxnreg for the main kernel (comma-separated sweep).", + ) + parser.add_argument( + "--num-ctas", + type=str, + default=None, + help="Override num_ctas for the main kernel (comma-separated sweep).", + ) + parser.add_argument( + "--cta-per-sm", + type=str, + default=None, + help="CTAs per SM in the 1D persistent grid for mode=persistent_main " + "(comma-separated sweep). num_persistent = cta_per_sm × NUM_SMS. " + "Default = 1 (one CTA per SM). Replaces the old --num-persistent. " + "Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--num-loop-stages", + type=str, + default=None, + help="num_stages on the inner tl.range(...) persistent loop for " + "mode=persistent_main (comma-separated sweep). Default = 2. Note: " + "this is loop-level, NOT the kernel-arg num_stages (which only " + "pipelines dot-feeding loads). Watch Triton issue #8259 — " + "num_stages>1 + flatten=True can corrupt stores in non-dot kernels. " + "Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--flatten", + type=str, + default=None, + help="`flatten` arg on tl.range(...) for mode=persistent_main " + "(comma-separated 0/1 sweep). Default = 1. Ignored for " + "non-persistent_main modes.", + ) + parser.add_argument( + "--warp-specialize", + type=str, + default=None, + help="`warp_specialize` arg on tl.range(...) for mode=persistent_main " + "(comma-separated 0/1 sweep). Default = 0. Triton 3.6 only " + "supports it on simple matmul loops; our scan loop probably won't " + "pattern-match — exposed as a knob for sweep experiments. Requires " + "num_warps >= 4 if 1. Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--sr-modes", + type=str, + default="RN", + help="Comma-separated rounding modes to sweep: any combination of " + "{RN, SR}. SR (stochastic rounding) is silently skipped for state " + "dtypes that don't support it (bf16, fp32). Default 'RN' matches " + "legacy --philox-rounding=False behavior.", + ) + parser.add_argument( + "--rectangle-for-nowrite", + type=str, + default="0", + help="Comma-separated 0/1 values: 0 = replay-style nowrite kernel, " + "1 = dedicated rectangle nowrite kernel. Sweep both with '0,1' to " + "compare in one invocation. Silently no-op for write cells (the " + "write path always uses replay-style). Only applies to the " + "checkpointing variant.", + ) + parser.add_argument( + "--use-tma-rect-load", + type=str, + default=None, + help="Comma-separated 0/1 sweep. Use TMA (host-built tensor " + "descriptor) for state load in the rectangle nowrite path. " + "Cells where the rect path isn't reachable (e.g. mode=monolithic " + "+ WC=True) skip the value=1 case as a dupe.", + ) + parser.add_argument( + "--use-tma-replay-write-load", + type=str, + default=None, + help="Comma-separated 0/1 sweep. TMA state LOAD in replay main " + "when WC=True. Independent from nowrite-load and rect TMA — see " + "CHECKPOINTING_DESIGN.md item #17 for measured perf.", + ) + parser.add_argument( + "--use-tma-replay-nowrite-load", + type=str, + default=None, + help="Comma-separated 0/1 sweep. TMA state LOAD in replay main " + "when WC=False. Design doc reports the largest win on this path " + "(int8 b>=64: -8 to -12%%).", + ) + parser.add_argument( + "--use-tma-replay-write-store", + type=str, + default=None, + help="Comma-separated 0/1 sweep. TMA state STORE in replay main " + "(WC=True path only — no-op for WC=False). Independent from all " + "load TMA flags.", + ) + parser.add_argument( + "--modes", + type=str, + default="monolithic", + help="Comma-separated dispatch modes to sweep, any of " + "{monolithic,dynamic,doublelaunch}. monolithic = today's behavior " + "(one kernel pair, write_checkpoint applied to whole batch); " + "dynamic = single kernel pair that dispatches per-slot at runtime " + "based on PNAT (rectangle_for_nowrite picks RECTANGLE constexpr); " + "doublelaunch = two kernel pairs launched in sequence with " + "EARLY_OUT=True, each handling slots whose mode matches it. " + "Only applies to the checkpointing variant; non-monolithic modes " + "ignore --write-modes (per-slot from PNAT).", + ) + parser.add_argument( + "--mix-csv", + type=str, + default=None, + help="Path to AL histogram CSV (cols: AL, count). When set, an " + "additional 'mix' cell is emitted per (batch, mtp, dtype, sr, " + "mode, RECT, M, W, ...) combo where prev_tokens varies per iter, " + "drawn from the steady-state PNAT distribution induced by the " + "AL histogram. Mix cells run only on dynamic and doublelaunch " + "modes (mono on a mixed batch corrupts wrong-mode slots). " + "Each iteration of the captured CUDA graph has a different " + "pre-baked prev_tokens vector; warmup iters use distinct samples " + "from the timed iters so nsys-included warmup leaks don't bias.", + ) + parser.add_argument( + "--mix-csv-column", + type=int, + default=1, + help="Column index (0-based) in the AL histogram CSV for the " + "count/probability column. Default 1 (second column).", + ) + parser.add_argument( + "--mix-seed", + type=int, + default=42, + help="RNG seed for the steady-state PNAT sampler. Same seed " + "across runs => same per-slot samples for reproducible " + "comparisons.", + ) + parser.add_argument( + "--sort-slots", + type=str, + default="0", + help="Comma-separated 0/1. When 1, mix scenarios pre-sort slots " + "write-first (write slots at the head of slot_perm, nowrite at the " + "tail) and the dl-family kernels read pid_b through that perm — " + "clusters early-outs at one end of the grid. Only meaningful for " + "doublelaunch/dlgrouped/maindl with mix scenarios; mono/dynamic " + "and pure-batch cells skip sort=1.", + ) + parser.add_argument( + "--reverse-nowrite", + type=str, + default="0", + help="Comma-separated 0/1. When 1 (and --sort-slots 1), the " + "nowrite-side kernels in dlgrouped/doublelaunch/maindl walk the " + "perm in reverse so both halves of the dl chain front-load real " + "work. reverse=1 with sort=0 is skipped (no perm to reverse).", + ) + parser.add_argument( + "--hardcode-sort", + type=str, + default="0", + help="Comma-separated 0/1. When 1, the per-iter prev_tokens " + "samples are pre-sorted write-first OFFLINE (CPU-side) before " + "the timed region — kernel runs unchanged (USE_PERM=False) but " + "the EO gate sees sorted PNAT so early-outs cluster naturally. " + "Zero per-program load cost vs --sort-slots; output is " + "scrambled (we don't permute x/B/C/dt) but timing is meaningful. " + "Used to isolate whether clustering helps independent of the " + "perm-load overhead in the sort-slots path.", + ) + parser.add_argument( + "--mix-iters", + type=int, + default=None, + help="Iteration count override for mix scenarios (each iter is a " + "different per-slot prev_tokens draw). Default (None) uses " + "--iters. Mix scenarios benefit from more iters since each " + "iter samples a different mix; pure scenarios don't.", + ) + parser.add_argument( + "--philox-rounding", + action="store_true", + help="DEPRECATED — equivalent to --sr-modes SR. Retained for " + "backward compatibility; use --sr-modes for new scripts. fp16 SR " + "and fp8 SR require sm_100a (Blackwell B200+).", + ) + parser.add_argument( + "--philox-rounds", + type=int, + default=5, + help="Number of Philox PRNG rounds. Default 5 matches the " + "Nemotron-3-Super-120B production config (mamba_ssm_philox_rounds=5 " + "in examples/configs and tests/integration/perf configs). The " + "wrapper's generic fallback default is 10; callers without explicit " + "config see 10. Only consulted when --philox-rounding is enabled.", + ) + parser.add_argument( + "--variant", + choices=["replay", "checkpointing"], + default="replay", + help="Which kernel to time as the 'replay' row. 'replay' = today's " + "kernel (selective_state_update.py:replay). 'checkpointing' = " + "checkpointing_state_update.py. Both share the same wrapper signature.", + ) + parser.add_argument( + "--full-import", + action="store_true", + help="Use standard tensorrt_llm import path instead of fast direct " + "module loading. Slower (~40s startup) but guaranteed correct " + "if the fast path breaks due to package changes.", + ) + args = parser.parse_args() + + # Backward-compat: --philox-rounding implies --sr-modes SR if --sr-modes + # was left at the default. If both are set explicitly, error. + sr_modes_default = (args.sr_modes == "RN") + if args.philox_rounding: + if not sr_modes_default and args.sr_modes != "SR": + parser.error( + "--philox-rounding (deprecated) is incompatible with explicit " + f"--sr-modes={args.sr_modes!r}. Use --sr-modes SR (or " + "RN,SR) instead and drop --philox-rounding." + ) + args.sr_modes = "SR" + + sr_modes = [m.strip() for m in args.sr_modes.split(",") if m.strip()] + for m in sr_modes: + if m not in ("RN", "SR"): + parser.error(f"--sr-modes value must be RN or SR, got {m!r}") + args.sr_modes_list = sr_modes + + rect_modes = [v.strip() for v in args.rectangle_for_nowrite.split(",") if v.strip()] + rect_list = [] + for v in rect_modes: + if v not in ("0", "1"): + parser.error(f"--rectangle-for-nowrite value must be 0 or 1, got {v!r}") + rect_list.append(v == "1") + args.rectangle_for_nowrite_list = rect_list + + sort_modes = [v.strip() for v in (args.sort_slots or "0").split(",") if v.strip()] + sort_list = [] + for v in sort_modes: + if v not in ("0", "1"): + parser.error(f"--sort-slots value must be 0 or 1, got {v!r}") + sort_list.append(v == "1") + args.sort_slots_list = sort_list + + rev_modes = [v.strip() for v in (args.reverse_nowrite or "0").split(",") if v.strip()] + rev_list = [] + for v in rev_modes: + if v not in ("0", "1"): + parser.error(f"--reverse-nowrite value must be 0 or 1, got {v!r}") + rev_list.append(v == "1") + args.reverse_nowrite_list = rev_list + + hsort_modes = [v.strip() for v in (args.hardcode_sort or "0").split(",") if v.strip()] + hsort_list = [] + for v in hsort_modes: + if v not in ("0", "1"): + parser.error(f"--hardcode-sort value must be 0 or 1, got {v!r}") + hsort_list.append(v == "1") + args.hardcode_sort_list = hsort_list + + if args.write_modes is not None: + wm = [v.strip() for v in args.write_modes.split(",") if v.strip()] + write_list = [] + for v in wm: + if v not in ("0", "1"): + parser.error(f"--write-modes value must be 0 or 1, got {v!r}") + write_list.append(v == "1") + args.write_modes_list = write_list + else: + args.write_modes_list = [args.write_checkpoint] + + modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] + valid_modes = { + "monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", "persistent_main", "persistent_dynamic", + } + for m in modes_raw: + if m not in valid_modes: + parser.error( + f"--modes value must be one of {sorted(valid_modes)}, got {m!r}" + ) + args.modes_list = modes_raw or ["monolithic"] + return args + + +class _Tee: + """Write to both stdout and a file simultaneously.""" + + def __init__(self, path: str): + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + self._file = open(path, "w") # noqa: SIM115 + self._stdout = sys.stdout + + def write(self, data): + self._stdout.write(data) + self._file.write(data) + + def flush(self): + self._stdout.flush() + self._file.flush() + + def close(self): + self._file.close() + + +if __name__ == "__main__": + _args = _parse_args() + + _out_path = None + if _args.output != "-": + _ts = datetime.now().strftime("%Y%m%d_%H%M%S") + _fname = f"benchmark_replay_{_ts}.txt" + if _args.output is None: + _out_path = os.path.expanduser(f"~/nemo_logs/{_fname}") + elif os.path.isdir(_args.output) or _args.output.endswith("/"): + _out_path = os.path.join(_args.output, _fname) + else: + _out_path = _args.output + + if _out_path: + _tee = _Tee(_out_path) + sys.stdout = _tee + print(f"# benchmark_replay_selective_state_update {datetime.now().isoformat()}") + print(f"# cmd: {' '.join(sys.argv)}") + + try: + _run_benchmark(_args) + finally: + if _out_path: + sys.stdout = _tee._stdout + _tee.close() + print(f"\nResults saved to: {_out_path}") From 5789041a32f9561f308fdc070c686652557060cc Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 13:11:41 -0700 Subject: [PATCH 33/89] bench: speed up mamba replay benchmark CUPTI parsing Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 768 ++++++++++++++---- 1 file changed, 630 insertions(+), 138 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 3baeb1247ae2..f0919e9a3a54 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -39,9 +39,10 @@ Three modes: --cupti --cuda-graph (default) - Capture one CUDA graph per cell (warmup + timed iters inlined), - replay once, read kernel start/end from CUPTI. ~20× faster than - nsys-wrapped capture and matches it to within ~1% / noise floor. + Capture a small CUDA graph for the cell, replay it for warmup + timed + iterations, and read kernel start/end from CUPTI. Raw CUPTI buffers are + parsed out-of-process on the timed path, with a cached ordinal plan used + to keep only the kernels we care about. --cupti --no-cuda-graph Eager loop with CUPTI. Per-kernel timestamps are still accurate, but @@ -118,14 +119,20 @@ """ import argparse +import atexit +import ctypes import importlib import itertools import json +import multiprocessing as mp import os +import queue import statistics import sys +import threading import time from datetime import datetime +from multiprocessing import shared_memory from pathlib import Path import numpy as np @@ -558,7 +565,7 @@ def _build_tensors( # CUPTI in-process kernel timing # # Self-contained module-in-a-file. Reads kernel start/end timestamps directly -# from the GPU profiling fabric via NVIDIA's cupti-python bindings (1 ns +# from the GPU profiling fabric via CUPTI's Activity API (1 ns # resolution), avoiding two pitfalls of the cuda-events path: # # 1. cudaEvent.elapsed_time() resolution (~0.5 us) is too coarse for the @@ -568,9 +575,8 @@ def _build_tensors( # profile-export-sqlite-parse pipeline is heavy and out-of-process. # # This is functionally equivalent to wrapping each cell in nsys, except it -# runs in the same Python process with no serialization. When this proves -# out, lift `CuptiKernelTimer` and `_time_kernel_cuda_graph_cupti` into a -# proper TRT-LLM utility module — there is no benchmark-specific code below. +# runs in the same benchmark process and sends raw activity buffers to a +# parser process instead of materializing Python objects in the CUPTI callback. # ============================================================================= @@ -643,36 +649,248 @@ def _kernels_per_iter_baseline(with_conv1d: bool) -> int: return 2 if with_conv1d else 1 +_LIBCUPTI_CANDIDATES = ( + os.environ.get("CUPTI_LIBRARY_PATH"), + "/usr/local/lib/python3.12/dist-packages/nvidia/cu13/lib/libcupti.so.13", + "libcupti.so.13", + "libcupti.so", +) +_CUPTI_SUCCESS = 0 +_CUPTI_ERROR_MAX_LIMIT_REACHED = 12 +_CUPTI_ERROR_INVALID_KIND = 21 +_CUPTI_ACTIVITY_KIND_KERNEL = 3 +_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL = 10 +_CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER = 5 +_CUPTI_HOST_BUFFER_BYTES = 1024 * 1024 +_CUPTI_HOST_BUFFER_COUNT = 16 +_CUDA_GRAPH_GROUP_ITERS = 4 + + +def _load_libcupti() -> ctypes.CDLL: + errors = [] + for candidate in _LIBCUPTI_CANDIDATES: + if not candidate: + continue + try: + return ctypes.CDLL(candidate) + except OSError as exc: + errors.append(f"{candidate}: {exc}") + raise ImportError("Unable to load libcupti: " + "; ".join(errors)) + + +class _CuptiActivityKernel11Prefix(ctypes.Structure): + _pack_ = 1 + _fields_ = [ + ("kind", ctypes.c_int), + ("cache_config", ctypes.c_uint8), + ("shared_memory_config", ctypes.c_uint8), + ("registers_per_thread", ctypes.c_uint16), + ("partitioned_global_cache_requested", ctypes.c_int), + ("partitioned_global_cache_executed", ctypes.c_int), + ("start", ctypes.c_uint64), + ("end", ctypes.c_uint64), + ("completed", ctypes.c_uint64), + ("device_id", ctypes.c_uint32), + ("context_id", ctypes.c_uint32), + ("stream_id", ctypes.c_uint32), + ("grid_x", ctypes.c_int32), + ("grid_y", ctypes.c_int32), + ("grid_z", ctypes.c_int32), + ("block_x", ctypes.c_int32), + ("block_y", ctypes.c_int32), + ("block_z", ctypes.c_int32), + ("static_shared_memory", ctypes.c_int32), + ("dynamic_shared_memory", ctypes.c_int32), + ("local_memory_per_thread", ctypes.c_uint32), + ("local_memory_total", ctypes.c_uint32), + ("correlation_id", ctypes.c_uint32), + ("grid_id", ctypes.c_int64), + ("name", ctypes.c_void_p), + ("reserved0", ctypes.c_void_p), + ("queued", ctypes.c_uint64), + ("submitted", ctypes.c_uint64), + ("launch_type", ctypes.c_uint8), + ("is_shared_memory_carveout_requested", ctypes.c_uint8), + ("shared_memory_carveout_requested", ctypes.c_uint8), + ("padding", ctypes.c_uint8), + ("shared_memory_executed", ctypes.c_uint32), + ("graph_node_id", ctypes.c_uint64), + ] + + +def _configure_cupti_get_next_record(libcupti) -> None: + libcupti.cuptiActivityGetNextRecord.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_void_p), + ] + libcupti.cuptiActivityGetNextRecord.restype = ctypes.c_int + + +def _parse_cupti_buffer_ptr(libcupti, buffer_ptr: int, valid_size: int, *, include_names: bool): + records = [] + zero_ts_count = 0 + zero_ts_names: dict[str, int] = {} + record_ptr = ctypes.c_void_p(None) + while True: + result = libcupti.cuptiActivityGetNextRecord( + ctypes.c_void_p(buffer_ptr), + valid_size, + ctypes.byref(record_ptr), + ) + if result == _CUPTI_SUCCESS: + kind = ctypes.cast(record_ptr, ctypes.POINTER(ctypes.c_int)).contents.value + if kind not in (_CUPTI_ACTIVITY_KIND_KERNEL, _CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL): + continue + kernel = ctypes.cast(record_ptr, ctypes.POINTER(_CuptiActivityKernel11Prefix)).contents + name = None + if include_names: + if kernel.name: + name = ctypes.string_at(kernel.name).decode("utf-8", errors="replace") + else: + name = "?" + if kernel.start == 0 or kernel.end == 0: + zero_ts_count += 1 + if name is not None: + zero_ts_names[name] = zero_ts_names.get(name, 0) + 1 + continue + if include_names: + records.append(( + name, + int(kernel.start), + int(kernel.end), + int(kernel.correlation_id), + 0, + int(kernel.graph_node_id), + int(kernel.stream_id), + )) + else: + records.append(( + int(kernel.start), + int(kernel.end), + int(kernel.correlation_id), + int(kernel.graph_node_id), + int(kernel.stream_id), + )) + elif result == _CUPTI_ERROR_MAX_LIMIT_REACHED: + break + elif result == _CUPTI_ERROR_INVALID_KIND: + break + else: + raise RuntimeError(f"cuptiActivityGetNextRecord failed with CUptiResult={result}") + return records, zero_ts_count, zero_ts_names + + +def _apply_cupti_filter_plan(numeric_records, filter_plan): + if not filter_plan: + return [ + (None, start, end, corr, 0, graph_node_id, stream_id) + for start, end, corr, graph_node_id, stream_id in sorted(numeric_records) + ] + + filtered = [] + replay_idx = 0 + record_idx = 0 + for start, end, corr, graph_node_id, stream_id in sorted(numeric_records): + if replay_idx >= len(filter_plan): + break + records_per_replay, ordinal_names = filter_plan[replay_idx] + if record_idx < len(ordinal_names): + name = ordinal_names[record_idx] + if name is not None: + filtered.append((name, start, end, corr, 0, graph_node_id, stream_id)) + record_idx += 1 + if record_idx >= records_per_replay: + replay_idx += 1 + record_idx = 0 + return filtered + + +def _cupti_parser_worker(input_queue, output_queue, ready_event) -> None: + libcupti = _load_libcupti() + _configure_cupti_get_next_record(libcupti) + shared_blocks: dict[str, shared_memory.SharedMemory] = {} + records_by_generation: dict[int, list[tuple[int, int, int, int, int]]] = {} + zero_ts_by_generation: dict[int, int] = {} + ready_event.set() + while True: + item = input_queue.get() + if item is None: + break + kind = item[0] + if kind == "buffer": + _, generation, buffer_id, name, valid_size = item + shm = shared_blocks.get(name) + if shm is None: + shm = shared_memory.SharedMemory(name=name) + shared_blocks[name] = shm + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + parser_ptr = ctypes.addressof(shared_char) + records, zero_ts_count, _ = _parse_cupti_buffer_ptr( + libcupti, + parser_ptr, + valid_size, + include_names=False, + ) + records_by_generation.setdefault(generation, []).extend(records) + zero_ts_by_generation[generation] = zero_ts_by_generation.get(generation, 0) + zero_ts_count + ctypes.memset(parser_ptr, 0, len(shm.buf)) + except Exception as exc: # pragma: no cover - diagnostic worker path + output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) + finally: + del shared_char + output_queue.put({"kind": "buffer_done", "generation": generation, "buffer_id": buffer_id}) + elif kind == "finish": + _, generation, filter_plan = item + try: + raw_records = records_by_generation.pop(generation, []) + zero_ts_count = zero_ts_by_generation.pop(generation, 0) + filtered_records = _apply_cupti_filter_plan(raw_records, filter_plan) + output_queue.put({ + "kind": "finish_done", + "generation": generation, + "records": filtered_records, + "zero_ts_count": zero_ts_count, + "zero_ts_names": {}, + "raw_record_count": len(raw_records), + }) + except Exception as exc: # pragma: no cover - diagnostic worker path + output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) + else: + output_queue.put({"kind": "error", "generation": -1, "error": f"unknown parser message {kind!r}"}) + for shm in shared_blocks.values(): + shm.close() + + class CuptiKernelTimer: - """Process-singleton wrapper around CUPTI's CONCURRENT_KERNEL activity. - - CUPTI's callbacks are global (one subscriber per process), so the timer - is constructed lazily once via `CuptiKernelTimer.get()`. cupti-python - parses the activity buffer for us — `buffer_completed` receives a Python - list of typed activity objects, not a raw byte buffer — so no FFI is - needed. - - Usage: - timer = CuptiKernelTimer.get() - timer.start() # arms; drops any stale records - - records, zero_ts_count, zero_ts_names = timer.stop() - # records: list of tuples per kernel - # (name, start_ns, end_ns, corr, - # graph_id, graph_node_id, stream) - # zero_ts_count: kernel records CUPTI - # delivered with start=0 or end=0 - # (couldn't timestamp); zero_ts_names: - # name → count breakdown. - - The callback fires from a CUPTI worker thread, so a lock guards the - record buffer. Records are kept tiny (tuple of ints + str) to minimize - Python overhead in the hot path of the callback. + """Raw CUPTI Activity timer with out-of-process parsing for timed runs. + + CUPTI's callback gives us raw activity buffers. The callback only hands + shared-memory buffer metadata to a parser process, so the main process + avoids the cupti-python per-record object creation cost during the timed + path. A single local calibration replay may parse names in-process to + build an ordinal filter plan for a just-captured CUDA graph. """ _instance = None _import_error = None + _request_callback_type = ctypes.CFUNCTYPE( + None, + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t), + ) + _complete_callback_type = ctypes.CFUNCTYPE( + None, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_size_t, + ) + @classmethod def get(cls) -> "CuptiKernelTimer": if cls._instance is not None: @@ -680,93 +898,241 @@ def get(cls) -> "CuptiKernelTimer": if cls._import_error is not None: raise cls._import_error try: - from cupti import cupti as _c - except ImportError as e: # pragma: no cover — env-dependent - cls._import_error = e + cls._instance = cls() + return cls._instance + except ImportError as exc: # pragma: no cover - env-dependent + cls._import_error = exc raise - cls._instance = cls._init(_c) - return cls._instance - @classmethod - def _init(cls, _c) -> "CuptiKernelTimer": - import threading - - self = object.__new__(cls) - self._c = _c - self._records: list[tuple] = [] - self._zero_ts_count = 0 # how many kernel records were dropped due to start/end==0 - self._zero_ts_names: dict = {} # name -> count of zero-ts drops (for diagnostics) + def __init__(self) -> None: + self._libcupti = _load_libcupti() + self._configure_functions() self._lock = threading.Lock() + self._shared_buffers: dict[int, shared_memory.SharedMemory] = {} + self._buffer_id_by_ptr: dict[int, int] = {} + self._free_buffer_ids: list[int] = [] + self._local_completed: list[tuple[int, int]] = [] + self._mode = "drop" + self._generation = 0 + self._finish_results: dict[int, dict] = {} + self._parser_errors: list[str] = [] + self._filter_plan = () + self._mp_ctx = mp.get_context("spawn") + self._parse_input_queue = self._mp_ctx.Queue() + self._parse_output_queue = self._mp_ctx.Queue() + ready_event = self._mp_ctx.Event() + self._parse_process = self._mp_ctx.Process( + target=_cupti_parser_worker, + args=(self._parse_input_queue, self._parse_output_queue, ready_event), + ) + self._parse_process.start() + if not ready_event.wait(timeout=10.0): + raise RuntimeError("CUPTI parser process did not initialize") + + self._set_zeroed_host_buffer_attr() + for _ in range(_CUPTI_HOST_BUFFER_COUNT): + self._free_buffer_ids.append(self._allocate_shared_buffer()) + + self._request_callback = self._request_callback_type(self._request_buffer) + self._complete_callback = self._complete_callback_type(self._complete_buffer) + self._check(self._libcupti.cuptiActivityRegisterCallbacks( + self._request_callback, + self._complete_callback, + )) + self._check(self._libcupti.cuptiActivityEnable(_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL)) + atexit.register(self.close) + + def _configure_functions(self) -> None: + self._libcupti.cuptiActivityRegisterCallbacks.argtypes = [ + self._request_callback_type, + self._complete_callback_type, + ] + self._libcupti.cuptiActivityRegisterCallbacks.restype = ctypes.c_int + self._libcupti.cuptiActivityEnable.argtypes = [ctypes.c_int] + self._libcupti.cuptiActivityEnable.restype = ctypes.c_int + self._libcupti.cuptiActivityFlushAll.argtypes = [ctypes.c_uint32] + self._libcupti.cuptiActivityFlushAll.restype = ctypes.c_int + self._libcupti.cuptiActivitySetAttribute.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_void_p, + ] + self._libcupti.cuptiActivitySetAttribute.restype = ctypes.c_int + _configure_cupti_get_next_record(self._libcupti) + + def _set_zeroed_host_buffer_attr(self) -> None: + value_obj = ctypes.c_uint8(1) + size_obj = ctypes.c_size_t(ctypes.sizeof(value_obj)) + result = self._libcupti.cuptiActivitySetAttribute( + _CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER, + ctypes.byref(size_obj), + ctypes.byref(value_obj), + ) + if result != _CUPTI_SUCCESS: + print( + "[WARN] CUPTI zeroed host-buffer attribute failed; " + f"continuing with default CUPTI buffer handling (CUptiResult={result}).", + file=sys.stderr, + ) - # CUPTI callback contract (from cupti-python-samples/cupti_common.py): - # buffer_requested() -> (buffer_size, max_num_records) - # buffer_completed(activities: list) - # Setting max_num_records=0 (unbounded) avoids spurious buffer - # requests. 8 MiB matches the sample defaults. - def _buf_req(): - return (8 * 1024 * 1024, 0) - - kernel_kinds = (_c.ActivityKind.CONCURRENT_KERNEL, _c.ActivityKind.KERNEL) - - def _buf_done(activities): - recs = [] - zero_drops_local = 0 - zero_names_local: dict = {} - for a in activities: - if a.kind not in kernel_kinds: - continue - # start/end == 0 means CUPTI couldn't time this kernel. - # Track these instead of silently dropping — they're a sign - # CUPTI is failing to record kernels we DID launch. - if a.start == 0 or a.end == 0: - zero_drops_local += 1 - name = getattr(a, "name", "?") - zero_names_local[name] = zero_names_local.get(name, 0) + 1 - continue - recs.append(( - a.name, - int(a.start), - int(a.end), - int(a.correlation_id), - int(a.graph_id), - int(a.graph_node_id), - int(a.stream_id), - )) - if recs or zero_drops_local: - with self._lock: - if recs: - self._records.extend(recs) - if zero_drops_local: - self._zero_ts_count += zero_drops_local - for k, v in zero_names_local.items(): - self._zero_ts_names[k] = self._zero_ts_names.get(k, 0) + v - - # Hold strong refs so the C side never sees GC'd Python callables. - self._buf_req = _buf_req - self._buf_done = _buf_done - - _c.activity_register_callbacks(_buf_req, _buf_done) - _c.activity_enable(_c.ActivityKind.CONCURRENT_KERNEL) - return self - - def start(self) -> None: - """Arm capture: flush any stale records, then clear the buffer.""" - self._c.activity_flush_all(1) + def _check(self, result: int) -> None: + if result != _CUPTI_SUCCESS: + raise RuntimeError(f"CUPTI call failed with CUptiResult={result}") + + def _allocate_shared_buffer(self) -> int: + buffer_id = len(self._shared_buffers) + shm = shared_memory.SharedMemory(create=True, size=_CUPTI_HOST_BUFFER_BYTES) + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + ptr = ctypes.addressof(shared_char) + finally: + del shared_char + if ptr % 8 != 0: + shm.close() + shm.unlink() + raise RuntimeError("CUPTI shared-memory activity buffer was not 8-byte aligned") + self._shared_buffers[buffer_id] = shm + self._buffer_id_by_ptr[ptr] = buffer_id + return buffer_id + + def _buffer_ptr(self, buffer_id: int) -> int: + shm = self._shared_buffers[buffer_id] + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + return ctypes.addressof(shared_char) + finally: + del shared_char + + def _request_buffer(self, buffer, size, max_num_records) -> None: with self._lock: - self._records.clear() - self._zero_ts_count = 0 - self._zero_ts_names = {} + if self._free_buffer_ids: + buffer_id = self._free_buffer_ids.pop() + else: + buffer_id = self._allocate_shared_buffer() + ptr = self._buffer_ptr(buffer_id) + buffer[0] = ptr + size[0] = _CUPTI_HOST_BUFFER_BYTES + max_num_records[0] = 0 + + def _complete_buffer(self, context, stream_id, buffer, size, valid_size) -> None: + del context, stream_id, size + buffer_ptr = int(buffer) + valid_size_int = int(valid_size) + with self._lock: + mode = self._mode + generation = self._generation + buffer_id = self._buffer_id_by_ptr[buffer_ptr] + if valid_size_int == 0 or mode == "drop": + self._free_buffer_ids.append(buffer_id) + return + if mode == "local": + self._local_completed.append((buffer_id, valid_size_int)) + return + shm = self._shared_buffers[buffer_id] + self._parse_input_queue.put(("buffer", generation, buffer_id, shm.name, valid_size_int)) + + def _drain_parser_results(self) -> None: + while True: + try: + result = self._parse_output_queue.get_nowait() + except queue.Empty: + break + kind = result.get("kind") + if kind == "buffer_done": + with self._lock: + self._free_buffer_ids.append(int(result["buffer_id"])) + elif kind == "finish_done": + self._finish_results[int(result["generation"])] = result + elif kind == "error": + self._parser_errors.append(str(result.get("error"))) - def stop(self) -> tuple[list[tuple], int, dict]: - """Flush and return all kernel records + count of zero-timestamp drops. + def _flush(self, flag: int) -> None: + self._check(self._libcupti.cuptiActivityFlushAll(flag)) - Returns (records, zero_ts_count, zero_ts_names_dict). The latter two - are diagnostic: nonzero values mean CUPTI delivered records with - start=0 or end=0, indicating it failed to timestamp the kernel. - """ - self._c.activity_flush_all(1) + def _begin(self, mode: str, filter_plan=()) -> int: + with self._lock: + self._mode = "drop" + self._flush(1) + self._drain_parser_results() + with self._lock: + self._generation += 1 + generation = self._generation + self._mode = mode + self._local_completed = [] + self._filter_plan = filter_plan + return generation + + def capture_names(self, replay_fn) -> tuple[list[tuple], int, dict]: + """Run a small calibration replay and parse kernel names locally.""" + self._begin("local") + replay_fn() + torch.cuda.synchronize() + self._flush(0) + records: list[tuple] = [] + zero_ts_count = 0 + zero_ts_names: dict[str, int] = {} with self._lock: - return (list(self._records), self._zero_ts_count, dict(self._zero_ts_names)) + completed = list(self._local_completed) + self._local_completed = [] + self._mode = "drop" + for buffer_id, valid_size in completed: + ptr = self._buffer_ptr(buffer_id) + recs, zeros, zero_names = _parse_cupti_buffer_ptr( + self._libcupti, + ptr, + valid_size, + include_names=True, + ) + records.extend(recs) + zero_ts_count += zeros + for name, count in zero_names.items(): + zero_ts_names[name] = zero_ts_names.get(name, 0) + count + ctypes.memset(ptr, 0, _CUPTI_HOST_BUFFER_BYTES) + with self._lock: + self._free_buffer_ids.append(buffer_id) + records.sort(key=lambda r: r[1]) + return records, zero_ts_count, zero_ts_names + + def start(self, filter_plan=()) -> None: + self._begin("parser", filter_plan) + + def stop(self) -> tuple[list[tuple], int, dict, int]: + generation = self._generation + self._flush(0) + with self._lock: + self._mode = "drop" + filter_plan = self._filter_plan + self._parse_input_queue.put(("finish", generation, filter_plan)) + deadline = time.perf_counter() + 10.0 + while time.perf_counter() < deadline: + self._drain_parser_results() + if self._parser_errors: + raise RuntimeError("CUPTI parser process failed: " + "; ".join(self._parser_errors)) + result = self._finish_results.pop(generation, None) + if result is not None: + return ( + list(result["records"]), + int(result["zero_ts_count"]), + dict(result["zero_ts_names"]), + int(result["raw_record_count"]), + ) + time.sleep(0.001) + raise TimeoutError("Timed out waiting for CUPTI parser process") + + def close(self) -> None: + parse_process = getattr(self, "_parse_process", None) + if parse_process is not None and parse_process.is_alive(): + self._parse_input_queue.put(None) + parse_process.join(timeout=5.0) + if parse_process.is_alive(): + parse_process.terminate() + parse_process.join(timeout=1.0) + for shm in getattr(self, "_shared_buffers", {}).values(): + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass # ============================================================================= @@ -801,7 +1167,7 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, """ records = [ r for r in records - if any(s in r[0] for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS) + if r[0] is not None and any(s in r[0] for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS) ] records.sort(key=lambda r: r[1]) # by start_ns @@ -871,6 +1237,61 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, _PRE_GRAPH_WARMUP_ITERS = 3 # standard practice; see commit msg / design doc +_CUPTI_FILTER_PLAN_CACHE: dict[tuple, tuple[int, tuple[str | None, ...]]] = {} + + +def _target_name_or_none(name: str | None) -> str | None: + if name is None: + return None + if any(s in name for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS): + return name + return None + + +def _capture_group_graph(args, run_fn, reset_fn, group_iters: int) -> torch.cuda.CUDAGraph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(group_iters): + reset_fn() + if args.l2_flush: + _l2_flush.fill_(0.0) + run_fn() + torch.cuda.synchronize() + return graph + + +def _graph_group_iters(total_iters: int, pre_iter_fn) -> int: + if pre_iter_fn is not None: + return 1 + for group_iters in (_CUDA_GRAPH_GROUP_ITERS, 2): + if total_iters % group_iters == 0: + return group_iters + return 1 + + +def _get_cupti_filter_plan(timer: CuptiKernelTimer, graph, cache_key: tuple | None, + group_iters: int) -> tuple[int, tuple[str | None, ...]]: + full_cache_key = None if cache_key is None else (cache_key, group_iters) + if full_cache_key is not None: + cached = _CUPTI_FILTER_PLAN_CACHE.get(full_cache_key) + if cached is not None: + return cached + + records, zero_ts_count, zero_ts_names = timer.capture_names(graph.replay) + if zero_ts_count: + print( + f"[WARN] CUPTI calibration saw {zero_ts_count} zero-timestamp records " + f"(breakdown {zero_ts_names}); continuing with nonzero records.", + file=sys.stderr, + ) + ordinal_names = tuple(_target_name_or_none(r[0]) for r in records) + target_count = sum(name is not None for name in ordinal_names) + if target_count == 0: + raise RuntimeError("CUPTI calibration did not find any target kernel records") + plan = (len(records), ordinal_names) + if full_cache_key is not None: + _CUPTI_FILTER_PLAN_CACHE[full_cache_key] = plan + return plan def _time_kernel_cuda_graph( @@ -882,6 +1303,7 @@ def _time_kernel_cuda_graph( expected_K: int, pre_iter_fn=None, iters_override: int | None = None, + cupti_plan_key: tuple | None = None, ) -> dict: """CUDA-graph CUPTI timer (graph-per-iter design). @@ -926,32 +1348,59 @@ def _time_kernel_cuda_graph( run_fn() torch.cuda.synchronize() + total_iters = warmup + iters + group_iters = _graph_group_iters(total_iters, pre_iter_fn) + # Reset just before capture so warmup state changes don't bleed in. reset_fn() torch.cuda.synchronize() - # Capture ONE iter. pre_iter_fn deliberately not in here. - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - reset_fn() - if args.l2_flush: - _l2_flush.fill_(0.0) - run_fn() + # Capture a small group of identical logical iterations. Mix/pre_iter + # cells stay at one iter per replay because their per-iter input update + # still runs outside the graph. + g = _capture_group_graph(args, run_fn, reset_fn, group_iters) torch.cuda.synchronize() + records_per_replay, ordinal_names = _get_cupti_filter_plan( + timer, + g, + cupti_plan_key, + group_iters, + ) + target_count = sum(name is not None for name in ordinal_names) + expected_targets_per_replay = expected_K * group_iters + if target_count != expected_targets_per_replay: + print( + f"[WARN] CUPTI calibration mismatch for {tag!r}: expected " + f"{expected_targets_per_replay} target records in a {group_iters}-iter graph replay, " + f"got {target_count} target records out of {records_per_replay} total records.", + file=sys.stderr, + ) + # Time: replay the per-iter graph `warmup + iters` times, with # pre_iter_fn called between replays on the same stream. CUPTI # records every kernel launch; _stats_from_cupti_records validates # against expected_K and slices warmup off the front. - timer.start() + graph_replays = total_iters // group_iters + filter_plan = ((records_per_replay, ordinal_names),) * graph_replays + timer.start(filter_plan) torch.cuda.nvtx.range_push(tag) - for i in range(warmup + iters): + for i in range(graph_replays): if pre_iter_fn is not None: pre_iter_fn(i) g.replay() torch.cuda.synchronize() torch.cuda.nvtx.range_pop() - records, zero_ts_count, zero_ts_names = timer.stop() + records, zero_ts_count, zero_ts_names, raw_record_count = timer.stop() + expected_raw_record_count = records_per_replay * graph_replays + if raw_record_count != expected_raw_record_count: + print( + f"[WARN] CUPTI raw-record mismatch for {tag!r}: expected " + f"{records_per_replay} total records/replay × {graph_replays} replays " + f"= {expected_raw_record_count}, got {raw_record_count}. SKIPPING cell.", + file=sys.stderr, + ) + return None return _stats_from_cupti_records(records, warmup, iters, tag, expected_K, zero_ts_count=zero_ts_count, @@ -967,6 +1416,7 @@ def _time_kernel_eager( expected_K: int, pre_iter_fn=None, iters_override: int | None = None, + cupti_plan_key: tuple | None = None, ) -> dict: """Non-graph CUPTI timer (for ncu wrapping, debugging, etc.). @@ -978,19 +1428,22 @@ def _time_kernel_eager( warmup = args.warmup iters = iters_override if iters_override is not None else args.iters - timer.start() - torch.cuda.nvtx.range_push(tag) - # Unified warmup+iters loop; CUPTI filters by warmup count internally. - for i in range(warmup + iters): - reset_fn() - if args.l2_flush: - _flush_l2() # includes synchronize - if pre_iter_fn is not None: - pre_iter_fn(i) - run_fn() - torch.cuda.synchronize() - torch.cuda.nvtx.range_pop() - records, zero_ts_count, zero_ts_names = timer.stop() + del cupti_plan_key + + def _run_eager_loop(): + torch.cuda.nvtx.range_push(tag) + # Unified warmup+iters loop; CUPTI filters by warmup count internally. + for i in range(warmup + iters): + reset_fn() + if args.l2_flush: + _flush_l2() # includes synchronize + if pre_iter_fn is not None: + pre_iter_fn(i) + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + + records, zero_ts_count, zero_ts_names = timer.capture_names(_run_eager_loop) return _stats_from_cupti_records(records, warmup, iters, tag, expected_K, zero_ts_count=zero_ts_count, @@ -1047,6 +1500,7 @@ def _time_kernel( expected_K: int, pre_iter_fn=None, iters_override: int | None = None, + cupti_plan_key: tuple | None = None, ) -> dict: """Dispatch to graph-CUPTI / eager-CUPTI / no-timer path. @@ -1072,12 +1526,14 @@ def _time_kernel( expected_K=expected_K, pre_iter_fn=pre_iter_fn, iters_override=iters_override, + cupti_plan_key=cupti_plan_key, ) return _time_kernel_eager( args, run_fn, reset_fn, tag, expected_K=expected_K, pre_iter_fn=pre_iter_fn, iters_override=iters_override, + cupti_plan_key=cupti_plan_key, ) @@ -1568,6 +2024,19 @@ def _run_baseline(): stats = _time_kernel( args, _run_baseline, reset_fn, tag, expected_K=_kernels_per_iter_baseline(with_conv1d), + cupti_plan_key=( + "baseline", + args.baseline, + batch, + mtp_len, + state_dtype_name, + act_dtype_name, + with_conv1d, + bool(args.l2_flush), + bool(args.external_pdl), + bool(use_philox), + _kernels_per_iter_baseline(with_conv1d), + ), ) if stats is not None: @@ -2168,15 +2637,38 @@ def _emit_split(name_w, name_nw, val_w, val_nw): # the skipped list for an external rerun in a fresh process. retry_budget = max(0, getattr(args, "cupti_retry", 1)) stats = None + expected_K = _kernels_per_iter_incremental( + mode, with_conv1d=with_conv1d, + persistent_skip_empty=scenario_skip_empty, + ) + plan_key = ( + "incremental", + args.variant, + mode, + batch, + mtp_len, + state_dtype_name, + act_dtype_name, + with_conv1d, + bool(args.l2_flush), + bool(args.external_pdl), + bool(args.internal_pdl), + bool(use_philox), + bool(rectangle_for_nowrite), + bool(write_checkpoint), + bool(sort_slots), + bool(reverse_nowrite), + bool(hardcode_sort), + scenario_pre_iter is not None, + expected_K, + ) for attempt in range(retry_budget + 1): stats = _time_kernel( args, _run_incr, reset_fn, sweep_tag, - expected_K=_kernels_per_iter_incremental( - mode, with_conv1d=with_conv1d, - persistent_skip_empty=scenario_skip_empty, - ), + expected_K=expected_K, pre_iter_fn=scenario_pre_iter, iters_override=scenario_iters, + cupti_plan_key=plan_key, ) if stats is not None: break From a6604e51b27c70f10c0a5e7f937d0d4bbae3cc84 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 15:13:16 -0700 Subject: [PATCH 34/89] bench: tune mamba replay benchmark host overhead Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 228 ++++++++++++++++-- 1 file changed, 202 insertions(+), 26 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index f0919e9a3a54..3f13650c79ae 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -663,7 +663,7 @@ def _kernels_per_iter_baseline(with_conv1d: bool) -> int: _CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER = 5 _CUPTI_HOST_BUFFER_BYTES = 1024 * 1024 _CUPTI_HOST_BUFFER_COUNT = 16 -_CUDA_GRAPH_GROUP_ITERS = 4 +_DEFAULT_CUDA_GRAPH_GROUP_ITERS = 2 def _load_libcupti() -> ctypes.CDLL: @@ -917,6 +917,9 @@ def __init__(self) -> None: self._finish_results: dict[int, dict] = {} self._parser_errors: list[str] = [] self._filter_plan = () + self._last_start_timing: dict[str, float] = {} + self._last_stop_timing: dict[str, float] = {} + self._current_flush_period_ms = 0 self._mp_ctx = mp.get_context("spawn") self._parse_input_queue = self._mp_ctx.Queue() self._parse_output_queue = self._mp_ctx.Queue() @@ -952,6 +955,8 @@ def _configure_functions(self) -> None: self._libcupti.cuptiActivityEnable.restype = ctypes.c_int self._libcupti.cuptiActivityFlushAll.argtypes = [ctypes.c_uint32] self._libcupti.cuptiActivityFlushAll.restype = ctypes.c_int + self._libcupti.cuptiActivityFlushPeriod.argtypes = [ctypes.c_uint32] + self._libcupti.cuptiActivityFlushPeriod.restype = ctypes.c_int self._libcupti.cuptiActivitySetAttribute.argtypes = [ ctypes.c_int, ctypes.POINTER(ctypes.c_size_t), @@ -1031,35 +1036,65 @@ def _complete_buffer(self, context, stream_id, buffer, size, valid_size) -> None shm = self._shared_buffers[buffer_id] self._parse_input_queue.put(("buffer", generation, buffer_id, shm.name, valid_size_int)) + def _handle_parser_result(self, result: dict) -> None: + kind = result.get("kind") + if kind == "buffer_done": + with self._lock: + self._free_buffer_ids.append(int(result["buffer_id"])) + elif kind == "finish_done": + self._finish_results[int(result["generation"])] = result + elif kind == "error": + self._parser_errors.append(str(result.get("error"))) + def _drain_parser_results(self) -> None: while True: try: result = self._parse_output_queue.get_nowait() except queue.Empty: break - kind = result.get("kind") - if kind == "buffer_done": - with self._lock: - self._free_buffer_ids.append(int(result["buffer_id"])) - elif kind == "finish_done": - self._finish_results[int(result["generation"])] = result - elif kind == "error": - self._parser_errors.append(str(result.get("error"))) + self._handle_parser_result(result) def _flush(self, flag: int) -> None: self._check(self._libcupti.cuptiActivityFlushAll(flag)) - def _begin(self, mode: str, filter_plan=()) -> int: + def _set_flush_period_ms(self, period_ms: int) -> None: + if period_ms == self._current_flush_period_ms: + return + self._check(self._libcupti.cuptiActivityFlushPeriod(period_ms)) + self._current_flush_period_ms = period_ms + + def _begin( + self, + mode: str, + filter_plan=(), + flush_period_ms: int = 0, + collect_timing: bool = False, + ) -> int: + start_timing: dict[str, float] = {} with self._lock: self._mode = "drop" + phase_start_s = time.perf_counter() if collect_timing else 0.0 self._flush(1) + if collect_timing: + start_timing["forced_flush_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) + phase_start_s = time.perf_counter() if collect_timing else 0.0 self._drain_parser_results() + if collect_timing: + start_timing["drain_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) with self._lock: self._generation += 1 generation = self._generation self._mode = mode self._local_completed = [] self._filter_plan = filter_plan + if flush_period_ms > 0: + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._set_flush_period_ms(flush_period_ms) + if collect_timing: + start_timing["period_enable_ms"] = 1000.0 * ( + time.perf_counter() - phase_start_s + ) + self._last_start_timing = start_timing return generation def capture_names(self, replay_fn) -> tuple[list[tuple], int, dict]: @@ -1093,23 +1128,51 @@ def capture_names(self, replay_fn) -> tuple[list[tuple], int, dict]: records.sort(key=lambda r: r[1]) return records, zero_ts_count, zero_ts_names - def start(self, filter_plan=()) -> None: - self._begin("parser", filter_plan) - - def stop(self) -> tuple[list[tuple], int, dict, int]: + def start( + self, + filter_plan=(), + flush_period_ms: int = 0, + collect_timing: bool = False, + ) -> None: + self._begin("parser", filter_plan, flush_period_ms, collect_timing) + + def stop(self, collect_timing: bool = False) -> tuple[list[tuple], int, dict, int]: + stop_timing: dict[str, float] = {} + stop_start_s = time.perf_counter() if collect_timing else 0.0 generation = self._generation + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._set_flush_period_ms(0) + if collect_timing: + stop_timing["period_disable_ms"] = 1000.0 * ( + time.perf_counter() - phase_start_s + ) + phase_start_s = time.perf_counter() if collect_timing else 0.0 self._flush(0) + if collect_timing: + stop_timing["flush_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) with self._lock: self._mode = "drop" filter_plan = self._filter_plan self._parse_input_queue.put(("finish", generation, filter_plan)) + phase_start_s = time.perf_counter() if collect_timing else 0.0 deadline = time.perf_counter() + 10.0 while time.perf_counter() < deadline: - self._drain_parser_results() + timeout_s = max(0.0, min(0.01, deadline - time.perf_counter())) + try: + result = self._parse_output_queue.get(timeout=timeout_s) + except queue.Empty: + continue + self._handle_parser_result(result) if self._parser_errors: raise RuntimeError("CUPTI parser process failed: " + "; ".join(self._parser_errors)) result = self._finish_results.pop(generation, None) if result is not None: + if collect_timing: + stop_timing["parser_wait_ms"] = 1000.0 * ( + time.perf_counter() - phase_start_s + ) + stop_timing["total_ms"] = 1000.0 * (time.perf_counter() - stop_start_s) + self._last_stop_timing = stop_timing return ( list(result["records"]), int(result["zero_ts_count"]), @@ -1119,6 +1182,12 @@ def stop(self) -> tuple[list[tuple], int, dict, int]: time.sleep(0.001) raise TimeoutError("Timed out waiting for CUPTI parser process") + def last_start_timing(self) -> dict[str, float]: + return dict(self._last_start_timing) + + def last_stop_timing(self) -> dict[str, float]: + return dict(self._last_stop_timing) + def close(self) -> None: parse_process = getattr(self, "_parse_process", None) if parse_process is not None and parse_process.is_alive(): @@ -1240,6 +1309,34 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, _CUPTI_FILTER_PLAN_CACHE: dict[tuple, tuple[int, tuple[str | None, ...]]] = {} +class _HostTiming: + def __init__(self, enabled: bool) -> None: + self.enabled = enabled + self.values: dict[str, float | int | bool] = {} + self._total_start_s = time.perf_counter() if enabled else 0.0 + self._phase_start_s = 0.0 + + def start(self) -> None: + if self.enabled: + self._phase_start_s = time.perf_counter() + + def stop(self, key: str) -> None: + if self.enabled: + self.values[key] = 1000.0 * (time.perf_counter() - self._phase_start_s) + + def add(self, key: str, value: float | int | bool) -> None: + if self.enabled: + self.values[key] = value + + def stop_total(self) -> None: + if self.enabled: + self.values["total_ms"] = 1000.0 * (time.perf_counter() - self._total_start_s) + + def attach(self, stats: dict | None) -> None: + if self.enabled and stats is not None: + stats["host_timing"] = self.values + + def _target_name_or_none(name: str | None) -> str | None: if name is None: return None @@ -1260,10 +1357,11 @@ def _capture_group_graph(args, run_fn, reset_fn, group_iters: int) -> torch.cuda return graph -def _graph_group_iters(total_iters: int, pre_iter_fn) -> int: +def _graph_group_iters(args, total_iters: int, pre_iter_fn) -> int: if pre_iter_fn is not None: return 1 - for group_iters in (_CUDA_GRAPH_GROUP_ITERS, 2): + requested_group_iters = max(1, int(args.cuda_graph_group_iters)) + for group_iters in (requested_group_iters, 2): if total_iters % group_iters == 0: return group_iters return 1 @@ -1334,6 +1432,7 @@ def _time_kernel_cuda_graph( call. Used to give mix scenarios a higher iter count than pure (more iters = more independent mix draws averaged in). """ + host_timing = _HostTiming(bool(getattr(args, "host_timing", False))) timer = CuptiKernelTimer.get() warmup = args.warmup iters = iters_override if iters_override is not None else args.iters @@ -1341,32 +1440,44 @@ def _time_kernel_cuda_graph( # Pre-graph eager warmup: full per-iter chain × N. Settles allocator # + autotune; pre_iter_fn included so any side effects it has are # exercised before capture. + host_timing.start() for _ in range(_PRE_GRAPH_WARMUP_ITERS): reset_fn() if pre_iter_fn is not None: pre_iter_fn(0) run_fn() torch.cuda.synchronize() + host_timing.stop("pre_graph_warmup_ms") total_iters = warmup + iters - group_iters = _graph_group_iters(total_iters, pre_iter_fn) + group_iters = _graph_group_iters(args, total_iters, pre_iter_fn) # Reset just before capture so warmup state changes don't bleed in. + host_timing.start() reset_fn() torch.cuda.synchronize() + host_timing.stop("pre_capture_reset_ms") # Capture a small group of identical logical iterations. Mix/pre_iter # cells stay at one iter per replay because their per-iter input update # still runs outside the graph. + host_timing.start() g = _capture_group_graph(args, run_fn, reset_fn, group_iters) torch.cuda.synchronize() + host_timing.stop("graph_capture_ms") + plan_cache_key = None if cupti_plan_key is None else (cupti_plan_key, group_iters) + host_timing.add("cupti_plan_cached", ( + plan_cache_key is not None and plan_cache_key in _CUPTI_FILTER_PLAN_CACHE + )) + host_timing.start() records_per_replay, ordinal_names = _get_cupti_filter_plan( timer, g, cupti_plan_key, group_iters, ) + host_timing.stop("cupti_plan_ms") target_count = sum(name is not None for name in ordinal_names) expected_targets_per_replay = expected_K * group_iters if target_count != expected_targets_per_replay: @@ -1383,15 +1494,34 @@ def _time_kernel_cuda_graph( # against expected_K and slices warmup off the front. graph_replays = total_iters // group_iters filter_plan = ((records_per_replay, ordinal_names),) * graph_replays - timer.start(filter_plan) + cupti_flush_period_ms = max(0, int(getattr(args, "cupti_flush_period_ms", 0))) + host_timing.start() + timer.start( + filter_plan, + flush_period_ms=cupti_flush_period_ms, + collect_timing=host_timing.enabled, + ) + host_timing.stop("cupti_start_ms") + for key, value in timer.last_start_timing().items(): + host_timing.add(f"cupti_start_{key}", value) torch.cuda.nvtx.range_push(tag) + host_timing.start() for i in range(graph_replays): if pre_iter_fn is not None: pre_iter_fn(i) g.replay() + host_timing.stop("graph_enqueue_ms") + host_timing.start() torch.cuda.synchronize() + host_timing.stop("graph_sync_ms") torch.cuda.nvtx.range_pop() - records, zero_ts_count, zero_ts_names, raw_record_count = timer.stop() + host_timing.start() + records, zero_ts_count, zero_ts_names, raw_record_count = timer.stop( + collect_timing=host_timing.enabled, + ) + host_timing.stop("cupti_stop_ms") + for key, value in timer.last_stop_timing().items(): + host_timing.add(f"cupti_stop_{key}", value) expected_raw_record_count = records_per_replay * graph_replays if raw_record_count != expected_raw_record_count: print( @@ -1402,9 +1532,20 @@ def _time_kernel_cuda_graph( ) return None - return _stats_from_cupti_records(records, warmup, iters, tag, expected_K, - zero_ts_count=zero_ts_count, - zero_ts_names=zero_ts_names) + host_timing.start() + stats = _stats_from_cupti_records(records, warmup, iters, tag, expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names) + host_timing.stop("stats_ms") + host_timing.stop_total() + host_timing.add("graph_group_iters", group_iters) + host_timing.add("graph_replays", graph_replays) + host_timing.add("cupti_records_per_replay", records_per_replay) + host_timing.add("cupti_target_records_per_replay", target_count) + host_timing.add("cupti_raw_records", raw_record_count) + host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) + host_timing.attach(stats) + return stats def _time_kernel_eager( @@ -1424,6 +1565,7 @@ def _time_kernel_eager( come from CUPTI — same accuracy as the graph path, just slower per-iter (extra Python + sync overhead). """ + host_timing = _HostTiming(bool(getattr(args, "host_timing", False))) timer = CuptiKernelTimer.get() warmup = args.warmup iters = iters_override if iters_override is not None else args.iters @@ -1443,11 +1585,18 @@ def _run_eager_loop(): torch.cuda.synchronize() torch.cuda.nvtx.range_pop() + host_timing.start() records, zero_ts_count, zero_ts_names = timer.capture_names(_run_eager_loop) + host_timing.stop("timed_loop_and_cupti_parse_ms") - return _stats_from_cupti_records(records, warmup, iters, tag, expected_K, - zero_ts_count=zero_ts_count, - zero_ts_names=zero_ts_names) + host_timing.start() + stats = _stats_from_cupti_records(records, warmup, iters, tag, expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names) + host_timing.stop("stats_ms") + host_timing.stop_total() + host_timing.attach(stats) + return stats def _run_kernel_untimed(args, run_fn, reset_fn, tag: str) -> dict: @@ -2809,6 +2958,8 @@ def _print_row( k: stats[k] for k in ("median", "p95", "p99", "n") if k in stats } + if "host_timing" in stats: + row_stats["host_timing"] = stats["host_timing"] json_results[key] = row_stats # Append to JSONL sidecar if a path is set (incremental persistence). # Open per-write because args is pickled to ProcessPoolExecutor @@ -3288,6 +3439,15 @@ def _parse_args() -> argparse.Namespace: "single CUDA graph with per-iteration events " "inside the graph, eliminating all host overhead.", ) + parser.add_argument( + "--cuda-graph-group-iters", + type=int, + default=_DEFAULT_CUDA_GRAPH_GROUP_ITERS, + help="For pure CUDA-graph cells, capture this many logical benchmark " + "iterations per graph replay when warmup + iters is divisible by this " + "value. Mix cells stay at one logical iteration per replay because " + "their per-iteration input update runs outside the graph.", + ) parser.add_argument( "--cupti", action=argparse.BooleanOptionalAction, @@ -3299,6 +3459,14 @@ def _parse_args() -> argparse.Namespace: "— use when wrapping the bench in nsys/ncu, where the external " "profiler provides timings and our CUPTI subscriber would conflict.", ) + parser.add_argument( + "--cupti-flush-period-ms", + type=int, + default=0, + help="If >0, ask CUPTI to periodically flush activity buffers during " + "the timed CUDA-graph region. This can overlap raw-buffer parsing with " + "long timed cells; 0 leaves flushing explicit at the end of each cell.", + ) parser.add_argument( "--json-output", default=None, @@ -3315,6 +3483,14 @@ def _parse_args() -> argparse.Namespace: "each kernel) — useful for PDL overlap analysis but adds ~4 KB/cell. " "Default off keeps records to ~40 bytes (median/p95/p99/n only).", ) + parser.add_argument( + "--host-timing", + action=argparse.BooleanOptionalAction, + default=False, + help="Attach benchmark host-side phase timings to JSON/JSONL results. " + "Useful for diagnosing benchmark overhead, but it adds roughly 1 KB " + "per compact JSONL row and several perf_counter calls per cell.", + ) parser.add_argument( "--cupti-retry", type=int, From 00d78f6445b3a5d5fb1ae2affd6f1bf674a14007 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 15:33:10 -0700 Subject: [PATCH 35/89] bench: defer mamba replay CUPTI parsing Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 295 +++++++++++++++--- 1 file changed, 249 insertions(+), 46 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 3f13650c79ae..6903d259700c 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -1054,6 +1054,10 @@ def _drain_parser_results(self) -> None: break self._handle_parser_result(result) + def is_generation_ready(self, generation: int) -> bool: + self._drain_parser_results() + return generation in self._finish_results or bool(self._parser_errors) + def _flush(self, flag: int) -> None: self._check(self._libcupti.cuptiActivityFlushAll(flag)) @@ -1136,9 +1140,8 @@ def start( ) -> None: self._begin("parser", filter_plan, flush_period_ms, collect_timing) - def stop(self, collect_timing: bool = False) -> tuple[list[tuple], int, dict, int]: + def stop_async(self, collect_timing: bool = False) -> tuple[int, dict[str, float]]: stop_timing: dict[str, float] = {} - stop_start_s = time.perf_counter() if collect_timing else 0.0 generation = self._generation phase_start_s = time.perf_counter() if collect_timing else 0.0 self._set_flush_period_ms(0) @@ -1154,24 +1157,31 @@ def stop(self, collect_timing: bool = False) -> tuple[list[tuple], int, dict, in self._mode = "drop" filter_plan = self._filter_plan self._parse_input_queue.put(("finish", generation, filter_plan)) + self._last_stop_timing = stop_timing + return generation, stop_timing + + def wait_for_generation( + self, + generation: int, + stop_timing: dict[str, float] | None = None, + collect_timing: bool = False, + ) -> tuple[list[tuple], int, dict, int]: + if stop_timing is None: + stop_timing = {} phase_start_s = time.perf_counter() if collect_timing else 0.0 deadline = time.perf_counter() + 10.0 while time.perf_counter() < deadline: - timeout_s = max(0.0, min(0.01, deadline - time.perf_counter())) - try: - result = self._parse_output_queue.get(timeout=timeout_s) - except queue.Empty: - continue - self._handle_parser_result(result) - if self._parser_errors: - raise RuntimeError("CUPTI parser process failed: " + "; ".join(self._parser_errors)) result = self._finish_results.pop(generation, None) if result is not None: if collect_timing: stop_timing["parser_wait_ms"] = 1000.0 * ( time.perf_counter() - phase_start_s ) - stop_timing["total_ms"] = 1000.0 * (time.perf_counter() - stop_start_s) + stop_timing["total_ms"] = ( + stop_timing.get("period_disable_ms", 0.0) + + stop_timing.get("flush_ms", 0.0) + + stop_timing["parser_wait_ms"] + ) self._last_stop_timing = stop_timing return ( list(result["records"]), @@ -1179,9 +1189,20 @@ def stop(self, collect_timing: bool = False) -> tuple[list[tuple], int, dict, in dict(result["zero_ts_names"]), int(result["raw_record_count"]), ) - time.sleep(0.001) + timeout_s = max(0.0, min(0.01, deadline - time.perf_counter())) + try: + parser_result = self._parse_output_queue.get(timeout=timeout_s) + except queue.Empty: + continue + self._handle_parser_result(parser_result) + if self._parser_errors: + raise RuntimeError("CUPTI parser process failed: " + "; ".join(self._parser_errors)) raise TimeoutError("Timed out waiting for CUPTI parser process") + def stop(self, collect_timing: bool = False) -> tuple[list[tuple], int, dict, int]: + generation, stop_timing = self.stop_async(collect_timing) + return self.wait_for_generation(generation, stop_timing, collect_timing) + def last_start_timing(self) -> dict[str, float]: return dict(self._last_start_timing) @@ -1337,6 +1358,65 @@ def attach(self, stats: dict | None) -> None: stats["host_timing"] = self.values +class _PendingCuptiStats: + + def __init__( + self, + timer: CuptiKernelTimer, + generation: int, + stop_timing: dict[str, float], + host_timing: _HostTiming, + *, + warmup: int, + iters: int, + tag: str, + expected_K: int, + expected_raw_record_count: int, + ) -> None: + self._timer = timer + self._generation = generation + self._stop_timing = stop_timing + self._host_timing = host_timing + self._warmup = warmup + self._iters = iters + self._tag = tag + self._expected_K = expected_K + self._expected_raw_record_count = expected_raw_record_count + + def is_ready(self) -> bool: + return self._timer.is_generation_ready(self._generation) + + def resolve(self) -> dict | None: + records, zero_ts_count, zero_ts_names, raw_record_count = self._timer.wait_for_generation( + self._generation, + self._stop_timing, + collect_timing=self._host_timing.enabled, + ) + for key, value in self._timer.last_stop_timing().items(): + self._host_timing.add(f"cupti_stop_{key}", value) + if raw_record_count != self._expected_raw_record_count: + print( + f"[WARN] CUPTI raw-record mismatch for {self._tag!r}: expected " + f"{self._expected_raw_record_count}, got {raw_record_count}. SKIPPING cell.", + file=sys.stderr, + ) + return None + + self._host_timing.start() + stats = _stats_from_cupti_records( + records, + self._warmup, + self._iters, + self._tag, + self._expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names, + ) + self._host_timing.stop("stats_ms") + self._host_timing.attach(stats) + return stats + + def _target_name_or_none(name: str | None) -> str | None: if name is None: return None @@ -1515,14 +1595,38 @@ def _time_kernel_cuda_graph( torch.cuda.synchronize() host_timing.stop("graph_sync_ms") torch.cuda.nvtx.range_pop() + expected_raw_record_count = records_per_replay * graph_replays host_timing.start() + if int(getattr(args, "cupti_defer_depth", 1)) > 1: + generation, stop_timing = timer.stop_async(collect_timing=host_timing.enabled) + host_timing.stop("cupti_stop_ms") + for key, value in timer.last_stop_timing().items(): + host_timing.add(f"cupti_stop_{key}", value) + host_timing.stop_total() + host_timing.add("graph_group_iters", group_iters) + host_timing.add("graph_replays", graph_replays) + host_timing.add("cupti_records_per_replay", records_per_replay) + host_timing.add("cupti_target_records_per_replay", target_count) + host_timing.add("cupti_raw_records_expected", expected_raw_record_count) + host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) + return _PendingCuptiStats( + timer, + generation, + stop_timing, + host_timing, + warmup=warmup, + iters=iters, + tag=tag, + expected_K=expected_K, + expected_raw_record_count=expected_raw_record_count, + ) + records, zero_ts_count, zero_ts_names, raw_record_count = timer.stop( collect_timing=host_timing.enabled, ) host_timing.stop("cupti_stop_ms") for key, value in timer.last_stop_timing().items(): host_timing.add(f"cupti_stop_{key}", value) - expected_raw_record_count = records_per_replay * graph_replays if raw_record_count != expected_raw_record_count: print( f"[WARN] CUPTI raw-record mismatch for {tag!r}: expected " @@ -1543,6 +1647,7 @@ def _time_kernel_cuda_graph( host_timing.add("cupti_records_per_replay", records_per_replay) host_timing.add("cupti_target_records_per_replay", target_count) host_timing.add("cupti_raw_records", raw_record_count) + host_timing.add("cupti_raw_records_expected", expected_raw_record_count) host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) host_timing.attach(stats) return stats @@ -2188,22 +2293,18 @@ def _run_baseline(): ), ) - if stats is not None: - _print_row( - show_kernel_col, - args.baseline, - batch, - mtp_len, - "N/A", - state_dtype_name, - act_dtype_name, - stats, - json_results=getattr(args, "_json_results", None), - tp_size=args.tp_size, - json_detailed=getattr(args, "json_detailed", False), - jsonl_path=getattr(args, "_jsonl_path", None), - jsonl_host=getattr(args, "_jsonl_host", None), - ) + _submit_result_job( + args, + stats, + show_kernel_col=show_kernel_col, + kernel_name=args.baseline, + batch=batch, + mtp_len=mtp_len, + prev_k="N/A", + state_dtype_name=state_dtype_name, + act_dtype_name=act_dtype_name, + skipped_tag=tag, + ) # --- Sweep parameter parsing (invariant across prev_k) --- def _parse_sweep(val): @@ -2784,7 +2885,12 @@ def _emit_split(name_w, name_nw, val_w, val_nw): # because the failure is transient at the kernel-launch level. # Per --cupti-retry budget. On final failure, append tag to # the skipped list for an external rerun in a fresh process. - retry_budget = max(0, getattr(args, "cupti_retry", 1)) + defer_results = ( + args.cuda_graph + and getattr(args, "cupti", True) + and int(getattr(args, "cupti_defer_depth", 1)) > 1 + ) + retry_budget = 0 if defer_results else max(0, getattr(args, "cupti_retry", 1)) stats = None expected_K = _kernels_per_iter_incremental( mode, with_conv1d=with_conv1d, @@ -2834,6 +2940,7 @@ def _emit_split(name_w, name_nw, val_w, val_nw): # For pure scenarios, n_writes is constant: 0 (nowrite) or # batch (write), determined by scn["fill"] + mtp_len > max_window. # For mix, scn carries the precomputed per-iter array. + per_iter_nw = None if stats is not None and getattr(args, "json_detailed", False): eff_iters = scenario_iters if scenario_iters is not None else args.iters if scn["fill"] is not None: @@ -2847,25 +2954,21 @@ def _emit_split(name_w, name_nw, val_w, val_nw): per_iter_nw = nw_full[args.warmup:args.warmup + eff_iters].tolist() else: per_iter_nw = None - if per_iter_nw is not None: - stats["n_writes_per_iter"] = per_iter_nw if stats is not None: - _print_row( - show_kernel_col, - args.variant, - batch, - mtp_len, - prev_k_for_print, - state_dtype_name, - act_dtype_name, + _submit_result_job( + args, stats, - sweep_suffix, - json_results=getattr(args, "_json_results", None), - tp_size=args.tp_size, - json_detailed=getattr(args, "json_detailed", False), - jsonl_path=getattr(args, "_jsonl_path", None), - jsonl_host=getattr(args, "_jsonl_host", None), + show_kernel_col=show_kernel_col, + kernel_name=args.variant, + batch=batch, + mtp_len=mtp_len, + prev_k=prev_k_for_print, + state_dtype_name=state_dtype_name, + act_dtype_name=act_dtype_name, + sweep_suffix=sweep_suffix, + per_iter_nw=per_iter_nw, + skipped_tag=sweep_tag, ) @@ -2975,6 +3078,95 @@ def _print_row( f.write(json.dumps(rec) + "\n") +def _finish_result_job(args, job: dict) -> None: + result = job["result"] + if isinstance(result, _PendingCuptiStats): + stats = result.resolve() + else: + stats = result + + if stats is None: + skipped_tag = job.get("skipped_tag") + if skipped_tag is not None: + args._skipped_cells.append(skipped_tag) + return + + per_iter_nw = job.get("per_iter_nw") + if per_iter_nw is not None and getattr(args, "json_detailed", False): + stats["n_writes_per_iter"] = per_iter_nw + + _print_row( + job["show_kernel_col"], + job["kernel_name"], + job["batch"], + job["mtp_len"], + job["prev_k"], + job["state_dtype_name"], + job["act_dtype_name"], + stats, + job.get("sweep_suffix", ""), + json_results=getattr(args, "_json_results", None), + tp_size=args.tp_size, + json_detailed=getattr(args, "json_detailed", False), + jsonl_path=getattr(args, "_jsonl_path", None), + jsonl_host=getattr(args, "_jsonl_host", None), + ) + + +def _drain_pending_results(args, *, force: bool = False) -> None: + pending_results = getattr(args, "_pending_results", None) + if not pending_results: + return + + max_pending = max(1, int(getattr(args, "cupti_defer_depth", 1))) + while pending_results: + first_result = pending_results[0]["result"] + should_block = force or len(pending_results) >= max_pending + if ( + not should_block + and isinstance(first_result, _PendingCuptiStats) + and not first_result.is_ready() + ): + break + job = pending_results.pop(0) + _finish_result_job(args, job) + + +def _submit_result_job( + args, + result, + *, + show_kernel_col, + kernel_name, + batch, + mtp_len, + prev_k, + state_dtype_name, + act_dtype_name, + sweep_suffix="", + per_iter_nw=None, + skipped_tag=None, +) -> None: + job = { + "result": result, + "show_kernel_col": show_kernel_col, + "kernel_name": kernel_name, + "batch": batch, + "mtp_len": mtp_len, + "prev_k": prev_k, + "state_dtype_name": state_dtype_name, + "act_dtype_name": act_dtype_name, + "sweep_suffix": sweep_suffix, + "per_iter_nw": per_iter_nw, + "skipped_tag": skipped_tag, + } + if isinstance(result, _PendingCuptiStats): + args._pending_results.append(job) + _drain_pending_results(args) + else: + _finish_result_job(args, job) + + # Main benchmark loop @@ -2982,6 +3174,7 @@ def _run_benchmark(args) -> None: # JSON accumulator — populated by _print_row when --json-output is set. # Stash on args so we don't need to thread a dict through every helper. args._json_results = {} if getattr(args, "json_output", None) else None + args._pending_results = [] # JSONL incremental sidecar. Path = `.jsonl`. Each completed # cell appends one line `{"key": , "stats": {...}}` to this file @@ -3306,6 +3499,8 @@ def _run_benchmark(args) -> None: mix_samples_sorted_cpu=mix_samples_sorted_cpu, ) + _drain_pending_results(args, force=True) + if args.profile: torch.cuda.cudart().cudaProfilerStop() @@ -3467,6 +3662,14 @@ def _parse_args() -> argparse.Namespace: "the timed CUDA-graph region. This can overlap raw-buffer parsing with " "long timed cells; 0 leaves flushing explicit at the end of each cell.", ) + parser.add_argument( + "--cupti-defer-depth", + type=int, + default=4, + help="Maximum number of CUDA-graph CUPTI timing results that may be " + "left for the parser process while the main process starts later cells. " + "1 preserves synchronous per-cell parsing and inline retry behavior.", + ) parser.add_argument( "--json-output", default=None, From fc0dce3a6d303f7cc75bc2486500df35d2b0cc6f Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 16:43:57 -0700 Subject: [PATCH 36/89] bench: compute compact CUPTI stats in parser process Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 145 +++++++++++++----- 1 file changed, 107 insertions(+), 38 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 6903d259700c..8bae887676ef 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -842,11 +842,32 @@ def _cupti_parser_worker(input_queue, output_queue, ready_event) -> None: del shared_char output_queue.put({"kind": "buffer_done", "generation": generation, "buffer_id": buffer_id}) elif kind == "finish": - _, generation, filter_plan = item + if len(item) == 4: + _, generation, filter_plan, stats_request = item + else: + _, generation, filter_plan = item + stats_request = None try: raw_records = records_by_generation.pop(generation, []) zero_ts_count = zero_ts_by_generation.pop(generation, 0) filtered_records = _apply_cupti_filter_plan(raw_records, filter_plan) + stats = None + parser_stats_ms = 0.0 + stats_ready = stats_request is not None + if stats_request is not None: + stats_start_s = time.perf_counter() + stats = _stats_from_cupti_records( + filtered_records, + int(stats_request["warmup"]), + int(stats_request["iters"]), + str(stats_request["tag"]), + int(stats_request["expected_K"]), + zero_ts_count=zero_ts_count, + zero_ts_names={}, + include_details=bool(stats_request.get("include_details", True)), + ) + parser_stats_ms = 1000.0 * (time.perf_counter() - stats_start_s) + filtered_records = [] output_queue.put({ "kind": "finish_done", "generation": generation, @@ -854,6 +875,9 @@ def _cupti_parser_worker(input_queue, output_queue, ready_event) -> None: "zero_ts_count": zero_ts_count, "zero_ts_names": {}, "raw_record_count": len(raw_records), + "stats": stats, + "stats_ready": stats_ready, + "parser_stats_ms": parser_stats_ms, }) except Exception as exc: # pragma: no cover - diagnostic worker path output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) @@ -1140,7 +1164,11 @@ def start( ) -> None: self._begin("parser", filter_plan, flush_period_ms, collect_timing) - def stop_async(self, collect_timing: bool = False) -> tuple[int, dict[str, float]]: + def stop_async( + self, + collect_timing: bool = False, + stats_request: dict | None = None, + ) -> tuple[int, dict[str, float]]: stop_timing: dict[str, float] = {} generation = self._generation phase_start_s = time.perf_counter() if collect_timing else 0.0 @@ -1156,16 +1184,16 @@ def stop_async(self, collect_timing: bool = False) -> tuple[int, dict[str, float with self._lock: self._mode = "drop" filter_plan = self._filter_plan - self._parse_input_queue.put(("finish", generation, filter_plan)) + self._parse_input_queue.put(("finish", generation, filter_plan, stats_request)) self._last_stop_timing = stop_timing return generation, stop_timing - def wait_for_generation( + def wait_for_generation_result( self, generation: int, stop_timing: dict[str, float] | None = None, collect_timing: bool = False, - ) -> tuple[list[tuple], int, dict, int]: + ) -> dict: if stop_timing is None: stop_timing = {} phase_start_s = time.perf_counter() if collect_timing else 0.0 @@ -1183,12 +1211,7 @@ def wait_for_generation( + stop_timing["parser_wait_ms"] ) self._last_stop_timing = stop_timing - return ( - list(result["records"]), - int(result["zero_ts_count"]), - dict(result["zero_ts_names"]), - int(result["raw_record_count"]), - ) + return result timeout_s = max(0.0, min(0.01, deadline - time.perf_counter())) try: parser_result = self._parse_output_queue.get(timeout=timeout_s) @@ -1199,6 +1222,20 @@ def wait_for_generation( raise RuntimeError("CUPTI parser process failed: " + "; ".join(self._parser_errors)) raise TimeoutError("Timed out waiting for CUPTI parser process") + def wait_for_generation( + self, + generation: int, + stop_timing: dict[str, float] | None = None, + collect_timing: bool = False, + ) -> tuple[list[tuple], int, dict, int]: + result = self.wait_for_generation_result(generation, stop_timing, collect_timing) + return ( + list(result["records"]), + int(result["zero_ts_count"]), + dict(result["zero_ts_names"]), + int(result["raw_record_count"]), + ) + def stop(self, collect_timing: bool = False) -> tuple[list[tuple], int, dict, int]: generation, stop_timing = self.stop_async(collect_timing) return self.wait_for_generation(generation, stop_timing, collect_timing) @@ -1243,7 +1280,8 @@ def _stats_from_spans(spans_us: list[float]) -> dict: def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, zero_ts_count: int = 0, - zero_ts_names: dict | None = None): + zero_ts_names: dict | None = None, + include_details: bool = True): """Bin a flat CUPTI kernel record stream into per-iter spans + per-kernel relative timestamps. Used by both graph and eager CUPTI paths. @@ -1314,15 +1352,17 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, iter_start_ns = min(r[1] for r in chunk) iter_end_ns = max(r[2] for r in chunk) spans_us.append((iter_end_ns - iter_start_ns) / 1000.0) - for r in chunk: - name = r[0] - slot = per_kernel.setdefault(name, {"start_us": [], "end_us": []}) - slot["start_us"].append((r[1] - iter_start_ns) / 1000.0) - slot["end_us"].append((r[2] - iter_start_ns) / 1000.0) + if include_details: + for r in chunk: + name = r[0] + slot = per_kernel.setdefault(name, {"start_us": [], "end_us": []}) + slot["start_us"].append((r[1] - iter_start_ns) / 1000.0) + slot["end_us"].append((r[2] - iter_start_ns) / 1000.0) out = _stats_from_spans(spans_us) - out["iters_us"] = spans_us - out["per_kernel"] = per_kernel + if include_details: + out["iters_us"] = spans_us + out["per_kernel"] = per_kernel return out @@ -1387,13 +1427,14 @@ def is_ready(self) -> bool: return self._timer.is_generation_ready(self._generation) def resolve(self) -> dict | None: - records, zero_ts_count, zero_ts_names, raw_record_count = self._timer.wait_for_generation( + result = self._timer.wait_for_generation_result( self._generation, self._stop_timing, collect_timing=self._host_timing.enabled, ) for key, value in self._timer.last_stop_timing().items(): self._host_timing.add(f"cupti_stop_{key}", value) + raw_record_count = int(result["raw_record_count"]) if raw_record_count != self._expected_raw_record_count: print( f"[WARN] CUPTI raw-record mismatch for {self._tag!r}: expected " @@ -1402,17 +1443,22 @@ def resolve(self) -> dict | None: ) return None - self._host_timing.start() - stats = _stats_from_cupti_records( - records, - self._warmup, - self._iters, - self._tag, - self._expected_K, - zero_ts_count=zero_ts_count, - zero_ts_names=zero_ts_names, - ) - self._host_timing.stop("stats_ms") + if result.get("stats_ready"): + stats = result.get("stats") + self._host_timing.add("stats_ms", 0.0) + self._host_timing.add("parser_stats_ms", float(result.get("parser_stats_ms", 0.0))) + else: + self._host_timing.start() + stats = _stats_from_cupti_records( + list(result["records"]), + self._warmup, + self._iters, + self._tag, + self._expected_K, + zero_ts_count=int(result["zero_ts_count"]), + zero_ts_names=dict(result["zero_ts_names"]), + ) + self._host_timing.stop("stats_ms") self._host_timing.attach(stats) return stats @@ -1598,7 +1644,16 @@ def _time_kernel_cuda_graph( expected_raw_record_count = records_per_replay * graph_replays host_timing.start() if int(getattr(args, "cupti_defer_depth", 1)) > 1: - generation, stop_timing = timer.stop_async(collect_timing=host_timing.enabled) + generation, stop_timing = timer.stop_async( + collect_timing=host_timing.enabled, + stats_request={ + "warmup": warmup, + "iters": iters, + "tag": tag, + "expected_K": expected_K, + "include_details": bool(getattr(args, "json_detailed", False)), + }, + ) host_timing.stop("cupti_stop_ms") for key, value in timer.last_stop_timing().items(): host_timing.add(f"cupti_stop_{key}", value) @@ -1637,9 +1692,16 @@ def _time_kernel_cuda_graph( return None host_timing.start() - stats = _stats_from_cupti_records(records, warmup, iters, tag, expected_K, - zero_ts_count=zero_ts_count, - zero_ts_names=zero_ts_names) + stats = _stats_from_cupti_records( + records, + warmup, + iters, + tag, + expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names, + include_details=bool(getattr(args, "json_detailed", False)), + ) host_timing.stop("stats_ms") host_timing.stop_total() host_timing.add("graph_group_iters", group_iters) @@ -1695,9 +1757,16 @@ def _run_eager_loop(): host_timing.stop("timed_loop_and_cupti_parse_ms") host_timing.start() - stats = _stats_from_cupti_records(records, warmup, iters, tag, expected_K, - zero_ts_count=zero_ts_count, - zero_ts_names=zero_ts_names) + stats = _stats_from_cupti_records( + records, + warmup, + iters, + tag, + expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names, + include_details=bool(getattr(args, "json_detailed", False)), + ) host_timing.stop("stats_ms") host_timing.stop_total() host_timing.attach(stats) From ccdc79982007100492c024f53f374733244706fa Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 17:02:10 -0700 Subject: [PATCH 37/89] bench: trim cuda graph capture host overhead Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../modules/mamba/benchmark_replay_selective_state_update.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 8bae887676ef..82bc414d38b1 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -1479,7 +1479,6 @@ def _capture_group_graph(args, run_fn, reset_fn, group_iters: int) -> torch.cuda if args.l2_flush: _l2_flush.fill_(0.0) run_fn() - torch.cuda.synchronize() return graph @@ -1589,7 +1588,6 @@ def _time_kernel_cuda_graph( # still runs outside the graph. host_timing.start() g = _capture_group_graph(args, run_fn, reset_fn, group_iters) - torch.cuda.synchronize() host_timing.stop("graph_capture_ms") plan_cache_key = None if cupti_plan_key is None else (cupti_plan_key, group_iters) From 80e57109baba5744beb7bcb9f7452d77d600ac3a Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 17:04:18 -0700 Subject: [PATCH 38/89] bench: group mixed cuda graph replays Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 145 +++++++++++++----- 1 file changed, 108 insertions(+), 37 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 82bc414d38b1..574db578a827 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -663,7 +663,8 @@ def _kernels_per_iter_baseline(with_conv1d: bool) -> int: _CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER = 5 _CUPTI_HOST_BUFFER_BYTES = 1024 * 1024 _CUPTI_HOST_BUFFER_COUNT = 16 -_DEFAULT_CUDA_GRAPH_GROUP_ITERS = 2 +_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE = 2 +_DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX = 4 def _load_libcupti() -> ctypes.CDLL: @@ -1471,10 +1472,18 @@ def _target_name_or_none(name: str | None) -> str | None: return None -def _capture_group_graph(args, run_fn, reset_fn, group_iters: int) -> torch.cuda.CUDAGraph: +def _capture_group_graph( + args, + run_fn, + reset_fn, + group_iters: int, + graph_pre_iter_fn=None, +) -> torch.cuda.CUDAGraph: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - for _ in range(group_iters): + for j in range(group_iters): + if graph_pre_iter_fn is not None: + graph_pre_iter_fn(j) reset_fn() if args.l2_flush: _l2_flush.fill_(0.0) @@ -1482,10 +1491,18 @@ def _capture_group_graph(args, run_fn, reset_fn, group_iters: int) -> torch.cuda return graph -def _graph_group_iters(args, total_iters: int, pre_iter_fn) -> int: - if pre_iter_fn is not None: +def _graph_group_iters(args, total_iters: int, pre_iter_fn, pre_iter_group_factory) -> int: + if pre_iter_fn is not None and pre_iter_group_factory is None: return 1 - requested_group_iters = max(1, int(args.cuda_graph_group_iters)) + requested = getattr(args, "cuda_graph_group_iters", None) + if requested is None: + requested_group_iters = ( + _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX + if pre_iter_group_factory is not None + else _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE + ) + else: + requested_group_iters = max(1, int(requested)) for group_iters in (requested_group_iters, 2): if total_iters % group_iters == 0: return group_iters @@ -1525,16 +1542,15 @@ def _time_kernel_cuda_graph( *, expected_K: int, pre_iter_fn=None, + pre_iter_group_factory=None, iters_override: int | None = None, cupti_plan_key: tuple | None = None, ) -> dict: """CUDA-graph CUPTI timer (graph-per-iter design). - Captures one CUDA graph holding a single iter's worth of work (reset - + l2_flush + run_fn) and replays it `warmup + iters` times. Per-iter - setup (`pre_iter_fn`, e.g. mix-mode PNAT/n_writes copies) runs OUTSIDE - the graph, on the same CUDA stream so the order - `pre_iter_fn → l2_flush → run_fn` is preserved on every replay. + Captures one CUDA graph holding a small group of logical iterations + (per-iter setup + reset + l2_flush + run_fn) and replays it enough + times to cover `warmup + iters`. Why graph-per-iter (vs the older "one giant graph holding all iters" design): instantiating a CUDA graph is expensive — proportional to @@ -1542,12 +1558,9 @@ def _time_kernel_cuda_graph( cheaper than one big graph instantiated for each cell of a sweep. Replays are cheap regardless. - Why pre_iter_fn outside the graph: it depends on the iter index `i` - (different sample per iter), but graph capture would bake in the - capture-time `i`. Putting the per-iter copy outside the graph also - has a useful side effect: PNAT (the per-iter input) is loaded into - L2 by the copy, then evicted by the in-graph L2 flush, so the kernel - reads PNAT cold — closer to production behavior than the old design. + Mix cells use a per-replay device window: an outside-graph copy loads + the next group of PNAT/n_writes samples, then graph-captured per-iter + copies update kernel inputs before each reset + L2 flush + run. Pre-graph eager warmup (3 iters): forces PyTorch's caching allocator + Triton's autotune cache to settle before capture so the graph @@ -1575,7 +1588,11 @@ def _time_kernel_cuda_graph( host_timing.stop("pre_graph_warmup_ms") total_iters = warmup + iters - group_iters = _graph_group_iters(args, total_iters, pre_iter_fn) + group_iters = _graph_group_iters(args, total_iters, pre_iter_fn, pre_iter_group_factory) + pre_replay_fn = None + graph_pre_iter_fn = None + if pre_iter_group_factory is not None and group_iters > 1: + pre_replay_fn, graph_pre_iter_fn = pre_iter_group_factory(group_iters) # Reset just before capture so warmup state changes don't bleed in. host_timing.start() @@ -1584,12 +1601,18 @@ def _time_kernel_cuda_graph( host_timing.stop("pre_capture_reset_ms") # Capture a small group of identical logical iterations. Mix/pre_iter - # cells stay at one iter per replay because their per-iter input update - # still runs outside the graph. + # cells can group when they provide a graph-side pre-iter updater backed + # by a per-replay device window. host_timing.start() - g = _capture_group_graph(args, run_fn, reset_fn, group_iters) + g = _capture_group_graph(args, run_fn, reset_fn, group_iters, graph_pre_iter_fn) host_timing.stop("graph_capture_ms") + if pre_replay_fn is not None: + host_timing.start() + pre_replay_fn(0) + torch.cuda.synchronize() + host_timing.stop("graph_preload_ms") + plan_cache_key = None if cupti_plan_key is None else (cupti_plan_key, group_iters) host_timing.add("cupti_plan_cached", ( plan_cache_key is not None and plan_cache_key in _CUPTI_FILTER_PLAN_CACHE @@ -1612,10 +1635,10 @@ def _time_kernel_cuda_graph( file=sys.stderr, ) - # Time: replay the per-iter graph `warmup + iters` times, with - # pre_iter_fn called between replays on the same stream. CUPTI - # records every kernel launch; _stats_from_cupti_records validates - # against expected_K and slices warmup off the front. + # Time: replay the grouped graph enough times to cover warmup+iters. + # Mix cells preload one device window per replay on the same stream. + # CUPTI records every kernel launch; _stats_from_cupti_records + # validates against expected_K and slices warmup off the front. graph_replays = total_iters // group_iters filter_plan = ((records_per_replay, ordinal_names),) * graph_replays cupti_flush_period_ms = max(0, int(getattr(args, "cupti_flush_period_ms", 0))) @@ -1631,7 +1654,9 @@ def _time_kernel_cuda_graph( torch.cuda.nvtx.range_push(tag) host_timing.start() for i in range(graph_replays): - if pre_iter_fn is not None: + if pre_replay_fn is not None: + pre_replay_fn(i) + elif pre_iter_fn is not None: pre_iter_fn(i) g.replay() host_timing.stop("graph_enqueue_ms") @@ -1820,6 +1845,7 @@ def _time_kernel( *, expected_K: int, pre_iter_fn=None, + pre_iter_group_factory=None, iters_override: int | None = None, cupti_plan_key: tuple | None = None, ) -> dict: @@ -1846,6 +1872,7 @@ def _time_kernel( args, run_fn, reset_fn, tag, expected_K=expected_K, pre_iter_fn=pre_iter_fn, + pre_iter_group_factory=pre_iter_group_factory, iters_override=iters_override, cupti_plan_key=cupti_plan_key, ) @@ -2454,10 +2481,9 @@ def _split_or_share(split_csv, shared_values): # Mix scenario: skip on monolithic (mono on mixed PNAT corrupts the # wrong-mode slots). Persistent_main + mix is now supported: bench # pre-bakes both a per-iter PNAT samples tensor and a per-iter - # n_writes samples tensor; pre_iter_fn copies row i of each into the - # kernel-input tensors (PNAT and n_writes_dev) on the same stream as - # the captured CUDA graph, so they're cold w.r.t. the in-graph L2 - # flush. + # n_writes samples tensor; grouped graph capture copies window rows + # into kernel-input tensors (PNAT and n_writes_dev) before each + # in-graph L2 flush, so the timed kernels read PNAT cold. if mix_samples_cpu is not None and mode != "monolithic": device = state_work.device # Hardcode-sort: per-iter prev_tokens are CPU-sorted write-first. @@ -2494,6 +2520,7 @@ def _split_or_share(split_csv, shared_values): # Build _mix_pre_iter — the closure that runs OUTSIDE the captured # graph between replays. Updates: prev_tokens (always), # slot_perm_buf (when sort_slots), n_writes_dev_mix (persistent). + perm_samples_gpu = None if sort_slots and perm_samples_cpu is not None: perm_samples_gpu = torch.from_numpy(perm_samples_cpu).to( device=device, dtype=torch.int32 @@ -2520,6 +2547,45 @@ def _mix_pre_iter(i, _s=samples_gpu, _ns=n_writes_samples_gpu, def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): _pt.copy_(_s[i]) + def _mix_pre_iter_group_factory( + group_iters, + _s=samples_gpu, + _ps=perm_samples_gpu, + _ns=n_writes_samples_gpu, + _pt=prev_tokens, + _pm=slot_perm_buf, + _nw=n_writes_dev_mix, + ): + sample_window = torch.empty( + (group_iters, _s.shape[1]), device=_s.device, dtype=_s.dtype, + ) + perm_window = ( + torch.empty((group_iters, _ps.shape[1]), device=_ps.device, dtype=_ps.dtype) + if _ps is not None else None + ) + nw_window = ( + torch.empty((group_iters,), device=_ns.device, dtype=_ns.dtype) + if _ns is not None else None + ) + + def _pre_replay(replay_idx): + start = replay_idx * group_iters + end = start + group_iters + sample_window.copy_(_s[start:end]) + if perm_window is not None: + perm_window.copy_(_ps[start:end]) + if nw_window is not None: + nw_window.copy_(_ns[start:end]) + + def _graph_pre_iter(j): + _pt.copy_(sample_window[j]) + if perm_window is not None: + _pm.copy_(perm_window[j]) + if nw_window is not None: + _nw.copy_(nw_window[j:j + 1]) + + return _pre_replay, _graph_pre_iter + # Mix iters override: if --mix-iters set, use it; else use args.iters. mix_iters = getattr(args, "mix_iters", None) scenarios.append({ @@ -2527,6 +2593,7 @@ def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): "print_label": "mix", "fill": None, "pre_iter": _mix_pre_iter, + "pre_iter_group_factory": _mix_pre_iter_group_factory, "iters": mix_iters, # None => use args.iters # Pass through to _run_incr so the wrapper receives _n_writes_dev # (mix scenarios) instead of _n_writes (pure scenarios). @@ -2548,6 +2615,7 @@ def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): prev_tokens.fill_(scn["fill"]) prev_k_for_print = scn["print_label"] scenario_pre_iter = scn["pre_iter"] + scenario_pre_iter_group_factory = scn.get("pre_iter_group_factory") scenario_iters = scn.get("iters") # None => use args.iters tag = f"incr_b{batch}_mtp{mtp_len}_{scn['label']}_s{state_dtype_name}_a{act_dtype_name}" @@ -2788,8 +2856,8 @@ def _run_incr( ) # n_writes plumbing: pure scenarios pass an int # (host knows the value, can host-skip empty halves); - # mix scenarios pass a (1,) device tensor that the - # bench's pre_iter_fn updates per replay. + # mix scenarios pass a (1,) device tensor updated + # per iter by the benchmark pre-iter path. # _persistent_skip_empty_halves=False on mix so both # halves always launch (kernel uses device n_writes # to derive its slot range). @@ -2989,6 +3057,7 @@ def _emit_split(name_w, name_nw, val_w, val_nw): args, _run_incr, reset_fn, sweep_tag, expected_K=expected_K, pre_iter_fn=scenario_pre_iter, + pre_iter_group_factory=scenario_pre_iter_group_factory, iters_override=scenario_iters, cupti_plan_key=plan_key, ) @@ -3704,11 +3773,13 @@ def _parse_args() -> argparse.Namespace: parser.add_argument( "--cuda-graph-group-iters", type=int, - default=_DEFAULT_CUDA_GRAPH_GROUP_ITERS, - help="For pure CUDA-graph cells, capture this many logical benchmark " - "iterations per graph replay when warmup + iters is divisible by this " - "value. Mix cells stay at one logical iteration per replay because " - "their per-iteration input update runs outside the graph.", + default=None, + help="Capture this many logical benchmark iterations per graph " + "replay when warmup + iters is divisible by this value. Default " + f"auto-selects {_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE} for pure " + f"cells and {_DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX} for mix cells. " + "Mix cells use a per-replay device window so they can group " + "iterations too.", ) parser.add_argument( "--cupti", From c69ea445b813d49417e9a073fd134d33482fae94 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 17:16:23 -0700 Subject: [PATCH 39/89] bench: reduce cuda graph pre-capture warmup Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../benchmark_replay_selective_state_update.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 574db578a827..7e4ba2c889cc 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -1367,7 +1367,7 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, return out -_PRE_GRAPH_WARMUP_ITERS = 3 # standard practice; see commit msg / design doc +_PRE_GRAPH_WARMUP_ITERS = 1 _CUPTI_FILTER_PLAN_CACHE: dict[tuple, tuple[int, tuple[str | None, ...]]] = {} @@ -1562,7 +1562,7 @@ def _time_kernel_cuda_graph( the next group of PNAT/n_writes samples, then graph-captured per-iter copies update kernel inputs before each reset + L2 flush + run. - Pre-graph eager warmup (3 iters): forces PyTorch's caching allocator + Pre-graph eager warmup: forces PyTorch's caching allocator + Triton's autotune cache to settle before capture so the graph doesn't bake in init-only allocations. @@ -1575,16 +1575,19 @@ def _time_kernel_cuda_graph( warmup = args.warmup iters = iters_override if iters_override is not None else args.iters - # Pre-graph eager warmup: full per-iter chain × N. Settles allocator - # + autotune; pre_iter_fn included so any side effects it has are - # exercised before capture. + # Pre-graph eager warmup: full per-iter chain once. This settles + # Triton/PyTorch setup and wrapper-side intermediate allocations; + # skipping it risks lazy work leaking into graph capture. + warmup_iters = _PRE_GRAPH_WARMUP_ITERS + host_timing.add("pre_graph_warmup_iters", warmup_iters) host_timing.start() - for _ in range(_PRE_GRAPH_WARMUP_ITERS): + for _ in range(warmup_iters): reset_fn() if pre_iter_fn is not None: pre_iter_fn(0) run_fn() - torch.cuda.synchronize() + if warmup_iters > 0: + torch.cuda.synchronize() host_timing.stop("pre_graph_warmup_ms") total_iters = warmup + iters From 73a438a40f2ffa8ebc5c3310b06b3977eed879f5 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 17:28:52 -0700 Subject: [PATCH 40/89] bench: add mix-only replay benchmark mode Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 38 ++++++++++++------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 7e4ba2c889cc..fe6dc36dadf9 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -2467,20 +2467,21 @@ def _split_or_share(split_csv, shared_values): # tensor, with the per-iter copy captured inside the CUDA graph. Pure # and mix can coexist in one call so a single nsys trace covers both. scenarios = [] - for prev_k in prev_ks: - # On the nowrite path, new tokens append at [prev_k, prev_k+T) of - # the active buffer, so prev_k+T must fit within max_window. - # mode != monolithic dispatches per-slot from PNAT, so any - # prev_k <= max_window is valid for those modes. - if mode == "monolithic" and not write_checkpoint and prev_k + mtp_len > max_window: - continue - scenarios.append({ - "label": f"k{prev_k}", - "print_label": prev_k, - "fill": prev_k, - "pre_iter": None, - "iters": None, # use args.iters - }) + if not (getattr(args, "mix_only", False) and mix_samples_cpu is not None): + for prev_k in prev_ks: + # On the nowrite path, new tokens append at [prev_k, prev_k+T) of + # the active buffer, so prev_k+T must fit within max_window. + # mode != monolithic dispatches per-slot from PNAT, so any + # prev_k <= max_window is valid for those modes. + if mode == "monolithic" and not write_checkpoint and prev_k + mtp_len > max_window: + continue + scenarios.append({ + "label": f"k{prev_k}", + "print_label": prev_k, + "fill": prev_k, + "pre_iter": None, + "iters": None, # use args.iters + }) # Mix scenario: skip on monolithic (mono on mixed PNAT corrupts the # wrong-mode slots). Persistent_main + mix is now supported: bench # pre-bakes both a per-iter PNAT samples tensor and a per-iter @@ -4214,6 +4215,13 @@ def _parse_args() -> argparse.Namespace: "--iters. Mix scenarios benefit from more iters since each " "iter samples a different mix; pure scenarios don't.", ) + parser.add_argument( + "--mix-only", + action=argparse.BooleanOptionalAction, + default=False, + help="When --mix-csv is set, emit only mix scenarios and skip the " + "pure prev_k sibling scenarios. Default: false.", + ) parser.add_argument( "--philox-rounding", action="store_true", @@ -4247,6 +4255,8 @@ def _parse_args() -> argparse.Namespace: "if the fast path breaks due to package changes.", ) args = parser.parse_args() + if args.mix_only and args.mix_csv is None: + parser.error("--mix-only requires --mix-csv") # Backward-compat: --philox-rounding implies --sr-modes SR if --sr-modes # was left at the default. If both are set explicitly, error. From eb762147ae45836793074eb84a9d21ed04455854 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 17:34:02 -0700 Subject: [PATCH 41/89] bench: keep compact per-iter mix stats Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 40 +++++++++++-------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index fe6dc36dadf9..b984886150a2 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -68,7 +68,7 @@ { "metadata": {timestamp, cmd, tp_size, warmup, iters, variant, cupti}, "results": { - "": {median, p95, p99, n, [iters_us], [per_kernel]} + "": {median, p95, p99, n, iters_us, [n_writes_per_iter], [per_kernel]} } } @@ -87,7 +87,9 @@ min(kernel_start_ns) across the iter's kernels — same convention as nsys-derived collect.py used to use. - n: number of timed iters that contributed. - - iters_us: list of length n, raw per-iter spans (only with --json-detailed). + - iters_us: list of length n, raw per-iter spans. + - n_writes_per_iter: for mix rows, list of length n with the number of + write-path slots in each timed iteration. - per_kernel: {: {start_us: [...], end_us: [...]}} where timestamps are RELATIVE to that iter's first kernel start, in us. Lets you see PDL overlap directly without an external profiler. Only with @@ -1361,8 +1363,8 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, slot["end_us"].append((r[2] - iter_start_ns) / 1000.0) out = _stats_from_spans(spans_us) + out["iters_us"] = spans_us if include_details: - out["iters_us"] = spans_us out["per_kernel"] = per_kernel return out @@ -3076,17 +3078,20 @@ def _emit_split(name_w, name_nw, val_w, val_nw): if stats is None: args._skipped_cells.append(sweep_tag) - # Attach n_writes_per_iter for --json-detailed bucketing. + # Attach n_writes_per_iter when it is needed for scoring. # For pure scenarios, n_writes is constant: 0 (nowrite) or # batch (write), determined by scn["fill"] + mtp_len > max_window. # For mix, scn carries the precomputed per-iter array. per_iter_nw = None - if stats is not None and getattr(args, "json_detailed", False): + if stats is not None and ( + getattr(args, "json_detailed", False) or scn["fill"] is None + ): eff_iters = scenario_iters if scenario_iters is not None else args.iters if scn["fill"] is not None: - # Pure scenario: constant n_writes for every iter. - is_write = (scn["fill"] + mtp_len > max_window) - per_iter_nw = [batch if is_write else 0] * eff_iters + if getattr(args, "json_detailed", False): + # Pure scenario: constant n_writes for every iter. + is_write = (scn["fill"] + mtp_len > max_window) + per_iter_nw = [batch if is_write else 0] * eff_iters else: # Mix scenario: slice off warmup, keep timed iters. nw_full = scn.get("n_writes_per_iter") @@ -3172,9 +3177,9 @@ def _print_row( """Print one summary row and optionally accumulate stats for JSON output. `stats` is a dict from _time_kernel: {median, p95, p99, n, iters_us, - [per_kernel]}. The summary table only shows the headline percentiles. - JSON output captures median/p95/p99/n by default; with json_detailed=True - it also captures the per-iter and per-kernel data. + [n_writes_per_iter], [per_kernel]}. The summary table only shows the + headline percentiles. JSON output captures compact per-iter spans by + default; with json_detailed=True it also captures per-kernel data. When `jsonl_path` is provided, also appends one JSON line per row to the JSONL sidecar (crash-safe incremental persistence; lets a killed sweep @@ -3198,7 +3203,8 @@ def _print_row( row_stats = stats else: row_stats = { - k: stats[k] for k in ("median", "p95", "p99", "n") + k: stats[k] + for k in ("median", "p95", "p99", "n", "iters_us", "n_writes_per_iter") if k in stats } if "host_timing" in stats: @@ -3232,7 +3238,7 @@ def _finish_result_job(args, job: dict) -> None: return per_iter_nw = job.get("per_iter_nw") - if per_iter_nw is not None and getattr(args, "json_detailed", False): + if per_iter_nw is not None: stats["n_writes_per_iter"] = per_iter_nw _print_row( @@ -3823,10 +3829,10 @@ def _parse_args() -> argparse.Namespace: "--json-detailed", action=argparse.BooleanOptionalAction, default=False, - help="When --json-output is set, also include iters_us (raw per-iter " - "spans) and per_kernel (per-iter relative start/end timestamps for " - "each kernel) — useful for PDL overlap analysis but adds ~4 KB/cell. " - "Default off keeps records to ~40 bytes (median/p95/p99/n only).", + help="When --json-output is set, also include per_kernel " + "(per-iter relative start/end timestamps for each kernel). Compact " + "JSON always includes iters_us, and mix rows include " + "n_writes_per_iter. Default off keeps records compact.", ) parser.add_argument( "--host-timing", From 25836d843a72d85d5e4a9cd744a097750fc431df Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 16:33:19 -0700 Subject: [PATCH 42/89] bench: keep CUPTI flush probe scratch helper Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../_torch/modules/mamba/cupti_flush_probe.py | 1029 +++++++++++++++++ 1 file changed, 1029 insertions(+) create mode 100644 tests/unittest/_torch/modules/mamba/cupti_flush_probe.py diff --git a/tests/unittest/_torch/modules/mamba/cupti_flush_probe.py b/tests/unittest/_torch/modules/mamba/cupti_flush_probe.py new file mode 100644 index 000000000000..b05bbfded754 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/cupti_flush_probe.py @@ -0,0 +1,1029 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Microbenchmarks for CUPTI Activity flush behavior. + +This is intentionally independent of the Mamba replay benchmark. It answers +questions about CUPTI overheads with a one-kernel CUDA graph: + +* Does flush cost scale with completed kernel records? +* Can a background flush overlap Python execution and CUDA graph enqueue? +* If work B is already enqueued after an event following work A, can a + non-forced flush deliver A's completed records while B is still running? +* Does CUPTI periodic flushing reduce the explicit end-of-cell flush cost? +* How much of flush time is Python callback work? + +TODO: Once the raw parser path is validated in the full benchmark, test +capturing 2/4/8 benchmark iterations per CUDA graph replay. Larger graphs +should reduce Python graph.replay() calls while keeping graph instantiation +cheap enough for sweeps. +""" + +from __future__ import annotations + +import argparse +import ctypes +import json +import multiprocessing as mp +import queue +import sys +import threading +import time +from collections.abc import Iterable +from multiprocessing import shared_memory + +import torch + + +_LIBCUPTI_PATH = "/usr/local/lib/python3.12/dist-packages/nvidia/cu13/lib/libcupti.so.13" +_CUPTI_SUCCESS = 0 +_CUPTI_ATTR_DEVICE_BUFFER_SIZE = 0 +_CUPTI_ATTR_DEVICE_BUFFER_POOL_LIMIT = 2 +_CUPTI_ATTR_ZEROED_OUT_ACTIVITY_BUFFER = 5 +_CUPTI_ATTR_DEVICE_BUFFER_PRE_ALLOCATE_VALUE = 6 +_CUPTI_ATTR_MEM_ALLOCATION_TYPE_HOST_PINNED = 8 +_CUPTI_ERROR_MAX_LIMIT_REACHED = 12 +_CUPTI_ERROR_INVALID_KIND = 21 +_CUPTI_ACTIVITY_KIND_KERNEL = 3 +_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL = 10 + + +def set_cupti_size_t_attribute(attr: int, value: int | None) -> None: + if value is None: + return + libcupti = ctypes.CDLL(_LIBCUPTI_PATH) + set_attribute = libcupti.cuptiActivitySetAttribute + set_attribute.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_void_p, + ] + set_attribute.restype = ctypes.c_int + value_obj = ctypes.c_size_t(value) + size_obj = ctypes.c_size_t(ctypes.sizeof(value_obj)) + result = set_attribute(attr, ctypes.byref(size_obj), ctypes.byref(value_obj)) + if result != _CUPTI_SUCCESS: + raise RuntimeError(f"cuptiActivitySetAttribute({attr}, {value}) failed with CUptiResult={result}") + + +def set_cupti_uint8_attribute(attr: int, value: bool | None) -> None: + if value is None: + return + libcupti = ctypes.CDLL(_LIBCUPTI_PATH) + set_attribute = libcupti.cuptiActivitySetAttribute + set_attribute.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_void_p, + ] + set_attribute.restype = ctypes.c_int + value_obj = ctypes.c_uint8(1 if value else 0) + size_obj = ctypes.c_size_t(ctypes.sizeof(value_obj)) + result = set_attribute(attr, ctypes.byref(size_obj), ctypes.byref(value_obj)) + if result != _CUPTI_SUCCESS: + raise RuntimeError(f"cuptiActivitySetAttribute({attr}, {int(value)}) failed with CUptiResult={result}") + + +class _CuptiActivityKernel11Prefix(ctypes.Structure): + _pack_ = 1 + _fields_ = [ + ("kind", ctypes.c_int), + ("cache_config", ctypes.c_uint8), + ("shared_memory_config", ctypes.c_uint8), + ("registers_per_thread", ctypes.c_uint16), + ("partitioned_global_cache_requested", ctypes.c_int), + ("partitioned_global_cache_executed", ctypes.c_int), + ("start", ctypes.c_uint64), + ("end", ctypes.c_uint64), + ("completed", ctypes.c_uint64), + ("device_id", ctypes.c_uint32), + ("context_id", ctypes.c_uint32), + ("stream_id", ctypes.c_uint32), + ("grid_x", ctypes.c_int32), + ("grid_y", ctypes.c_int32), + ("grid_z", ctypes.c_int32), + ("block_x", ctypes.c_int32), + ("block_y", ctypes.c_int32), + ("block_z", ctypes.c_int32), + ("static_shared_memory", ctypes.c_int32), + ("dynamic_shared_memory", ctypes.c_int32), + ("local_memory_per_thread", ctypes.c_uint32), + ("local_memory_total", ctypes.c_uint32), + ("correlation_id", ctypes.c_uint32), + ("grid_id", ctypes.c_int64), + ("name", ctypes.c_void_p), + ("reserved0", ctypes.c_void_p), + ("queued", ctypes.c_uint64), + ("submitted", ctypes.c_uint64), + ("launch_type", ctypes.c_uint8), + ("is_shared_memory_carveout_requested", ctypes.c_uint8), + ("shared_memory_carveout_requested", ctypes.c_uint8), + ("padding", ctypes.c_uint8), + ("shared_memory_executed", ctypes.c_uint32), + ("graph_node_id", ctypes.c_uint64), + ] + + +def parse_cupti_buffer_ptr(buffer_ptr: int, valid_size: int) -> dict[str, float | int]: + libcupti = ctypes.CDLL(_LIBCUPTI_PATH) + get_next_record = libcupti.cuptiActivityGetNextRecord + get_next_record.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_void_p), + ] + get_next_record.restype = ctypes.c_int + record_ptr = ctypes.c_void_p(None) + records = 0 + kernel_records = 0 + zero_ts = 0 + invalid_kind = 0 + min_start = None + max_end = 0 + start_s = time.perf_counter() + while True: + result = get_next_record(ctypes.c_void_p(buffer_ptr), valid_size, ctypes.byref(record_ptr)) + if result == _CUPTI_SUCCESS: + records += 1 + kind = ctypes.cast(record_ptr, ctypes.POINTER(ctypes.c_int)).contents.value + if kind in (_CUPTI_ACTIVITY_KIND_KERNEL, _CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL): + kernel = ctypes.cast(record_ptr, ctypes.POINTER(_CuptiActivityKernel11Prefix)).contents + if kernel.start == 0 or kernel.end == 0: + zero_ts += 1 + else: + kernel_records += 1 + min_start = kernel.start if min_start is None else min(min_start, kernel.start) + max_end = max(max_end, kernel.end) + elif result == _CUPTI_ERROR_MAX_LIMIT_REACHED: + break + elif result == _CUPTI_ERROR_INVALID_KIND: + invalid_kind += 1 + break + else: + raise RuntimeError(f"cuptiActivityGetNextRecord failed with CUptiResult={result}") + parse_s = time.perf_counter() - start_s + span_us = 0.0 if min_start is None else (max_end - min_start) / 1000.0 + return { + "records": records, + "kernel_records": kernel_records, + "zero_ts": zero_ts, + "invalid_kind": invalid_kind, + "parse_ms": 1000.0 * parse_s, + "span_us": span_us, + } + + +def parse_cupti_payload(payload: bytes) -> dict[str, float | int]: + raw_buffer = ctypes.create_string_buffer(payload) + return parse_cupti_buffer_ptr(ctypes.addressof(raw_buffer), len(payload)) + + +def cupti_parser_worker(input_queue, output_queue, ready_event, zero_shm_before_ack: bool) -> None: + ctypes.CDLL(_LIBCUPTI_PATH) + ready_event.set() + shared_blocks: dict[str, shared_memory.SharedMemory] = {} + while True: + item = input_queue.get() + if item is None: + break + generation = item[0] + try: + transport = item[1] + if transport == "bytes": + result = parse_cupti_payload(item[2]) + elif transport == "shm": + _, _, name, buffer_ptr, valid_size = item + shm = shared_blocks.get(name) + if shm is None: + shm = shared_memory.SharedMemory(name=name) + shared_blocks[name] = shm + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + parser_ptr = ctypes.addressof(shared_char) + result = parse_cupti_buffer_ptr(parser_ptr, valid_size) + if zero_shm_before_ack: + ctypes.memset(parser_ptr, 0, len(shm.buf)) + finally: + del shared_char + result["buffer_ptr"] = buffer_ptr + else: + raise ValueError(f"unknown parser transport {transport!r}") + result["generation"] = generation + output_queue.put(result) + except Exception as exc: # pragma: no cover - diagnostic worker path + output_queue.put({"generation": generation, "error": repr(exc)}) + for shm in shared_blocks.values(): + shm.close() + + +class CuptiActivityProbe: + """Tiny CUPTI Activity wrapper with selectable callback cost.""" + + def __init__(self, callback_mode: str, host_buffer_bytes: int, max_records_per_buffer: int) -> None: + from cupti import cupti as cupti + + self._cupti = cupti + self._callback_mode = callback_mode + self._host_buffer_bytes = host_buffer_bytes + self._max_records_per_buffer = max_records_per_buffer + self._lock = threading.Lock() + self._record_count = 0 + self._zero_ts_count = 0 + self._callback_count = 0 + self._name_count: dict[str, int] = {} + self._records: list[tuple[int, int, int]] = [] + self._kernel_kinds = ( + cupti.ActivityKind.CONCURRENT_KERNEL, + cupti.ActivityKind.KERNEL, + ) + + def buffer_requested() -> tuple[int, int]: + return (self._host_buffer_bytes, self._max_records_per_buffer) + + def buffer_completed(activities) -> None: + local_count = 0 + local_zero_ts = 0 + local_names: dict[str, int] = {} + local_records: list[tuple[int, int, int]] = [] + for activity in activities: + if activity.kind not in self._kernel_kinds: + continue + if activity.start == 0 or activity.end == 0: + local_zero_ts += 1 + continue + local_count += 1 + if self._callback_mode in ("names", "records"): + name = getattr(activity, "name", "?") + local_names[name] = local_names.get(name, 0) + 1 + if self._callback_mode in ("numeric", "records"): + local_records.append(( + int(activity.start), + int(activity.end), + int(activity.graph_node_id), + )) + with self._lock: + self._record_count += local_count + self._zero_ts_count += local_zero_ts + self._callback_count += 1 + for name, count in local_names.items(): + self._name_count[name] = self._name_count.get(name, 0) + count + self._records.extend(local_records) + + self._buffer_requested = buffer_requested + self._buffer_completed = buffer_completed + cupti.activity_register_callbacks(buffer_requested, buffer_completed) + cupti.activity_enable(cupti.ActivityKind.CONCURRENT_KERNEL) + + def clear(self) -> None: + """Force-drain stale records, then clear local counters.""" + self.flush(flag=1) + with self._lock: + self._record_count = 0 + self._zero_ts_count = 0 + self._callback_count = 0 + self._name_count = {} + self._records = [] + + def flush(self, flag: int) -> float: + start_s = time.perf_counter() + self._cupti.activity_flush_all(flag) + return time.perf_counter() - start_s + + def set_flush_period(self, period_ms: int) -> None: + self._cupti.activity_flush_period(period_ms) + + def snapshot(self) -> dict[str, object]: + with self._lock: + return { + "records": self._record_count, + "zero_ts": self._zero_ts_count, + "callbacks": self._callback_count, + "unique_names": len(self._name_count), + } + + +class RawCuptiActivityProbe: + """CUPTI Activity wrapper using raw ctypes callbacks. + + Unlike cupti-python's high-level callback, this callback receives the raw + CUPTI activity buffer pointer. It only records the pointer and valid byte + count, so flush does not materialize one Python object per activity record. + """ + + _request_callback_type = ctypes.CFUNCTYPE( + None, + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t), + ) + _complete_callback_type = ctypes.CFUNCTYPE( + None, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_size_t, + ) + + def __init__( + self, + host_buffer_bytes: int, + max_records_per_buffer: int, + parse_process: bool, + parse_transport: str, + host_buffer_count: int, + parser_zero_shm: bool, + ) -> None: + self._libcupti = ctypes.CDLL(_LIBCUPTI_PATH) + self._host_buffer_bytes = host_buffer_bytes + self._max_records_per_buffer = max_records_per_buffer + self._parse_transport = parse_transport + self._lock = threading.Lock() + self._buffers: dict[int, ctypes.Array] = {} + self._shared_buffers: dict[int, shared_memory.SharedMemory] = {} + self._free_ptrs: list[int] = [] + self._completed: list[tuple[int, int, int]] = [] + self._request_count = 0 + self._complete_count = 0 + self._parse_process_enabled = parse_process + self._parser_records = 0 + self._parser_kernel_records = 0 + self._parser_zero_ts = 0 + self._parser_buffers = 0 + self._parser_errors = 0 + self._parser_parse_ms = 0.0 + self._parser_span_us = 0.0 + self._mp_ctx = None + self._parse_input_queue = None + self._parse_output_queue = None + self._parse_process = None + self._generation = 0 + + self._configure_functions() + if parse_process: + self._mp_ctx = mp.get_context("spawn") + self._parse_input_queue = self._mp_ctx.Queue() + self._parse_output_queue = self._mp_ctx.Queue() + ready_event = self._mp_ctx.Event() + self._parse_process = self._mp_ctx.Process( + target=cupti_parser_worker, + args=(self._parse_input_queue, self._parse_output_queue, ready_event, parser_zero_shm), + ) + self._parse_process.start() + if not ready_event.wait(timeout=5.0): + raise RuntimeError("CUPTI parser process did not initialize") + if parse_transport == "shm": + for _ in range(host_buffer_count): + self._free_ptrs.append(self._allocate_shared_buffer()) + self._request_callback = self._request_callback_type(self._request_buffer) + self._complete_callback = self._complete_callback_type(self._complete_buffer) + self._check(self._libcupti.cuptiActivityRegisterCallbacks( + self._request_callback, + self._complete_callback, + )) + self._check(self._libcupti.cuptiActivityEnable(10)) + + def _configure_functions(self) -> None: + self._libcupti.cuptiActivityRegisterCallbacks.argtypes = [ + self._request_callback_type, + self._complete_callback_type, + ] + self._libcupti.cuptiActivityRegisterCallbacks.restype = ctypes.c_int + self._libcupti.cuptiActivityEnable.argtypes = [ctypes.c_int] + self._libcupti.cuptiActivityEnable.restype = ctypes.c_int + self._libcupti.cuptiActivityFlushAll.argtypes = [ctypes.c_uint32] + self._libcupti.cuptiActivityFlushAll.restype = ctypes.c_int + self._libcupti.cuptiActivityFlushPeriod.argtypes = [ctypes.c_uint32] + self._libcupti.cuptiActivityFlushPeriod.restype = ctypes.c_int + + def _check(self, result: int) -> None: + if result != _CUPTI_SUCCESS: + raise RuntimeError(f"CUPTI call failed with CUptiResult={result}") + + def _allocate_buffer(self) -> int: + raw_buffer = ctypes.create_string_buffer(self._host_buffer_bytes + 8) + raw_ptr = ctypes.addressof(raw_buffer) + aligned_ptr = (raw_ptr + 7) & ~7 + self._buffers[aligned_ptr] = raw_buffer + return aligned_ptr + + def _allocate_shared_buffer(self) -> int: + shm = shared_memory.SharedMemory(create=True, size=self._host_buffer_bytes + 8) + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + raw_ptr = ctypes.addressof(shared_char) + finally: + del shared_char + aligned_ptr = (raw_ptr + 7) & ~7 + if aligned_ptr != raw_ptr: + raise RuntimeError("shared memory buffer was not 8-byte aligned") + self._shared_buffers[aligned_ptr] = shm + return aligned_ptr + + def _request_buffer(self, buffer, size, max_num_records) -> None: + with self._lock: + if self._free_ptrs: + ptr = self._free_ptrs.pop() + else: + if self._parse_process_enabled and self._parse_transport == "shm": + ptr = self._allocate_shared_buffer() + else: + ptr = self._allocate_buffer() + self._request_count += 1 + buffer[0] = ptr + size[0] = self._host_buffer_bytes + max_num_records[0] = self._max_records_per_buffer + + def _complete_buffer(self, context, stream_id, buffer, size, valid_size) -> None: + del context, stream_id + buffer_ptr = int(buffer) + valid_size_int = int(valid_size) + size_int = int(size) + with self._lock: + generation = self._generation + if self._parse_process_enabled: + if self._parse_transport == "shm": + shm = self._shared_buffers[buffer_ptr] + self._parse_input_queue.put(( + generation, + "shm", + shm.name, + buffer_ptr, + valid_size_int, + )) + else: + payload = ctypes.string_at(buffer, valid_size_int) + self._parse_input_queue.put((generation, "bytes", payload)) + with self._lock: + if self._parse_process_enabled and self._parse_transport == "bytes": + self._free_ptrs.append(buffer_ptr) + elif self._parse_process_enabled: + pass + else: + self._completed.append((buffer_ptr, valid_size_int, size_int)) + self._complete_count += 1 + + def clear(self) -> None: + self.flush(flag=1) + self.wait_for_parser(timeout_s=1.0) + with self._lock: + self._free_ptrs.extend(ptr for ptr, _, _ in self._completed) + self._completed = [] + self._request_count = 0 + self._complete_count = 0 + self._generation += 1 + self._parser_records = 0 + self._parser_kernel_records = 0 + self._parser_zero_ts = 0 + self._parser_buffers = 0 + self._parser_errors = 0 + self._parser_parse_ms = 0.0 + self._parser_span_us = 0.0 + + def flush(self, flag: int) -> float: + start_s = time.perf_counter() + self._check(self._libcupti.cuptiActivityFlushAll(flag)) + return time.perf_counter() - start_s + + def set_flush_period(self, period_ms: int) -> None: + self._check(self._libcupti.cuptiActivityFlushPeriod(period_ms)) + + def _drain_parser_results(self) -> None: + if not self._parse_process_enabled: + return + while True: + try: + result = self._parse_output_queue.get_nowait() + except queue.Empty: + break + if result.get("generation") != self._generation: + continue + if "error" in result: + self._parser_errors += 1 + continue + buffer_ptr = result.get("buffer_ptr") + if buffer_ptr is not None: + with self._lock: + self._free_ptrs.append(int(buffer_ptr)) + self._parser_buffers += 1 + self._parser_records += int(result["records"]) + self._parser_kernel_records += int(result["kernel_records"]) + self._parser_zero_ts += int(result["zero_ts"]) + self._parser_parse_ms += float(result["parse_ms"]) + self._parser_span_us = max(self._parser_span_us, float(result["span_us"])) + + def wait_for_parser(self, timeout_s: float = 0.5) -> None: + if not self._parse_process_enabled: + return + deadline_s = time.perf_counter() + timeout_s + while True: + self._drain_parser_results() + with self._lock: + complete_count = self._complete_count + parser_buffers = self._parser_buffers + self._parser_errors + if parser_buffers >= complete_count or time.perf_counter() >= deadline_s: + break + time.sleep(0.001) + + def snapshot(self) -> dict[str, object]: + self._drain_parser_results() + with self._lock: + snapshot = { + "records": None, + "zero_ts": None, + "callbacks": self._complete_count, + "requests": self._request_count, + "valid_bytes": sum(valid_size for _, valid_size, _ in self._completed), + "completed_buffers": len(self._completed), + "unique_names": None, + } + if self._parse_process_enabled: + snapshot.update({ + "parser_buffers": self._parser_buffers, + "parser_records": self._parser_records, + "parser_kernel_records": self._parser_kernel_records, + "parser_zero_ts": self._parser_zero_ts, + "parser_errors": self._parser_errors, + "parser_parse_ms": self._parser_parse_ms, + "parser_span_us": self._parser_span_us, + }) + return snapshot + + def close(self) -> None: + if not self._parse_process_enabled: + return + self._parse_input_queue.put(None) + self._parse_process.join(timeout=5.0) + if self._parse_process.is_alive(): + self._parse_process.terminate() + self._parse_process.join(timeout=1.0) + for shm in self._shared_buffers.values(): + shm.close() + shm.unlink() + + +def emit(case: str, **fields: object) -> None: + print(json.dumps({"case": case, **fields}, sort_keys=True), flush=True) + + +def parse_ints(text: str) -> list[int]: + return [int(part.strip()) for part in text.split(",") if part.strip()] + + +def make_one_kernel_graph(elements: int, kernels_per_graph: int) -> torch.cuda.CUDAGraph: + tensor = torch.ones((elements,), device="cuda") + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(kernels_per_graph): + tensor.add_(1.0) + torch.cuda.synchronize() + return graph + + +def replay(graph: torch.cuda.CUDAGraph, count: int) -> float: + start_s = time.perf_counter() + for _ in range(count): + graph.replay() + return time.perf_counter() - start_s + + +def graph_setup_case(elements: int, kernels_per_graph_values: Iterable[int], replays: int, trials: int) -> None: + warm_tensor = torch.ones((elements,), device="cuda") + for _ in range(3): + warm_tensor.add_(1.0) + torch.cuda.synchronize() + warm_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(warm_graph): + warm_tensor.add_(1.0) + torch.cuda.synchronize() + warm_graph.replay() + torch.cuda.synchronize() + + for kernels_per_graph in kernels_per_graph_values: + capture_times_ms = [] + replay_times_ms = [] + warmup_times_ms = [] + for trial in range(trials): + tensor = torch.ones((elements,), device="cuda") + torch.cuda.synchronize() + + warm_start_s = time.perf_counter() + for _ in range(3): + for _ in range(kernels_per_graph): + tensor.add_(1.0) + torch.cuda.synchronize() + warmup_s = time.perf_counter() - warm_start_s + + capture_start_s = time.perf_counter() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(kernels_per_graph): + tensor.add_(1.0) + torch.cuda.synchronize() + capture_s = time.perf_counter() - capture_start_s + + replay_s = replay(graph, replays) + torch.cuda.synchronize() + capture_times_ms.append(1000.0 * capture_s) + replay_times_ms.append(1000.0 * replay_s) + warmup_times_ms.append(1000.0 * warmup_s) + emit( + "graph_setup_trial", + kernels_per_graph=kernels_per_graph, + trial=trial, + warmup_ms=1000.0 * warmup_s, + capture_instantiate_ms=1000.0 * capture_s, + replay_count=replays, + replay_enqueue_ms=1000.0 * replay_s, + replay_enqueue_us_per_replay=1_000_000.0 * replay_s / replays, + replay_enqueue_us_per_kernel=1_000_000.0 * replay_s / (replays * kernels_per_graph), + ) + capture_sorted = sorted(capture_times_ms) + replay_sorted = sorted(replay_times_ms) + warmup_sorted = sorted(warmup_times_ms) + mid = len(capture_sorted) // 2 + emit( + "graph_setup", + kernels_per_graph=kernels_per_graph, + trials=trials, + warmup_ms_median=warmup_sorted[mid], + capture_instantiate_ms_median=capture_sorted[mid], + capture_instantiate_ms_min=min(capture_times_ms), + capture_instantiate_ms_max=max(capture_times_ms), + replay_count=replays, + replay_enqueue_ms_median=replay_sorted[mid], + replay_enqueue_us_per_replay=1000.0 * replay_sorted[mid] / replays, + replay_enqueue_us_per_kernel=1000.0 * replay_sorted[mid] / (replays * kernels_per_graph), + ) + + +def wait_for_parser_if_present(probe, timeout_s: float = 0.5) -> None: + if hasattr(probe, "wait_for_parser"): + probe.wait_for_parser(timeout_s=timeout_s) + + +def completed_scaling( + probe: CuptiActivityProbe, + graph: torch.cuda.CUDAGraph, + counts: Iterable[int], + flags: Iterable[int], + kernels_per_graph: int, +) -> None: + for flag in flags: + for count in counts: + probe.clear() + enqueue_s = replay(graph, count) + sync_start_s = time.perf_counter() + torch.cuda.synchronize() + sync_s = time.perf_counter() - sync_start_s + flush_s = probe.flush(flag) + wait_for_parser_if_present(probe) + emit( + "completed_scaling", + count=count, + graph_replays=count, + kernels_per_graph=kernels_per_graph, + expected_kernel_records=count * kernels_per_graph, + flag=flag, + enqueue_ms=1000.0 * enqueue_s, + sync_ms=1000.0 * sync_s, + flush_ms=1000.0 * flush_s, + **probe.snapshot(), + ) + + +def overlap_enqueue( + probe: CuptiActivityProbe, + graph: torch.cuda.CUDAGraph, + count: int, + flag: int, + yield_before_overlap: bool, + kernels_per_graph: int, +) -> None: + probe.clear() + replay(graph, 20) + torch.cuda.synchronize() + + baseline_enqueue_s = replay(graph, count) + torch.cuda.synchronize() + + probe.clear() + replay(graph, count) + torch.cuda.synchronize() + + ready = threading.Event() + flush_result: dict[str, float] = {} + + def flush_worker() -> None: + ready.set() + flush_result["flush_s"] = probe.flush(flag) + + thread = threading.Thread(target=flush_worker) + thread.start() + ready.wait() + if yield_before_overlap: + time.sleep(0) + overlap_enqueue_s = replay(graph, count) + thread.join() + torch.cuda.synchronize() + tail_flush_s = probe.flush(flag) + wait_for_parser_if_present(probe) + emit( + "overlap_enqueue", + count=count, + graph_replays=count, + kernels_per_graph=kernels_per_graph, + expected_kernel_records=count * kernels_per_graph, + flag=flag, + baseline_enqueue_ms=1000.0 * baseline_enqueue_s, + overlap_enqueue_ms=1000.0 * overlap_enqueue_s, + background_flush_ms=1000.0 * flush_result["flush_s"], + tail_flush_ms=1000.0 * tail_flush_s, + switch_interval_ms=1000.0 * sys.getswitchinterval(), + yield_before_overlap=yield_before_overlap, + **probe.snapshot(), + ) + + +def gil_gap(probe: CuptiActivityProbe, graph: torch.cuda.CUDAGraph, count: int, flag: int) -> None: + probe.clear() + replay(graph, count) + torch.cuda.synchronize() + + ready = threading.Event() + flush_result: dict[str, float] = {} + + def flush_worker() -> None: + ready.set() + flush_result["flush_s"] = probe.flush(flag) + + thread = threading.Thread(target=flush_worker) + thread.start() + ready.wait() + loop_start_s = time.perf_counter() + previous_s = loop_start_s + max_gap_s = 0.0 + iterations = 0 + while thread.is_alive(): + now_s = time.perf_counter() + max_gap_s = max(max_gap_s, now_s - previous_s) + previous_s = now_s + iterations += 1 + thread.join() + loop_s = time.perf_counter() - loop_start_s + wait_for_parser_if_present(probe) + emit( + "gil_gap", + count=count, + flag=flag, + background_flush_ms=1000.0 * flush_result["flush_s"], + main_loop_ms=1000.0 * loop_s, + max_main_thread_gap_ms=1000.0 * max_gap_s, + main_loop_iterations=iterations, + **probe.snapshot(), + ) + + +def prefix_event_flush( + probe: CuptiActivityProbe, + graph: torch.cuda.CUDAGraph, + prefix_count: int, + tail_count: int, + flag: int, + kernels_per_graph: int, +) -> None: + probe.clear() + replay(graph, 20) + torch.cuda.synchronize() + probe.clear() + + enqueue_start_s = time.perf_counter() + replay(graph, prefix_count) + event = torch.cuda.Event() + event.record() + replay(graph, tail_count) + enqueue_s = time.perf_counter() - enqueue_start_s + + event_sync_start_s = time.perf_counter() + event.synchronize() + event_sync_s = time.perf_counter() - event_sync_start_s + + flush1_s = probe.flush(flag) + snapshot_after_flush1 = probe.snapshot() + + tail_sync_start_s = time.perf_counter() + torch.cuda.synchronize() + tail_sync_s = time.perf_counter() - tail_sync_start_s + + flush2_s = probe.flush(flag) + wait_for_parser_if_present(probe) + emit( + "prefix_event_flush", + prefix_count=prefix_count, + tail_count=tail_count, + kernels_per_graph=kernels_per_graph, + expected_kernel_records=(prefix_count + tail_count) * kernels_per_graph, + flag=flag, + enqueue_ms=1000.0 * enqueue_s, + event_sync_ms=1000.0 * event_sync_s, + flush1_ms=1000.0 * flush1_s, + records_after_flush1=snapshot_after_flush1["records"], + callbacks_after_flush1=snapshot_after_flush1["callbacks"], + tail_sync_ms=1000.0 * tail_sync_s, + flush2_ms=1000.0 * flush2_s, + **probe.snapshot(), + ) + + +def periodic_flush( + probe: CuptiActivityProbe, + graph: torch.cuda.CUDAGraph, + count: int, + period_ms: int, + flag: int, + kernels_per_graph: int, +) -> None: + probe.clear() + probe.set_flush_period(period_ms) + try: + enqueue_s = replay(graph, count) + sync_start_s = time.perf_counter() + torch.cuda.synchronize() + sync_s = time.perf_counter() - sync_start_s + snapshot_before_explicit = probe.snapshot() + flush_s = probe.flush(flag) + wait_for_parser_if_present(probe) + emit( + "periodic_flush", + count=count, + graph_replays=count, + kernels_per_graph=kernels_per_graph, + expected_kernel_records=count * kernels_per_graph, + period_ms=period_ms, + flag=flag, + enqueue_ms=1000.0 * enqueue_s, + sync_ms=1000.0 * sync_s, + records_before_explicit_flush=snapshot_before_explicit["records"], + callbacks_before_explicit_flush=snapshot_before_explicit["callbacks"], + explicit_flush_ms=1000.0 * flush_s, + **probe.snapshot(), + ) + finally: + probe.set_flush_period(0) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--case", + choices=("all", "scaling", "overlap", "gil", "prefix", "periodic", "graph-setup"), + default="all", + ) + parser.add_argument("--collector", choices=("python", "raw"), default="python") + parser.add_argument( + "--parse-process", + action="store_true", + help="With --collector raw, copy completed buffers to a parser process.", + ) + parser.add_argument( + "--parse-transport", + choices=("bytes", "shm"), + default="bytes", + help="How the raw collector sends completed buffers to the parser process.", + ) + parser.add_argument( + "--parser-zero-shm", + action="store_true", + help="With --parse-transport shm, zero the shared buffer in the parser process before acking reuse.", + ) + parser.add_argument("--callback-mode", choices=("count", "numeric", "names", "records"), default="count") + parser.add_argument("--host-buffer-bytes", type=int, default=8 * 1024 * 1024) + parser.add_argument("--host-buffer-count", type=int, default=8) + parser.add_argument("--max-records-per-buffer", type=int, default=0) + parser.add_argument( + "--device-buffer-size", + type=int, + default=None, + help="CUPTI_ACTIVITY_ATTR_DEVICE_BUFFER_SIZE, set before CUDA init.", + ) + parser.add_argument( + "--device-buffer-pool-limit", + type=int, + default=None, + help="CUPTI_ACTIVITY_ATTR_DEVICE_BUFFER_POOL_LIMIT, set before CUDA init.", + ) + parser.add_argument( + "--device-buffer-prealloc", + type=int, + default=None, + help="CUPTI_ACTIVITY_ATTR_DEVICE_BUFFER_PRE_ALLOCATE_VALUE, set before CUDA init.", + ) + parser.add_argument( + "--device-buffer-host-pinned", + action=argparse.BooleanOptionalAction, + default=None, + help="CUPTI_ACTIVITY_ATTR_MEM_ALLOCATION_TYPE_HOST_PINNED, set before CUDA init.", + ) + parser.add_argument( + "--zeroed-out-host-buffer", + action=argparse.BooleanOptionalAction, + default=None, + help="CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER. Experimental with cupti-python buffers.", + ) + parser.add_argument("--elements", type=int, default=1024 * 1024) + parser.add_argument("--kernels-per-graph", type=int, default=1) + parser.add_argument("--kernels-per-graph-list", default="1,2,4,8,16") + parser.add_argument("--graph-setup-replays", type=int, default=300) + parser.add_argument("--graph-setup-trials", type=int, default=5) + parser.add_argument("--counts", default="10,50,100,300,1000") + parser.add_argument("--flags", default="0,1") + parser.add_argument("--count", type=int, default=300) + parser.add_argument("--switch-interval-ms", type=float, default=None) + parser.add_argument("--yield-before-overlap", action="store_true") + parser.add_argument("--prefix-count", type=int, default=100) + parser.add_argument("--tail-count", type=int, default=4000) + parser.add_argument("--period-ms", type=int, default=1) + args = parser.parse_args() + + if args.switch_interval_ms is not None: + sys.setswitchinterval(args.switch_interval_ms / 1000.0) + + set_cupti_size_t_attribute(_CUPTI_ATTR_DEVICE_BUFFER_SIZE, args.device_buffer_size) + set_cupti_size_t_attribute(_CUPTI_ATTR_DEVICE_BUFFER_POOL_LIMIT, args.device_buffer_pool_limit) + set_cupti_size_t_attribute(_CUPTI_ATTR_DEVICE_BUFFER_PRE_ALLOCATE_VALUE, args.device_buffer_prealloc) + set_cupti_uint8_attribute(_CUPTI_ATTR_MEM_ALLOCATION_TYPE_HOST_PINNED, args.device_buffer_host_pinned) + set_cupti_uint8_attribute(_CUPTI_ATTR_ZEROED_OUT_ACTIVITY_BUFFER, args.zeroed_out_host_buffer) + + torch.cuda.init() + if args.case == "graph-setup": + graph_setup_case( + args.elements, + parse_ints(args.kernels_per_graph_list), + args.graph_setup_replays, + args.graph_setup_trials, + ) + return + + graph = make_one_kernel_graph(args.elements, args.kernels_per_graph) + if args.collector == "raw": + probe = RawCuptiActivityProbe( + args.host_buffer_bytes, + args.max_records_per_buffer, + args.parse_process, + args.parse_transport, + args.host_buffer_count, + args.parser_zero_shm, + ) + else: + probe = CuptiActivityProbe(args.callback_mode, args.host_buffer_bytes, args.max_records_per_buffer) + counts = parse_ints(args.counts) + flags = parse_ints(args.flags) + + try: + if args.case in ("all", "scaling"): + completed_scaling(probe, graph, counts, flags, args.kernels_per_graph) + if args.case in ("all", "overlap"): + for flag in flags: + overlap_enqueue( + probe, + graph, + args.count, + flag, + args.yield_before_overlap, + args.kernels_per_graph, + ) + if args.case in ("all", "gil"): + for flag in flags: + gil_gap(probe, graph, args.count, flag) + if args.case in ("all", "prefix"): + for flag in flags: + prefix_event_flush( + probe, + graph, + args.prefix_count, + args.tail_count, + flag, + args.kernels_per_graph, + ) + if args.case in ("all", "periodic"): + for flag in flags: + periodic_flush(probe, graph, args.count, args.period_ms, flag, args.kernels_per_graph) + finally: + if hasattr(probe, "close"): + probe.close() + + +if __name__ == "__main__": + main() From f5ccc2845f754d58a30cfd4a9fb7164cdc2ac6f9 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 12 May 2026 15:30:22 -0700 Subject: [PATCH 43/89] bench: JSONL-canonical artifact + host-blind resume + cell-list mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JSONL is the only data artifact now — every cell host-stamped and wall-clock timestamped ("t": time.time() in seconds) per row, so post-hoc efficiency analysis can derive per-cell wall time and identify startup-bound vs steady-state segments without separate profiling. Drop the clean-exit '.json' write. Resume on startup is host-blind: a bench restarted on a different node loads every prior key into _done_keys and only times the cells that weren't covered elsewhere. Sidecar '.meta.json' appended across runs. mkdir -p the output directory if the caller (e.g. a search driver) didn't create it. --cell-list replaces --retry-cells. Each JSON entry is a dict of canonical knob keys (Mw, Mnw, Ww, Wnw, Sw, Snw, pW, pS, H, R, CT, CPSw, CPSnw, LSw, LSnw, FL, WS, TMARL, TMAWL, TMANL, TMAWS, SR, RECT, WC, MODE, SORT, REVN, HSORT — tied forms M/W/S/CPS/LS auto-expanded). Bench auto-covers the per-knob value range from the list, builds the canonical knob-tuple-set for inner-loop filtering, and times only listed cells. Dict-based matching survives bench gaining new knobs (old cell-list files keep working — knobs not in the dict retain CLI defaults). Cell-list is loaded in main() BEFORE the args.*_list derivations so modes_list / sr_modes_list / rectangle_for_nowrite_list etc. respect the override. Use case: search_driver dumps cfg dicts directly; no tag-string drift, no cartesian-cover derivation in the orchestrator, no silent skip when a new knob is added. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 434 ++++++++++++++---- 1 file changed, 334 insertions(+), 100 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index b984886150a2..05e0138eb2ce 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -2119,10 +2119,21 @@ def _ps(val): errors.append((futures[fut], e)) if errors: - for cfg, e in errors: - print(f"[compile-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", + # Don't abort the whole bench just because some (M, W, S, TMA, ...) + # combinations refuse to compile. Driver-style invocations + # (search loops, sparse cell-lists) often hit a few bad combos in + # otherwise-valid sweeps; aborting wastes the rest of the bench's + # work. Log them once each and continue — the inner-loop runs + # below will trip their compile again per-cell, fail safely + # (per-cell is caught), and the cells just don't get timed. + print(f"[compile-warmup] {len(errors)} configs failed to compile " + f"(continuing; cells will be skipped):", file=sys.stderr) + for cfg, e in errors[:5]: + first_line = str(e).strip().split("\n")[0][:120] + print(f" FAILED {cfg}: {type(e).__name__}: {first_line}", file=sys.stderr) - raise errors[0][1] + if len(errors) > 5: + print(f" ... and {len(errors) - 5} more", file=sys.stderr) print(f"[compile-warmup] done in {time.perf_counter() - t0:.1f}s") @@ -2994,11 +3005,14 @@ def _emit_split(name_w, name_nw, val_w, val_nw): sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") reset_fn = _reset_conv1d_realistic if with_conv1d else _reset - # --retry-cells filter: only time cells whose tag is in the - # retry set. Cheaper than re-enumerating cells in the orchestrator. - retry_set = getattr(args, "_retry_cells_set", None) - if retry_set and sweep_tag not in retry_set: - continue + # --cell-list filter: only time cells whose (canonical-knob-values) + # tuple is in the loaded set. Robust to bench gaining new knobs + # (old cell-list files keep working: any keys they don't list + # become wildcards that retain CLI defaults). + if args._cell_list_keys: + _tup = _current_cell_tuple(args, locals()) + if _tup is None or _tup not in args._cell_list_set: + continue # Resume from JSONL: skip cells already recorded. Built the same # way _print_row builds JSON keys; must stay in sync. done_keys = getattr(args, "_done_keys", None) @@ -3168,24 +3182,26 @@ def _print_row( act_dtype_name, stats, sweep_suffix="", - json_results=None, tp_size=None, json_detailed=False, jsonl_path=None, jsonl_host=None, ): - """Print one summary row and optionally accumulate stats for JSON output. + """Print one summary row and append the result to the JSONL sidecar. `stats` is a dict from _time_kernel: {median, p95, p99, n, iters_us, [n_writes_per_iter], [per_kernel]}. The summary table only shows the - headline percentiles. JSON output captures compact per-iter spans by - default; with json_detailed=True it also captures per-kernel data. + headline percentiles. JSONL captures the compact per-iter spans + + n_writes_per_iter by default; with json_detailed=True it also captures + per-kernel data. - When `jsonl_path` is provided, also appends one JSON line per row to the + When `jsonl_path` is provided, appends one JSON line per row to the JSONL sidecar (crash-safe incremental persistence; lets a killed sweep - resume from the last completed cell on rerun). Open-per-write because - `args` is pickled to ProcessPoolExecutor workers and file handles aren't - picklable. + resume from the last completed cell on rerun, even across hosts). Open + per-write because `args` is pickled to ProcessPoolExecutor workers and + file handles aren't picklable. JSONL is the canonical artifact — the + bench no longer writes a final `.json` summary; use `jsonl_to_json.py` + if a one-shot `.json` snapshot is needed. """ kernel_col = f"{kernel_name:>11} | " if show_kernel_col else "" print( @@ -3194,7 +3210,7 @@ def _print_row( f"{stats['median']:>9.2f} | {stats['p95']:>7.2f} | {stats['p99']:>7.2f} |" f"{sweep_suffix}" ) - if json_results is not None: + if jsonl_path is not None: key = _build_json_key( kernel_name, batch, mtp_len, prev_k, state_dtype_name, sweep_suffix, tp_size, @@ -3207,9 +3223,8 @@ def _print_row( for k in ("median", "p95", "p99", "n", "iters_us", "n_writes_per_iter") if k in stats } - if "host_timing" in stats: - row_stats["host_timing"] = stats["host_timing"] - json_results[key] = row_stats + if "host_timing" in stats: + row_stats["host_timing"] = stats["host_timing"] # Append to JSONL sidecar if a path is set (incremental persistence). # Open per-write because args is pickled to ProcessPoolExecutor # workers, and file handles aren't picklable. A clean SIGTERM or @@ -3217,7 +3232,13 @@ def _print_row( # newline; catastrophic kills can leave a partial last line, which # the resume reader tolerates via json.JSONDecodeError pass. if jsonl_path is not None: - rec = {"key": key, "stats": row_stats} + # Wall-clock timestamp (float seconds since UNIX epoch) at write + # time. Lets post-hoc analysis diff consecutive rows to derive + # per-cell wall budget and identify startup-bound vs steady-state + # segments (cells/sec, downtime between bench invocations) without + # needing to instrument the bench's outer loops separately. + import time as _time + rec = {"key": key, "stats": row_stats, "t": _time.time()} if jsonl_host is not None: rec["host"] = jsonl_host with open(jsonl_path, "a") as f: @@ -3251,7 +3272,6 @@ def _finish_result_job(args, job: dict) -> None: job["act_dtype_name"], stats, job.get("sweep_suffix", ""), - json_results=getattr(args, "_json_results", None), tp_size=args.tp_size, json_detailed=getattr(args, "json_detailed", False), jsonl_path=getattr(args, "_jsonl_path", None), @@ -3313,21 +3333,201 @@ def _submit_result_job( _finish_result_job(args, job) +# Cell-list mode — canonical knob-key mapping to argparse args + local +# loop variable. See _load_cell_list_into_args / inner-loop filter. +# +# Each entry: cell-key → (args attribute name, comma-separated string flag) +# For split (write/nowrite) knobs, we use Xw / Xnw keys. Tied forms (M, W, +# S, CPS, LS) accepted on load and expanded to their w/nw variants. +_CELL_LIST_KEY_TO_ARG = { + "Mw": "block_size_m_write", + "Mnw": "block_size_m_nowrite", + "Ww": "num_warps_write", + "Wnw": "num_warps_nowrite", + "Sw": "num_stages_write", + "Snw": "num_stages_nowrite", + "CPSw": "cta_per_sm_write", + "CPSnw": "cta_per_sm_nowrite", + "LSw": "num_loop_stages_write", + "LSnw": "num_loop_stages_nowrite", + "pW": "precompute_num_warps", + "pS": "precompute_num_stages", + "H": "heads_per_block", + "R": "maxnreg", + "CT": "num_ctas", + "FL": "flatten", + "WS": "warp_specialize", + "TMARL": "use_tma_rect_load", + "TMAWL": "use_tma_replay_write_load", + "TMANL": "use_tma_replay_nowrite_load", + "TMAWS": "use_tma_replay_write_store", + "RECT": "rectangle_for_nowrite", + "WC": "write_modes", + "SORT": "sort_slots", + "REVN": "reverse_nowrite", + "HSORT": "hardcode_sort", + # MODE and SR get special handling (string values): + # MODE → args.modes (single mode name) + # SR → args.sr_modes ("RN" if 0, "SR" if 1) +} + +# Split-knob tied form: "M" expands to both "Mw" and "Mnw". +_CELL_LIST_TIED_EXPANSIONS = { + "M": ("Mw", "Mnw"), + "W": ("Ww", "Wnw"), + "S": ("Sw", "Snw"), + "CPS": ("CPSw", "CPSnw"), + "LS": ("LSw", "LSnw"), +} + + +def _normalize_cell(cell: dict) -> dict: + """Expand tied-form keys (M, W, S, CPS, LS) to their w/nw variants. + Returns a new dict with only canonical split-or-plain keys. + """ + out = dict(cell) + for tied, (w_key, nw_key) in _CELL_LIST_TIED_EXPANSIONS.items(): + if tied in out: + v = out.pop(tied) + out.setdefault(w_key, v) + out.setdefault(nw_key, v) + return out + + +def _load_cell_list_into_args(args) -> None: + """Read --cell-list JSON, normalize, override args.* knob ranges, and + populate args._cell_list_keys + args._cell_list_set for the inner-loop + filter. Errors out if cells aren't uniform (different key sets). + """ + with open(args.cell_list) as f: + raw = json.load(f) + if not isinstance(raw, list): + sys.exit(f"--cell-list: expected JSON list, got {type(raw).__name__}") + cells = [_normalize_cell(c) for c in raw] + if not cells: + print("[cell-list] empty list — nothing to time", file=sys.stderr) + return + # All cells must share the same key set (uniform schema) + keys0 = frozenset(cells[0].keys()) + for i, c in enumerate(cells[1:], start=1): + if frozenset(c.keys()) != keys0: + sys.exit( + f"--cell-list: cells must have uniform key sets; cell[0] " + f"has {sorted(keys0)} but cell[{i}] has {sorted(c.keys())}" + ) + + # Auto-cover: collect per-knob value set across all cells + cover: dict = {} + for c in cells: + for k, v in c.items(): + cover.setdefault(k, set()).add(v) + # Apply overrides + for key, vals in cover.items(): + if key in _CELL_LIST_KEY_TO_ARG: + arg_name = _CELL_LIST_KEY_TO_ARG[key] + vals_str = ",".join(str(v) for v in sorted(vals)) + setattr(args, arg_name, vals_str) + elif key == "MODE": + args.modes = ",".join(sorted({str(v) for v in vals})) + elif key == "SR": + args.sr_modes = ",".join(sorted({"SR" if v else "RN" for v in vals})) + else: + print(f"[cell-list] WARNING: unknown key {key!r} in cells; " + f"will not override any args.* attribute (the value will " + f"still be matched in the filter if a matching local var " + f"is in scope)", file=sys.stderr) + + # Canonical key order (sorted) for tuple matching in the inner loop + args._cell_list_keys = tuple(sorted(keys0)) + args._cell_list_set = { + tuple(c[k] for k in args._cell_list_keys) for c in cells + } + print(f"[cell-list] loaded {len(cells)} cells with keys " + f"{list(args._cell_list_keys)}; overrode args.* to auto-cover", + file=sys.stderr) + + +# Maps cell-list key → name of the local variable in _bench_config's inner +# loop. Used to extract the "current cell" tuple for the filter check. +# Keep in sync with the loop-variable names; the filter is lenient about +# missing names (it picks them up from the inner scope at runtime). +_CELL_LIST_KEY_TO_LOCAL = { + "Mw": "block_size_m_w", + "Mnw": "block_size_m_nw", + "Ww": "num_warps_w", + "Wnw": "num_warps_nw", + "Sw": "num_stages_w", + "Snw": "num_stages_nw", + "CPSw": "cta_per_sm_w", + "CPSnw": "cta_per_sm_nw", + "LSw": "num_loop_stages_w", + "LSnw": "num_loop_stages_nw", + "pW": "precompute_num_warps", + "pS": "precompute_num_stages", + "H": "heads_per_block", + "R": "maxnreg", + "CT": "num_ctas", + "FL": "flatten", + "WS": "warp_specialize", + "TMARL": "use_tma_rect_load", + "TMAWL": "use_tma_replay_write_load", + "TMANL": "use_tma_replay_nowrite_load", + "TMAWS": "use_tma_replay_write_store", + "RECT": "rectangle_for_nowrite", + "WC": "write_checkpoint", + "MODE": "mode", + "SORT": "sort_slots", + "REVN": "reverse_nowrite", + "HSORT": "hardcode_sort", + "SR": "use_philox", +} + + +def _current_cell_tuple(args, locals_dict: dict) -> tuple | None: + """Build the (key1=val1, key2=val2, ...) tuple for the current inner-loop + iteration, matching args._cell_list_keys' order. Used by the inner-loop + filter to check membership in args._cell_list_set. Returns None if any + expected local is missing (the bench evolved a knob name — caller skips). + """ + if not args._cell_list_keys: + return None + vals = [] + for k in args._cell_list_keys: + local_name = _CELL_LIST_KEY_TO_LOCAL.get(k, k) + if local_name not in locals_dict: + return None + v = locals_dict[local_name] + # Coerce bools to ints to match cell-list JSON (1/0) + if isinstance(v, bool): + v = int(v) + vals.append(v) + return tuple(vals) + + # Main benchmark loop def _run_benchmark(args) -> None: - # JSON accumulator — populated by _print_row when --json-output is set. - # Stash on args so we don't need to thread a dict through every helper. - args._json_results = {} if getattr(args, "json_output", None) else None + # Pending-results FIFO for srxl's deferred CUPTI parsing pipeline. Each + # entry holds a _PendingCuptiStats handle; _drain_pending_results pulls + # ready entries and routes them to _print_row (which appends to JSONL). args._pending_results = [] # JSONL incremental sidecar. Path = `.jsonl`. Each completed - # cell appends one line `{"key": , "stats": {...}}` to this file - # as it finishes timing. On startup we read this sidecar (if present) and - # populate _json_results + _done_keys so a killed bench can resume without + # cell appends one line `{"key": , "stats": {...}, "host": }` + # to this file as it finishes timing. On startup we read this sidecar (if + # present) and populate _done_keys so a killed bench can resume without # redoing already-timed cells. Crash-safe by construction: append-only # writes survive SIGTERM/SIGKILL/reboot mid-sweep. + # + # Resume is host-blind: _done_keys includes records from any host, so a + # bench restarted on a different node fills in the missing cells without + # redoing cells already covered elsewhere. Cross-host *timings* aren't + # directly comparable, but each JSONL record carries its `host` stamp so + # the analyzer can group/compare per host. This bench no longer writes a + # final `.json` summary — the JSONL is the canonical artifact; use the + # `jsonl_to_json.py` helper if a one-shot `.json` snapshot is needed. + # # Note: we store only paths/strings on `args` because args is pickled to # ProcessPoolExecutor workers during compile-warmup, and file handles # (TextIOWrapper) aren't picklable. _print_row open-appends per cell. @@ -3338,15 +3538,11 @@ def _run_benchmark(args) -> None: import socket args._jsonl_host = socket.gethostname() args._jsonl_path = args.json_output + ".jsonl" - # Read existing JSONL if present; build skip set IFF hostname matches. - # Cross-node timings aren't directly comparable (CPU/GPU clock/topology - # differ), so resuming on a new host would mix incompatible numbers. - # If the JSONL was recorded on a different host, log + skip resume so - # the user can decide (rename / archive the old file). + # Read existing JSONL if present: load every record's key into the + # skip set regardless of host (gap-fill on a new node). if os.path.exists(args._jsonl_path): n_loaded = 0 - n_skipped_host = 0 - recorded_hosts = set() + host_counts: dict[str, int] = {} with open(args._jsonl_path) as f: for line in f: line = line.strip() @@ -3357,58 +3553,95 @@ def _run_benchmark(args) -> None: except json.JSONDecodeError: # Tolerate partial last line from a crash mid-write. continue - rec_host = rec.get("host") - if rec_host is not None: - recorded_hosts.add(rec_host) - if rec_host is not None and rec_host != args._jsonl_host: - n_skipped_host += 1 - continue k = rec.get("key") if k is None: continue - args._json_results[k] = rec.get("stats", {}) args._done_keys.add(k) n_loaded += 1 + rec_host = rec.get("host") + if rec_host: + host_counts[rec_host] = host_counts.get(rec_host, 0) + 1 if n_loaded: + host_summary = ", ".join( + f"{h}={n}" for h, n in sorted(host_counts.items()) + ) if host_counts else "(no host stamps)" print( f"[resume] {args._jsonl_path}: loaded {n_loaded} prior " - f"cell results from host={args._jsonl_host}; sweep will " - f"skip them.", - file=sys.stderr, - ) - if n_skipped_host: - print( - f"[resume] WARNING: {args._jsonl_path} also contains " - f"{n_skipped_host} records from other hosts " - f"{sorted(recorded_hosts - {args._jsonl_host})!r}; ignoring " - f"them. If you want a clean run, remove or archive the " - f"jsonl file first.", + f"cell results across hosts [{host_summary}]; sweep will " + f"skip them. New cells stamp host={args._jsonl_host}.", file=sys.stderr, ) + # Sidecar metadata: cmd, host, tp_size, variant, cupti, etc. Written + # once at startup; helps later analysis identify how this JSONL was + # produced even though there's no top-level .json wrapper anymore. + meta_path = args.json_output + ".meta.json" + meta_payload = { + "timestamp": datetime.now().isoformat(), + "host": args._jsonl_host, + "cmd": " ".join(sys.argv), + "tp_size": getattr(args, "tp_size", None), + "warmup": getattr(args, "warmup", None), + "iters": getattr(args, "iters", None), + "variant": getattr(args, "variant", None), + "cupti": getattr(args, "cupti", False), + } + # Append to a list so successive runs (gap-fill, retry) keep history. + existing_meta = [] + if os.path.exists(meta_path): + try: + with open(meta_path) as f: + existing_meta = json.load(f) + if not isinstance(existing_meta, list): + existing_meta = [existing_meta] + except (OSError, json.JSONDecodeError): + existing_meta = [] + existing_meta.append(meta_payload) + # Bench is sometimes invoked with --json-output pointing into a dir + # the caller hasn't created (subprocess driver, search loop, etc.). + # Ensure the dir exists before writing the meta sidecar OR the JSONL. + os.makedirs(os.path.dirname(os.path.abspath(meta_path)), exist_ok=True) + tmp = meta_path + ".tmp" + with open(tmp, "w") as f: + json.dump(existing_meta, f, indent=2) + os.replace(tmp, meta_path) + # Skipped cells accumulator — populated by _bench_config when CUPTI capture # mismatch causes a cell to be skipped. Written to args.skipped_output - # (or derived from json_output) at end of run. Pair with --retry-cells - # to re-time only the skipped cells in a fresh process. + # (or derived from json_output) at end of run. args._skipped_cells = [] - # Retry-cells filter set. When non-empty, only cells whose tag is in this - # set will be timed; all others are silently skipped. Cells are matched - # against the sweep_tag string built in _bench_config. - args._retry_cells_set: set[str] = set() - if getattr(args, "retry_cells", None): - with open(args.retry_cells) as f: - data = json.load(f) - # Accept either a list of tag strings, or the same shape we write - # (dict with "skipped" key). Tolerant of both for hand-edited files. - if isinstance(data, dict) and "skipped" in data: - data = data["skipped"] - args._retry_cells_set = set(data) - print( - f"[retry] --retry-cells loaded {len(args._retry_cells_set)} tags " - f"from {args.retry_cells}; sweep will skip all other cells.", - file=sys.stderr, - ) + # Cell-list filter (replaces the old --retry-cells tag-string filter). + # When set, the sweep iterates ONLY the cells described in the list. + # + # Each entry in the JSON file is a dict of canonical knob keys → values, + # using the same names that appear in the sweep_tag (Mw/Mnw, Ww/Wnw, + # Sw/Snw, pW, pS, H, R, CT, CPSw/CPSnw, LSw/LSnw, FL, WS, TMARL, + # TMAWL, TMANL, TMAWS, SR, RECT, WC, MODE, SORT, REVN, HSORT). Each + # cell may also use the tied forms M / W / S / CPS / LS (single value + # applied to both write and nowrite halves). + # + # On load we: + # - Override the bench's CLI knob args (`args.block_size_m_write`, + # etc.) with the union of values present across all cells per knob, + # so the cartesian iteration auto-covers the list. + # - Build `args._cell_list_keys` (the canonical key order used by + # every cell — must be uniform across the list) and + # `args._cell_list_set` (frozen tuples for O(1) membership check + # inside the inner loop). + # + # In the inner loop, we build the current iteration's tuple and skip + # cells not in the set. Dict-matching is robust to bench gaining new + # knobs (old cell-list files keep working — newly-added knobs simply + # aren't matched on, so they retain CLI defaults). + # Cell-list state may already have been populated by main() (so that + # the args.*_list derivations downstream see the override). Default to + # empty if not. + if not hasattr(args, "_cell_list_keys"): + args._cell_list_keys: tuple = () + args._cell_list_set: set = set() + if getattr(args, "cell_list", None): + _load_cell_list_into_args(args) assert args.nheads % args.tp_size == 0, ( f"nheads ({args.nheads}) must be divisible by tp_size ({args.tp_size})" @@ -3650,28 +3883,16 @@ def _run_benchmark(args) -> None: if args.profile: torch.cuda.cudart().cudaProfilerStop() - if args.json_output and args._json_results is not None: - payload = { - "metadata": { - "timestamp": datetime.now().isoformat(), - "cmd": " ".join(sys.argv), - "tp_size": args.tp_size, - "warmup": args.warmup, - "iters": args.iters, - "variant": args.variant, - "cupti": getattr(args, "cupti", False), - }, - "results": args._json_results, - } - tmp = args.json_output + ".tmp" - with open(tmp, "w") as f: - json.dump(payload, f, indent=2) - os.replace(tmp, args.json_output) - print(f"\nJSON results written to: {args.json_output} " - f"({len(args._json_results)} entries)") + # JSONL is the canonical artifact (written incrementally per cell with + # host stamps). No clean-exit `.json` write — use `jsonl_to_json.py` to + # materialize a snapshot when an analyzer wants one. + if args.json_output and args._jsonl_path is not None: + print(f"\nJSONL results: {args._jsonl_path} " + f"(meta sidecar: {args.json_output}.meta.json)") - # Write the skipped-cells sidecar. Pair with --retry-cells in a separate - # invocation to re-time the failed cells in a fresh process. + # Write the skipped-cells sidecar. Caller can convert this list to a + # --cell-list JSON (one dict per skipped cell) to drive a retry pass in + # a fresh process. skipped_path = getattr(args, "skipped_output", None) if skipped_path is None and args.json_output: # Derive default: foo.json -> foo.skipped.json @@ -3858,16 +4079,20 @@ def _parse_args() -> argparse.Namespace: help="Path to write the list of cells that failed CUPTI capture even " "after --cupti-retry retries (JSON list of sweep_tag strings). " "Default: derived from --json-output by replacing .json with " - ".skipped.json. Pair with --retry-cells in a separate invocation " - "(fresh process = fresh CUPTI subscriber) to re-time these cells.", + ".skipped.json.", ) parser.add_argument( - "--retry-cells", + "--cell-list", default=None, - help="Path to a JSON list of cell tags (as written by --skipped-output " - "in a prior invocation). When set, the sweep iterates as normal but " - "skips any cell whose tag is NOT in the listed set. Lets collect.py " - "drive a retry pass over only the cells that failed the first time.", + help="Path to a JSON list of cell dicts (one per cell to time). " + "Each dict has canonical knob keys → values: Mw, Mnw, Ww, Wnw, Sw, " + "Snw, pW, pS, H, R, CT, CPSw, CPSnw, LSw, LSnw, FL, WS, TMARL, " + "TMAWL, TMANL, TMAWS, SR, RECT, WC, MODE, SORT, REVN, HSORT (tied " + "forms M / W / S / CPS / LS are also accepted and auto-expanded). " + "When set, bench's CLI knob ranges are auto-overridden to the " + "per-knob union across all cells, and the inner-loop filter skips " + "any iteration whose knob-value tuple isn't in the list. All cells " + "must share the same key set (uniform schema).", ) parser.add_argument( "--prev-tokens-fracs", @@ -4264,6 +4489,15 @@ def _parse_args() -> argparse.Namespace: if args.mix_only and args.mix_csv is None: parser.error("--mix-only requires --mix-csv") + # Cell-list (if any) must be applied BEFORE the post-argparse string→list + # derivations below — those build args.*_list from args.* strings, so a + # cell-list override of e.g. args.modes='maindl' needs to land before + # args.modes_list is computed. The function populates args._cell_list_keys + # and args._cell_list_set, plus overrides args.* knob strings to the + # per-knob union of values across the listed cells. + if getattr(args, "cell_list", None): + _load_cell_list_into_args(args) + # Backward-compat: --philox-rounding implies --sr-modes SR if --sr-modes # was left at the default. If both are set explicitly, error. sr_modes_default = (args.sr_modes == "RN") From 0abdb909572aa233dae64c6e277ffad50fe6663e Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 13 May 2026 13:10:12 -0700 Subject: [PATCH 44/89] bench: compile-warmup CPS-grouping + cell-list direct enum + phase markers - _compile_warmup_phase: switch from inner-cartesian (read args.block_size_m etc which is None when --cell-list is used) to two paths: * cell-list mode: enumerate args._cell_list_set directly * sweep-args mode: split-aware (read both shared and write/nowrite args) - CPS-grouping: cells differing only in CPSw/CPSnw share a worker so in-process Triton cache hits across the div_by_16 spec-bucket boundary - _warm_one_config: accept list-of-dicts for inner_overrides (CPS-grouped task) - Add timestamped phase markers ([phase], [compile-warmup]) to stderr with flush - Add flush=True to existing [WARN]/[retry]/rec[N] prints for live Popen view Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 333 +++++++++++++----- 1 file changed, 242 insertions(+), 91 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 05e0138eb2ce..a16097bf45d6 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -665,7 +665,7 @@ def _kernels_per_iter_baseline(with_conv1d: bool) -> int: _CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER = 5 _CUPTI_HOST_BUFFER_BYTES = 1024 * 1024 _CUPTI_HOST_BUFFER_COUNT = 16 -_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE = 2 +_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE = 1 _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX = 4 @@ -1327,6 +1327,7 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, f"(warmup+iters) = {expected_total} records, got {total}. " f"Kernel record counts: {name_counts}.{zero_msg} SKIPPING cell.", file=sys.stderr, + flush=True, ) # Per-record dump: (name, start_ns_rel, end_ns_rel, corr_id, graph_id, stream_id). # Times relative to first record so absolute ns isn't drowning output. @@ -1341,9 +1342,11 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, f" rec[{i:3d}] name={r[0]!r} start={rel_start:.2f}us " f"end={rel_end:.2f}us corr={r[3]} graph={r[4]} stream={r[6]}", file=sys.stderr, + flush=True, ) if len(records) > 30: - print(f" ... ({len(records) - 30} more records elided)", file=sys.stderr) + print(f" ... ({len(records) - 30} more records elided)", + file=sys.stderr, flush=True) return None K = expected_K timed = records[warmup * K:] @@ -1494,21 +1497,23 @@ def _capture_group_graph( def _graph_group_iters(args, total_iters: int, pre_iter_fn, pre_iter_group_factory) -> int: + """Pick the graph-group size unconditionally; the caller is expected to + round total_iters up to a multiple of this so all iters fit in clean + replays. Sample arrays are pre-padded at allocation (see _sample_pnat + call site) so the per-replay window can index past the user-requested + iter count by up to group_iters-1 extra samples. + """ if pre_iter_fn is not None and pre_iter_group_factory is None: + # Per-iter callback without a group-factory: can't batch. return 1 requested = getattr(args, "cuda_graph_group_iters", None) if requested is None: - requested_group_iters = ( + return ( _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX if pre_iter_group_factory is not None else _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE ) - else: - requested_group_iters = max(1, int(requested)) - for group_iters in (requested_group_iters, 2): - if total_iters % group_iters == 0: - return group_iters - return 1 + return max(1, int(requested)) def _get_cupti_filter_plan(timer: CuptiKernelTimer, graph, cache_key: tuple | None, @@ -1594,6 +1599,13 @@ def _time_kernel_cuda_graph( total_iters = warmup + iters group_iters = _graph_group_iters(args, total_iters, pre_iter_fn, pre_iter_group_factory) + # Args are rounded at argparse-time so warmup+iters/mix_iters are already + # multiples of the relevant group_iters. Assert here to catch any caller + # bypassing argparse. + assert total_iters % group_iters == 0, ( + f"total_iters={total_iters} not a multiple of group_iters={group_iters}; " + f"args.warmup/iters/mix_iters should be rounded post-argparse." + ) pre_replay_fn = None graph_pre_iter_fn = None if pre_iter_group_factory is not None and group_iters > 1: @@ -1900,45 +1912,53 @@ def _warm_one_config(args, cfg, baseline_fn) -> None: aren't picklable). Each worker process holds its own GIL → no serialization between concurrent compiles. - ``cfg`` is a tuple of (outer_cfg, inner_overrides): + ``cfg`` is a tuple of (outer_cfg, inner_overrides_or_list): * outer_cfg = (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, rect, write_ckpt, mode, sort_slots, reverse_nowrite, hardcode_sort) - * inner_overrides = dict of args attribute name -> single-value string - to clamp the per-cell inner-knob sweep to ONE - combination. Triggers exactly one Triton compile - per worker invocation, so N workers achieve - N-way concurrency regardless of outer config - count. (Prior design fanned out only at outer - granularity, capping concurrency at ~10 even - with --compile-threads 50.) + * inner_overrides_or_list = dict of args attribute name -> value-string, + OR a list of such dicts. In the list form (CPS-grouped task) the + worker compiles each entry sequentially within the same process so + Triton's in-process kernel cache catches value-spec hits across + related entries (e.g. CPS={1,2} and {4,8} each form a `div_by_16` + spec bucket; the second compile in a bucket short-circuits). ``baseline_fn`` is optional — when ``None``, only the checkpointing kernel is warmed (the baseline-selection kernel can be warmed once in the parent if needed). This lets us avoid pickling C-extension function references across processes. """ - outer_cfg, inner_overrides = cfg + outer_cfg, inner_overrides_or_list = cfg + overrides_list = (inner_overrides_or_list + if isinstance(inner_overrides_or_list, list) + else [inner_overrides_or_list]) (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, rect, write_ckpt, mode, sort_slots, reverse_nowrite, hardcode_sort) = outer_cfg - # Clone args and override inner-knob sweep lists to single values. - # _bench_config then iterates a 1×1×...×1 cartesian inside. import argparse as _ap - args_copy = _ap.Namespace(**vars(args)) - for k, v in inner_overrides.items(): - setattr(args_copy, k, v) - _bench_config( - args_copy, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, - sr_mode=sr_mode, rectangle_for_nowrite=rect, - write_checkpoint=write_ckpt, mode=mode, - sort_slots=sort_slots, reverse_nowrite=reverse_nowrite, - hardcode_sort=hardcode_sort, - warmup_only=True, - ) + for inner_overrides in overrides_list: + # Fresh clone per entry: prevents knob-value leakage between + # consecutive cells in a CPS-grouped task (entries may set + # different non-CPS knobs in degenerate edge cases). + args_copy = _ap.Namespace(**vars(args)) + for k, v in inner_overrides.items(): + setattr(args_copy, k, v) + _bench_config( + args_copy, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, + sr_mode=sr_mode, rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, mode=mode, + sort_slots=sort_slots, reverse_nowrite=reverse_nowrite, + hardcode_sort=hardcode_sort, + warmup_only=True, + ) def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, baseline_fn, max_workers: int) -> None: + _cw_t0 = time.perf_counter() + def _cw(label: str) -> None: + dt = time.perf_counter() - _cw_t0 + print(f"[compile-warmup] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) + _cw("entered _compile_warmup_phase") """Parallel compile-warmup using a ProcessPoolExecutor with `spawn` start method. @@ -2040,10 +2060,27 @@ def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtyp hardcode_sort, )) - # Enumerate inner-knob cartesian — same axes _bench_config iterates - # internally. Each (outer × inner) tuple becomes one task; workers - # then trigger exactly one Triton compile per task, giving true - # N-way concurrency with --compile-threads N. + # Enumerate inner-knob signatures. Two paths: + # (1) --cell-list mode (preferred when set): pull exactly the cells + # that will be timed from args._cell_list_set. No synthetic + # cartesian — we only pre-compile what will run. + # (2) Sweep-args mode: cartesian over split-aware axes that read + # BOTH unsplit (args.X) and per-half (args.X_write/_nowrite) + # knob settings. Older code read only args.X and silently + # enumerated 1 inner combo when callers set only the per-half + # versions (all cell-list usage, plus any --block-size-m-write/ + # _nowrite CLI invocation), causing massive in-process JIT + # compile tax for persistent_main especially. + # + # In BOTH paths we GROUP tasks by non-CPS signature so each worker + # process compiles all CPS values for its group sequentially. + # NUM_PERSISTENT = CPS * num_sms is a runtime int but Triton auto- + # specializes on `div_by_16`, partitioning {CPS=1,2} (132,264) from + # {CPS=4,8} (528,1056) into two distinct compiled variants. By + # keeping all CPS variants for one (M,W,S,LS,TMA,...) signature in + # the same worker, the second compile in each spec bucket hits the + # in-process Triton cache (no disk-cache round trip). + def _ps(val): if val is None or (isinstance(val, str) and not val): return [None] @@ -2051,60 +2088,127 @@ def _ps(val): return [v.strip() for v in val.split(",") if v.strip()] return [val] - # CPS (cta_per_sm) collapsed to first value: NUM_PERSISTENT is runtime now, - # so different CPS values share the same compiled kernel. Collapsing here - # avoids enumerating 4-8x redundant tasks that would each pay ~50-100ms - # bench setup overhead for a cache hit on the same kernel hash. - _cps_for_compile = _ps(args.cta_per_sm)[:1] - knob_axes = [ - ("block_size_m", _ps(args.block_size_m)), - ("num_warps", _ps(args.num_warps)), - ("num_stages", _ps(args.num_stages)), - ("precompute_num_warps", _ps(args.precompute_num_warps)), - ("precompute_num_stages", _ps(args.precompute_num_stages)), - ("heads_per_block", _ps(args.heads_per_block)), - ("maxnreg", _ps(args.maxnreg)), - ("num_ctas", _ps(args.num_ctas)), - ("cta_per_sm", _cps_for_compile), # collapsed (runtime) - ("num_loop_stages", _ps(args.num_loop_stages)), - ("flatten", _ps(args.flatten)), - ("warp_specialize", _ps(args.warp_specialize)), - ("use_tma_rect_load", _ps(args.use_tma_rect_load)), - ("use_tma_replay_write_load", _ps(args.use_tma_replay_write_load)), - ("use_tma_replay_nowrite_load", _ps(args.use_tma_replay_nowrite_load)), - ("use_tma_replay_write_store", _ps(args.use_tma_replay_write_store)), - ] - import itertools as _it - inner_combos = list(_it.product(*(values for _, values in knob_axes))) + def _split_or_pair(shared_attr, w_attr, nw_attr): + """Return list of (write_val, nowrite_val) strings. + + Reads shared (args.X), write-side (args.X_write), and nowrite-side + (args.X_nowrite) values. If both per-half attrs are None, emits + tied pairs (v,v) over the shared values. If either per-half is + set, cartesian-iterates per-half values, falling back to shared + for whichever side is None. + """ + w = _ps(getattr(args, w_attr, None)) + nw = _ps(getattr(args, nw_attr, None)) + s = _ps(getattr(args, shared_attr, None)) + if w == [None] and nw == [None]: + return [(v, v) for v in s] + if w == [None]: + w = s + if nw == [None]: + nw = s + return [(a, b) for a in w for b in nw] + + _cw(f"built {len(configs)} outer configs") + # Build per-cell inner-overrides dicts. + cell_set = getattr(args, "_cell_list_set", set()) + cell_keys = getattr(args, "_cell_list_keys", ()) + if cell_set: + inner_dicts = [] + for tup in cell_set: + d = dict(zip(cell_keys, tup)) + inner = {} + for k, v in d.items(): + if k in _CELL_LIST_KEY_TO_ARG: + inner[_CELL_LIST_KEY_TO_ARG[k]] = str(v) + inner_dicts.append(inner) + else: + m_pairs = _split_or_pair("block_size_m", "block_size_m_write", "block_size_m_nowrite") + w_pairs = _split_or_pair("num_warps", "num_warps_write", "num_warps_nowrite") + ns_pairs = _split_or_pair("num_stages", "num_stages_write", "num_stages_nowrite") + cps_pairs = _split_or_pair("cta_per_sm", "cta_per_sm_write", "cta_per_sm_nowrite") + ls_pairs = _split_or_pair("num_loop_stages", "num_loop_stages_write", "num_loop_stages_nowrite") + pw_vals = _ps(args.precompute_num_warps) + ps_vals = _ps(args.precompute_num_stages) + h_vals = _ps(args.heads_per_block) + mr_vals = _ps(args.maxnreg) + ct_vals = _ps(args.num_ctas) + fl_vals = _ps(args.flatten) + wsp_vals = _ps(args.warp_specialize) + trl_vals = _ps(args.use_tma_rect_load) + twl_vals = _ps(args.use_tma_replay_write_load) + tnl_vals = _ps(args.use_tma_replay_nowrite_load) + tws_vals = _ps(args.use_tma_replay_write_store) + import itertools as _it + inner_dicts = [] + for ((mw, mnw), (ww, wnw), (sw, snw), (cw, cnw), (lw, lnw), + pw, ps_, h, mr, ct, fl, wsp, + trl, twl, tnl, tws) in _it.product( + m_pairs, w_pairs, ns_pairs, cps_pairs, ls_pairs, + pw_vals, ps_vals, h_vals, mr_vals, ct_vals, + fl_vals, wsp_vals, + trl_vals, twl_vals, tnl_vals, tws_vals): + d = {} + for k, v in ( + ("block_size_m_write", mw), + ("block_size_m_nowrite", mnw), + ("num_warps_write", ww), + ("num_warps_nowrite", wnw), + ("num_stages_write", sw), + ("num_stages_nowrite", snw), + ("cta_per_sm_write", cw), + ("cta_per_sm_nowrite", cnw), + ("num_loop_stages_write", lw), + ("num_loop_stages_nowrite", lnw), + ("precompute_num_warps", pw), + ("precompute_num_stages", ps_), + ("heads_per_block", h), + ("maxnreg", mr), + ("num_ctas", ct), + ("flatten", fl), + ("warp_specialize", wsp), + ("use_tma_rect_load", trl), + ("use_tma_replay_write_load", twl), + ("use_tma_replay_nowrite_load", tnl), + ("use_tma_replay_write_store", tws), + ): + if v is not None: + d[k] = str(v) + inner_dicts.append(d) + + # Group by non-CPS signature. All inner_dicts that match on every + # key except CPSw/CPSnw/cta_per_sm land in the same worker task; the + # worker will compile them in sequence, sharing Triton's in-process + # kernel cache across the divides bucket boundary at most once per + # bucket (CPS in {1,2} vs {4,8}). + _cps_keys = ("cta_per_sm_write", "cta_per_sm_nowrite", "cta_per_sm") + groups: dict = {} + for d in inner_dicts: + sig = tuple(sorted((k, v) for k, v in d.items() if k not in _cps_keys)) + groups.setdefault(sig, []).append(d) - # Cross product outer × inner. Override only knobs that have an - # explicit value (skip None — those leave args. at its CLI default, - # which _bench_config handles via its own _parse_sweep). tasks = [] for outer in configs: - for inner_tuple in inner_combos: - inner_overrides = { - name: str(val) - for (name, _), val in zip(knob_axes, inner_tuple) - if val is not None - } - tasks.append((outer, inner_overrides)) + for sig, group in groups.items(): + tasks.append((outer, group)) - # Shuffle to reduce cross-worker race on the same kernel hash. Two - # workers picking adjacent tasks (same mode, neighboring knob value) - # could both miss + compile the same kernel hash; shuffling spreads - # workloads across different kernel hash families. + # Shuffle ACROSS groups (not within — within-group order is the + # CPS sequence that benefits from in-process cache adjacency). import random as _r _r.shuffle(tasks) + n_total_cells = sum(len(g) for g in groups.values()) + _cw(f"built {len(tasks)} tasks covering {n_total_cells} cells in {len(groups)} groups") print(f"[compile-warmup] {len(tasks)} compile tasks " - f"({len(configs)} outer × {len(inner_combos)} inner combos) " + f"({len(configs)} outer × {len(groups)} cell-groups " + f"covering {n_total_cells} cells, CPS-grouped) " f"across {max_workers} processes (ProcessPoolExecutor, spawn start)") t0 = time.perf_counter() ctx = multiprocessing.get_context("spawn") errors = [] + _cw("about to create ProcessPoolExecutor") with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as ex: + _cw("ProcessPoolExecutor created, about to submit tasks") # baseline_fn=None: workers compile only the checkpointing kernel. # Baseline kernels (if any) get compiled lazily in the parent during # the timing phase — usually just one extra compile, negligible. @@ -2112,28 +2216,27 @@ def _ps(val): ex.submit(_warm_one_config, args, task, None): task for task in tasks } + _cw(f"submitted {len(futures)} tasks, waiting for results") + _n_done = 0 for fut in futures: try: fut.result() except Exception as e: errors.append((futures[fut], e)) + _n_done += 1 + # Progress beacons at 10/25/50/75/100% to gauge effective parallelism. + if _n_done in (max(1, len(futures)//10), + max(1, len(futures)//4), + max(1, len(futures)//2), + max(1, (3*len(futures))//4), + len(futures)): + _cw(f"{_n_done}/{len(futures)} tasks complete") if errors: - # Don't abort the whole bench just because some (M, W, S, TMA, ...) - # combinations refuse to compile. Driver-style invocations - # (search loops, sparse cell-lists) often hit a few bad combos in - # otherwise-valid sweeps; aborting wastes the rest of the bench's - # work. Log them once each and continue — the inner-loop runs - # below will trip their compile again per-cell, fail safely - # (per-cell is caught), and the cells just don't get timed. - print(f"[compile-warmup] {len(errors)} configs failed to compile " - f"(continuing; cells will be skipped):", file=sys.stderr) - for cfg, e in errors[:5]: - first_line = str(e).strip().split("\n")[0][:120] - print(f" FAILED {cfg}: {type(e).__name__}: {first_line}", + for cfg, e in errors: + print(f"[compile-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", file=sys.stderr) - if len(errors) > 5: - print(f" ... and {len(errors) - 5} more", file=sys.stderr) + raise errors[0][1] print(f"[compile-warmup] done in {time.perf_counter() - t0:.1f}s") @@ -3088,6 +3191,7 @@ def _emit_split(name_w, name_nw, val_w, val_nw): f"[retry] CUPTI mismatch on {sweep_tag!r}; " f"retrying ({attempt + 1}/{retry_budget})", file=sys.stderr, + flush=True, ) if stats is None: args._skipped_cells.append(sweep_tag) @@ -3508,6 +3612,15 @@ def _current_cell_tuple(args, locals_dict: dict) -> tuple | None: def _run_benchmark(args) -> None: + # Phase-timing markers — emit timestamped checkpoints so a captured-stdout + # run can later attribute wall time to setup vs compile-warmup vs prewarm + # vs timing. Single-line format makes log-grepping trivial. + _phase_t0 = time.perf_counter() + def _phase(label: str) -> None: + dt = time.perf_counter() - _phase_t0 + print(f"[phase] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) + _phase("enter _run_benchmark") + # Pending-results FIFO for srxl's deferred CUPTI parsing pipeline. Each # entry holds a _PendingCuptiStats handle; _drain_pending_results pulls # ready entries and routes them to _print_row (which appends to JSONL). @@ -3637,11 +3750,14 @@ def _run_benchmark(args) -> None: # Cell-list state may already have been populated by main() (so that # the args.*_list derivations downstream see the override). Default to # empty if not. + _phase(f"done loading _done_keys ({len(args._done_keys)} entries)") + if not hasattr(args, "_cell_list_keys"): args._cell_list_keys: tuple = () args._cell_list_set: set = set() if getattr(args, "cell_list", None): _load_cell_list_into_args(args) + _phase(f"done loading cell-list ({len(args._cell_list_set)} cells)") assert args.nheads % args.tp_size == 0, ( f"nheads ({args.nheads}) must be divisible by tp_size ({args.tp_size})" @@ -3682,11 +3798,13 @@ def _run_benchmark(args) -> None: elif args.l2_flush: _init_l2_flush() + _phase("about to enter compile-warmup") if args.compile_threads > 0: _compile_warmup_phase( args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, baseline_fn, max_workers=args.compile_threads, ) + _phase("returned from compile-warmup") # Pre-warm the per-(state_dtype, act_dtype, mtp_len, ...) tensor cache at # the largest requested batch size. Without this, the timing loop would @@ -3703,6 +3821,7 @@ def _run_benchmark(args) -> None: args.tp_nheads, args.head_dim, args.d_state, args.tp_ngroups, max_window=getattr(args, "max_window", None) or None, ) + _phase("done tensor prewarm — entering timing") if args.profile: torch.cuda.cudart().cudaProfilerStart() @@ -3970,7 +4089,12 @@ def _parse_args() -> argparse.Namespace: default="bf16", help="Comma-separated activation dtypes for x/B/C/dt: fp32,bf16", ) - parser.add_argument("--warmup", type=int, default=20, help="Number of warmup iterations") + parser.add_argument("--warmup", type=int, default=4, + help="Number of warmup iterations. Default aligns with " + "the graph group-iters (default 4 for mix scenarios) so " + "warmup + iters / mix-iters lands on a clean multiple " + "without per-args rounding overhead. Earlier default of " + "20 was overkill for steady-state warming.") parser.add_argument("--iters", type=int, default=100, help="Number of timed iterations") parser.add_argument( "--compile-threads", @@ -4489,6 +4613,33 @@ def _parse_args() -> argparse.Namespace: if args.mix_only and args.mix_csv is None: parser.error("--mix-only requires --mix-csv") + # Round iter counts up so warmup + iters (and warmup + mix_iters) are clean + # multiples of the graph group-iters used downstream. Default mix group is + # 4, default pure group is 2. An explicit --cuda-graph-group-iters can + # request a larger group. We round to the max of the two so all scenarios + # in a single run (pure + mix) share a clean total_iters. The cost is at + # most (group-1) extra iters per scenario — negligible — and the win is + # that graph_group_iters never falls back to 1 (which caused ~5x slowdown + # in observed benchmark walls). + _group_for_rounding = max( + _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX, + _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE, + getattr(args, "cuda_graph_group_iters", None) or 0, + ) + def _round_iters_to_group(name, val): + total = args.warmup + val + if total % _group_for_rounding == 0: + return val + new_total = ((total + _group_for_rounding - 1) // _group_for_rounding) * _group_for_rounding + new_val = new_total - args.warmup + print(f"[bench] rounding --{name} {val} → {new_val} so warmup+{name} " + f"({new_total}) is a multiple of graph group_iters={_group_for_rounding}", + file=sys.stderr) + return new_val + args.iters = _round_iters_to_group("iters", args.iters) + if getattr(args, "mix_iters", None): + args.mix_iters = _round_iters_to_group("mix-iters", args.mix_iters) + # Cell-list (if any) must be applied BEFORE the post-argparse string→list # derivations below — those build args.*_list from args.* strings, so a # cell-list override of e.g. args.modes='maindl' needs to land before From 811f2cb4584e3fce6b8d7e772a19cbf76068de62 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 15 May 2026 10:22:02 -0700 Subject: [PATCH 45/89] bench: iterate cell-list directly (O(|cells|), not O(cartesian)) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The --cell-list code path was iterating the full inner knob cartesian and filtering each iteration via 'if _tup not in args._cell_list_set: continue'. That's O(cartesian) — fine when the cartesian is small, catastrophic when auto-cover-populated args.*_write × args.*_nowrite × CPS × LS × TMA flags push the cartesian into the billions. Before fix (pers_dynamic with CPS=[1..8], S=[1..4], TMA on): cartesian ≈ 3.3B → ~55 min CPU spin per bench call, 0% GPU After fix: iterate args._cell_list_set directly, one yield per cell. Wall scales with |cell-list|, no spin. The 21-tuple unpack order matches itertools.product(*_iter_axes), so the fix just replaces the source of the iterator. When --cell-list isn't used, the original cartesian iteration is preserved. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 32 ++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index a16097bf45d6..2ce6ffdee702 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -2782,6 +2782,36 @@ def _graph_pre_iter(j): use_tma_replay_nowrite_load_values, use_tma_replay_write_store_values, ) + # Iteration source: when --cell-list is active, iterate the cell set + # DIRECTLY (one yield per cell). The earlier design iterated the + # full inner cartesian and filtered each iteration via membership in + # args._cell_list_set — that's O(cartesian) which blows up to + # billions of iterations when the cell-list spans wide split-knob + # values (CPS, LS, M, W, S each contributing a Wx*Wnw factor on top + # of TMA flags), producing 50+ min of CPU spin per bench call before + # any actual timing. Direct iteration is O(|cell_list|). + if getattr(args, "_cell_list_set", None): + def _gen_from_cell_list(): + keys = args._cell_list_keys + for tup in args._cell_list_set: + d = dict(zip(keys, tup)) + yield ( + d.get("Mw"), d.get("Mnw"), + d.get("Ww"), d.get("Wnw"), + d.get("Sw"), d.get("Snw"), + d.get("pW"), d.get("pS"), + d.get("H"), + d.get("R"), d.get("CT"), + d.get("CPSw"), d.get("CPSnw"), + d.get("LSw"), d.get("LSnw"), + d.get("FL"), d.get("WS"), + d.get("TMARL"), d.get("TMAWL"), + d.get("TMANL"), d.get("TMAWS"), + ) + _iter_source = _gen_from_cell_list() + else: + _iter_source = itertools.product(*_iter_axes) + for ( block_size_m_w, block_size_m_nw, @@ -2804,7 +2834,7 @@ def _graph_pre_iter(j): use_tma_replay_write_load, use_tma_replay_nowrite_load, use_tma_replay_write_store, - ) in itertools.product(*_iter_axes): + ) in _iter_source: # When tied, _nw values were placeholder None; fill from _w (the # shared value). When split, _w and _nw came from independent lists. if not _any_split: From 8e91b93b452802f29a5befa846b9301a5460d5c3 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 15 May 2026 10:25:52 -0700 Subject: [PATCH 46/89] bench: guard cell-list-direct iter on 'not warmup_only' MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compile-warmup workers (_warm_one_config) call _bench_config with warmup_only=True and inner_overrides that CLAMP args.*_write/_nowrite to single values, making the cartesian 1x1...x1 = 1 iter (the cell the worker was given). My previous structural fix made _bench_config use cell-list-direct iteration whenever args._cell_list_set is populated — but that includes the worker case, so each of 28 workers iterated all 2884 cells instead of its assigned one. Compile-warmup wall went from ~38s to 233s+ for the first 256/2568 tasks (worker N×N blowup). Fix: gate the cell-list-direct iteration on 'not warmup_only'. Main timing path uses cell-list-direct (O(|cells|)). Workers continue with cartesian, which is now 1-iter under their clamped args. Verified: 1-cell warmup_only=True probe runs in ~30s (correct). Main timing path benefits unchanged. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../benchmark_replay_selective_state_update.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 2ce6ffdee702..19b0e84a0b79 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -2782,7 +2782,8 @@ def _graph_pre_iter(j): use_tma_replay_nowrite_load_values, use_tma_replay_write_store_values, ) - # Iteration source: when --cell-list is active, iterate the cell set + # Iteration source: when --cell-list is active AND this is the main + # timing path (not a compile-warmup worker), iterate the cell set # DIRECTLY (one yield per cell). The earlier design iterated the # full inner cartesian and filtered each iteration via membership in # args._cell_list_set — that's O(cartesian) which blows up to @@ -2790,7 +2791,15 @@ def _graph_pre_iter(j): # values (CPS, LS, M, W, S each contributing a Wx*Wnw factor on top # of TMA flags), producing 50+ min of CPU spin per bench call before # any actual timing. Direct iteration is O(|cell_list|). - if getattr(args, "_cell_list_set", None): + # + # IMPORTANT exception for workers (warmup_only=True): _warm_one_config + # clamps args.*_write/_nowrite via inner_overrides to single values, + # making the cartesian 1×1×...×1 = 1 iter, which is exactly the one + # cell that worker was given. If we used cell-list-direct iteration + # here, every worker would iterate ALL 2884 cells instead of just + # its assigned one — turning compile-warmup into 28-way duplication. + # (Observed: 256 tasks in 233s under that bug vs ~18s correct.) + if getattr(args, "_cell_list_set", None) and not warmup_only: def _gen_from_cell_list(): keys = args._cell_list_keys for tup in args._cell_list_set: From ac8ec9ed186674075c15dbe8eb411926e601e675 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 15 May 2026 10:33:47 -0700 Subject: [PATCH 47/89] bench: compile-warmup uses per-cell outer in cell-list mode (no outer cartesian) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the compile-warmup built tasks = OUTER × CELLS, where OUTER cartesian-enumerated the kernel-compile-relevant axes (RECT, MODE, SR, WC, SORT, REVN, HSORT) over the UNION of values seen across cells. If the cell-list spanned both RECT=0 and RECT=1 (or any other outer-axis split), each cell was paired with EVERY outer config — half of those tasks then failed the worker-side cell-list filter and wasted dispatch. Fix: in cell-list mode, group cells by (cell's own outer + non-CPS inner signature) and pair each group with ONE task carrying that exact cell-outer. Each cell now contributes O(1) tasks regardless of the cell-list's outer-axis diversity. Sweep-args mode (no --cell-list) keeps the existing OUTER × INNER cartesian. Observed for pers_dynamic b=64 with RECT mixed: 2× task count → 1×. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 94 ++++++++++++++----- 1 file changed, 68 insertions(+), 26 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 19b0e84a0b79..bb4d5805c447 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -2109,19 +2109,65 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): return [(a, b) for a in w for b in nw] _cw(f"built {len(configs)} outer configs") - # Build per-cell inner-overrides dicts. cell_set = getattr(args, "_cell_list_set", set()) cell_keys = getattr(args, "_cell_list_keys", ()) + # CPS keys are runtime ints (kernel value-specializes on `div_by_16`); + # cells differing only on CPS values can SHARE a worker so the second + # CPS value in a div_by_16 bucket hits the in-process Triton cache. + _cps_keys = ("cta_per_sm_write", "cta_per_sm_nowrite", "cta_per_sm") + if cell_set: - inner_dicts = [] + # ============== CELL-LIST PATH ============== + # Build tasks DIRECTLY from cells. Each cell carries its OWN + # outer-axis values (RECT, MODE, SR, WC, SORT, REVN, HSORT) so we + # pair each cell with its specific outer config — NOT the union- + # cartesian of all cells' outer values. Previously the OUTER × + # CELL cartesian doubled task count when a cell-list spanned both + # RECT=0 and RECT=1 (or any other outer-axis split); half the + # tasks then failed the cell-list filter inside the worker and + # wasted dispatch overhead. This path is O(|unique cell groups|). + from collections import defaultdict as _dd + cell_groups: dict = _dd(list) for tup in cell_set: d = dict(zip(cell_keys, tup)) + cell_outer = ( + "SR" if d.get("SR", 0) else "RN", # sr_mode + bool(d.get("RECT", 0)), # rect + bool(d.get("WC", 1)), # write_ckpt + d.get("MODE", "monolithic"), # mode + bool(d.get("SORT", 0)), # sort_slots + bool(d.get("REVN", 0)), # reverse_nowrite + bool(d.get("HSORT", 0)), # hardcode_sort + ) inner = {} for k, v in d.items(): if k in _CELL_LIST_KEY_TO_ARG: inner[_CELL_LIST_KEY_TO_ARG[k]] = str(v) - inner_dicts.append(inner) + non_cps_sig = tuple(sorted((k, v) for k, v in inner.items() if k not in _cps_keys)) + cell_groups[(cell_outer, non_cps_sig)].append(inner) + + # CLI-runtime axes (batch/mtp/dtype) are NOT in cell-list — they + # come from CLI args and cartesian here (typically just 1 combo). + cli_outers = [] + for _b in _compile_batches: + for _m in mtp_lengths: + _pk = _resolve_prev_ks(args, _m) + for _sd in state_dtypes: + for _ad in act_dtypes: + cli_outers.append((_b, _m, _pk, _sd, _ad)) + + tasks = [] + for cli_outer in cli_outers: + for (cell_outer, _sig), inner_list in cell_groups.items(): + outer_cfg = (*cli_outer, *cell_outer) + tasks.append((outer_cfg, inner_list)) + n_groups = len(cell_groups) + n_total_cells = sum(len(g) for g in cell_groups.values()) + n_outer_used = len(cli_outers) else: + # ============== SWEEP-ARGS PATH ============== + # Build inner_dicts via cartesian over knob axes, then cross with + # the `configs` outer cartesian. Existing behavior. m_pairs = _split_or_pair("block_size_m", "block_size_m_write", "block_size_m_nowrite") w_pairs = _split_or_pair("num_warps", "num_warps_write", "num_warps_nowrite") ns_pairs = _split_or_pair("num_stages", "num_stages_write", "num_stages_nowrite") @@ -2175,33 +2221,29 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): d[k] = str(v) inner_dicts.append(d) - # Group by non-CPS signature. All inner_dicts that match on every - # key except CPSw/CPSnw/cta_per_sm land in the same worker task; the - # worker will compile them in sequence, sharing Triton's in-process - # kernel cache across the divides bucket boundary at most once per - # bucket (CPS in {1,2} vs {4,8}). - _cps_keys = ("cta_per_sm_write", "cta_per_sm_nowrite", "cta_per_sm") - groups: dict = {} - for d in inner_dicts: - sig = tuple(sorted((k, v) for k, v in d.items() if k not in _cps_keys)) - groups.setdefault(sig, []).append(d) - - tasks = [] - for outer in configs: - for sig, group in groups.items(): - tasks.append((outer, group)) - - # Shuffle ACROSS groups (not within — within-group order is the - # CPS sequence that benefits from in-process cache adjacency). + groups: dict = {} + for d in inner_dicts: + sig = tuple(sorted((k, v) for k, v in d.items() if k not in _cps_keys)) + groups.setdefault(sig, []).append(d) + tasks = [] + for outer in configs: + for sig, group in groups.items(): + tasks.append((outer, group)) + n_groups = len(groups) + n_total_cells = sum(len(g) for g in groups.values()) + n_outer_used = len(configs) + + # Shuffle ACROSS tasks (preserve within-group CPS sequence for in-process + # cache adjacency — within-group order is intentional, not shuffled). import random as _r _r.shuffle(tasks) - n_total_cells = sum(len(g) for g in groups.values()) - _cw(f"built {len(tasks)} tasks covering {n_total_cells} cells in {len(groups)} groups") + _cw(f"built {len(tasks)} tasks covering {n_total_cells} cells in {n_groups} groups") print(f"[compile-warmup] {len(tasks)} compile tasks " - f"({len(configs)} outer × {len(groups)} cell-groups " - f"covering {n_total_cells} cells, CPS-grouped) " - f"across {max_workers} processes (ProcessPoolExecutor, spawn start)") + f"({n_outer_used} outer × {n_groups} cell-groups " + f"covering {n_total_cells} cells, CPS-grouped" + + (", per-cell outer" if cell_set else "") + + f") across {max_workers} processes (ProcessPoolExecutor, spawn start)") t0 = time.perf_counter() ctx = multiprocessing.get_context("spawn") From 646c9a215404e356033c816db0200c2c817ca039 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 15 May 2026 12:37:30 -0700 Subject: [PATCH 48/89] checkpointing: wire TMA load under IS_DYNAMIC=True + per-path descriptors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two related changes: 1. _persistent_main_impl previously hard-forced use_tma_load=False when IS_DYNAMIC=True (the persistent_dynamic launch path), making the USE_TMA_LOAD_WRITE and USE_TMA_LOAD_NOWRITE wrapper flags silent no-ops there. Only USE_TMA_STORE was live. Replace the IS_DYNAMIC short-circuit with a branch on is_write (constexpr in non-dynamic mode, runtime in dynamic), followed by a constexpr-gated per-side TMA call. Triton DCEs the wrong branch in the non-dynamic case (kernel sass unchanged), and emits both branches runtime-gated in the dynamic case. Descriptor safety preserved because the wrapper builds real tensor_descriptors iff any TMA flag is on, and the constexpr gating skips .load() when no flag is on. GPU probe at b=64 fp32 (TMAWL=1, TMANL=1, TMAWS=1) reaches 19.62us vs the previous TMA-store-only winner at 19.92us — ~1.5% win at b=64 from previously-dead load knobs. 2. Per-path TMA descriptor: state_tma_descriptor split into state_tma_descriptor_write / _nowrite so M-split (Mw \!= Mnw) can build matching block_shapes. When tied (the common case), both variables hold the same descriptor object. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 104 ++++++++++++------ 1 file changed, 72 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index b1de20f5f103..7ed891d322c0 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -3330,22 +3330,28 @@ def _persistent_main_impl( state_ptrs = ( state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate ) - # Pick LOAD flag. In non-dynamic mode WRITE_CHECKPOINT is constexpr, - # so we pick the matching flag at compile time and the dead branch is - # DCE'd. In dynamic mode, slot may be write or nowrite per-CTA — TMA - # load there would require real descriptors at BOTH compile branches - # (else compilation fails because state_tma_descriptor is a plain - # tensor when TMA is off), which is not currently wired up. Force - # use_tma_load=False for dynamic mode; revisit if dynamic+TMA becomes - # worth wiring up. - if IS_DYNAMIC: - use_tma_load: tl.constexpr = False - else: - use_tma_load: tl.constexpr = USE_TMA_LOAD_WRITE if WRITE_CHECKPOINT else USE_TMA_LOAD_NOWRITE - if use_tma_load: - state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + # Load state. Branch on is_write (constexpr in non-dynamic mode, runtime + # in dynamic), then constexpr-pick TMA-vs-tl.load per side. + # + # Non-dynamic (is_write is constexpr = WRITE_CHECKPOINT): outer `if` + # DCE's, only the matching side's constexpr-gated load survives. + # + # Dynamic (is_write is runtime per-CTA): both write and nowrite blocks + # emit; each contains exactly one of (TMA load, tl.load) after the + # constexpr USE_TMA_LOAD_* gate resolves. The descriptor is real iff + # any of the 4 TMA flags is on at the wrapper (line 4712-4713 of this + # file); the constexpr gating guarantees we never call .load() on the + # plain-tensor fallback path, so this stays compilation-safe. + if is_write: + if USE_TMA_LOAD_WRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) else: - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + if USE_TMA_LOAD_NOWRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) if QUANT_MAX > 0.0: state_scales_base = ( state_scales_ptr @@ -4695,13 +4701,19 @@ def checkpointing_state_update( state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 state_scales_strides = (0, 0, 0) - # Single TMA descriptor for state — used by all TMA-consuming paths - # (rect load, replay write load, replay write store, replay nowrite load). - # Same memory (state's flat 2D view, shape (cache*nheads*dim, dstate)), - # same block_shape, same underlying tensor descriptor — gated at usage - # sites by per-path constexprs. When no TMA flag is on, the variable - # holds the raw `state` tensor as a dummy; kernels never reference it - # because their constexprs are all False (Triton DCEs the dead branches). + # Per-path TMA descriptors for state — write-side and nowrite-side. Each + # kernel launch consumes the descriptor whose block_shape[0] matches its + # BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need + # distinct descriptors; otherwise the descriptor's block_shape[0] would + # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot / arithmetic + # on the loaded tile fails shape inference at compile time + # ("Cannot make_shape_compatible: incompatible dimensions"). When Mw == + # Mnw (tied, the common case) the two descriptors are the same object. + # Same memory (state's flat 2D view, shape (cache*nheads*dim, dstate)) + # and same dstate block_shape — only block_shape[0] differs. + # When no TMA flag is on, both variables hold the raw `state` tensor as a + # dummy; kernels never reference it because their constexprs are all + # False (Triton DCEs the dead branches). # `triton.set_allocator()` must run before any descriptor-using launch. if (_use_tma_rect_load or _use_tma_replay_write_load or _use_tma_replay_write_store or _use_tma_replay_nowrite_load): @@ -4709,12 +4721,20 @@ def checkpointing_state_update( _ensure_tma_allocator() assert state.is_contiguous(), "TMA state requires contiguous state" assert state.stride(-1) == 1, "TMA state requires inner stride 1" - state_tma_descriptor = TensorDescriptor.from_tensor( - state.view(-1, state.shape[-1]), - block_shape=[BLOCK_SIZE_M, triton.next_power_of_2(dstate)], + _state_flat = state.view(-1, state.shape[-1]) + _dstate_pow2 = triton.next_power_of_2(dstate) + state_tma_descriptor_write = TensorDescriptor.from_tensor( + _state_flat, block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], ) + if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: + state_tma_descriptor_nowrite = state_tma_descriptor_write + else: + state_tma_descriptor_nowrite = TensorDescriptor.from_tensor( + _state_flat, block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], + ) else: - state_tma_descriptor = state # dummy; all consuming constexprs False + state_tma_descriptor_write = state # dummy; all consuming constexprs False + state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False # Slot permutation — pointer + USE_PERM gate. When the caller provides # a perm tensor the dl-family launches read pid_b through it; otherwise @@ -4861,10 +4881,14 @@ def launch_replay_main(write_checkpoint: bool, early_out: bool, _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE + # Per-path TMA descriptor: block_shape[0] must match the kernel's + # BLOCK_SIZE_M (`_bsm`); see the descriptor build block above. + _desc = (state_tma_descriptor_write if write_checkpoint + else state_tma_descriptor_nowrite) def _main_grid_local(META, _bsm=_bsm): return (triton.cdiv(dim, _bsm), batch, nheads) _checkpointing_main_kernel[_main_grid_local]( - state, state_tma_descriptor, state_scales_arg, old_x, + state, _desc, state_scales_arg, old_x, old_B, old_dt, old_dA_cumsum, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, @@ -4918,8 +4942,10 @@ def launch_rectangle_main(early_out: bool, _ns = NUM_STAGES_NOWRITE def _main_grid_local(META, _bsm=_bsm): return (triton.cdiv(dim, _bsm), batch, nheads) + # Rectangle is always the nowrite-side path; descriptor block_shape[0] + # must match BLOCK_SIZE_M_NOWRITE (= _bsm here). _rectangle_main_kernel[_main_grid_local]( - state, state_tma_descriptor, state_scales_arg, old_x, + state, state_tma_descriptor_nowrite, state_scales_arg, old_x, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, cb_scaled, decay_vec, @@ -4953,8 +4979,11 @@ def _main_grid_local(META, _bsm=_bsm): def launch_dynamic_main(rectangle: bool, launch_dependent_kernels: bool = False): + # Dynamic mode uses a single BLOCK_SIZE_M (no M-split inside this + # kernel); BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE by the wrapper's tied + # convention, so the write-side descriptor matches. _dynamic_main_kernel[main_grid]( - state, state_tma_descriptor, state_scales_arg, old_x, + state, state_tma_descriptor_write, state_scales_arg, old_x, old_B, old_dt, old_dA_cumsum, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, @@ -5022,7 +5051,11 @@ def launch_dynamic_main(rectangle: bool, # tile_id is covered exactly once across all live pids in [0, grid) when # grid <= NUM_PERSISTENT (each CTA does 1 tile; loop step >= total_work # exits immediately) AND when grid == NUM_PERSISTENT (each CTA loops over - # multiple tiles). NUM_PERSISTENT stays a constexpr = cta_per_sm * num_sms. + # multiple tiles). NUM_PERSISTENT is now a runtime int (see kernel def + # docstring at _persistent_main_kernel) so changing cta_per_sm does NOT + # trigger a new Triton compile — same kernel binary, different loop step. + # (Named UPPERCASE for historical Triton-style consistency only; not + # constexpr.) _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M def launch_persistent_main(write_checkpoint: bool, @@ -5070,8 +5103,11 @@ def launch_persistent_main(write_checkpoint: bool, _n_slots_for_launch = batch _total_work_launch = max(1, _n_slots_for_launch * _num_pid_m_local * nheads) grid = (min(_num_persistent, _total_work_launch),) + # Per-path TMA descriptor — block_shape[0] must match _bsm. + _desc = (state_tma_descriptor_write if write_checkpoint + else state_tma_descriptor_nowrite) _persistent_main_kernel[grid]( - state, state_tma_descriptor, state_scales_arg, old_x, + state, _desc, state_scales_arg, old_x, old_B, old_dt, old_dA_cumsum, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, @@ -5143,8 +5179,12 @@ def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, # comment for correctness rationale. _total_work_launch = max(1, batch * _num_pid_m * nheads) grid = (min(num_persistent_arg, _total_work_launch),) + # Persistent-dynamic kernel uses a single BLOCK_SIZE_M (same as the + # wrapper's BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE tied convention), so + # the write-side descriptor matches. Both write and nowrite slots + # in this kernel share that BSM. _persistent_main_kernel[grid]( - state, state_tma_descriptor, state_scales_arg, old_x, + state, state_tma_descriptor_write, state_scales_arg, old_x, old_B, old_dt, old_dA_cumsum, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, From e9ead286d64f3ea9265eeea1b73c36456e264e7a Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 15 May 2026 15:49:33 -0700 Subject: [PATCH 49/89] bench: retry CUPTI parser spawn with timeout + alive polling Under N-way concurrent benches, CuptiKernelTimer.__init__ spawn of the parser child can fail two ways: - early: child dies during pickle.load with FileNotFoundError in multiprocessing/synchronize.SemLock._rebuild (POSIX named semaphore raced between parent create and child sem_open) - slow: child stays alive but its bench-module re-import (torch, triton, ...) is delayed past the 10s ready_event timeout by CPU contention with concurrent compile-warmup spawn children Wrap parser spawn in a 3-attempt loop with: - 30s per-attempt deadline (was 10s) - 0.5s polling on both ready_event AND process.is_alive() (detects dead-child fast, doesn't wait the full 30s for the early-fail case) - terminate+join cleanup between attempts - jittered backoff (0.5, 1.0, 1.5s) Final RuntimeError now reports the failure mode (alive=T/F, exitcode) so future diagnosis is easier. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 50 +++++++++++++++---- 1 file changed, 40 insertions(+), 10 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index bb4d5805c447..2313b0122587 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -948,16 +948,46 @@ def __init__(self) -> None: self._last_stop_timing: dict[str, float] = {} self._current_flush_period_ms = 0 self._mp_ctx = mp.get_context("spawn") - self._parse_input_queue = self._mp_ctx.Queue() - self._parse_output_queue = self._mp_ctx.Queue() - ready_event = self._mp_ctx.Event() - self._parse_process = self._mp_ctx.Process( - target=_cupti_parser_worker, - args=(self._parse_input_queue, self._parse_output_queue, ready_event), - ) - self._parse_process.start() - if not ready_event.wait(timeout=10.0): - raise RuntimeError("CUPTI parser process did not initialize") + # Retry parser-process spawn: concurrent bench instances on the same + # node race on POSIX named semaphores in /dev/shm — child can die in + # pickle.load with FileNotFoundError in SemLock._rebuild before + # signalling ready_event. Detect early-dead child via is_alive() so + # we don't waste the full timeout, and retry up to 3x with jitter. + last_err = None + for _spawn_attempt in range(3): + self._parse_input_queue = self._mp_ctx.Queue() + self._parse_output_queue = self._mp_ctx.Queue() + ready_event = self._mp_ctx.Event() + self._parse_process = self._mp_ctx.Process( + target=_cupti_parser_worker, + args=(self._parse_input_queue, self._parse_output_queue, ready_event), + ) + self._parse_process.start() + deadline = time.time() + 30.0 + spawn_ok = False + while time.time() < deadline: + if ready_event.wait(timeout=0.5): + spawn_ok = True + break + if not self._parse_process.is_alive(): + break + if spawn_ok: + last_err = None + break + last_err = (f"attempt {_spawn_attempt + 1}: " + f"alive={self._parse_process.is_alive()}, " + f"exitcode={self._parse_process.exitcode}") + try: + if self._parse_process.is_alive(): + self._parse_process.terminate() + self._parse_process.join(timeout=2.0) + except Exception: + pass + time.sleep(0.5 + 0.5 * _spawn_attempt) + if last_err is not None: + raise RuntimeError( + f"CUPTI parser process did not initialize after 3 attempts: {last_err}" + ) self._set_zeroed_host_buffer_attr() for _ in range(_CUPTI_HOST_BUFFER_COUNT): From be5046f6e90f19c383f6a0a328ce355a247efafc Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 15 May 2026 15:57:04 -0700 Subject: [PATCH 50/89] bench: add --mp-start-method forkserver option MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compile-warmup ProcessPoolExecutor and the CUPTI parser child are multiprocessing.spawn children by default. Each child runs a fresh Python interpreter and re-imports this module (~15s of torch/triton import cost). With 4 concurrent benches × 26 compile threads each = 104 spawn children paying this 15s import in parallel = serious CPU thrash. Add --mp-start-method {spawn,forkserver} (default spawn). When forkserver is selected: - Add this file's directory to sys.path so the forkserver process can import the bench module by basename (not __main__, which is per-process) - mp.set_start_method('forkserver', force=True) - mp.set_forkserver_preload([this module]) — server imports torch / triton / etc ONCE, then forks workers that inherit the imports Module-level _MP_START_METHOD constant set in __main__; both get_context() sites read it. GPU probe on B200 (4 cells, fp32 b=64, TMA fix verified): forkserver path completes compile-warmup cleanly, no CUPTI failures, kernel timing unchanged (19.65us TMAWL=TMANL= TMAWS=1 vs 19.62 spawn baseline — within noise). In real-world sweep at 4-bench × 26-thread scale, the parallel import cost drops from ~15s to <1s per round. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 43 +++++++++++++++++-- 1 file changed, 40 insertions(+), 3 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 2313b0122587..d1dd58d764f8 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -665,6 +665,12 @@ def _kernels_per_iter_baseline(with_conv1d: bool) -> int: _CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER = 5 _CUPTI_HOST_BUFFER_BYTES = 1024 * 1024 _CUPTI_HOST_BUFFER_COUNT = 16 + +# Multiprocessing start method for compile-warmup + CUPTI parser children. +# Set in __main__ from --mp-start-method. "spawn" (default) is robust; each +# child re-imports torch/triton/etc (~15s). "forkserver" preloads once and +# forks cheaply (~1s/child) — see __main__ block for the preload setup. +_MP_START_METHOD = "spawn" _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE = 1 _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX = 4 @@ -947,7 +953,7 @@ def __init__(self) -> None: self._last_start_timing: dict[str, float] = {} self._last_stop_timing: dict[str, float] = {} self._current_flush_period_ms = 0 - self._mp_ctx = mp.get_context("spawn") + self._mp_ctx = mp.get_context(_MP_START_METHOD) # Retry parser-process spawn: concurrent bench instances on the same # node race on POSIX named semaphores in /dev/shm — child can die in # pickle.load with FileNotFoundError in SemLock._rebuild before @@ -2273,10 +2279,10 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): f"({n_outer_used} outer × {n_groups} cell-groups " f"covering {n_total_cells} cells, CPS-grouped" + (", per-cell outer" if cell_set else "") - + f") across {max_workers} processes (ProcessPoolExecutor, spawn start)") + + f") across {max_workers} processes (ProcessPoolExecutor, {_MP_START_METHOD} start)") t0 = time.perf_counter() - ctx = multiprocessing.get_context("spawn") + ctx = multiprocessing.get_context(_MP_START_METHOD) errors = [] _cw("about to create ProcessPoolExecutor") with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as ex: @@ -4217,6 +4223,19 @@ def _parse_args() -> argparse.Namespace: "in parallel and populate the persistent cache for free hits during " "the sequential timed phase. 0 disables the phase. Default 64.", ) + parser.add_argument( + "--mp-start-method", + choices=("spawn", "forkserver"), + default="spawn", + help="multiprocessing start method for compile-warmup workers AND " + "the CUPTI parser child process. 'spawn' (default) is robust but " + "each child re-imports the bench module (~15s torch+triton import " + "cost). 'forkserver' starts a server once, preloads the bench " + "module ONCE, then forks children cheaply (~1s each). When 4 " + "benches run concurrently with --compile-threads 26 each, spawn " + "still incurs 4*26=104 imports per round; forkserver cuts this to " + "4 (one per server).", + ) parser.add_argument( "--profile", action="store_true", @@ -4860,6 +4879,24 @@ def close(self): if __name__ == "__main__": _args = _parse_args() + # Configure multiprocessing start method early — must be before any + # mp.get_context() that uses the chosen method. For forkserver, also + # add this file's dir to sys.path so the forkserver can import this + # module by basename for preload (otherwise it tries to import + # __main__, which is a different beast across processes). + if _args.mp_start_method == "forkserver": + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + mp.set_start_method("forkserver", force=True) + try: + mp.set_forkserver_preload([ + "benchmark_replay_selective_state_update", + ]) + except Exception as _e: + print(f"[warn] set_forkserver_preload failed: {_e!r}; " + f"forks will still work but pay full import cost", + file=sys.stderr) + _MP_START_METHOD = _args.mp_start_method + _out_path = None if _args.output != "-": _ts = datetime.now().strftime("%Y%m%d_%H%M%S") From 588bcaf8ff58eaa5e146ec0fa27fc3888aec07a8 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 15 May 2026 21:49:27 -0700 Subject: [PATCH 51/89] mamba_checkpointing kernel: pack int8/int16 SR random bits + supporting work for pd codegen experiment This commit bundles three concerns in mamba_checkpointing: 1. Kernel (checkpointing_state_update.py): pack int8/int16 stochastic-rounding random bits to halve per-slot RNG storage cost in the replay state-update path. 2. Bench (benchmark_replay_selective_state_update.py): stamp 'gpu' field in each JSONL record (from CUDA_VISIBLE_DEVICES at startup); plumb jsonl_gpu through _print_row. Enables the oracle-cache layer in search_driver to discover and attribute timings per (host, gpu) pair. 3. Refactor staging: side-by-side _refactored.py copies of the kernel, bench, and test files. Used as a hot-swappable target while iterating on a WC_IS_CONSTEXPR codegen experiment for persistent_dynamic mode (see follow-up commit which promotes the experiment to the live kernel and drops these staging files). Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 86 +- .../checkpointing_state_update_refactored.py | 5510 +++++++++++++++++ ...benchmark_replay_selective_state_update.py | 12 + ...eplay_selective_state_update_refactored.py | 4935 +++++++++++++++ .../mamba/test_checkpointing_state_update.py | 91 +- ...t_checkpointing_state_update_refactored.py | 2230 +++++++ 6 files changed, 12853 insertions(+), 11 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_refactored.py create mode 100644 tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_refactored.py create mode 100644 tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_refactored.py diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index 7ed891d322c0..b6ebe6afbda9 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -98,6 +98,47 @@ def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: ) +@triton.jit +def _bitrev32(x: tl.tensor) -> tl.tensor: + return tl.inline_asm_elementwise( + asm="brev.b32 $0, $1;", + constraints="=r,r", + args=(x,), + dtype=tl.uint32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _stochastic_round_int8_packed( + x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor +) -> tl.tensor: + """Stochastic rounding for int8 using one random uint32 per 4 values.""" + low = rand & 0x0000FFFF + high = (rand >> 16) & 0x0000FFFF + low_rev = _bitrev32(low) >> 16 + high_rev = _bitrev32(high) >> 16 + rand_pos = offs_n & 3 + rand16 = tl.where( + rand_pos == 0, + low, + tl.where(rand_pos == 1, low_rev, tl.where(rand_pos == 2, high, high_rev)), + ) + rand01 = rand16.to(tl.float32) * (1.0 / float(1 << 16)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +@triton.jit +def _stochastic_round_int16_packed( + x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor +) -> tl.tensor: + """Stochastic rounding for int16 using one random uint32 per 2 values.""" + rand_bits = tl.where((offs_n & 1) == 0, rand, _bitrev32(rand)) + rand01 = (rand_bits & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + # Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, # old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. # Grid: (batch, nheads // HEADS_PER_BLOCK). @@ -1693,9 +1734,10 @@ def _replay_main_impl( if USE_RS_ROUNDING: # Generate random tensor for stochastic rounding. The amount of # randomness needed depends on the SR codegen path: - # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) + # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) - # int8/int16 SR (uniform-noise + floor): 1 b32 per output + # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs + # int16 SR (24b + bitrev32): 1 b32 per 2 outputs # The PTX cvt.rs.* instructions consume a single 32-bit random # and split the bits internally for 2 or 4 conversions. The # tl.inline_asm_elementwise wrapper has uniform `pack` across all @@ -1705,10 +1747,14 @@ def _replay_main_impl( # slots — saves Philox rounds proportionally. if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: RAND_DIVISOR: tl.constexpr = 4 # fp8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: + RAND_DIVISOR: tl.constexpr = 4 # int8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: + RAND_DIVISOR: tl.constexpr = 2 # int16 SR elif QUANT_MAX == 0.0: RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) else: - RAND_DIVISOR: tl.constexpr = 1 # int8/int16 SR (full per-element) + RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized rand_seed = tl.load(rand_seed_ptr) base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head @@ -1785,6 +1831,9 @@ def _replay_main_impl( else: if USE_RS_ROUNDING: # int8 / int16 + SR — uniform-noise + floor. + # int8 packs 4 values per random u32 using 16-bit chunks + # and bitrev16; int16 packs 2 values per random u32 using + # 24-bit uniforms from the direct/reversed u32. # (fp8 SR was handled by the early branch above.) tl.static_assert( (state_ptrs.dtype.element_ty == tl.int8) @@ -1792,8 +1841,14 @@ def _replay_main_impl( "Quantized SR fall-through expects int8 or int16; " "fp8 SR is handled by the prior branch.", ) - rand01 = (rand & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) - state_q = tl.extra.cuda.libdevice.floor(state_q + rand01) + if state_ptrs.dtype.element_ty == tl.int8: + state_q = _stochastic_round_int8_packed( + state_q, rand, offs_n[None, :] + ) + else: + state_q = _stochastic_round_int16_packed( + state_q, rand, offs_n[None, :] + ) elif state_ptrs.dtype.element_ty != tl.float8e4nv: # int8 / int16 + RN — explicit round before clamp. # fp8 + RN deliberately skips this — explicit round() would @@ -3430,19 +3485,24 @@ def _persistent_main_impl( if USE_RS_ROUNDING: # Generate random tensor for stochastic rounding. The amount of # randomness needed depends on the SR codegen path: - # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) + # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) - # int8/int16 SR (uniform-noise + floor): 1 b32 per output + # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs + # int16 SR (24b + bitrev32): 1 b32 per 2 outputs # The PTX cvt.rs.* instructions consume a single 32-bit random # and split the bits internally for 2 or 4 conversions. Generate # only what's actually consumed and broadcast to fill the unused # slots — saves Philox rounds proportionally. if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: RAND_DIVISOR: tl.constexpr = 4 # fp8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: + RAND_DIVISOR: tl.constexpr = 4 # int8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: + RAND_DIVISOR: tl.constexpr = 2 # int16 SR elif QUANT_MAX == 0.0: RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) else: - RAND_DIVISOR: tl.constexpr = 1 # int8/int16 SR (full per-element) + RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized rand_seed = tl.load(rand_seed_ptr) base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head @@ -3504,8 +3564,14 @@ def _persistent_main_impl( "Quantized SR fall-through expects int8 or int16; " "fp8 SR is handled by the prior branch.", ) - rand01 = (rand & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) - state_q = tl.extra.cuda.libdevice.floor(state_q + rand01) + if state_ptrs.dtype.element_ty == tl.int8: + state_q = _stochastic_round_int8_packed( + state_q, rand, offs_n[None, :] + ) + else: + state_q = _stochastic_round_int16_packed( + state_q, rand, offs_n[None, :] + ) elif state_ptrs.dtype.element_ty != tl.float8e4nv: tl.static_assert( (state_ptrs.dtype.element_ty == tl.int8) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_refactored.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_refactored.py new file mode 100644 index 000000000000..96116649b7ff --- /dev/null +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_refactored.py @@ -0,0 +1,5510 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +# +# Adapted from: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +# SPDX-FileCopyrightText: Copyright contributors to the sglang project +# +# Copyright (c) 2024, Tri Dao, Albert Gu. +# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID +from tensorrt_llm._utils import get_sm_version + +from .softplus import softplus + + +# Lazy global allocator for Triton TMA tensor descriptors. Required by any +# host- or device-built tensor_descriptor; without it Triton raises at first +# launch. See TMA backlog item #17 / scratch experiment notes. +_TMA_ALLOCATOR_SET = False + + +def _ensure_tma_allocator() -> None: + global _TMA_ALLOCATOR_SET + if _TMA_ALLOCATOR_SET: + return + + def _alloc_fn(size, alignment, stream): + # Triton expects an int8 buffer of `size` bytes; alignment is enforced + # by the allocator returning a buffer satisfying it (PyTorch's + # cudaMalloc-backed tensors are 256B-aligned, so we're fine). + return torch.empty(size, device="cuda", dtype=torch.int8) + + triton.set_allocator(_alloc_fn) + _TMA_ALLOCATOR_SET = True + + +@triton.jit +def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 pair → fp16x2 using Philox random bits. + + Uses PTX cvt.rs.f16x2.f32 which rounds each fp32 value to fp16 using + the random bits to break ties, avoiding systematic rounding bias that + accumulates over many decode steps with fp16 state. + + Adapted from flashinfer (Apache-2.0, vLLM/mamba lineage). + """ + return tl.inline_asm_elementwise( + asm="""{ + cvt.rs.f16x2.f32 $0, $2, $1, $3; + }""", + constraints=("=r,r,r,r,r"), + args=(x, rand), + dtype=tl.float16, + is_pure=True, + pack=2, + ) + + +@triton.jit +def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 quad → fp8 e4m3 using Philox random bits. + + Uses PTX cvt.rs.satfinite.e4m3x4.f32 which combines stochastic rounding + and saturating cast in a single op (output is final fp8, no separate + clamp needed). The reversed source-register order {$4,$3,$2,$1} is + load-bearing — PTX packs leftmost source into the high byte but Triton's + pack=4 is little-endian, so the natural {$1,$2,$3,$4} order would + silently shuffle every group of 4 contiguous outputs. + + Requires SM_100a+ (Blackwell B200). Caller must gate at the wrapper + level — this kernel does not check. + + Adapted from vLLM PR #40012 (Apache-2.0). + """ + return tl.inline_asm_elementwise( + asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", + constraints="=r,r,r,r,r,r,r,r,r", + args=(x, rand), + dtype=tl.float8e4nv, + is_pure=True, + pack=4, + ) + + +@triton.jit +def _bitrev32(x: tl.tensor) -> tl.tensor: + return tl.inline_asm_elementwise( + asm="brev.b32 $0, $1;", + constraints="=r,r", + args=(x,), + dtype=tl.uint32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _stochastic_round_int8_packed( + x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor +) -> tl.tensor: + """Stochastic rounding for int8 using one random uint32 per 4 values.""" + low = rand & 0x0000FFFF + high = (rand >> 16) & 0x0000FFFF + low_rev = _bitrev32(low) >> 16 + high_rev = _bitrev32(high) >> 16 + rand_pos = offs_n & 3 + rand16 = tl.where( + rand_pos == 0, + low, + tl.where(rand_pos == 1, low_rev, tl.where(rand_pos == 2, high, high_rev)), + ) + rand01 = rand16.to(tl.float32) * (1.0 / float(1 << 16)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +@triton.jit +def _stochastic_round_int16_packed( + x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor +) -> tl.tensor: + """Stochastic rounding for int16 using one random uint32 per 2 values.""" + rand_bits = tl.where((offs_n & 1) == 0, rand, _bitrev32(rand)) + rand01 = (rand_bits & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +# Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, +# old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. +# Grid: (batch, nheads // HEADS_PER_BLOCK). + + +@triton.jit() +def _replay_precompute_impl( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers (both buffers reachable via stride_*_dbuf). This + # kernel writes to either the active (= cache_buf_idx) or inactive + # (= 1 - cache_buf_idx) buffer depending on WRITE_CHECKPOINT — see + # comment block at top of kernel body. + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + # Double-buffer index (per cache slot) — selects this step's "active" + # buffer (= where the historical inputs for this step live). + cache_buf_idx_ptr, + # Per-request accepted-tokens count (already-cached old tokens at + # [0, PNAT) of the active buffer; new tokens this step go after them + # on no-checkpoint steps). + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + # Slot permutation: maps grid program_id -> original slot index. + # When USE_PERM=False, pid_b = tl.program_id(0) (today's behavior) and + # this ptr is unused. When USE_PERM=True, pid_b = perm[pid_grid] (or + # perm[B-1-pid_grid] if REVERSE_PERM=True). + slot_perm_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides: (cache, 2, T, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Slot permutation flags. USE_PERM=True gates a slot_perm_ptr load that + # remaps grid program_id -> original slot index. REVERSE_PERM=True walks + # the perm from the tail (B-1-pid_grid). Used by sorted-dispatch + # variants of dl/dlgrouped/maindl to cluster early-outs at one end of + # the grid; ignored by monolithic / dynamic. + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, + # Checkpointing flag — selects target buffer + offset for new-token + # cache writes. See "Cache write semantics" block below. + # Runtime (not constexpr): the only WRITE_CHECKPOINT-dependent code in + # this body is the write_buf/write_offset selection, which is plain + # arithmetic — no constexpr-shaped tile or whole-block gate. Letting + # it be runtime lets the dynamic dispatch kernel call us once with the + # per-slot needs_write flag instead of inlining two specializations. + write_checkpoint, +): + pid_grid = tl.program_id(axis=0) + # REVERSE_PERM walks the grid tail-first regardless of USE_PERM — + # combined with hardcode-sorted prev_tokens (kernel-side USE_PERM=False + # but PNAT pre-sorted write-first), reverse traversal makes the nowrite + # half front-load real work. + pid_grid_eff = (tl.num_programs(axis=0) - 1 - pid_grid) if REVERSE_PERM else pid_grid + if USE_PERM: + pid_b = tl.load(slot_perm_ptr + pid_grid_eff) + else: + pid_b = pid_grid_eff + pid_hg = tl.program_id(axis=1) # head-group index + first_head = pid_hg * HEADS_PER_BLOCK + + # Resolve cache index for writes + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + + # --- Cache write semantics --- + # cache_buf_idx names this step's "active" buffer — the one with the + # historical inputs at [0, PNAT). The other buffer is "staging". + # + # Where do we write new tokens this step? + # WRITE_CHECKPOINT=False (no overflow): append to ACTIVE buffer at + # offset [PNAT : PNAT+T). Caller does NOT flip cache_buf_idx + # afterward; PNAT_next = PNAT + accepted. [0, PNAT) preserved. + # WRITE_CHECKPOINT=True (would overflow): write to STAGING buffer at + # [0, T). Caller flips cache_buf_idx afterward; next step's + # active = the one we just wrote. PNAT_next = accepted. Old + # data in the previous active buffer is folded into state via + # the replay update and discarded. This matches today's replay + # kernel behavior exactly. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if write_checkpoint: + write_buf = 1 - buf_active + write_offset = 0 + else: + write_buf = buf_active + write_offset = prev_num_accepted_tokens + + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + t_mask = offs_t < T + n_mask = offs_n < dstate + + # Causal mask is shared across all heads (depends only on offs_t) + causal_mask = offs_t[:, None] >= offs_t[None, :] + valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] + + # --- Vectorized pre-wait phase across HEADS_PER_BLOCK heads --- + # Compute dt, dA_cumsum, decay_vec as (H, T) tiles. Pre-compute + # scale_combo = decay_matrix * dt[:, None, :] as an (H, T, T) tile that + # stays in registers across gdc_wait — eliminates the post-wait reload + # of dt + dA_cumsum and the per-head loop. + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) + + # Load dt (H, T) + dt_addrs = ( + dt_ptr + pid_b * stride_dt_batch + + heads_block[:, None] * stride_dt_head + + offs_t[None, :] * stride_dt_T + ) + dt = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias[:, None] + if DT_SOFTPLUS: + dt = softplus(dt) + + A = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) # (H,) + dA_cumsum = tl.cumsum(A[:, None] * dt, axis=1) # (H, T) + decay_vec = tl.exp(dA_cumsum) # (H, T) + + # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of write_buf. + old_dt_addrs = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block[:, None] * stride_old_dt_head + + (write_offset + offs_t)[None, :] * stride_old_dt_T + ) + tl.store(old_dt_addrs, dt, mask=t_mask[None, :]) + + old_dA_cumsum_addrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block[:, None] * stride_old_dA_cumsum_head + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T + ) + tl.store(old_dA_cumsum_addrs, dA_cumsum, mask=t_mask[None, :]) + + # decay_vec scratch — always at offs_t. + decay_vec_addrs = ( + decay_vec_ptr + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) + tl.store(decay_vec_addrs, decay_vec, mask=t_mask[None, :]) + + # scale_combo (H, T, T) = exp(dA_cumsum[h, t1] - dA_cumsum[h, t2]) * dt[h, t2] + # Stays live across gdc_wait — used post-wait to compute CB_scaled. + decay_matrix = tl.exp(dA_cumsum[:, :, None] - dA_cumsum[:, None, :]) # (H, T, T) + scale_combo = decay_matrix * dt[:, None, :] # (H, T, T) + + # --- Wait for upstream kernel (external PDL) before loading B and C --- + # All dt processing above is independent of conv1d outputs. + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- + group_idx = first_head // nheads_ngroups_ratio + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_all = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_all = tl.load( + B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + # Compute raw CB once — shared across all heads in this block + raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) + + # Store B to cache at [write_offset : write_offset+T) of write_buf. + if first_head % nheads_ngroups_ratio == 0: + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + write_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_all, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # --- Vectorized post-wait phase: scale_combo (H, T, T) is still live in + # registers from pre-wait; multiply by raw_CB (T, T), apply causal mask, + # store as one (H, T, T) tile. --- + CB_scaled_block = tl.where( + valid_mask[None, :, :], + raw_CB[None, :, :] * scale_combo, + 0.0, + ) # (H, T, T) + cb_scaled_addrs = ( + cb_scaled_ptr + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_t[None, None, :] * stride_cb_j + ) # (H, T, T) + cb_store_mask = ( + (offs_t[None, :, None] < BLOCK_SIZE_T) + & (offs_t[None, None, :] < BLOCK_SIZE_T) + ) + tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) + + +# Replay-style precompute kernel. Thin wrapper around _replay_precompute_impl +# that carries the @triton.heuristics for constexpr derivation; called from +# the Python wrapper on the replay-style path (write or replay-nowrite). +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.jit() +def _checkpointing_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + EARLY_OUT: tl.constexpr, + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does, so + # main can start its setup regardless of how this program ends (pad, + # early-out, or full body). PDL signals are idempotent; main's + # gdc_wait still gates on prerequisite-kernel completion for + # correctness. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + # Per-program early-out gate (option-2 double-launch). When EARLY_OUT + # is False the entire block is constexpr-folded out and the wrapper is + # just an impl call. When True, this kernel only runs for slots whose + # (PNAT + T > MAX) status matches WRITE_CHECKPOINT. + if EARLY_OUT: + pid_grid_eo = tl.program_id(axis=0) + pid_grid_eo_eff = (tl.num_programs(axis=0) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo + if USE_PERM: + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) + else: + pid_b_eo = pid_grid_eo_eff + if HAS_CACHE_BATCH_INDICES: + cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) + if cbi_eo == pad_slot_id: + return + else: + cbi_eo = pid_b_eo.to(tl.int64) + pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) + if (pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH) != WRITE_CHECKPOINT: + return + _replay_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + pad_slot_id, + T, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + USE_PERM, + REVERSE_PERM, + WRITE_CHECKPOINT, + ) + + +# Rectangle precompute kernel: produces a (T, K) CB rectangle that combines +# old-token (B from cache, k ∈ [0, PNAT)) and new-token (B from input, k ∈ +# [MAX-T, MAX) at compile-time-static shift) contributions in a single matmul. +# Used only on no-checkpoint steps (nowrite path); pairs with +# `_rectangle_main_kernel`. K-axis size = max(np2(MAX_REPLAY_BUFFER_LENGTH), +# 16); the static layout is sound because nowrite implies PNAT + T <= +# MAX_REPLAY_BUFFER_LENGTH, so old [0, PNAT) and new [MAX-T, MAX) never +# overlap. Also folds total_decay into decay_vec at precomp time so main +# can skip materializing a state_prev_decayed (M, dstate) tile. + + +@triton.jit() +def _rectangle_precompute_impl( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle + decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) + # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite + # path: read from buf_active at [0, PNAT), write new tokens at + # [PNAT, PNAT+T) of buf_active (same buffer). + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + # Slot permutation: see _replay_precompute_impl for semantics. + slot_perm_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # rectangle K-axis bound + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides (rectangle: (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides: (cache, 2, T_max, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T_max) + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T_max) + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Slot permutation flags — see _replay_precompute_impl. + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, +): + pid_grid = tl.program_id(axis=0) + # REVERSE_PERM walks the grid tail-first regardless of USE_PERM — + # combined with hardcode-sorted prev_tokens (kernel-side USE_PERM=False + # but PNAT pre-sorted write-first), reverse traversal makes the nowrite + # half front-load real work. + pid_grid_eff = (tl.num_programs(axis=0) - 1 - pid_grid) if REVERSE_PERM else pid_grid + if USE_PERM: + pid_b = tl.load(slot_perm_ptr + pid_grid_eff) + else: + pid_b = pid_grid_eff + pid_hg = tl.program_id(axis=1) + first_head = pid_hg * HEADS_PER_BLOCK + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_buf = buf_active + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout: old at [0, PNAT) (mask is_old_k); + # new at [MAX-T, MAX) at compile-time shift K_NEW_SHIFT = MAX - T. + # PNAT + T <= MAX is guaranteed on the nowrite path → no overlap. + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis (output rows) + offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis (rectangle input cols) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + t_mask = offs_t < T + n_mask = offs_n < dstate + + # K-axis masks (approach C: runtime PNAT-offset instead of K_NEW_SHIFT) + # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Loop 1: per-head dt processing. dt → dt_processed → dA_cumsum → + # decay_vec_new (= exp(cumAdt_new)). Stored to write_buf for next step. + # decay_vec_full (= total_decay * decay_vec_new) is finalized in loop 2 + # once total_decay is loaded; loop 1 stores raw decay_vec_new to scratch. + for h_local in range(HEADS_PER_BLOCK): + head_idx = first_head + h_local + + dt_base = dt_ptr + pid_b * stride_dt_batch + head_idx * stride_dt_head + dt = tl.load(dt_base + offs_t * stride_dt_T, mask=t_mask, other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + head_idx * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias + if DT_SOFTPLUS: + dt = softplus(dt) + + A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) + dA_cumsum = tl.cumsum(A * dt, axis=0) + + # Store dt and dA_cumsum to write_buf at [write_offset, write_offset+T) + # for next step's replay/rectangle use. + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + tl.store( + old_dt_base + (write_offset + offs_t) * stride_old_dt_T, + dt, + mask=t_mask, + ) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + tl.store( + old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, + dA_cumsum, + mask=t_mask, + ) + + # ---- Hoisted: cache-only loads independent of conv1d ---- + # old_B (group-level, BLOCK_K × BLOCK_DSTATE = ~8KB tile) and the + # decay_vec_full per-head pre-compute (which writes to DRAM and doesn't + # need cross-gdc_wait variables) are issued BEFORE gdc_wait so their + # HBM latency overlaps with conv1d. Per-head factor_dt/exp_diff stay + # below gdc_wait — they need cross-iteration spans, which Triton can't + # express without a DRAM round-trip; the per-head LOADS in the post- + # wait loop are small and cheap, so leave them. + group_idx = first_head // nheads_ngroups_ratio + + # Group-level: old B from active buffer at [0, PNAT) of the K-axis. + old_B_read_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + buf_active * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_load = tl.load( + old_B_read_base + + safe_old_k[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + mask=is_old_k[:, None] & n_mask[None, :], + other=0.0, + ) + + # Pre-wait: vectorized across HEADS_PER_BLOCK heads. Compute decay_vec_full + # (H, T) and combo = factor_dt * exp_diff (H, T, K). Store decay_vec_full; + # combo_block stays in registers across gdc_wait — used directly post-wait + # to compute rect_CB_scaled without a global memory roundtrip. + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) + + # Per-head bases (H,) — broadcast with offs_k or offs_t for 2D loads. + old_dt_read_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_active * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_read_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + ) + old_dt_write_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_write_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + ) + + # (H, K) loads at [0, PNAT) — old data from previous step. + hk_mask = is_old_k[None, :] # (1, K) + old_dt_all = tl.load( + old_dt_read_h[:, None] + safe_old_k[None, :] * stride_old_dt_T, + mask=hk_mask, other=0.0, + ).to(tl.float32) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, + mask=hk_mask, other=0.0, + ).to(tl.float32) + # (H,) scalar-per-head: total_dA_cumsum at prev_k_idx. + total_dA_cumsum = tl.load( + old_dA_cumsum_read_h + prev_k_idx * stride_old_dA_cumsum_T + ).to(tl.float32) + # (H, T) loads at [PNAT, PNAT+T) — this step's dA_cumsum_new from loop 1. + ht_mask = t_mask[None, :] # (1, T) + dA_cumsum_new = tl.load( + old_dA_cumsum_write_h[:, None] + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T, + mask=ht_mask, other=0.0, + ).to(tl.float32) + # (H, K) loads at K_NEW_SHIFT-shifted positions for new tokens. + hkn_mask = is_new_k[None, :] + dt_at_kn = tl.load( + old_dt_write_h[:, None] + + (write_offset + safe_k_new)[None, :] * stride_old_dt_T, + mask=hkn_mask, other=0.0, + ).to(tl.float32) + dA_cumsum_at_kn = tl.load( + old_dA_cumsum_write_h[:, None] + + (write_offset + safe_k_new)[None, :] * stride_old_dA_cumsum_T, + mask=hkn_mask, other=0.0, + ).to(tl.float32) + + # decay_vec_full = total_decay * exp(cumAdt_new). (H, T). + total_decay = tl.where( + prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0 + ) # (H,) + decay_vec_full_block = total_decay[:, None] * tl.exp(dA_cumsum_new) # (H, T) + decay_vec_addrs = ( + decay_vec_ptr + + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) # (H, T) + tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) + + # combo_block = factor_dt * exp_diff — (H, T, K). Stays in registers + # across gdc_wait. + factor_dt = tl.where(is_old_k[None, :], old_dt_all, dt_at_kn) # (H, K) + s_k = tl.where( + is_old_k[None, :], + total_dA_cumsum[:, None] - old_dA_cumsum_all, + -dA_cumsum_at_kn, + ) # (H, K) + # exp_diff (H, T, K) = exp(s_k (H, 1, K) + dA_cumsum_new (H, T, 1)). + exp_diff = tl.exp(s_k[:, None, :] + dA_cumsum_new[:, :, None]) + combo_block = factor_dt[:, None, :] * exp_diff # (H, T, K) + + # ---- gdc_wait: from here on we depend on conv1d's outputs ---- + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # Conv1d outputs: B and C + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + B_new_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_all = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_new_orig = tl.load( + B_new_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_new_shifted = tl.load( + B_new_base + safe_k_new[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=is_new_k[:, None] & n_mask[None, :], + other=0.0, + ) + # Disjoint masks: old at [0, PNAT), new at [K_NEW_SHIFT, K_NEW_SHIFT+T). + B_combined = old_B_load + B_new_shifted + raw_rect_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_combined).to(tl.bfloat16)) + + # Append new B to cache at [PNAT, PNAT+T) of write_buf (once per group). + if first_head % nheads_ngroups_ratio == 0: + old_B_write_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + write_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_write_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_new_orig, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # Causal mask (BLOCK_SIZE_T × BLOCK_SIZE_K, shared across heads). + # Approach C: new tokens at runtime [PNAT, PNAT+T) instead of K_NEW_SHIFT. + t_idx_2d = offs_t[:, None] + k_idx_2d = offs_k[None, :] + is_old_k_2d = k_idx_2d < prev_num_accepted_tokens + k_new_idx_2d = k_idx_2d - prev_num_accepted_tokens + is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) + causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] + + # Post-wait vectorized: combo_block (H, T, K) is still live in registers. + # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store as + # one (H, T, K) tile. + rect_CB_scaled_block = tl.where( + causal_combined[None, :, :], + raw_rect_CB[None, :, :] * combo_block, + 0.0, + ) # (H, T, K) + cb_scaled_addrs = ( + cb_scaled_ptr + + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_k[None, None, :] * stride_cb_j + ) # (H, T, K) + cb_store_mask_3d = ( + (offs_t[None, :, None] < BLOCK_SIZE_T) + & (offs_k[None, None, :] < BLOCK_SIZE_K) + ) # (1, T, K) → broadcasts to (H, T, K) + tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) + + +# Rectangle precompute kernel. Thin wrapper around _rectangle_precompute_impl +# that carries the @triton.heuristics for constexpr derivation; called from +# the Python wrapper on the rectangle nowrite path. +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _rectangle_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + EARLY_OUT: tl.constexpr, + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does, so + # main can start its setup regardless of how this program ends. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + # Per-program early-out gate. Rectangle is nowrite-only, so EARLY_OUT + # skips slots whose PNAT + T > MAX (slots that would need write). + if EARLY_OUT: + pid_grid_eo = tl.program_id(axis=0) + pid_grid_eo_eff = (tl.num_programs(axis=0) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo + if USE_PERM: + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) + else: + pid_b_eo = pid_grid_eo_eff + if HAS_CACHE_BATCH_INDICES: + cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) + if cbi_eo == pad_slot_id: + return + else: + cbi_eo = pid_b_eo.to(tl.int64) + pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) + if pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH: + return + _rectangle_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + USE_PERM, + REVERSE_PERM, + ) + + +# Dynamic precompute kernel. Single launchable kernel that, per program, +# reads PNAT and dispatches to one of the existing impls: +# +# if pnat + T > MAX: replay_precompute_impl(WRITE_CHECKPOINT=True) +# else if RECTANGLE: rectangle_precompute_impl +# else: replay_precompute_impl(WRITE_CHECKPOINT=False) +# +# RECTANGLE is constexpr (compile-time tuning param); the inner branch +# is folded so only one of the two nowrite paths is emitted per +# specialization. Reg envelope = max(replay_write, X) where X depends +# on RECTANGLE. cb_scaled is allocated (T, K) by the wrapper regardless; +# replay paths write to the first T columns, rectangle writes the full K. +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _dynamic_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. + RECTANGLE: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + pid_b = tl.program_id(axis=0) + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH + # write_checkpoint is now runtime in replay precompute, so a single + # call site handles both write and nowrite for the replay branch. + # Take rectangle only when RECTANGLE is True AND this slot doesn't + # need write; everything else funnels into replay. + if needs_write_runtime or not RECTANGLE: + _replay_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) + pad_slot_id, + T, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + False, # USE_PERM + False, # REVERSE_PERM + needs_write_runtime, + ) + else: + _rectangle_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + False, # USE_PERM + False, # REVERSE_PERM + ) + + +# Main kernel: tl.dot replay + precomputed CB output. +# Grid: (cdiv(dim, M), batch, nheads). + + +@triton.jit() +def _replay_main_impl( + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or + # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor + # USE_TMA_STORE is enabled (kernel ignores it via constexpr). + state_tma_descriptor, + # Per-(cache, head, dim) decode scale, fp32, only consulted when QUANT_MAX>0. + # Layout (cache, nheads, dim) — broadcast over dstate at load/store. + state_scales_ptr, + # Cache READ pointers (read-buffer from previous step) + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + # Cache WRITE pointer (write-buffer for old_x only; B/dt/dA_cumsum written by precompute) + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + # New input pointers + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + # Precomputed pointers + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + # Slot permutation: see _replay_precompute_impl for semantics. + slot_perm_ptr, + # Stochastic rounding + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # cache T-axis capacity (= max_window) + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides: (cache, nheads, dim) — only used when QUANT_MAX>0 + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides: (cache, T, nheads, dim) — single-buffered + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides: (cache, 2, T, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + # State quantization: 0.0 means non-quantized (fp16/bf16/fp32); >0 means + # quantized (int8=127, int16=32767, fp8_e4m3fn=448). Single in-kernel + # switch for the dequant-on-load and encode-on-store paths. Wrapper sets + # this from state.dtype; kernel-entry static_assert below pins the + # invariant that it must coincide with int8/int16/float8e4nv state dtype. + QUANT_MAX: tl.constexpr, + # Checkpointing flag + WRITE_CHECKPOINT: tl.constexpr, # When True: quantize+write post-replay state to HBM (checkpoint step). + # When False: skip state write entirely (non-checkpoint step). + # The rectangle non-checkpoint path is implemented in + # _rectangle_main_kernel (separate kernel pair, picked by + # the wrapper via rectangle_for_nowrite=True). + # When True: signal PDL dependents at the very top of every program + # (including pad/early-out programs). Used by doublelaunch and maindl + # so the next kernel (the second main, or the second precompute) can + # start its setup while this main is still computing. Default False + # for monolithic / dynamic / the LAST main in dl/maindl chains. + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + # Slot permutation flags — see _replay_precompute_impl. + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, + # TMA toggles — gated per-path inside the body (write-load picks + # USE_TMA_LOAD_WRITE; nowrite-load picks USE_TMA_LOAD_NOWRITE; store + # only fires on the write path and uses USE_TMA_STORE). The wrapper + # passes write_load_value when WC=True and nowrite_load_value when + # WC=False; the unused flag is dummy False. Both are constexpr; + # is_write here is constexpr (= WRITE_CHECKPOINT), so use_tma_load + # constexpr-folds. + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # Hoisted PDL signal: fire as the first thing every program does, so + # downstream kernels can start setup regardless of how this program + # ends (pad / early-out / full body). + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized + # state dtype (int8 / int16 / float8e4nv) and only those. Cheap + # insurance against a wrapper bug that desynchronizes the two. + tl.static_assert( + (QUANT_MAX > 0.0) + == ( + (state_ptr.dtype.element_ty == tl.int8) + or (state_ptr.dtype.element_ty == tl.int16) + or (state_ptr.dtype.element_ty == tl.float8e4nv) + ), + "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", + ) + + pid_m = tl.program_id(axis=0) + pid_grid_b = tl.program_id(axis=1) + pid_grid_b_eff = (tl.num_programs(axis=1) - 1 - pid_grid_b) if REVERSE_PERM else pid_grid_b + if USE_PERM: + pid_b = tl.load(slot_perm_ptr + pid_grid_b_eff) + else: + pid_b = pid_grid_b_eff + pid_h = tl.program_id(axis=2) + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Active buffer (= cache_buf_idx) holds the historical inputs for this + # step at [0, PNAT). The replay phase reads from there. The new-tokens + # write target depends on WRITE_CHECKPOINT — see Cache write semantics + # block in the precompute kernel for the full rationale. + active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if WRITE_CHECKPOINT: + write_buf = 1 - active_buf # noqa: F841 — old_x is single-buffered (no use here) + write_offset = 0 + else: + write_buf = active_buf # noqa: F841 — old_x is single-buffered (no use here) + write_offset = prev_num_accepted_tokens + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + # Replay axis: separate from offs_t. Spans [0, BLOCK_SIZE_WINDOW) ⊇ + # [0, MAX_REPLAY_BUFFER_LENGTH); used for old-token loads (mask: offs_window < PNAT). + offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # Load state. state_tma_descriptor is a host-built tensor_descriptor + # over a flat (cache*nheads*dim, dstate) view of state when any TMA + # path is enabled; raw `state_ptr` is the underlying tensor and is + # always passed. state_ptrs / state_ptr_raw are the raw-pointer view + # used for !TMA load and store paths. offs_y is the flat row index + # for TMA load/store; computed unconditionally (cheap int math; DCE'd + # when no TMA path is reachable). + state_mask = m_mask[:, None] & n_mask[None, :] + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + # Pick LOAD flag based on WRITE_CHECKPOINT (constexpr). Use an if/else + # with tl.constexpr annotations on each branch — a plain ternary binds + # the result to a Python (non-constexpr) name and the `if use_tma_load` + # gate below becomes a runtime branch, which forces BOTH the + # `state_tma_descriptor.load(...)` and `tl.load(state_ptrs, ...)` paths + # to compile. When TMA is off, the wrapper passes the plain state + # tensor as state_tma_descriptor (no `.load()` method) → compile fails. + # Wrapper passes USE_TMA_LOAD_WRITE = write-load value when WC=True + # (NOWRITE flag is dummy False then), and the converse when WC=False. + if WRITE_CHECKPOINT: + use_tma_load: tl.constexpr = USE_TMA_LOAD_WRITE + else: + use_tma_load: tl.constexpr = USE_TMA_LOAD_NOWRITE + if use_tma_load: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + # Dequantize on load (per-(head, dim) decode scale, broadcast over dstate). + # Only consulted when QUANT_MAX>0 — non-quantized paths skip entirely. + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, + other=1.0, + ).to(tl.float32) + state = state * decode_scale[:, None] + + # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) + group_idx = pid_h // nheads_ngroups_ratio + + # Old-token mask along the WINDOW (replay) axis. PNAT ≤ MAX ≤ + # BLOCK_SIZE_WINDOW, so this enables all valid old-cache positions. + # (Distinct from t_mask = offs_t < T which gates output T-rows only.) + old_window_mask = offs_window < prev_num_accepted_tokens + + # Load precomputed dt and dA_cumsum from READ buffer at [0, PNAT). + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + active_buf * stride_old_dt_dbuf + + pid_h * stride_old_dt_head + ) + old_dt_all = tl.load( + old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 + ).to(tl.float32) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + active_buf * stride_old_dA_cumsum_dbuf + + pid_h * stride_old_dA_cumsum_head + ) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, + mask=old_window_mask, other=0.0, + ).to(tl.float32) + + # Load dA_cumsum at prev_k-1 directly via pointer math (avoids masked reduction). + # Clamp to [0, MAX-1] defensively — caller contract gives PNAT ≤ MAX. + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( + tl.float32 + ) + + # Step 0 invariant: PNAT=0 means `state` is already last step's state (not + # two back). coeff is all-zero (old_window_mask all-false), total_decay + # is 1.0, so the replay leaves `state` unchanged — cache contents don't matter. + coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all + + # Load old_x at [0, PNAT) of the WINDOW axis (single-buffered cache). + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + old_x_all = tl.load( + old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + mask=old_window_mask[:, None] & m_mask[None, :], + other=0.0, + ) + + # Load old_B from READ buffer at [0, PNAT) of the WINDOW axis. + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + active_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_all = tl.load( + old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + mask=old_window_mask[:, None] & n_mask[None, :], + other=0.0, + ).to(tl.float32) + + # Scale B by coefficients + dB_scaled = coeff[:, None] * old_B_all + + # Apply total decay to initial state FIRST, then add contributions + total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) + state *= total_decay + + # tl.dot fast-forward: old_x^T @ dB_scaled → (M, dstate) + state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) + + # Write post-replay state — only on checkpoint steps. When + # WRITE_CHECKPOINT is False, the replay computed `state` is local-only and + # discarded; skipping the HBM store + Philox path is the main performance + # win of replay-style checkpointing on the common (non-checkpoint) step. + if WRITE_CHECKPOINT: + if USE_RS_ROUNDING: + # Generate random tensor for stochastic rounding. The amount of + # randomness needed depends on the SR codegen path: + # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) + # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) + # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs + # int16 SR (24b + bitrev32): 1 b32 per 2 outputs + # The PTX cvt.rs.* instructions consume a single 32-bit random + # and split the bits internally for 2 or 4 conversions. The + # tl.inline_asm_elementwise wrapper has uniform `pack` across all + # args, so it provides 2 (fp16) or 4 (fp8) rand inputs per asm + # call but only the first is read; the others are dead. Generate + # only what's actually consumed and broadcast to fill the unused + # slots — saves Philox rounds proportionally. + if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: + RAND_DIVISOR: tl.constexpr = 4 # fp8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: + RAND_DIVISOR: tl.constexpr = 4 # int8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: + RAND_DIVISOR: tl.constexpr = 2 # int16 SR + elif QUANT_MAX == 0.0: + RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) + else: + RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized + + rand_seed = tl.load(rand_seed_ptr) + base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head + # Number of unique randoms per row = dstate / RAND_DIVISOR. + # randint4x emits 4 randoms per offset, so use that / 4 offsets. + offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) + rand_offsets_q = ( + base_rand + + offs_m[:, None] * stride_state_dim + + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) + ) # (M, dstate / (4*RAND_DIVISOR)) + if PHILOX_ROUNDS > 0: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + else: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) + r01 = tl.join(r0, r1) + r23 = tl.join(r2, r3) + r0123 = tl.join(r01, r23) + rand_compact = tl.reshape( + r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) + ) + # Broadcast each unique rand to RAND_DIVISOR adjacent positions + # in the dstate axis. Pack-group (pack=2 fp16 / pack=4 fp8) + # consumes adjacent positions; the unique rand lands at the + # asm's read slot ($3 fp16 / $5 fp8); duplicates feed the dead + # slots ($4 fp16; $6/$7/$8 fp8). Triton's broadcast_to is + # stride-0 in IR. + # + # Tested zero-fill alternative (tl.join with zeros): essentially + # equivalent register count (95 vs 96 at one config) and same + # timing. ptxas does not use RZ for the dead asm input slots + # in either case; the extra ~15 regs vs pre-fix come from + # rand_compact's lifetime across the asm call, not from the + # fill pattern. Broadcast wins on simplicity. + if RAND_DIVISOR > 1: + rand_3d = rand_compact[:, :, None] + rand_3d = tl.broadcast_to( + rand_3d, + (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), + ) + rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + else: + rand = rand_compact + + if QUANT_MAX > 0.0: + # Quantized state path: int8 / int16 / fp8_e4m3fn (RN or SR). + # 1) Per-(head, dim) channel scale via amax over dstate. + amax = tl.max(tl.abs(state), axis=1) # (M,) + encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) # (M,) + decode_scale = 1.0 / encode_scale # (M,) + # 2) Store decode_scale (1/encode) so reads do a single multiply. + state_scales_ptrs = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + + offs_m * stride_state_scales_dim + ) + tl.store(state_scales_ptrs, decode_scale, mask=m_mask) + # 3) Scale state into quant range — into a NEW variable so the + # downstream output phase still sees the dequantized fp32 state. + state_q = state * encode_scale[:, None] + # 4) Round per dtype. Order matters: handle fp8 SR first (PTX + # combines round + saturating cast in one op, output is final fp8 + # so we store and finish on that branch). Other branches share + # the clamp + cast tail below. + if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): + # fp8_e4m3fn + SR — PTX cvt.rs.satfinite.e4m3x4.f32. Output + # is final fp8 (saturate included); store directly. + _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_fp8sr) + else: + tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) + else: + if USE_RS_ROUNDING: + # int8 / int16 + SR — uniform-noise + floor. + # int8 packs 4 values per random u32 using 16-bit chunks + # and bitrev16; int16 packs 2 values per random u32 using + # 24-bit uniforms from the direct/reversed u32. + # (fp8 SR was handled by the early branch above.) + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized SR fall-through expects int8 or int16; " + "fp8 SR is handled by the prior branch.", + ) + if state_ptrs.dtype.element_ty == tl.int8: + state_q = _stochastic_round_int8_packed( + state_q, rand, offs_n[None, :] + ) + else: + state_q = _stochastic_round_int16_packed( + state_q, rand, offs_n[None, :] + ) + elif state_ptrs.dtype.element_ty != tl.float8e4nv: + # int8 / int16 + RN — explicit round before clamp. + # fp8 + RN deliberately skips this — explicit round() would + # destroy fp8 sub-integer precision; native cast at store + # does RN at the fp8 grid resolution. + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized RN with explicit round() expects int8 or int16.", + ) + state_q = tl.extra.cuda.libdevice.round(state_q) + # Clamp + cast tail: int8/int16 (RN+SR) and fp8 RN. + # fp8 RN reaches here without prior round() — .to(float8e4nv) + # does native RN at the fp8 grid. + state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) + _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_cast) + else: + tl.store(state_ptrs, _state_q_cast, mask=state_mask) + elif USE_RS_ROUNDING: + # Non-quantized + SR: only fp16 (bf16 has no PTX SR cast; fp32 + # doesn't need rounding). + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "Non-quantized SR only supports fp16 state.", + ) + _state_sr = _stochastic_round_fp16x2(state, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_sr) + else: + tl.store(state_ptrs, _state_sr, mask=state_mask) + else: + # Non-quantized + RN: fp16 / bf16 / fp32 native cast. + _state_cast = state.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_cast) + else: + tl.store(state_ptrs, _state_cast, mask=state_mask) + + # Phase 2: Output using precomputed CB_scaled and decay_vec + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + # Wait for precompute kernel (PDL) before reading its outputs. + # With chained PDL (conv1d → precompute → main), gdc_wait() ensures + # precompute has completed — which transitively ensures conv1d has + # completed (precompute waited on conv1d via its own gdc_wait). + # All loads below (x, C from conv1d; CB_scaled, decay_vec from precompute) + # are safe after this point. + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # Load conv1d outputs: C_all and x_all + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + x_all = tl.load( + x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + # Store new x to old_x cache at [write_offset : write_offset+T). + # old_x is single-buffered: write goes to the active buffer regardless; + # replay already read positions [0, PNAT) so write_offset = PNAT (no + # overlap) on no-checkpoint steps. On checkpoint steps write_offset = 0 + # (cache reset; old data folded into state via replay update). + tl.store( + old_x_base + + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_all, + mask=t_mask[:, None] & m_mask[None, :], + ) + x_all = x_all.to(tl.float32) + + # Load precomputed CB_scaled and decay_vec + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + tl.float32 + ) + + # init_out = C_all @ state^T * decay_vec + init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] + + # cb_out = CB_scaled @ x_all + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) + + out_all = init_out + cb_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# Replay-style main kernel. Thin wrapper around _replay_main_impl that carries +# the @triton.heuristics for constexpr derivation; called from the Python +# wrapper on the replay-style path (write or replay-nowrite). +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _checkpointing_main_kernel( + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or + # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor + # USE_TMA_STORE is enabled (kernel ignores it via constexpr). + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + EARLY_OUT: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, + # 3 TMA flags passed through to _replay_main_impl (which picks + # USE_TMA_LOAD_WRITE vs NOWRITE based on WRITE_CHECKPOINT). Wrapper + # passes write_load_value when WC=True (NOWRITE flag dummy False), + # and the converse when WC=False. + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # Hoisted PDL signal: fire as the first thing every program does. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + # Per-program early-out gate. Signal-then-skip lets the next kernel + # in dl/maindl chains start regardless of early-out outcome. + if EARLY_OUT: + pid_grid_eo = tl.program_id(axis=1) + pid_grid_eo_eff = (tl.num_programs(axis=1) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo + if USE_PERM: + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) + else: + pid_b_eo = pid_grid_eo_eff + if HAS_CACHE_BATCH_INDICES: + cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) + if cbi_eo == pad_slot_id: + return + else: + cbi_eo = pid_b_eo.to(tl.int64) + pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) + if (pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH) != WRITE_CHECKPOINT: + return + _replay_main_impl( + state_ptr, + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + rand_seed_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + WRITE_CHECKPOINT, + LAUNCH_DEPENDENT_KERNELS, + USE_PERM, + REVERSE_PERM, + USE_TMA_LOAD_WRITE, + USE_TMA_LOAD_NOWRITE, + USE_TMA_STORE, + ) + + +# Rectangle main kernel (nowrite-only): no replay step, no state HBM write, +# no SR codegen. state_out is computed from state_prev directly using the +# precomp-folded decay_vec_full; token_out is a single rectangle matmul over +# the (T, K) CB rectangle and (K, M) x_combined. Pairs with +# `_rectangle_precompute_kernel`. + + +@triton.jit() +def _rectangle_main_impl( + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as + # replay path). Used when USE_TMA_LOAD; ignored otherwise (kernel + # branches via constexpr). Wrapper passes the same descriptor as + # for replay paths — single underlying memory, consumed by per-path + # constexpr gates. + state_tma_descriptor, + state_scales_ptr, # only consulted when QUANT_MAX > 0 + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, # rectangle (batch, nheads, T, K) + decay_vec_ptr, # folded (batch, nheads, T) — total_decay * exp(cumAdt_new[t]) + state_batch_indices_ptr, + # Slot permutation: see _replay_precompute_impl for semantics. + slot_perm_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides (no quant-store path; state read-only for state_out) + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides: (cache, nheads, dim) — fp32, broadcast over dstate + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides: (cache, T_max, nheads, dim) — single-buffered + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides (rectangle (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + QUANT_MAX: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + # Slot permutation flags — see _replay_precompute_impl. + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, + USE_TMA_LOAD: tl.constexpr = False, +): + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + pid_m = tl.program_id(axis=0) + pid_grid_b = tl.program_id(axis=1) + pid_grid_b_eff = (tl.num_programs(axis=1) - 1 - pid_grid_b) if REVERSE_PERM else pid_grid_b + if USE_PERM: + pid_b = tl.load(slot_perm_ptr + pid_grid_b_eff) + else: + pid_b = pid_grid_b_eff + pid_h = tl.program_id(axis=2) + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout (matches precompute). + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_k = tl.arange(0, BLOCK_SIZE_K) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # K-axis masks (approach C: PNAT-runtime offset, matches precompute). + # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Load state. Read-only — no HBM write on the nowrite path. + # Quant scale hoist (backlog #16): for QUANT_MAX > 0 paths, defer the + # `* decode_scale` to AFTER the C @ state dot — applied to the (T, M) + # dot output instead of broadcast-multiplied into the (M, dstate) state + # tile. Algebra-equivalent (decode_scale is per-M, commutes with the + # matmul over dstate). Saves M·dstate fp32 muls (replaced by T·M), + # but the bigger potential win is shorter register lifetime for state + # (kept as native int8/int16/fp8 until just before the dot, where Triton + # casts to bf16 — vs current fp32 tile across the whole kernel). Only + # applies in rectangle main (no `state += dot` here). + if USE_TMA_LOAD: + # TMA descriptor (host-built) over the flat 2D view of state: + # shape=[cache_size * nheads * dim, dstate], strides=[dstate, 1]. + # Convert (cache, head, m) → flat row index using existing strides: + # rows-per-cache-slot = stride_state_batch / stride_state_dim + # rows-per-head = stride_state_head / stride_state_dim = dim (constexpr) + # rows-per-m = 1 + # Diagnostic: prior in-kernel descriptor attempts emitted + # ttng.tensormap_create setup (divergent shared-mem write) which + # blows up branch count; host-built descriptors avoid that. + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state = state_tma_descriptor.load([offs_y, 0]) + else: + state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, other=1.0, + ).to(tl.float32) + # state stays in native quant dtype — cast happens inside the dot below. + else: + state = state.to(tl.float32) + + # Group / pointer offset setup + group_idx = pid_h // nheads_ngroups_ratio + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + # Hoist: old_x from cache doesn't depend on conv1d/precompute, so issue + # the load BEFORE gdc_wait so its HBM latency overlaps with conv1d. + old_x_load = tl.load( + old_x_base + + safe_old_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + mask=is_old_k[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + + # PDL gate: precompute outputs (cb_scaled, decay_vec_full) become safe + # after gdc_wait. conv1d outputs (x, C) also gated by the chained PDL. + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # Load C and x (conv1d outputs after PDL wait) + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + # Single (BLOCK_K, M) load at PNAT-offset positions (approach C): + # K-axis [PNAT, PNAT+T) gets new tokens directly from x[0:T, :] via + # safe_k_new = offs_k - PNAT. Cache layout matches K-axis layout, so + # one load serves both matmul and cache write. + x_K = tl.load( + x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=is_new_k[:, None] & m_mask[None, :], + other=0.0, + ) + # Cache write: store at offs_k directly (is_new_k mask makes offs_k land + # at [PNAT, PNAT+T) in the cache, which is exactly write_offset+0..T-1). + tl.store( + old_x_base + + offs_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_K, + mask=is_new_k[:, None] & m_mask[None, :], + ) + + x_K_f32 = x_K.to(tl.float32) + # Matmul side: K-axis aligned; sum with old_x_load. + x_combined = old_x_load + x_K_f32 + + # T-axis view for D feedthrough / Z-gating: extract via (T, K) selection. + # Only materialized when needed. + if HAS_D or HAS_Z: + sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) + x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) + else: + x_all = x_K_f32 # placeholder; unused + + # Load precomputed rectangle CB and folded decay_vec. + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec_full = tl.load( + decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 + ).to(tl.float32) + + # state_out: state_prev contribution to output, with decay folded post-matmul. + # No state_prev_decayed (M, dstate) materialization — state is consumed + # directly by the matmul, then decay_vec_full multiplies the (T, M) result. + # For QUANT_MAX > 0 (#16 hoist): decode_scale also applies post-matmul + # at (T, M) granularity instead of pre-multiplied into the (M, dstate) + # state tile. Triton's tl.dot(a.to(bf16), b.to(bf16)) handles the + # int8/fp8 → bf16 cast inside the dot's input prep. + state_out = ( + tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) + * decay_vec_full[:, None] + ) + if QUANT_MAX > 0.0: + state_out = state_out * decode_scale[None, :] + + # token_out: combined old + new tokens contribution via the rectangle. + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) + + out_all = state_out + token_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# Rectangle main kernel. Thin wrapper around _rectangle_main_impl that carries +# the @triton.heuristics for constexpr derivation; called from the Python +# wrapper on the rectangle nowrite path. +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _rectangle_main_kernel( + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor for state (used when + # USE_TMA_LOAD). Same descriptor as replay paths use; gate via + # constexpr. Wrapper passes the unified state_tma_descriptor. + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + QUANT_MAX: tl.constexpr, + EARLY_OUT: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + USE_PERM: tl.constexpr, + REVERSE_PERM: tl.constexpr, + USE_TMA_LOAD: tl.constexpr = False, +): + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + # Per-program early-out gate. Rectangle is nowrite-only. + if EARLY_OUT: + pid_grid_eo = tl.program_id(axis=1) + pid_grid_eo_eff = (tl.num_programs(axis=1) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo + if USE_PERM: + pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) + else: + pid_b_eo = pid_grid_eo_eff + if HAS_CACHE_BATCH_INDICES: + cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) + if cbi_eo == pad_slot_id: + return + else: + cbi_eo = pid_b_eo.to(tl.int64) + pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) + if pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH: + return + _rectangle_main_impl( + state_ptr, + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + QUANT_MAX, + LAUNCH_DEPENDENT_KERNELS, + USE_PERM, + REVERSE_PERM, + USE_TMA_LOAD, + ) + + +# Dynamic main kernel. Single launchable kernel that, per program, reads +# PNAT and dispatches to one of the existing impls: +# +# if pnat + T > MAX: replay_main_impl(WRITE_CHECKPOINT=True) +# elif RECTANGLE (constexpr): rectangle_main_impl +# else: replay_main_impl(WRITE_CHECKPOINT=False) +# +# Unlike precompute, WRITE_CHECKPOINT stays constexpr in the main impl — +# the body has constexpr-gated state-write code (quant + Philox + HBM +# store) where folding meaningfully shrinks the codegen. So this kernel +# has TWO replay call sites (one per WRITE_CHECKPOINT specialization) +# both inlined, with a runtime branch picking which runs. Reg envelope = +# max(replay_write, X) where X = rectangle_nowrite (RECTANGLE=True) or +# replay_nowrite (RECTANGLE=False). cb_scaled is allocated (T, K) by the +# wrapper regardless. +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _dynamic_main_kernel( + # Pointers — union of replay-main and rectangle-main pointer args. + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as + # other main kernels). Currently always passed as dummy `state_ptr` + # by launch_dynamic_main since dynamic doesn't expose TMA toggles + # yet — kept in the signature for uniformity with replay/rect/persistent. + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides (replay only; passed but unused on rectangle path) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides (replay only) + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides (replay only) + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, # for replay path + BLOCK_SIZE_K: tl.constexpr, # for rectangle path + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. + RECTANGLE: tl.constexpr, + # Default False — dynamic main is normally terminal in its chain. + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + # 3 TMA toggles (matching dual-path kernel scheme): + # USE_TMA_LOAD_WRITE — write path's state load + # USE_TMA_LOAD_NOWRITE — nowrite path's state load (rect when + # RECTANGLE, else replay-nowrite) + # USE_TMA_STORE — write path's state store (no-op for nowrite) + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + pid_b = tl.program_id(axis=1) + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if pnat_local + T > MAX_REPLAY_BUFFER_LENGTH: + # Write slot — replay-style write (WRITE_CHECKPOINT=True constexpr). + _replay_main_impl( + state_ptr, + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) + rand_seed_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + True, # WRITE_CHECKPOINT (constexpr) + False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top + False, # USE_PERM + False, # REVERSE_PERM + USE_TMA_LOAD_WRITE, # write-load fires here + USE_TMA_LOAD_NOWRITE, # nowrite-load: dummy at this site + USE_TMA_STORE, # store fires (WC=True) + ) + else: + if RECTANGLE: + _rectangle_main_impl( + state_ptr, + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + QUANT_MAX, + False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top + False, # USE_PERM + False, # REVERSE_PERM + USE_TMA_LOAD_NOWRITE, # rect-load TMA flag + ) + else: + # Replay-style nowrite (WRITE_CHECKPOINT=False constexpr). + _replay_main_impl( + state_ptr, + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) + rand_seed_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + False, # WRITE_CHECKPOINT (constexpr) + False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top + False, # USE_PERM + False, # REVERSE_PERM + USE_TMA_LOAD_WRITE, # write-load: dummy at this site + USE_TMA_LOAD_NOWRITE, # nowrite-load fires here + USE_TMA_STORE, # store: dummy (WC=False) + ) + + +# Python wrapper + + +_QUANT_MAX_BY_DTYPE = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, +} + + +# ============================================================================ +# Persistent main kernel — 1D grid, persistent CTA loop with tl.range +# ============================================================================ +# +# Design (see ~/dev/scripts/mamba_replay/kernel_microbenchmarks/PERSISTENT_KERNELS.md +# for the full strawman): +# +# * Outer 1D grid of `NUM_PERSISTENT` CTAs (start at NUM_SMS, sweep upward). +# * Inside the kernel, a `tl.range(pid, total_work, NUM_PERSISTENT, flatten=True, +# num_stages=NUM_STAGES)` loop iterates over (slot, M_tile, head) work units. +# * Hard-sort PNAT host-side and pass `n_writes` as a runtime int32 scalar: +# the launcher invokes the kernel twice — once with slot_offset=0, +# n_slots=n_writes, WRITE_CHECKPOINT=True, and once with +# slot_offset=n_writes, n_slots=B-n_writes, WRITE_CHECKPOINT=False. +# * `_persistent_main_impl` is a copy of `_replay_main_impl`'s body with the +# program_id reads replaced by parameters and the slot_perm logic moved into +# the persistent loop wrapper. No code shared with the existing kernels; +# easy to delete if the experiment is abandoned. +# +# Notes: +# * `flatten=True` is canonical for Triton 3.6 persistent kernels (matches the +# upstream `_p_matmul_ogs.py` and tutorial 09). Combined with `num_stages=2` +# it pipelines the loop body — but watch open issue triton-lang/triton#8259 +# which reports this combo can corrupt stores in non-dot loops. First run +# correctness check is critical. +# * Warp specialization (`warp_specialize=True`) is NOT enabled — Triton 3.6 +# only supports it for simple matmul loops and our scan won't pattern-match. +# * No 2CTA cluster mode — that's dot-only per the kernel-tileir-optimization +# skill classification. + + +@triton.jit() +def _persistent_main_impl( + # Per-work-unit indices (computed by the persistent wrapper). + # `pid_b` is the post-perm slot index (caller has already applied any + # slot permutation and slot_offset). + pid_m, + pid_b, + pid_h, + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or + # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor + # USE_TMA_STORE is enabled (kernel ignores it via constexpr). + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + # IS_DYNAMIC: kept in the signature for caller-side bookkeeping (the + # outer _persistent_main_kernel still inspects it to decide the slot- + # IS_DYNAMIC: when True (persistent_dynamic), is_write is per-slot from + # PNAT. When False (persistent_main), is_write is constexpr from + # WRITE_CHECKPOINT. See also WC_IS_CONSTEXPR below. + IS_DYNAMIC: tl.constexpr, + # WC_IS_CONSTEXPR: when True, force is_write = WRITE_CHECKPOINT (constexpr) + # regardless of IS_DYNAMIC. Callers in RECT=1 use this in the is_w=True + # arm of _persistent_main_kernel (we know all slots that reach this call + # need is_write=True because is_w was the PNAT-derived runtime check, and + # this arm only fires when is_w is True). Passing WRITE_CHECKPOINT=True + # as a literal at the call site + WC_IS_CONSTEXPR=True here lets the inner + # body DCE the nowrite path under IS_DYNAMIC=True too — same codegen + # quality as persistent_main mode (-3.7% measured at b=1024 dyn-shape). + # When False (RECT=0 callers, where both write and nowrite slots are + # dispatched to ONE call), use the original runtime is_write under + # IS_DYNAMIC=True; avoids the binary-doubling regression that two + # specialized calls would cause. + WC_IS_CONSTEXPR: tl.constexpr = False, + # TMA flags — picked inside body based on is_write. When is_write is + # constexpr (either IS_DYNAMIC=False or WC_IS_CONSTEXPR=True), the + # use_tma_load = USE_TMA_LOAD_WRITE if is_write else USE_TMA_LOAD_NOWRITE + # ternary constexpr-folds and only one TMA load form survives. + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # IS_DYNAMIC: kernel-mode label, used by the OUTER _persistent_main_kernel + # to decide slot-range derivation and outer is_w dispatch strategy + # (constexpr WC for persistent_main; runtime is_w split -> 2 specialized + # impl calls for persistent_dynamic). Inside this impl, IS_DYNAMIC is + # NOT consulted at runtime -- WRITE_CHECKPOINT is the only constexpr that + # gates the write/nowrite codegen, in BOTH modes. + + # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized + # state dtype (int8 / int16 / float8e4nv) and only those. + tl.static_assert( + (QUANT_MAX > 0.0) + == ( + (state_ptr.dtype.element_ty == tl.int8) + or (state_ptr.dtype.element_ty == tl.int16) + or (state_ptr.dtype.element_ty == tl.float8e4nv) + ), + "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", + ) + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + # Resolve is_write: see WC_IS_CONSTEXPR / IS_DYNAMIC docs in the param + # list above. Three cases: + # - WC_IS_CONSTEXPR=True (RECT=1 is_w=True arm callers): use WC + # constexpr. Caller knows the slot needs write; inner DCEs nowrite + # paths. Avoids the binary-doubling overhead that calling the impl + # twice would cause, while still constexpr-DCEing the nowrite half. + # - IS_DYNAMIC=True (RECT=0 caller, persistent_dynamic): runtime + # branch on PNAT. Both write and nowrite codegen live in one body + # (no bloat) — same as the pre-refactor behavior. + # - IS_DYNAMIC=False (persistent_main): WC constexpr from caller. + if WC_IS_CONSTEXPR: + is_write: tl.constexpr = WRITE_CHECKPOINT + elif IS_DYNAMIC: + is_write = (prev_num_accepted_tokens + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_write = WRITE_CHECKPOINT + if is_write: + write_buf = 1 - active_buf # noqa: F841 + write_offset = 0 + else: + write_buf = active_buf # noqa: F841 + write_offset = prev_num_accepted_tokens + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # Load state. state_tma_descriptor is a host-built tensor_descriptor + # over a flat (cache*nheads*dim, dstate) view of state when any TMA + # path is enabled; raw `state_ptr` is the underlying tensor and is + # always passed. state_ptrs / state_ptr_raw are the raw-pointer view + # used for !TMA load and store paths. offs_y is the flat row index + # for TMA load/store; computed unconditionally (cheap int math; DCE'd + # when no TMA path is reachable). + state_mask = m_mask[:, None] & n_mask[None, :] + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + # Load state. Branch on is_write (constexpr = WRITE_CHECKPOINT in BOTH + # modes after the outer-dispatch refactor), then constexpr-pick TMA-vs- + # tl.load per side. Outer `if` DCE's, only the matching side's + # constexpr-gated load survives -- same compile-time picking for both + # persistent_main and persistent_dynamic (the latter dispatches at the + # outer kernel level so each impl instance sees a constexpr WC). + if is_write: + if USE_TMA_LOAD_WRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + else: + if USE_TMA_LOAD_NOWRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, + other=1.0, + ).to(tl.float32) + state = state * decode_scale[:, None] + + # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) + group_idx = pid_h // nheads_ngroups_ratio + + old_window_mask = offs_window < prev_num_accepted_tokens + + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + active_buf * stride_old_dt_dbuf + + pid_h * stride_old_dt_head + ) + old_dt_all = tl.load( + old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 + ).to(tl.float32) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + active_buf * stride_old_dA_cumsum_dbuf + + pid_h * stride_old_dA_cumsum_head + ) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, + mask=old_window_mask, other=0.0, + ).to(tl.float32) + + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( + tl.float32 + ) + + coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all + + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + old_x_all = tl.load( + old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + mask=old_window_mask[:, None] & m_mask[None, :], + other=0.0, + ) + + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + active_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_all = tl.load( + old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + mask=old_window_mask[:, None] & n_mask[None, :], + other=0.0, + ).to(tl.float32) + + dB_scaled = coeff[:, None] * old_B_all + + total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) + state *= total_decay + + state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) + + if is_write: + if USE_RS_ROUNDING: + # Generate random tensor for stochastic rounding. The amount of + # randomness needed depends on the SR codegen path: + # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) + # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) + # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs + # int16 SR (24b + bitrev32): 1 b32 per 2 outputs + # The PTX cvt.rs.* instructions consume a single 32-bit random + # and split the bits internally for 2 or 4 conversions. Generate + # only what's actually consumed and broadcast to fill the unused + # slots — saves Philox rounds proportionally. + if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: + RAND_DIVISOR: tl.constexpr = 4 # fp8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: + RAND_DIVISOR: tl.constexpr = 4 # int8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: + RAND_DIVISOR: tl.constexpr = 2 # int16 SR + elif QUANT_MAX == 0.0: + RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) + else: + RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized + + rand_seed = tl.load(rand_seed_ptr) + base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head + # Number of unique randoms per row = dstate / RAND_DIVISOR. + # randint4x emits 4 randoms per offset, so use that / 4 offsets. + offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) + rand_offsets_q = ( + base_rand + + offs_m[:, None] * stride_state_dim + + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) + ) # (M, dstate / (4*RAND_DIVISOR)) + if PHILOX_ROUNDS > 0: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + else: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) + r01 = tl.join(r0, r1) + r23 = tl.join(r2, r3) + r0123 = tl.join(r01, r23) + rand_compact = tl.reshape( + r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) + ) + # Broadcast each unique rand to RAND_DIVISOR adjacent positions. + # Pack-group (pack=2 fp16 / pack=4 fp8) consumes adjacent positions; + # the unique rand lands at the asm's read slot; duplicates feed + # the dead slots. Triton's broadcast_to is stride-0 in IR. + if RAND_DIVISOR > 1: + rand_3d = rand_compact[:, :, None] + rand_3d = tl.broadcast_to( + rand_3d, + (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), + ) + rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + else: + rand = rand_compact + + if QUANT_MAX > 0.0: + amax = tl.max(tl.abs(state), axis=1) + encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) + decode_scale = 1.0 / encode_scale + state_scales_ptrs = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + + offs_m * stride_state_scales_dim + ) + tl.store(state_scales_ptrs, decode_scale, mask=m_mask) + state_q = state * encode_scale[:, None] + if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): + _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_fp8sr) + else: + tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) + else: + if USE_RS_ROUNDING: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized SR fall-through expects int8 or int16; " + "fp8 SR is handled by the prior branch.", + ) + if state_ptrs.dtype.element_ty == tl.int8: + state_q = _stochastic_round_int8_packed( + state_q, rand, offs_n[None, :] + ) + else: + state_q = _stochastic_round_int16_packed( + state_q, rand, offs_n[None, :] + ) + elif state_ptrs.dtype.element_ty != tl.float8e4nv: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized RN with explicit round() expects int8 or int16.", + ) + state_q = tl.extra.cuda.libdevice.round(state_q) + state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) + _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_cast) + else: + tl.store(state_ptrs, _state_q_cast, mask=state_mask) + elif USE_RS_ROUNDING: + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "Non-quantized SR only supports fp16 state.", + ) + _state_sr = _stochastic_round_fp16x2(state, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_sr) + else: + tl.store(state_ptrs, _state_sr, mask=state_mask) + else: + _state_cast = state.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_cast) + else: + tl.store(state_ptrs, _state_cast, mask=state_mask) + + # Phase 2: Output using precomputed CB_scaled and decay_vec + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + x_all = tl.load( + x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_base + + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_all, + mask=t_mask[:, None] & m_mask[None, :], + ) + x_all = x_all.to(tl.float32) + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + tl.float32 + ) + + init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) + out_all = init_out + cb_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# `_persistent_rectangle_impl`: rectangle nowrite path for the persistent +# kernel. Body is a copy of `_rectangle_main_impl` with `pid_m`/`pid_b`/`pid_h` +# lifted to args (same pattern as `_persistent_main_impl` vs `_replay_main_impl`). +# Called only for nowrite slots when the kernel runs with RECTANGLE=True. +# Dropped from the rect impl: LAUNCH_DEPENDENT_KERNELS / USE_PERM / REVERSE_PERM +# (kernel-level, signalled once at top); slot_perm_ptr (kernel resolves perm). +@triton.jit() +def _persistent_rectangle_impl( + # Per-work-unit indices (computed by the persistent wrapper). + pid_m, + pid_b, + pid_h, + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as + # replay path). Used when USE_TMA_LOAD; ignored otherwise. + state_tma_descriptor, + state_scales_ptr, # only consulted when QUANT_MAX > 0 + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides (rectangle (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + QUANT_MAX: tl.constexpr, + USE_TMA_LOAD: tl.constexpr = False, +): + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout (matches precompute). + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_k = tl.arange(0, BLOCK_SIZE_K) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # K-axis masks (approach C: PNAT-runtime offset, matches precompute). + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Load state. Quant scale hoist: defer `* decode_scale` post-matmul. + if USE_TMA_LOAD: + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state = state_tma_descriptor.load([offs_y, 0]) + else: + state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, other=1.0, + ).to(tl.float32) + else: + state = state.to(tl.float32) + + # Group / pointer offset setup + group_idx = pid_h // nheads_ngroups_ratio + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + # Hoist: old_x doesn't depend on conv1d/precompute; load before gdc_wait. + old_x_load = tl.load( + old_x_base + + safe_old_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + mask=is_old_k[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + x_K = tl.load( + x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=is_new_k[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_base + + offs_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_K, + mask=is_new_k[:, None] & m_mask[None, :], + ) + + x_K_f32 = x_K.to(tl.float32) + x_combined = old_x_load + x_K_f32 + + if HAS_D or HAS_Z: + sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) + x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) + else: + x_all = x_K_f32 # placeholder; unused + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec_full = tl.load( + decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 + ).to(tl.float32) + + state_out = ( + tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) + * decay_vec_full[:, None] + ) + if QUANT_MAX > 0.0: + state_out = state_out * decode_scale[None, :] + + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) + + out_all = state_out + token_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# Persistent main kernel: 1D grid, persistent CTA loop. +# Heuristics mirror those of `_checkpointing_main_kernel`. +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"NUM_PID_M_BLOCKS": lambda args: triton.cdiv(args["dim"], args["BLOCK_SIZE_M"])} +) +@triton.jit() +def _persistent_main_kernel( + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view. + # Shared across BOTH the replay path (consumed by _persistent_main_impl + # when USE_TMA_LOAD_*/STORE) AND the rectangle path (consumed by + # _persistent_rectangle_impl when USE_TMA_LOAD) — same descriptor, same + # block_shape, just gated by separate constexprs per impl. Wrapper sets + # this to a TensorDescriptor when ANY of the three TMA flags is on, else + # to `state_ptr` (raw); each impl ignores it via its own constexpr when + # not consuming it. + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + rand_seed_ptr, + pad_slot_id, + # Persistent-loop work-distribution scalars. Caller pre-sorts the batch + # write-first; the kernel uses (n_writes, batch_total, WRITE_CHECKPOINT) + # to derive its own slot range. Write half processes [0, n_writes), + # nowrite half processes [n_writes, batch_total). + # + # n_writes_ptr is a device pointer to a (1,) int32 tensor. Reading + # from device memory (rather than taking a Python int kernel arg) is + # required so mix-mode benchmarking can vary n_writes per iter inside + # a captured CUDA graph — the source tensor's contents change, the + # pointer doesn't. Cost: one int load per kernel launch (~negligible). + # When IS_DYNAMIC=True the value is unused (Triton DCEs the load). + n_writes_ptr, # int32 *: device-side count of write-mode slots + batch_total, # int32: total slot count + nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + USE_PERM: tl.constexpr, + # NUM_PERSISTENT: runtime int (not constexpr). Used ONLY as the loop + # stride in `tl.range(pid, total_work, NUM_PERSISTENT, ...)`. Making it + # runtime collapses the cta_per_sm tuning dim from the kernel's compile + # signature: 8 CPS values used to mean 8x recompiles; now they share one + # compiled kernel. Work decomposition (pid_m, pid_b_local, pid_h) does + # NOT depend on NUM_PERSISTENT — it uses constexpr NUM_PID_M_BLOCKS and + # runtime n_slots_local — so loop unrolling and flatten=/num_stages=/ + # warp_specialize= optimizations on `tl.range` operate independently of + # the stride value. + NUM_PERSISTENT, + NUM_LOOP_STAGES: tl.constexpr, + NUM_PID_M_BLOCKS: tl.constexpr, + FLATTEN: tl.constexpr, + WARP_SPECIALIZE: tl.constexpr, + IS_DYNAMIC: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr = 16, # rectangle K-axis (heuristic-derived) + RECTANGLE: tl.constexpr = False, # when True, dispatch nowrite slots to _persistent_rectangle_impl + # 3 TMA toggles per the 3 live paths per-compilation: + # USE_TMA_LOAD_WRITE — replay-style state load when is_write + # USE_TMA_LOAD_NOWRITE — nowrite-path state load (rect when RECTANGLE, + # else replay-nowrite) + # USE_TMA_STORE — replay-style state store (only fires on write + # path; no-op when not is_write) + # Wrapper picks USE_TMA_LOAD_NOWRITE = _use_tma_rect_load (if rectangle) + # or _use_tma_replay_nowrite_load (if not). + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # PDL signal: fire once at kernel entry (not per work unit). + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + # Load runtime n_writes from device memory. Read once at kernel entry; + # used only by the !IS_DYNAMIC slot-range derivation below. Triton + # DCEs the load when IS_DYNAMIC=True (n_writes is dead there). + n_writes = tl.load(n_writes_ptr) + + # Derive this kernel's slot range. Two modes: + # IS_DYNAMIC=False (persistent_main): caller pre-sorts and splits halves; + # slot range is [0, n_writes) when WRITE_CHECKPOINT else [n_writes, batch_total) + # IS_DYNAMIC=True (persistent_dynamic): single launch covers full batch; + # each work-item dispatches via runtime PNAT check inside the impl. + if IS_DYNAMIC: + slot_lo = 0 + slot_hi = batch_total + else: + if WRITE_CHECKPOINT: + slot_lo = 0 + slot_hi = n_writes + else: + slot_lo = n_writes + slot_hi = batch_total + n_slots_local = slot_hi - slot_lo + + pid = tl.program_id(axis=0) + total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads + + # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) + # with pid_m varying fastest (M-tile cache locality on state load), then + # slot, then head — mirrors the existing 3D grid's axis ordering + # (axis=0 fastest = pid_m). + for tile_id in tl.range( + pid, total_work, NUM_PERSISTENT, + flatten=FLATTEN, num_stages=NUM_LOOP_STAGES, warp_specialize=WARP_SPECIALIZE, + ): + pid_m = tile_id % NUM_PID_M_BLOCKS + pid_b_local = (tile_id // NUM_PID_M_BLOCKS) % n_slots_local + pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) + # Translate local slot index → global slot index. When USE_PERM is + # set, the caller-provided slot_perm gives the original slot index + # for the post-sort position. + pid_b_grid = pid_b_local + slot_lo + if USE_PERM: + pid_b = tl.load(slot_perm_ptr + pid_b_grid) + else: + pid_b = pid_b_grid + + # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle + # impl. Pre-load cache_batch_idx + PNAT here only for the RECTANGLE + # path's branch decision. Both impls re-load and handle pad_slot_id + # internally (Triton's L1 cache makes the duplicate loads ~free). + if RECTANGLE: + if HAS_CACHE_BATCH_INDICES: + cbi_pre = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + is_pad = cbi_pre == pad_slot_id + else: + cbi_pre = pid_b.to(tl.int64) + is_pad = False + if not is_pad: + pnat_pre = tl.load(prev_num_accepted_tokens_ptr + cbi_pre) + if IS_DYNAMIC: + is_w = (pnat_pre + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_w = WRITE_CHECKPOINT + if is_w: + # Pass WRITE_CHECKPOINT=True constexpr to specialize this + # impl call for the write path. Under IS_DYNAMIC=True, the + # kernel-level WRITE_CHECKPOINT is False (launcher default), + # but the OUTER is_w branch we are inside narrows the + # runtime path to writes-only, so we override to True here + # so the impl's constexpr-gated `if is_write:` blocks DCE + # to the write-only codegen. Under IS_DYNAMIC=False + # (persistent_main), the kernel-level WRITE_CHECKPOINT is + # itself True for this half (write half launches with + # WC=True), and the outer is_w = WRITE_CHECKPOINT = True + # constexpr-folds; passing literal True here is consistent + # and constexpr-equivalent. + _persistent_main_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, + True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) + True, # WC_IS_CONSTEXPR — force inner to use WC constexpr + # 3 TMA flags: write-load fires here (we're in the + # is_write branch), nowrite-load is dead (no slot + # reaches it), store fires (write path). + USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + ) + else: + # Rectangle nowrite: pass state_ptr (raw, always) + + # state_tma_descriptor (the single unified descriptor — + # same memory replay paths use). Rect impl gates use + # of the descriptor via its USE_TMA_LOAD constexpr. + _persistent_rectangle_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, + LAUNCH_WITH_PDL, QUANT_MAX, + USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle + ) + # else: pad slot — skip both impls (both would early-return anyway) + else: + # No rectangle path — single _persistent_main_impl call covers + # both write and nowrite slots via WC constexpr (non-dynamic) or + # runtime is_write (IS_DYNAMIC=True). Pass all 3 TMA flags; + # impl picks USE_TMA_LOAD_WRITE vs USE_TMA_LOAD_NOWRITE based on + # its computed is_write — constexpr-folds when is_write is + # constexpr (non-dyn), runtime branch when IS_DYNAMIC=True. + # (Reverted from outer two-call dispatch: that doubled the + # compiled body size under IS_DYNAMIC=True and regressed RECT=0 + # perf by ~+24%.) + _persistent_main_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, + WRITE_CHECKPOINT, IS_DYNAMIC, + False, # WC_IS_CONSTEXPR=False — RECT=0 has both write/nowrite slots in one call + USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + ) + + +# ============================================================================ +# Python wrapper +# ============================================================================ + + +def checkpointing_state_update( + state: torch.Tensor, + old_x: torch.Tensor, + old_B: torch.Tensor, + old_dt: torch.Tensor, + old_dA_cumsum: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + out: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + pad_slot_id: int = PAD_SLOT_ID, + rand_seed: torch.Tensor | None = None, + philox_rounds: int = 10, + state_scales: torch.Tensor | None = None, + launch_with_pdl=False, + use_internal_pdl=True, + write_checkpoint: bool = True, + rectangle_for_nowrite: bool = False, + mode: str = "monolithic", + # Slot permutation: int32 (batch,) tensor mapping grid program_id -> + # original slot index. When provided, dl-family kernels (doublelaunch / + # dlgrouped / maindl) read pid_b through this perm so callers can pre-sort + # slots (e.g. write-first) to cluster early-outs at one end of the grid. + # Ignored by monolithic / dynamic. None => identity (today's behavior). + slot_perm: torch.Tensor | None = None, + # When True and slot_perm is provided, the nowrite-side kernels in + # dlgrouped/doublelaunch traverse the perm in reverse (B-1-pid_grid). + # Combined with a write-first sort, this front-loads real work in BOTH + # halves of the dl chain (writes from the head, nowrites from the tail). + reverse_nowrite: bool = False, + _block_size_m: int | None = None, + _num_warps: int | None = None, + _num_stages: int | None = None, + _precompute_num_warps: int | None = None, + _precompute_num_stages: int | None = None, + _heads_per_block: int | None = None, + _maxnreg: int | None = None, + _num_ctas: int | None = None, + # Per-main knobs (override shared values for one half of the dl-family / + # persistent_main launches). Default None = tied to the shared value + # (backward compat). The two main kernels (write vs nowrite) have + # different per-slot work — write does a state shift + store, nowrite + # just appends — so the optimum (M, W, S, H) can differ. Precompute + # knobs are intentionally NOT split: shared precompute wins (cheaper + # launch, hotter precompute outputs in L2). Persistent CPS / LS knobs + # are also split per-main since the two persistent_main launches have + # different grid sizes. + _block_size_m_write: int | None = None, + _block_size_m_nowrite: int | None = None, + _num_warps_write: int | None = None, + _num_warps_nowrite: int | None = None, + _num_stages_write: int | None = None, + _num_stages_nowrite: int | None = None, + # Note: heads_per_block / precompute_num_warps are NOT split — they only + # affect the precompute kernel, which is shared across write/nowrite. + # TMA state-tensor toggles — 4 independent paths (see CHECKPOINTING_DESIGN.md + # item #17 for measured perf profiles). Each is False=raw load/store, True= + # use a host-built TMA tensor_descriptor for that path. + _use_tma_rect_load: bool = False, # rect kernel's state load (nowrite-only) + _use_tma_replay_write_load: bool = False, # replay-style state load when WC=True + _use_tma_replay_write_store: bool = False, # replay-style state store when WC=True + _use_tma_replay_nowrite_load: bool = False, # replay-style state load when WC=False + # Persistent-mode bench kwargs (only consulted when mode == "persistent_main"): + # _n_writes : int — count of write-mode slots in the (pre-sorted) batch. + # Required when mode == "persistent_main"; the persistent kernel uses + # it as a runtime int32 to compute total_work for write/nowrite halves. + # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally + # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. Default = 1. + # _num_loop_stages : int — `num_stages` arg on the inner `tl.range(...)` + # persistent loop. Note: this is loop-level, NOT the kernel-arg + # `num_stages` (which only pipelines dot-feeding loads). Default 2. + # _flatten : bool — `flatten` arg on `tl.range(...)`. Default True + # (the canonical Triton 3.6 persistent idiom). + # _warp_specialize : bool — `warp_specialize` arg on `tl.range(...)`. + # Default False. Triton 3.6 only supports it on simple matmul loops; + # our scan loop probably won't pattern-match — but exposed as a knob + # for sweep experiments. Requires num_warps >= 4 if True. + _n_writes: int | None = None, + # Optional pre-allocated (1,) int32 device tensor for the persistent + # kernel's n_writes input. Bench passes this in mix scenarios so the + # captured CUDA graph can read varying n_writes per iter without + # re-capture. When None and `_n_writes` is provided, we allocate a + # scratch tensor and fill from `_n_writes` (pure scenarios). + _n_writes_dev: torch.Tensor | None = None, + # When True, persistent_main host-skips empty-half launches (n_writes=0 + # or =batch in pure scenarios). Default True preserves today's behavior. + # Set False to always launch both halves — used by mix scenarios (where + # host can't cheaply read n_writes per iter) and for fair K-consistent + # comparisons. + _persistent_skip_empty_halves: bool = True, + _cta_per_sm: int | None = None, + _num_loop_stages: int | None = None, + _flatten: bool | None = None, + _warp_specialize: bool | None = None, + # Per-main persistent-specific knobs. Same rationale as the BLOCK_SIZE_M + # split above: the two persistent_main launches (write half vs nowrite + # half) have different grid sizes and per-work-item costs, so they may + # want different cta_per_sm / num_loop_stages. + _cta_per_sm_write: int | None = None, + _cta_per_sm_nowrite: int | None = None, + _num_loop_stages_write: int | None = None, + _num_loop_stages_nowrite: int | None = None, +): + """ + Replay SSM state update with precomputed CB and tl.dot fast-forward. + + Two-kernel architecture: + 1. Precompute kernel: computes CB_scaled and decay_vec from B, C, dt, A. + Writes processed dt/dA_cumsum/B to double-buffered cache for next step. + 2. Main kernel: replays old tokens via tl.dot fast-forward on cached data, + then computes output using precomputed CB_scaled and new x/C inputs. + + PDL (Programmatic Dependent Launch) chain: + conv1d → (external PDL) → precompute → (internal PDL) → main + External PDL: precompute starts while conv1d is running; gdc_wait() + in precompute blocks until conv1d completes before loading B/C. + Internal PDL: main starts while precompute is running; main's replay + phase uses only cached data from the previous step. gdc_wait() in + main blocks until precompute completes before loading conv1d outputs + (x, C) and precompute outputs (CB_scaled, decay_vec). + + Uses double-buffered cache tensors. cache_buf_idx[slot] indicates which + buffer (0 or 1) to READ from for replay. The WRITE buffer is 1 - read. + Caller must flip cache_buf_idx[slot] after each call. + + Arguments: + state: (cache, nheads, dim, dstate) in-place. After the call, contains + the state after replaying prev_num_accepted_tokens old tokens. + old_x: (cache, T, nheads, dim) bf16 — old x cache (single-buffered). + old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. + old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. + old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. + cache_buf_idx: (cache,) int32 — which buffer to read (0 or 1). + prev_num_accepted_tokens: (cache,) int32. + x: (batch, T, nheads, dim) new token inputs. + dt: (batch, T, nheads, dim) with stride(-1)==0 (tie_hdim). + A: (nheads, dim, dstate) with stride(-1)==0, stride(-2)==0 (tie_hdim). + B: (batch, T, ngroups, dstate). + C: (batch, T, ngroups, dstate). + out: (batch, T, nheads, dim) preallocated output. + D: (nheads, dim) optional feed-through parameter. + z: (batch, T, nheads, dim) optional silu gate. + dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). + state_batch_indices: (batch,) optional cache slot mapping. + rand_seed: optional single-element int64 CUDA tensor for Philox PRNG seed. + When provided, state is stochastically rounded on store. Supported + for state.dtype in (fp16, int8, int16, fp8_e4m3fn); other dtypes + silently use deterministic rounding. fp16+SR and fp8+SR both + require sm_100a (Blackwell B200+) — wrapper asserts this loudly. + philox_rounds: number of Philox PRNG rounds (default 10). + state_scales: required when state.dtype in (int8, int16, fp8_e4m3fn). + Shape (cache_size, nheads, dim), fp32. Per-(head, dim) channel + decode scale (= 1 / encode_scale). The kernel writes scales on + checkpoint steps and reads them on load (broadcast over dstate). + Ignored for non-quantized state dtypes. + launch_with_pdl: enable external PDL (conv1d → precompute chain). + Defaults False; caller opts in when the upstream chain is PDL-safe. + Ignored on hardware that doesn't support PDL (sm < 90). + use_internal_pdl: enable internal PDL (precompute → main overlap). + Defaults True; override for testing only. + Ignored on hardware that doesn't support PDL (sm < 90). + + _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, + _precompute_num_warps, _precompute_num_stages, _heads_per_block, + _maxnreg, _num_ctas) are benchmark-only overrides; production callers + should leave them None to use the heuristic-tuned defaults. + """ + # PDL needs sm >= 90. + if get_sm_version() < 90: + launch_with_pdl = False + use_internal_pdl = False + + # Mode selection: + # mode="monolithic" (default): today's behavior. write_checkpoint and + # rectangle_for_nowrite together pick a single kernel pair for the + # whole batch. Calls the corresponding kernel pair with EARLY_OUT=False. + # mode="dynamic": single kernel pair (_dynamic_*_kernel) that dispatches + # per-slot at runtime based on PNAT. RECTANGLE constexpr (= + # rectangle_for_nowrite) picks whether the nowrite path is rectangle + # or replay-nowrite. write_checkpoint is ignored (per-slot from PNAT). + # mode="doublelaunch": two kernel pairs launched in sequence, each with + # EARLY_OUT=True, partitioning the batch by PNAT-derived mode. + # Write half: replay-write. Nowrite half: rectangle if + # rectangle_for_nowrite else replay-nowrite. write_checkpoint ignored. + # mode="dlgrouped": same 4 kernels as doublelaunch, but reordered to + # launch both precomputes first, then both mains. Lets the GPU + # run precomp1 || precomp2 in parallel before the mains start. + # write_checkpoint ignored. + # mode="maindl": shared (dynamic) precompute + doublelaunched main. + # One precompute call (_dynamic_precompute_kernel) handles per-slot + # dispatch, then two main kernels with EARLY_OUT=True for the write + # and nowrite halves. Strictly fewer kernel launches than + # doublelaunch (3 vs 4) at the cost of dispatch precompute's wider + # reg envelope. write_checkpoint ignored. + assert mode in ( + "monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", "persistent_main", "persistent_dynamic", + ), ( + f"unknown mode {mode!r}; expected one of " + "'monolithic', 'dynamic', 'doublelaunch', 'dlgrouped', 'maindl', " + "'dl_write_only', 'persistent_main', or 'persistent_dynamic'" + ) + use_rectangle = rectangle_for_nowrite and not write_checkpoint + + # --- Hardware support gates --- + # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX + # instructions (Ada Lovelace introduced them; Hopper/Blackwell carry them). + if state.dtype == torch.float8_e4m3fn: + assert get_sm_version() >= 89, ( + "fp8_e4m3fn state requires SM 89+ (Ada Lovelace / Hopper / Blackwell) " + f"for fp32↔fp8 cvt PTX instructions; current SM is {get_sm_version()}." + ) + + # PTX cvt.rs.* (stochastic rounding) family lands on Blackwell only. + # Wrapper fails loud; framework decides fall-back (e.g. drop SR, use RN). + # int8 / int16 SR uses pure-Triton libdevice.floor + uniform noise — no + # PTX SR instruction needed, runs anywhere. + if rand_seed is not None: + if state.dtype == torch.float16: + assert get_sm_version() >= 100, ( + "fp16 stochastic rounding (PTX cvt.rs.f16x2.f32) requires " + f"sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + ) + elif state.dtype == torch.float8_e4m3fn: + assert get_sm_version() >= 100, ( + "fp8 stochastic rounding (PTX cvt.rs.satfinite.e4m3x4.f32) " + f"requires sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + ) + + # --- Unsqueeze inputs to canonical shapes --- + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if x.dim() == 3: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if dt.dim() == 3: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if B.dim() == 3: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if C.dim() == 3: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None: + if z.dim() == 2: + z = z.unsqueeze(1) + if z.dim() == 3: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + if out.dim() == 2: + out = out.unsqueeze(1) + if out.dim() == 3: + out = out.unsqueeze(1) + + cache_size, nheads, dim, dstate = state.shape + batch, T, _, _ = x.shape + ngroups = B.shape[2] + assert nheads % ngroups == 0 + + # --- Quantization plumbing --- + # QUANT_MAX > 0 ⇔ state is int8 / int16 / fp8_e4m3fn. Kernel-entry + # static_assert on the Triton side mirrors this invariant. + quant_max = _QUANT_MAX_BY_DTYPE.get(state.dtype, 0.0) + is_quantized = quant_max > 0.0 + if is_quantized: + assert state_scales is not None, ( + f"state.dtype={state.dtype} requires state_scales tensor " + "(shape (cache_size, nheads, dim), fp32)." + ) + assert state_scales.shape == (cache_size, nheads, dim), ( + f"state_scales shape mismatch: expected {(cache_size, nheads, dim)}, " + f"got {state_scales.shape}." + ) + assert state_scales.dtype == torch.float32, ( + f"state_scales must be fp32, got {state_scales.dtype}." + ) + assert state_scales.device == state.device + + # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the + # placeholder degenerate case max_window = T (every step is a checkpoint + # step). For real replay-style checkpointing, max_window > T and + # `prev_num_accepted_tokens` can be 0..max_window. + max_window = old_x.shape[1] + assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" + # Replay-style code path uses BLOCK_SIZE_T = max(np2(T), 16) for the + # combined T-axis (T_new tile size) and reuses it for window loads. Until + # the heuristic is generalized to track max_window separately, require + # max_window to fit within that tile. + block_size_t = max(triton.next_power_of_2(T), 16) + assert max_window <= block_size_t, ( + f"max_window={max_window} exceeds BLOCK_SIZE_T={block_size_t} " + f"derived from T={T}; extend the heuristic to include max_window." + ) + + assert x.shape == (batch, T, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + assert B.shape == (batch, T, ngroups, dstate) + assert C.shape == B.shape + assert old_x.shape == (cache_size, max_window, nheads, dim) + assert old_B.shape == (cache_size, 2, max_window, ngroups, dstate) + assert old_dt.shape == (cache_size, 2, nheads, max_window) + assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_window) + assert cache_buf_idx.shape == (cache_size,) + assert prev_num_accepted_tokens.shape == (cache_size,) + + tie_hdim = ( + A.stride(-1) == 0 + and A.stride(-2) == 0 + and dt.stride(-1) == 0 + and (dt_bias is None or dt_bias.stride(-1) == 0) + ) + assert tie_hdim + + device = x.device + BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) + # Rectangle K-axis bound = window (max_window). Computed unconditionally + # so the launch sites can refer to it; only used on the rectangle path. + BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) + + # Allocate precomputed intermediates (per-call, not cached). Always + # allocate (T, K) — the largest layout that any path uses. Replay-style + # paths only touch the first T columns; rectangle/dynamic use the full K. + # The few extra unused columns per row are negligible (~6KB per layer at + # production sizes) and let the dispatch helpers share one buffer. + cb_scaled = torch.empty( + batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 + ) + decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) + + z_strides = ( + (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) if z is not None else (0, 0, 0, 0) + ) + + # Kernel tuning: BLOCK_SIZE_M, num_warps, HEADS_PER_BLOCK, precompute_num_warps. + # Dtype-aware heuristic from B200 sweeps (batch 1-512, T=6/32, TP=8, conv1d + + # chained PDL). Keyed on total_heads, BLOCK_SIZE_T, and state dtype; 16-bit + # states prefer different tiles from fp32 due to lower bandwidth. Philox + # gets its own branch — stochastic rounding shifts compute toward CUDA cores, + # so small-batch configs want more warps to hide the extra work. + total_heads = batch * nheads + heads_per_group = nheads // ngroups + state_is_16bit = state.dtype in (torch.float16, torch.bfloat16) + use_philox = rand_seed is not None + if BLOCK_SIZE_T <= 16: + if use_philox and state_is_16bit: + # Philox: more warps at small batch to hide CUDA core work. + # At large batch, converges to non-Philox fp16 config. + if total_heads <= 16: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 4, 4, 4, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + elif state_is_16bit: + if total_heads <= 16: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 4, 1 + elif total_heads <= 64: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 2, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 1, + min(2, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + else: # fp32 state (no Philox — fp32 doesn't need stochastic rounding) + if total_heads <= 32: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 4, 1 + elif total_heads <= 64: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + elif total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 2, 2, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 1, 2, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 2, + min(2, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + else: # T > 16 + if state_is_16bit: + if total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 16, + 1, + 4, + min(2, heads_per_group), + ) + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 1, + min(4, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 4, + min(2, heads_per_group), + ) + else: # fp32 state + if total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 2, + 4, + min(2, heads_per_group), + ) + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 2, + min(4, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 4, + min(2, heads_per_group), + ) + if _block_size_m is not None: + BLOCK_SIZE_M = _block_size_m + if _num_warps is not None: + num_warps = _num_warps + if _heads_per_block is not None: + heads_per_block = _heads_per_block + if _precompute_num_warps is not None: + precompute_num_warps = _precompute_num_warps + + # Per-main knob resolution: each _*_{write,nowrite} arg, if not None, + # overrides the corresponding shared value for ONE main launch only. + # Default (None) = tied to shared value (current behavior). + BLOCK_SIZE_M_WRITE = _block_size_m_write if _block_size_m_write is not None else BLOCK_SIZE_M + BLOCK_SIZE_M_NOWRITE = _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M + NUM_WARPS_WRITE = _num_warps_write if _num_warps_write is not None else num_warps + NUM_WARPS_NOWRITE = _num_warps_nowrite if _num_warps_nowrite is not None else num_warps + NUM_STAGES_WRITE = _num_stages_write if _num_stages_write is not None else _num_stages + NUM_STAGES_NOWRITE = _num_stages_nowrite if _num_stages_nowrite is not None else _num_stages + # Persistent-only per-main: + CTA_PER_SM_WRITE = _cta_per_sm_write if _cta_per_sm_write is not None else _cta_per_sm + CTA_PER_SM_NOWRITE = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _cta_per_sm + NUM_LOOP_STAGES_WRITE = _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages + NUM_LOOP_STAGES_NOWRITE = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages + + HAS_CACHE_BATCH_INDICES = state_batch_indices is not None + + assert nheads % heads_per_block == 0, ( + f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" + ) + assert heads_per_block <= heads_per_group, ( + f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" + ) + + # state_scales pointer + strides: real tensor when quantized, otherwise + # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). + if is_quantized: + state_scales_arg = state_scales + state_scales_strides = ( + state_scales.stride(0), + state_scales.stride(1), + state_scales.stride(2), + ) + else: + state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 + state_scales_strides = (0, 0, 0) + + # Per-path TMA descriptors for state — write-side and nowrite-side. Each + # kernel launch consumes the descriptor whose block_shape[0] matches its + # BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need + # distinct descriptors; otherwise the descriptor's block_shape[0] would + # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot / arithmetic + # on the loaded tile fails shape inference at compile time + # ("Cannot make_shape_compatible: incompatible dimensions"). When Mw == + # Mnw (tied, the common case) the two descriptors are the same object. + # Same memory (state's flat 2D view, shape (cache*nheads*dim, dstate)) + # and same dstate block_shape — only block_shape[0] differs. + # When no TMA flag is on, both variables hold the raw `state` tensor as a + # dummy; kernels never reference it because their constexprs are all + # False (Triton DCEs the dead branches). + # `triton.set_allocator()` must run before any descriptor-using launch. + if (_use_tma_rect_load or _use_tma_replay_write_load + or _use_tma_replay_write_store or _use_tma_replay_nowrite_load): + from triton.tools.tensor_descriptor import TensorDescriptor + _ensure_tma_allocator() + assert state.is_contiguous(), "TMA state requires contiguous state" + assert state.stride(-1) == 1, "TMA state requires inner stride 1" + _state_flat = state.view(-1, state.shape[-1]) + _dstate_pow2 = triton.next_power_of_2(dstate) + state_tma_descriptor_write = TensorDescriptor.from_tensor( + _state_flat, block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], + ) + if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: + state_tma_descriptor_nowrite = state_tma_descriptor_write + else: + state_tma_descriptor_nowrite = TensorDescriptor.from_tensor( + _state_flat, block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], + ) + else: + state_tma_descriptor_write = state # dummy; all consuming constexprs False + state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False + + # Slot permutation — pointer + USE_PERM gate. When the caller provides + # a perm tensor the dl-family launches read pid_b through it; otherwise + # we pass any valid pointer (state_batch_indices) and USE_PERM=False so + # the kernel falls back to pid_grid. Sort-driven dispatch (write-first + # clustering) is opt-in per call; monolithic / dynamic ignore the flag. + if slot_perm is not None: + assert slot_perm.dtype in (torch.int32, torch.int64), ( + f"slot_perm must be int32/int64, got {slot_perm.dtype}" + ) + assert slot_perm.numel() >= batch, ( + f"slot_perm has {slot_perm.numel()} entries; need >= batch ({batch})" + ) + slot_perm_arg = slot_perm + use_perm = True + else: + # Any valid ptr — gated by USE_PERM=False at compile time. + slot_perm_arg = state_batch_indices if state_batch_indices is not None else state + use_perm = False + + # Grid for main kernels (M tiling × batch × nheads). + def main_grid(META): + return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + + precomp_grid = (batch, nheads // heads_per_block) + d_strides = (D.stride(0), D.stride(1)) if D is not None else (0, 0) + + # ---- Launch helpers (close over locals) ------------------------------- + # Each helper is a thin closure that calls one Triton kernel with the + # full positional + kwarg argument list. Mode-dependent constexprs + # (write_checkpoint, early_out, rectangle) are passed in. + + def launch_replay_precompute(write_checkpoint: bool, early_out: bool, + reverse_perm: bool = False): + _checkpointing_precompute_kernel[precomp_grid]( + dt, dt_bias, A, B, C, + cb_scaled, decay_vec, + old_B, old_dt, old_dA_cumsum, + cache_buf_idx, prev_num_accepted_tokens, + state_batch_indices, slot_perm_arg, pad_slot_id, + T, max_window, dstate, nheads // ngroups, + dt.stride(0), dt.stride(1), dt.stride(2), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), B.stride(1), B.stride(2), B.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + dt_softplus, + HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, + LAUNCH_WITH_PDL=launch_with_pdl, + LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, + HEADS_PER_BLOCK=heads_per_block, + WRITE_CHECKPOINT=write_checkpoint, + EARLY_OUT=early_out, + USE_PERM=use_perm, + REVERSE_PERM=reverse_perm, + num_warps=precompute_num_warps, + **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), + launch_pdl=launch_with_pdl, + ) + + def launch_rectangle_precompute(early_out: bool, reverse_perm: bool = False): + _rectangle_precompute_kernel[precomp_grid]( + dt, dt_bias, A, B, C, + cb_scaled, decay_vec, + old_B, old_dt, old_dA_cumsum, + cache_buf_idx, prev_num_accepted_tokens, + state_batch_indices, slot_perm_arg, pad_slot_id, + T, max_window, dstate, nheads // ngroups, + dt.stride(0), dt.stride(1), dt.stride(2), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), B.stride(1), B.stride(2), B.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + dt_softplus, + HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, + LAUNCH_WITH_PDL=launch_with_pdl, + LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, + HEADS_PER_BLOCK=heads_per_block, + EARLY_OUT=early_out, + USE_PERM=use_perm, + REVERSE_PERM=reverse_perm, + num_warps=precompute_num_warps, + **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), + launch_pdl=launch_with_pdl, + ) + + def launch_dynamic_precompute(rectangle: bool): + _dynamic_precompute_kernel[precomp_grid]( + dt, dt_bias, A, B, C, + cb_scaled, decay_vec, + old_B, old_dt, old_dA_cumsum, + cache_buf_idx, prev_num_accepted_tokens, + state_batch_indices, pad_slot_id, + T, max_window, dstate, nheads // ngroups, + dt.stride(0), dt.stride(1), dt.stride(2), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), B.stride(1), B.stride(2), B.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + dt_softplus, + HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, + LAUNCH_WITH_PDL=launch_with_pdl, + LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, + HEADS_PER_BLOCK=heads_per_block, + RECTANGLE=rectangle, + num_warps=precompute_num_warps, + **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), + launch_pdl=launch_with_pdl, + ) + + def launch_replay_main(write_checkpoint: bool, early_out: bool, + launch_dependent_kernels: bool = False, + reverse_perm: bool = False): + # Per-main knob selection: write vs nowrite branches use independent + # M / num_warps / num_stages / heads_per_block values. Grid is + # M-dependent so it must be a closure over the selected M. + _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE + _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE + _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE + # Per-path TMA descriptor: block_shape[0] must match the kernel's + # BLOCK_SIZE_M (`_bsm`); see the descriptor build block above. + _desc = (state_tma_descriptor_write if write_checkpoint + else state_tma_descriptor_nowrite) + def _main_grid_local(META, _bsm=_bsm): + return (triton.cdiv(dim, _bsm), batch, nheads) + _checkpointing_main_kernel[_main_grid_local]( + state, _desc, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + _bsm, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=write_checkpoint, + EARLY_OUT=early_out, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + REVERSE_PERM=reverse_perm, + # Per-launch WC fixes which LOAD flag is "live"; pass write-load + # value when WC=True (NOWRITE flag dummy False), else converse. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), + USE_TMA_LOAD_NOWRITE=bool(_use_tma_replay_nowrite_load and not write_checkpoint), + USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), + num_warps=_nw, + **({"num_stages": _ns} if _ns else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + def launch_rectangle_main(early_out: bool, + launch_dependent_kernels: bool = False, + reverse_perm: bool = False): + # Rectangle is the nowrite-side path; use the nowrite-main knobs. + _bsm = BLOCK_SIZE_M_NOWRITE + _nw = NUM_WARPS_NOWRITE + _ns = NUM_STAGES_NOWRITE + def _main_grid_local(META, _bsm=_bsm): + return (triton.cdiv(dim, _bsm), batch, nheads) + # Rectangle is always the nowrite-side path; descriptor block_shape[0] + # must match BLOCK_SIZE_M_NOWRITE (= _bsm here). + _rectangle_main_kernel[_main_grid_local]( + state, state_tma_descriptor_nowrite, state_scales_arg, old_x, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, pad_slot_id, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + _bsm, + LAUNCH_WITH_PDL=use_internal_pdl, + QUANT_MAX=quant_max, + EARLY_OUT=early_out, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + REVERSE_PERM=reverse_perm, + USE_TMA_LOAD=bool(_use_tma_rect_load), + num_warps=_nw, + **({"num_stages": _ns} if _ns else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + def launch_dynamic_main(rectangle: bool, + launch_dependent_kernels: bool = False): + # Dynamic mode uses a single BLOCK_SIZE_M (no M-split inside this + # kernel); BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE by the wrapper's tied + # convention, so the write-side descriptor matches. + _dynamic_main_kernel[main_grid]( + state, state_tma_descriptor_write, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, rand_seed, pad_slot_id, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + RECTANGLE=rectangle, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + # 3 TMA flags. NOWRITE_LOAD picks rect-load vs replay-nowrite-load + # based on RECTANGLE constexpr (only one is reachable per compile). + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load), + USE_TMA_LOAD_NOWRITE=bool(_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load), + USE_TMA_STORE=bool(_use_tma_replay_write_store), + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + # ---- launch_persistent_main ------------------------------------------ + # Persistent-CTA main kernel. Single launch covers `n_slots` slots + # starting at `slot_offset`. Caller invokes twice: once for the write + # half (slot_offset=0, n_slots=n_writes, write_checkpoint=True) and + # once for the nowrite half (slot_offset=n_writes, + # n_slots=batch-n_writes, write_checkpoint=False). Hard-sort + # contract: caller has pre-sorted slots so [0, n_writes) are writes + # and [n_writes, batch) are nowrites. + + # Resolve persistent-mode bench knobs. Defaults: cta_per_sm = 1 + # (one CTA per SM, matches upstream `_p_matmul_ogs.py`); num_loop_stages + # = 2 (matches in-tree `swiglu` precedent for non-dot persistent loops); + # flatten = True (canonical Triton 3.6 idiom); warp_specialize = False. + _num_sms = torch.cuda.get_device_properties(device).multi_processor_count + cta_per_sm_arg = _cta_per_sm if _cta_per_sm else 1 + num_persistent_arg = cta_per_sm_arg * _num_sms + num_loop_stages_arg = _num_loop_stages if _num_loop_stages else 2 + flatten_arg = True if _flatten is None else bool(_flatten) + warp_specialize_arg = False if _warp_specialize is None else bool(_warp_specialize) + # Per-launch work-item count. At small batch, total_work may be < the + # full persistent grid; capping `grid` at `min(NUM_PERSISTENT, total_work)` + # avoids launching empty CTAs that pay setup cost for no work. Correctness: + # the kernel's `tl.range(pid, total_work, NUM_PERSISTENT)` ensures each + # tile_id is covered exactly once across all live pids in [0, grid) when + # grid <= NUM_PERSISTENT (each CTA does 1 tile; loop step >= total_work + # exits immediately) AND when grid == NUM_PERSISTENT (each CTA loops over + # multiple tiles). NUM_PERSISTENT is now a runtime int (see kernel def + # docstring at _persistent_main_kernel) so changing cta_per_sm does NOT + # trigger a new Triton compile — same kernel binary, different loop step. + # (Named UPPERCASE for historical Triton-style consistency only; not + # constexpr.) + _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + + def launch_persistent_main(write_checkpoint: bool, + n_writes_dev: torch.Tensor, + *, + host_n_writes: int | None = None, + skip_empty_halves: bool = True, + launch_dependent_kernels: bool = False, + rectangle: bool = False): + # `n_writes_dev` is a (1,) int32 device tensor; the kernel reads + # the count from device memory. `host_n_writes` is the same value + # known host-side (when available — pure scenarios) and lets us + # skip the launch entirely if its half is empty. In mix scenarios + # the host doesn't know n_writes per iter without a sync, so + # `host_n_writes is None` and `skip_empty_halves` is forced False + # — both halves always launch and the kernel processes whatever + # range device-n_writes implies. + if skip_empty_halves and host_n_writes is not None: + n_slots_for_kernel = host_n_writes if write_checkpoint else (batch - host_n_writes) + if n_slots_for_kernel <= 0: + return + # Per-main knob selection. The two persistent_main launches (write + # half vs nowrite half) get independent BLOCK_SIZE_M / num_warps / + # num_stages / cta_per_sm / num_loop_stages. See the per-main args + # block in the wrapper signature. + _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE + _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE + _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE + _cps = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE + _cps = _cps if _cps else 1 + _nls = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE + _nls = _nls if _nls else 2 + _num_persistent = _cps * _num_sms + _num_pid_m_local = (dim + _bsm - 1) // _bsm + # Grid sizing: cap at min(full persistent grid, actual total_work). + # `n_slots` for this launch is `host_n_writes` (write half) / `batch - + # host_n_writes` (nowrite half) when host knows it (pure); else upper + # bound `batch` for mix scenarios where host can't read n_writes_dev + # without a sync. Upper-bound is fine — the kernel's runtime check + # only iterates actual work; the only cost of overcounting is a few + # extra CTAs. + if host_n_writes is not None: + _n_slots_for_launch = host_n_writes if write_checkpoint else (batch - host_n_writes) + else: + _n_slots_for_launch = batch + _total_work_launch = max(1, _n_slots_for_launch * _num_pid_m_local * nheads) + grid = (min(_num_persistent, _total_work_launch),) + # Per-path TMA descriptor — block_shape[0] must match _bsm. + _desc = (state_tma_descriptor_write if write_checkpoint + else state_tma_descriptor_nowrite) + _persistent_main_kernel[grid]( + state, _desc, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + n_writes_dev, batch, nheads, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + _bsm, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=write_checkpoint, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + NUM_PERSISTENT=_num_persistent, + NUM_LOOP_STAGES=_nls, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=False, + RECTANGLE=rectangle, + # 3 TMA flags. IS_DYNAMIC=False: WC fixed per launch; impl + # constexpr-folds the LOAD pick. When WC=True (write half), + # NOWRITE_LOAD is dummy False; when WC=False, WRITE_LOAD/STORE + # dummy False. NOWRITE_LOAD picks rect-load (RECTANGLE) or + # replay-nowrite-load. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), + USE_TMA_LOAD_NOWRITE=bool( + (_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load) + and not write_checkpoint + ), + USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), + num_warps=_nw, + **({"num_stages": _ns} if _ns else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, + launch_dependent_kernels: bool = False, + rectangle: bool = False): + # Single-launch persistent kernel covering the whole batch with + # runtime per-slot WRITE_CHECKPOINT branch. No half-split, no + # n_writes needed (the kernel ignores n_writes_dev when + # IS_DYNAMIC=True; Triton DCEs the load). is_write is computed + # at runtime per work-item from the loaded PNAT. + # We still pass `n_writes_dev` (the same tensor the persistent_main + # path uses) so the kernel signature is uniform; the value is + # immaterial. + # Grid sizing: cap at total_work (= batch * num_pid_m * nheads) for + # the dynamic case (full-batch coverage); see launch_persistent_main + # comment for correctness rationale. + _total_work_launch = max(1, batch * _num_pid_m * nheads) + grid = (min(num_persistent_arg, _total_work_launch),) + # Persistent-dynamic kernel uses a single BLOCK_SIZE_M (same as the + # wrapper's BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE tied convention), so + # the write-side descriptor matches. Both write and nowrite slots + # in this kernel share that BSM. + _persistent_main_kernel[grid]( + state, state_tma_descriptor_write, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + n_writes_dev, batch, nheads, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=False, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + NUM_PERSISTENT=num_persistent_arg, + NUM_LOOP_STAGES=num_loop_stages_arg, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=True, + RECTANGLE=rectangle, + # 3 TMA flags. IS_DYNAMIC=True: is_write is runtime per slot; + # impl's load TMA picks per-slot (constexpr ternary becomes a + # runtime branch — both load forms emitted, ~negligible cost). + # NOWRITE_LOAD picks rect-load when RECTANGLE, else + # replay-nowrite-load. STORE only fires on runtime is_write. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load), + USE_TMA_LOAD_NOWRITE=bool( + _use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load + ), + USE_TMA_STORE=bool(_use_tma_replay_write_store), + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + # ---- Mode dispatch ---------------------------------------------------- + with torch.cuda.device(device.index): + if mode == "monolithic": + if use_rectangle: + launch_rectangle_precompute(early_out=False) + launch_rectangle_main(early_out=False) + else: + launch_replay_precompute(write_checkpoint=write_checkpoint, early_out=False) + launch_replay_main(write_checkpoint=write_checkpoint, early_out=False) + elif mode == "dynamic": + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_dynamic_main(rectangle=rectangle_for_nowrite) + elif mode == "maindl": + # Shared dispatch precompute, doublelaunched main. Precompute + # runs once with per-slot dispatch (saves the second precomp + # empty-grid tax of doublelaunch). Mains stay split with + # EARLY_OUT so each retains its constexpr-specialized reg + # envelope. First main signals PDL dependents so the second + # main can start its setup while the first is still computing. + # Dynamic precompute doesn't support sort (no early-out to + # cluster), so the perm only flows into the two EARLY_OUT mains. + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_replay_main(write_checkpoint=True, early_out=True, + launch_dependent_kernels=True) + if rectangle_for_nowrite: + launch_rectangle_main(early_out=True, reverse_perm=reverse_nowrite) + else: + launch_replay_main(write_checkpoint=False, early_out=True, + reverse_perm=reverse_nowrite) + elif mode == "dlgrouped": + # Same 4 kernels as doublelaunch but reordered: both precomputes + # first, then both mains. Lets the GPU run precomp1 || precomp2 + # in parallel (they're tiny grids) before the mains start, vs + # doublelaunch's interleaved precomp1→main1→precomp2→main2. + # First main signals PDL so the second main's setup overlaps. + # When slot_perm + reverse_nowrite are set, the nowrite-side + # walks the perm in reverse so both kernels front-load real work. + launch_replay_precompute(write_checkpoint=True, early_out=True) + if rectangle_for_nowrite: + launch_rectangle_precompute(early_out=True, reverse_perm=reverse_nowrite) + else: + launch_replay_precompute(write_checkpoint=False, early_out=True, + reverse_perm=reverse_nowrite) + launch_replay_main(write_checkpoint=True, early_out=True, + launch_dependent_kernels=True) + if rectangle_for_nowrite: + launch_rectangle_main(early_out=True, reverse_perm=reverse_nowrite) + else: + launch_replay_main(write_checkpoint=False, early_out=True, + reverse_perm=reverse_nowrite) + elif mode == "dl_write_only": + # Debug-only: just the write half of doublelaunch. EARLY_OUT=True + # means nowrite slots still pay the EO-gate tax (PNAT load + branch), + # but no nowrite-side kernels run. Used to isolate "is the sort + # regression in the write-side kernels?". + launch_replay_precompute(write_checkpoint=True, early_out=True) + launch_replay_main(write_checkpoint=True, early_out=True, + launch_dependent_kernels=False) + elif mode == "persistent_dynamic": + # Single-launch persistent kernel covering the full batch. + # Each work-item dispatches via runtime PNAT check (is_write = + # (pnat + T) > MAX). No n_writes/half-split — kernel ignores + # n_writes_dev when IS_DYNAMIC=True (Triton DCEs the load). + # We still need a valid pointer to satisfy the kernel arg + # signature; allocate or reuse `_n_writes_dev`. + n_writes_dev_local = ( + _n_writes_dev if _n_writes_dev is not None + else torch.zeros(1, dtype=torch.int32, device=device) + ) + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_persistent_dynamic_main( + n_writes_dev_local, + launch_dependent_kernels=False, + rectangle=rectangle_for_nowrite, + ) + elif mode == "persistent_main": + # Persistent-CTA main kernel. Reuses maindl's precompute + # structure (one shared dynamic_precompute that dispatches + # per-slot at runtime based on PNAT) followed by two + # persistent_main launches (write half + nowrite half). + # + # Hard-sort contract: caller has pre-sorted slots host-side so + # PNAT is monotone (writes first). Pass the perm via + # slot_perm + USE_PERM. + # + # n_writes is read by the kernel from a (1,) int32 device + # tensor. The caller can provide: + # * _n_writes_dev only (mix): a pre-filled (1,) int32 tensor + # it updates per iter via pre_iter_fn outside the captured + # graph. host can't cheaply read it without a sync, so both + # halves always launch. + # * _n_writes only (non-graph callers, e.g. unit tests): host + # int. We allocate the scratch tensor on the fly. CANNOT + # be used inside CUDA-graph capture — alloc inside capture + # invalidates the stream. + # * Both (pure under graph capture): caller pre-allocates the + # tensor outside capture and tells us the host value too. + # We skip the internal allocation and apply host-skip when + # _persistent_skip_empty_halves=True. This is the + # production-equivalent path the bench's pure cells take. + if _n_writes_dev is not None: + n_writes_dev_local = _n_writes_dev # no allocation + if _n_writes is not None: + # Caller provided both: pure scenario with pre-allocated + # tensor. Use host_n_writes for the skip-empty fast path. + assert 0 <= _n_writes <= batch, ( + f"_n_writes={_n_writes} must be in [0, batch={batch}]" + ) + host_n_writes_local = _n_writes + skip_empty_local = _persistent_skip_empty_halves + else: + # Mix: host doesn't know n_writes without a sync. + host_n_writes_local = None + skip_empty_local = False + else: + # No pre-allocated tensor. Fall back to on-the-fly alloc + # from _n_writes (host int). NOT graph-capture-safe. + assert _n_writes is not None, ( + "mode='persistent_main' requires either _n_writes " + "(host int, non-graph callers) or _n_writes_dev (device " + "tensor, recommended for graph-capture callers)." + ) + assert 0 <= _n_writes <= batch, ( + f"_n_writes={_n_writes} must be in [0, batch={batch}]" + ) + n_writes_dev_local = torch.tensor( + [_n_writes], dtype=torch.int32, device=device, + ) + host_n_writes_local = _n_writes + skip_empty_local = _persistent_skip_empty_halves + # rectangle_for_nowrite=True: precompute populates cb_scaled + # for the rect path; nowrite half uses the rectangle impl; + # write half always replay-style (rect doesn't apply). + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_persistent_main( + write_checkpoint=True, + n_writes_dev=n_writes_dev_local, + host_n_writes=host_n_writes_local, + skip_empty_halves=skip_empty_local, + launch_dependent_kernels=True, + rectangle=False, # write always replay-style + ) + launch_persistent_main( + write_checkpoint=False, + n_writes_dev=n_writes_dev_local, + host_n_writes=host_n_writes_local, + skip_empty_halves=skip_empty_local, + launch_dependent_kernels=False, + rectangle=rectangle_for_nowrite, + ) + else: # mode == "doublelaunch" + # Write half: always replay-style write. First main signals + # PDL dependents so the second precompute can start its setup + # while the first main is still computing. + launch_replay_precompute(write_checkpoint=True, early_out=True) + launch_replay_main(write_checkpoint=True, early_out=True, + launch_dependent_kernels=True) + # Nowrite half: rectangle if asked, else replay-nowrite. + if rectangle_for_nowrite: + launch_rectangle_precompute(early_out=True, reverse_perm=reverse_nowrite) + launch_rectangle_main(early_out=True, reverse_perm=reverse_nowrite) + else: + launch_replay_precompute(write_checkpoint=False, early_out=True, + reverse_perm=reverse_nowrite) + launch_replay_main(write_checkpoint=False, early_out=True, + reverse_perm=reverse_nowrite) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index d1dd58d764f8..e99b33675fa2 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -3407,6 +3407,7 @@ def _print_row( json_detailed=False, jsonl_path=None, jsonl_host=None, + jsonl_gpu=None, ): """Print one summary row and append the result to the JSONL sidecar. @@ -3462,6 +3463,8 @@ def _print_row( rec = {"key": key, "stats": row_stats, "t": _time.time()} if jsonl_host is not None: rec["host"] = jsonl_host + if jsonl_gpu is not None: + rec["gpu"] = jsonl_gpu with open(jsonl_path, "a") as f: f.write(json.dumps(rec) + "\n") @@ -3497,6 +3500,7 @@ def _finish_result_job(args, job: dict) -> None: json_detailed=getattr(args, "json_detailed", False), jsonl_path=getattr(args, "_jsonl_path", None), jsonl_host=getattr(args, "_jsonl_host", None), + jsonl_gpu=getattr(args, "_jsonl_gpu", None), ) @@ -3764,9 +3768,17 @@ def _phase(label: str) -> None: args._jsonl_path = None args._done_keys: set[str] = set() args._jsonl_host = None # hostname stamp for the current run + args._jsonl_gpu = None # GPU device id stamp (current process visibility) if getattr(args, "json_output", None): import socket args._jsonl_host = socket.gethostname() + # Capture GPU id once at startup. Used by the oracle-cache layer in + # search_driver to attribute timings to a specific (host, gpu) pair + # for cross-process pruning. os.environ['CUDA_VISIBLE_DEVICES'] + # is the right source pre-torch-init (it's what the harness sets); + # post-init we could use torch.cuda.current_device() but we keep it + # to env to avoid forcing a CUDA init at this point in startup. + args._jsonl_gpu = os.environ.get("CUDA_VISIBLE_DEVICES", "") args._jsonl_path = args.json_output + ".jsonl" # Read existing JSONL if present: load every record's key into the # skip set regardless of host (gap-fill on a new node). diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_refactored.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_refactored.py new file mode 100644 index 000000000000..298de246193f --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_refactored.py @@ -0,0 +1,4935 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Standalone benchmark for replay_selective_state_update (Triton kernel). + +Suitable for nsight-compute (ncu) and nsight-systems (nsys) capture. + +Fixed model config: NVIDIA-Nemotron-3-Super-120B-A12B at TP=8 + nheads=16, head_dim=64, d_state=128, ngroups=1 + +mtp_len is the per-request sequence length processed by replay: in MTP it +equals num_draft_tokens + 1 target token, so --mtp-lengths 6 models 5 drafts ++ 1 target. + +Baseline kernel (--baseline [triton|flashinfer]): + Calls selective_state_update with T=mtp_len tokens and disable_state_update=True, + matching the MTP scoring pass in mamba2_mixer.py exactly. + +Timing methodology +================== + +All in-bench timing comes from CUPTI's Activity API (1 ns kernel +timestamps from the GPU profiling fabric). cudaEvent.elapsed_time() was +removed — its ~0.5 us resolution overshoots CUPTI by ~50% on short kernels +in graphs, and we have no other use for it here. See the CUPTI block +lower in this file for the timer source. + +Three modes: + + --cupti --cuda-graph (default) + Capture a small CUDA graph for the cell, replay it for warmup + timed + iterations, and read kernel start/end from CUPTI. Raw CUPTI buffers are + parsed out-of-process on the timed path, with a cached ordinal plan used + to keep only the kernels we care about. + + --cupti --no-cuda-graph + Eager loop with CUPTI. Per-kernel timestamps are still accurate, but + the per-iter SPAN (max(end) - min(start)) now includes the Python + launch latency BETWEEN consecutive kernels in run_fn (~100 µs on + Hopper/Blackwell). Graph capture and PDL hide that latency; eager + mode honestly reports it. For per-kernel timing in eager mode, look + at per_kernel.start_us/end_us in --json-detailed output rather than + the span percentiles. Useful when graph capture is undesirable. + + --no-cupti (with or without --cuda-graph) + No in-bench timing — just runs the kernels for an external profiler + (nsys / ncu) to time. In-process CUPTI conflicts with nsys's own + subscriber, so disable ours when wrapping in nsys. Bench output + reports zeros for median/p95/p99; trust the external trace. + +JSON output schema (--json-output PATH) +======================================= + +Designed to be parsed by collect.py / report.py without touching sqlite or +NVTX traces. Future agents: prefer reading this JSON over re-running nsys. + + { + "metadata": {timestamp, cmd, tp_size, warmup, iters, variant, cupti}, + "results": { + "": {median, p95, p99, n, iters_us, [n_writes_per_iter], [per_kernel]} + } + } + +Key format mirrors collect.py's kernel_data.json convention: + incremental/{batch}/{mtp}/{sd}/k{prev_k}/{sweep_parts}/tp{tp} + triton/{batch}/{mtp}/{sd}/tp{tp} + flashinfer/{batch}/{mtp}/{sd}/tp{tp} + + - is normalized: bf16 / fp16 / fp32 / int8 / int16 / fp8. + - is e.g. "M16_W1_S3_SR0_RECT0_WC1" — flags concatenated by + underscore in canonical (M, W, S, pW, pS, H, R, CT, SR, RECT, WC) order. + - All numeric values in microseconds (us). + +Per-record fields: + - median, p95, p99: span statistics (us). Span = max(kernel_end_ns) - + min(kernel_start_ns) across the iter's kernels — same convention as + nsys-derived collect.py used to use. + - n: number of timed iters that contributed. + - iters_us: list of length n, raw per-iter spans. + - n_writes_per_iter: for mix rows, list of length n with the number of + write-path slots in each timed iteration. + - per_kernel: {: {start_us: [...], end_us: [...]}} where + timestamps are RELATIVE to that iter's first kernel start, in us. Lets + you see PDL overlap directly without an external profiler. Only with + --json-detailed. + +Example usage: + # Basic sweep (default = --cupti, just summary stats) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 1,2,4 --mtp-lengths 1,4,8 --warmup 5 --iters 20 + + # JSON output, summary stats only (compact) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 16 --mtp-lengths 6 --json-output /tmp/out.json + + # JSON output, full per-iter / per-kernel data (for PDL analysis etc.) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 16 --mtp-lengths 6 \\ + --json-output /tmp/out.json --json-detailed + + # nsys capture (--no-cupti so our subscriber doesn't conflict) + nsys profile --capture-range=cudaProfilerApi \\ + python benchmark_replay_selective_state_update.py --profile --no-cupti + + # ncu capture (--no-cupti --no-cuda-graph: each kernel replayable solo) + ncu --target-processes all \\ + python benchmark_replay_selective_state_update.py --profile \\ + --no-cupti --no-cuda-graph \\ + --batch-sizes 1 --mtp-lengths 4 --warmup 5 --iters 5 +""" + +import argparse +import atexit +import ctypes +import importlib +import itertools +import json +import multiprocessing as mp +import os +import queue +import statistics +import sys +import threading +import time +from datetime import datetime +from multiprocessing import shared_memory +from pathlib import Path + +import numpy as np +import torch +from einops import repeat + + +def _import_mamba_kernels_fast(): + """Load kernel modules directly (~40s faster than a full tensorrt_llm init). + Use --full-import as the fallback if module dependencies change. + + Strategy: stub the parent packages (tensorrt_llm, tensorrt_llm._torch, + tensorrt_llm._torch.modules) in sys.modules with __path__ set, but do + NOT execute their __init__.py. Then load the leaf kernel modules. + When a kernel body imports e.g. tensorrt_llm._utils.get_sm_version, + Python's machinery resolves it against our stub's __path__ and loads + only _utils.py — skipping the heavy tensorrt_llm package init. + """ + import types + + repo_root = Path(__file__).resolve().parents[5] + trtllm_dir = repo_root / "tensorrt_llm" + mamba_pkg = "tensorrt_llm._torch.modules.mamba" + mamba_dir = trtllm_dir / "_torch" / "modules" / "mamba" + + def _stub_pkg(fqn: str, pkg_dir: Path): + """Register a stub package in sys.modules without running its + __init__.py. Sets __path__ so Python can resolve submodule imports + against the real directory on disk.""" + if fqn in sys.modules: + return + stub = types.ModuleType(fqn) + stub.__path__ = [str(pkg_dir)] + sys.modules[fqn] = stub + + # Stub the parent chain so `from tensorrt_llm._utils import ...` (and + # similar) work without triggering tensorrt_llm/__init__.py. + _stub_pkg("tensorrt_llm", trtllm_dir) + _stub_pkg("tensorrt_llm._torch", trtllm_dir / "_torch") + _stub_pkg("tensorrt_llm._torch.modules", trtllm_dir / "_torch" / "modules") + + def _load(mod_name: str, file_name: str): + fqn = f"{mamba_pkg}.{mod_name}" if mod_name else mamba_pkg + if fqn in sys.modules: + return sys.modules[fqn] + spec = importlib.util.spec_from_file_location( + fqn, + mamba_dir / file_name, + submodule_search_locations=[str(mamba_dir)] if file_name == "__init__.py" else [], + ) + mod = importlib.util.module_from_spec(spec) + sys.modules[fqn] = mod + spec.loader.exec_module(mod) + return mod + + # 1. Package __init__ (defines PAD_SLOT_ID = -1) + _load("", "__init__.py") + # 2. softplus helper (used by both kernel modules) + _load("softplus", "softplus.py") + # 3. The actual kernels + replay_mod = _load("replay_selective_state_update", "replay_selective_state_update.py") + checkpoint_mod = _load("checkpointing_state_update_refactored", "checkpointing_state_update_refactored.py") + base_mod = _load("selective_state_update", "selective_state_update.py") + conv1d_mod = _load("causal_conv1d_triton", "causal_conv1d_triton.py") + + return ( + replay_mod.replay_selective_state_update, + checkpoint_mod.checkpointing_state_update, + base_mod.selective_state_update, + conv1d_mod.causal_conv1d_update, + ) + + +def _import_mamba_kernels_full(): + """Import via the standard tensorrt_llm package (slow but safe).""" + from tensorrt_llm._torch.modules.mamba.causal_conv1d_triton import causal_conv1d_update + from tensorrt_llm._torch.modules.mamba.checkpointing_state_update_refactored import ( + checkpointing_state_update, + ) + from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( + replay_selective_state_update, + ) + from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update + + return ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) + + +# Use fast import by default; --full-import parsed later but we need the +# functions at module level. Check sys.argv early. +if "--full-import" in sys.argv: + ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) = _import_mamba_kernels_full() +else: + try: + ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) = _import_mamba_kernels_fast() + except Exception as e: # noqa: BLE001 - exit loudly; don't hide a fast-import regression + print( + f"ERROR: fast import failed ({type(e).__name__}: {e})\n" + "Re-run with --full-import for the slow but stable path, " + "then file a bug or fix _import_mamba_kernels_fast.", + file=sys.stderr, + ) + sys.exit(1) + + +_VARIANT_FNS = { + "replay": lambda: replay_selective_state_update, + "checkpointing": lambda: checkpointing_state_update, +} + +# Model config defaults (Nemotron-3-Super-120B full model). +# --tp-size divides nheads and ngroups to get the per-GPU slice. +# TP=1: nheads=128, ngroups=8 +# TP=4: nheads=32, ngroups=2 +# TP=8: nheads=16, ngroups=1 (default) +NHEADS = 128 +HEAD_DIM = 64 +D_STATE = 128 +NGROUPS = 8 +TP_SIZE = 8 # default; overridden by --tp-size + +# L2 flush buffer: ~128 MB — larger than L2 on A100/H100/B200 +_L2_FLUSH_SIZE = 32 * 1024 * 1024 # float32 elements → 128 MB +_l2_flush: torch.Tensor | None = None + + +def _init_l2_flush() -> None: + global _l2_flush + _l2_flush = torch.empty(_L2_FLUSH_SIZE, dtype=torch.float32, device="cuda") + + +def _flush_l2() -> None: + """Evict L2 by writing to a large buffer then synchronising.""" + assert _l2_flush is not None + _l2_flush.fill_(0.0) + torch.cuda.synchronize() + + +def _resolve_prev_ks(args, mtp_len: int) -> list[int]: + """Resolve prev_k values for one mtp_len cell. + + Two input modes (mutually exclusive in spirit; absolute wins if both given): + --prev-tokens-int "0,10,11,16" → use literal integers, clamped to + [0, max_window] (where max_window is the cache T-axis capacity). + --prev-tokens-fracs "0,0.5,1.0" → fractions of mtp_len, clamped to + [0, mtp_len] (current behavior). + + For replay-style checkpointing the cache holds up to max_window old + tokens, so absolute integers are the right knob. Fractions are kept + for back-compat with prior placeholder runs. + """ + upper = getattr(args, "max_window", 0) or mtp_len + if getattr(args, "prev_tokens_int", None): + return sorted(set(max(0, min(upper, int(v))) for v in args.prev_tokens_int)) + return sorted( + set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs) + ) + + +# Tensor construction helpers + +# Module-level cache for tensor buffers shared across cells. Keyed by all +# the "fixed" dimensions (state_dtype, act_dtype, max_window, mtp_len, +# nheads, head_dim, d_state, ngroups). Within a key, the batch dim grows +# in place: if a new cell requests a batch <= cached max_batch, we return +# views (slices) of the existing tensors; if batch > cached max_batch, we +# realloc at the new batch (which becomes the new max). Tensors never shrink. +# +# Rationale: torch.randn/zeros for these tensor shapes at b=512 takes +# ~10-30ms per call. At ~895 cells/min with 5 different batch sizes, +# we were re-allocating every cell. Caching saves the bulk of that per-cell +# overhead, raising GPU util in the timing phase. +# +# Reset state lives in caller (state_work = state0.copy_), so cached state0 +# is purely a reference whose contents stay fixed once allocated. This is +# fine: it's only read by the reset path. +_TENSOR_CACHE: dict = {} + + +def _build_tensors( + batch: int, + mtp_len: int, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + nheads: int, + head_dim: int, + d_state: int, + ngroups: int, + max_window: int | None = None, +): + """ + Build all tensors for one benchmark configuration. + + nheads/ngroups are already TP-split (i.e. full_nheads // tp_size). + + Returns: + state0 : (batch, nheads, head_dim, d_state) – initial SSM state + x, dt, B, C : (batch, mtp_len, ...) – token inputs for both kernels + A, dt_bias, D : SSM parameters (float32, tie_hdim strides) + prev_tokens : (batch,) + out_incr : pre-allocated output for replay kernel (batch, mtp_len, nheads, head_dim) + out_base : pre-allocated output for baseline kernel (batch, mtp_len, nheads, head_dim) + intermediate_states_buffer: for baseline kernel (batch, mtp_len, nheads, head_dim, d_state) + """ + device = "cuda" + + # Cache lookup — grow batch in place if needed; else return views. + cache_key = (state_dtype, act_dtype, max_window, mtp_len, + nheads, head_dim, d_state, ngroups) + cached = _TENSOR_CACHE.get(cache_key) + if cached is not None and cached["max_batch"] >= batch: + # Hit — return slices for current batch. + b = batch + return ( + cached["state0"][:b], + cached["state_scales0"][:b] if cached["state_scales0"] is not None else None, + cached["old_x"][:b], + cached["old_B"][:b], + cached["old_dt"][:b], + cached["old_dA_cumsum"][:b], + cached["cache_buf_idx"][:b], + cached["x"][:b], + cached["dt"][:b], + cached["B"][:b], + cached["C"][:b], + cached["A"], + cached["dt_bias"], + cached["D"], + cached["prev_tokens"][:b], + cached["slot_perm_buf"][:b], + cached["out_incr"][:b], + cached["out_base"][:b], + cached["intermediate_states_buffer"][:b], + cached["xbc_input"][:b], + cached["conv_state"][:b], + cached["conv_weight"], + cached["conv_bias"], + cached["d_inner"], + cached["conv_dim"], + ) + + # Miss or grow. Allocate at new max_batch (existing data, if any, is + # released — caller code re-fills via reset paths anyway). Rebind + # `batch` locally to alloc_batch so the existing allocation code below + # uses the larger size; keep request_batch for the final slice. + request_batch = batch + alloc_batch = batch if cached is None else max(batch, cached["max_batch"]) + batch = alloc_batch + + torch.manual_seed(42) + + # --- SSM parameters (float32, tie_hdim strides) --- + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) # stride(-1)=0, stride(-2)=0 + + dt_bias_base = torch.randn(nheads, device=device, dtype=torch.float32) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) # stride(-1)=0 + + D_base = torch.randn(nheads, device=device, dtype=torch.float32) + D = repeat(D_base, "h -> h p", p=head_dim) + + # --- SSM state --- + # Quantized dtypes need their own initializer (torch.randn doesn't accept + # int) and a parallel fp32 scales tensor (per-(head, dim) channel decode + # scale, broadcast over dstate). Quant state is filled with realistic- + # range values via fp32 → quant; scales are derived consistently so the + # initial state isn't garbage on dequant. + _QUANT_BENCH = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, + } + if state_dtype in _QUANT_BENCH: + quant_max = _QUANT_BENCH[state_dtype] + state_fp32 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + amax = state_fp32.abs().amax(dim=-1) # (batch, nheads, head_dim) + encode_scale = quant_max / amax.clamp(min=1e-30) + state_scales0 = (1.0 / encode_scale).to(torch.float32) # decode scale + scaled = state_fp32 * encode_scale.unsqueeze(-1) + if state_dtype == torch.float8_e4m3fn: + state0 = scaled.clamp(-quant_max, quant_max).to(state_dtype) + else: + state0 = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) + else: + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state_scales0 = None + + # --- Cache tensors for replay kernel --- + # max_window is the cache T-axis capacity; defaults to mtp_len (the + # placeholder/degenerate case where every step is a checkpoint step). + # For real replay-style checkpointing, max_window > mtp_len. + cache_T = max_window if max_window is not None else mtp_len + # old_x: single-buffered (cache, max_window, nheads, dim) + old_x = torch.randn(batch, cache_T, nheads, head_dim, device=device, dtype=act_dtype) + # old_B: double-buffered (cache, 2, max_window, ngroups, dstate) + old_B = torch.randn(batch, 2, cache_T, ngroups, d_state, device=device, dtype=act_dtype) + # old_dt: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous + old_dt = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) + # old_dA_cumsum: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous + old_dA_cumsum = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) + # cache_buf_idx: which buffer to read (0 or 1) + cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) + + # --- Token inputs (used by both replay and baseline kernels) --- + x = torch.randn(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) + # dt must match D's dtype (fp32) for flashinfer — force it for all paths. + dt_base = torch.randn(batch, mtp_len, nheads, device=device, dtype=torch.float32) + dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) # tie_hdim + B = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) + C = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) + + # prev_tokens placeholder — overwritten per-run + prev_tokens = torch.zeros(batch, device=device, dtype=torch.int32) + # slot_perm placeholder — overwritten per-run by mix pre_iter_fn when + # sort_slots is enabled. Identity by default so cells that don't sort + # (or pure-batch cells) get a meaningful identity perm if the kernel + # ends up reading it (USE_PERM=False makes this path unused). + slot_perm_buf = torch.arange(batch, device=device, dtype=torch.int32) + + out_incr = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) + out_base = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) + + # intermediate_states_buffer is only consumed by the fp/baseline path; + # for quantized state dtypes we'll skip baselines entirely, so the buffer + # dtype falls back to fp32 to keep selective_state_update happy. + int_buffer_dtype = state_dtype if state_dtype not in _QUANT_BENCH else torch.float32 + intermediate_states_buffer = torch.zeros( + batch, mtp_len, nheads, head_dim, d_state, device=device, dtype=int_buffer_dtype + ) + + # --- Conv1d tensors (for --with-conv1d mode) --- + d_inner = nheads * head_dim + conv_dim = d_inner + 2 * ngroups * d_state + d_conv = 4 # conv kernel width for Nemotron/Mamba2 + + # xbc_input: (batch, conv_dim, mtp_len) — "hot" input from in_proj. + # Match production layout: in_proj output is (batch*mtp_len, conv_dim) + # contiguous, then .view(batch, mtp_len, conv_dim).transpose(1, 2) + # gives strides (mtp_len*conv_dim, 1, conv_dim) — NOT the standard + # (conv_dim*mtp_len, mtp_len, 1) of a freshly allocated 3D tensor. + # Conv1d preserves input strides in its output, so downstream split + # + view inherits the correct layout without needing .contiguous(). + xbc_input_flat = torch.randn(batch * mtp_len, conv_dim, device=device, dtype=act_dtype) + xbc_input = xbc_input_flat.view(batch, mtp_len, conv_dim).transpose(1, 2) + # conv_state: (batch, conv_dim, d_conv) — "cold" cache + conv_state = torch.randn(batch, conv_dim, d_conv, device=device, dtype=act_dtype) + # conv_weight: (conv_dim, d_conv) — parameter + conv_weight = torch.randn(conv_dim, d_conv, device=device, dtype=act_dtype) + # conv_bias: (conv_dim,) — parameter + conv_bias = torch.randn(conv_dim, device=device, dtype=act_dtype) + + # Store full-batch buffers in cache and return slices at request_batch. + _TENSOR_CACHE[cache_key] = { + "max_batch": alloc_batch, + "state0": state0, + "state_scales0": state_scales0, + "old_x": old_x, + "old_B": old_B, + "old_dt": old_dt, + "old_dA_cumsum": old_dA_cumsum, + "cache_buf_idx": cache_buf_idx, + "x": x, + "dt": dt, + "B": B, + "C": C, + "A": A, + "dt_bias": dt_bias, + "D": D, + "prev_tokens": prev_tokens, + "slot_perm_buf": slot_perm_buf, + "out_incr": out_incr, + "out_base": out_base, + "intermediate_states_buffer": intermediate_states_buffer, + "xbc_input": xbc_input, + "conv_state": conv_state, + "conv_weight": conv_weight, + "conv_bias": conv_bias, + "d_inner": d_inner, + "conv_dim": conv_dim, + } + rb = request_batch + return ( + state0[:rb], + state_scales0[:rb] if state_scales0 is not None else None, + old_x[:rb], + old_B[:rb], + old_dt[:rb], + old_dA_cumsum[:rb], + cache_buf_idx[:rb], + x[:rb], + dt[:rb], + B[:rb], + C[:rb], + A, + dt_bias, + D, + prev_tokens[:rb], + slot_perm_buf[:rb], + out_incr[:rb], + out_base[:rb], + intermediate_states_buffer[:rb], + xbc_input[:rb], + conv_state[:rb], + conv_weight, + conv_bias, + d_inner, + conv_dim, + ) + + +# ============================================================================= +# CUPTI in-process kernel timing +# +# Self-contained module-in-a-file. Reads kernel start/end timestamps directly +# from the GPU profiling fabric via CUPTI's Activity API (1 ns +# resolution), avoiding two pitfalls of the cuda-events path: +# +# 1. cudaEvent.elapsed_time() resolution (~0.5 us) is too coarse for the +# short kernels we care about, especially with PDL + cuda graphs at +# small batch — events recorded inside a graph have proven noisy. +# 2. nsys is the only known accurate alternative, but the +# profile-export-sqlite-parse pipeline is heavy and out-of-process. +# +# This is functionally equivalent to wrapping each cell in nsys, except it +# runs in the same benchmark process and sends raw activity buffers to a +# parser process instead of materializing Python objects in the CUPTI callback. +# ============================================================================= + + +# Substring match: kernels run_fn launches that we want to time. Mirrors +# the parser in scripts/.../collect.py so cupti and nsys-based outputs agree. +_CUPTI_KEEP_KERNEL_SUBSTRINGS = ( + "_replay_precompute", + "_checkpointing_precompute", + "_rectangle_precompute", + "_dynamic_precompute", + "_replay_state_update", + "_checkpointing_main", + "_rectangle_main", + "_dynamic_main", + "_persistent_main", + "selective_scan_update", + "selective_state_update", + "causal_conv1d_update", +) + + +def _kernels_per_iter_incremental( + mode: str, + with_conv1d: bool, + *, + persistent_skip_empty: bool = True, +) -> int: + """Expected number of CUPTI-tracked kernels per iter for the incremental + kernel chain, given the dispatch mode and the conv1d flag. + + Used to validate CUPTI record counts (no auto-inference — silent + mis-timing is the failure mode we're guarding against). + + `persistent_skip_empty=True` (today's behavior): the + `mode='persistent_main'` launch helper host-early-outs when its half + is empty (n_writes=0 or n_writes=batch in pure scenarios), so only + one of the two persistent_main_kernel launches actually fires per + iter. With `persistent_skip_empty=False` (future no-eo mode), both + halves always launch and K bumps by 1. + + `persistent_dynamic` always launches 1 main; not affected by the flag. + """ + if mode == "monolithic": + k = 2 # precomp + main + elif mode == "dynamic": + k = 2 # dynamic_precomp + dynamic_main + elif mode == "maindl": + k = 3 # 1 dynamic_precomp + 2 mains (write + nowrite) + elif mode in ("doublelaunch", "dlgrouped"): + k = 4 # 2 precomp + 2 main + elif mode == "dl_write_only": + k = 2 # 1 precomp + 1 main (write only) + elif mode == "persistent_dynamic": + k = 2 # 1 dynamic_precomp + 1 persistent_main + elif mode == "persistent_main": + k = 2 if persistent_skip_empty else 3 # see docstring + else: + raise ValueError(f"_kernels_per_iter_incremental: unknown mode {mode!r}") + if with_conv1d: + k += 1 + return k + + +def _kernels_per_iter_baseline(with_conv1d: bool) -> int: + """Expected kernels per iter for triton / flashinfer baselines. + + Both baselines run a single state-update kernel; `--with-conv1d` + prepends one conv1d kernel. + """ + return 2 if with_conv1d else 1 + + +_LIBCUPTI_CANDIDATES = ( + os.environ.get("CUPTI_LIBRARY_PATH"), + "/usr/local/lib/python3.12/dist-packages/nvidia/cu13/lib/libcupti.so.13", + "libcupti.so.13", + "libcupti.so", +) +_CUPTI_SUCCESS = 0 +_CUPTI_ERROR_MAX_LIMIT_REACHED = 12 +_CUPTI_ERROR_INVALID_KIND = 21 +_CUPTI_ACTIVITY_KIND_KERNEL = 3 +_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL = 10 +_CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER = 5 +_CUPTI_HOST_BUFFER_BYTES = 1024 * 1024 +_CUPTI_HOST_BUFFER_COUNT = 16 + +# Multiprocessing start method for compile-warmup + CUPTI parser children. +# Set in __main__ from --mp-start-method. "spawn" (default) is robust; each +# child re-imports torch/triton/etc (~15s). "forkserver" preloads once and +# forks cheaply (~1s/child) — see __main__ block for the preload setup. +_MP_START_METHOD = "spawn" +_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE = 1 +_DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX = 4 + + +def _load_libcupti() -> ctypes.CDLL: + errors = [] + for candidate in _LIBCUPTI_CANDIDATES: + if not candidate: + continue + try: + return ctypes.CDLL(candidate) + except OSError as exc: + errors.append(f"{candidate}: {exc}") + raise ImportError("Unable to load libcupti: " + "; ".join(errors)) + + +class _CuptiActivityKernel11Prefix(ctypes.Structure): + _pack_ = 1 + _fields_ = [ + ("kind", ctypes.c_int), + ("cache_config", ctypes.c_uint8), + ("shared_memory_config", ctypes.c_uint8), + ("registers_per_thread", ctypes.c_uint16), + ("partitioned_global_cache_requested", ctypes.c_int), + ("partitioned_global_cache_executed", ctypes.c_int), + ("start", ctypes.c_uint64), + ("end", ctypes.c_uint64), + ("completed", ctypes.c_uint64), + ("device_id", ctypes.c_uint32), + ("context_id", ctypes.c_uint32), + ("stream_id", ctypes.c_uint32), + ("grid_x", ctypes.c_int32), + ("grid_y", ctypes.c_int32), + ("grid_z", ctypes.c_int32), + ("block_x", ctypes.c_int32), + ("block_y", ctypes.c_int32), + ("block_z", ctypes.c_int32), + ("static_shared_memory", ctypes.c_int32), + ("dynamic_shared_memory", ctypes.c_int32), + ("local_memory_per_thread", ctypes.c_uint32), + ("local_memory_total", ctypes.c_uint32), + ("correlation_id", ctypes.c_uint32), + ("grid_id", ctypes.c_int64), + ("name", ctypes.c_void_p), + ("reserved0", ctypes.c_void_p), + ("queued", ctypes.c_uint64), + ("submitted", ctypes.c_uint64), + ("launch_type", ctypes.c_uint8), + ("is_shared_memory_carveout_requested", ctypes.c_uint8), + ("shared_memory_carveout_requested", ctypes.c_uint8), + ("padding", ctypes.c_uint8), + ("shared_memory_executed", ctypes.c_uint32), + ("graph_node_id", ctypes.c_uint64), + ] + + +def _configure_cupti_get_next_record(libcupti) -> None: + libcupti.cuptiActivityGetNextRecord.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_void_p), + ] + libcupti.cuptiActivityGetNextRecord.restype = ctypes.c_int + + +def _parse_cupti_buffer_ptr(libcupti, buffer_ptr: int, valid_size: int, *, include_names: bool): + records = [] + zero_ts_count = 0 + zero_ts_names: dict[str, int] = {} + record_ptr = ctypes.c_void_p(None) + while True: + result = libcupti.cuptiActivityGetNextRecord( + ctypes.c_void_p(buffer_ptr), + valid_size, + ctypes.byref(record_ptr), + ) + if result == _CUPTI_SUCCESS: + kind = ctypes.cast(record_ptr, ctypes.POINTER(ctypes.c_int)).contents.value + if kind not in (_CUPTI_ACTIVITY_KIND_KERNEL, _CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL): + continue + kernel = ctypes.cast(record_ptr, ctypes.POINTER(_CuptiActivityKernel11Prefix)).contents + name = None + if include_names: + if kernel.name: + name = ctypes.string_at(kernel.name).decode("utf-8", errors="replace") + else: + name = "?" + if kernel.start == 0 or kernel.end == 0: + zero_ts_count += 1 + if name is not None: + zero_ts_names[name] = zero_ts_names.get(name, 0) + 1 + continue + if include_names: + records.append(( + name, + int(kernel.start), + int(kernel.end), + int(kernel.correlation_id), + 0, + int(kernel.graph_node_id), + int(kernel.stream_id), + )) + else: + records.append(( + int(kernel.start), + int(kernel.end), + int(kernel.correlation_id), + int(kernel.graph_node_id), + int(kernel.stream_id), + )) + elif result == _CUPTI_ERROR_MAX_LIMIT_REACHED: + break + elif result == _CUPTI_ERROR_INVALID_KIND: + break + else: + raise RuntimeError(f"cuptiActivityGetNextRecord failed with CUptiResult={result}") + return records, zero_ts_count, zero_ts_names + + +def _apply_cupti_filter_plan(numeric_records, filter_plan): + if not filter_plan: + return [ + (None, start, end, corr, 0, graph_node_id, stream_id) + for start, end, corr, graph_node_id, stream_id in sorted(numeric_records) + ] + + filtered = [] + replay_idx = 0 + record_idx = 0 + for start, end, corr, graph_node_id, stream_id in sorted(numeric_records): + if replay_idx >= len(filter_plan): + break + records_per_replay, ordinal_names = filter_plan[replay_idx] + if record_idx < len(ordinal_names): + name = ordinal_names[record_idx] + if name is not None: + filtered.append((name, start, end, corr, 0, graph_node_id, stream_id)) + record_idx += 1 + if record_idx >= records_per_replay: + replay_idx += 1 + record_idx = 0 + return filtered + + +def _cupti_parser_worker(input_queue, output_queue, ready_event) -> None: + libcupti = _load_libcupti() + _configure_cupti_get_next_record(libcupti) + shared_blocks: dict[str, shared_memory.SharedMemory] = {} + records_by_generation: dict[int, list[tuple[int, int, int, int, int]]] = {} + zero_ts_by_generation: dict[int, int] = {} + ready_event.set() + while True: + item = input_queue.get() + if item is None: + break + kind = item[0] + if kind == "buffer": + _, generation, buffer_id, name, valid_size = item + shm = shared_blocks.get(name) + if shm is None: + shm = shared_memory.SharedMemory(name=name) + shared_blocks[name] = shm + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + parser_ptr = ctypes.addressof(shared_char) + records, zero_ts_count, _ = _parse_cupti_buffer_ptr( + libcupti, + parser_ptr, + valid_size, + include_names=False, + ) + records_by_generation.setdefault(generation, []).extend(records) + zero_ts_by_generation[generation] = zero_ts_by_generation.get(generation, 0) + zero_ts_count + ctypes.memset(parser_ptr, 0, len(shm.buf)) + except Exception as exc: # pragma: no cover - diagnostic worker path + output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) + finally: + del shared_char + output_queue.put({"kind": "buffer_done", "generation": generation, "buffer_id": buffer_id}) + elif kind == "finish": + if len(item) == 4: + _, generation, filter_plan, stats_request = item + else: + _, generation, filter_plan = item + stats_request = None + try: + raw_records = records_by_generation.pop(generation, []) + zero_ts_count = zero_ts_by_generation.pop(generation, 0) + filtered_records = _apply_cupti_filter_plan(raw_records, filter_plan) + stats = None + parser_stats_ms = 0.0 + stats_ready = stats_request is not None + if stats_request is not None: + stats_start_s = time.perf_counter() + stats = _stats_from_cupti_records( + filtered_records, + int(stats_request["warmup"]), + int(stats_request["iters"]), + str(stats_request["tag"]), + int(stats_request["expected_K"]), + zero_ts_count=zero_ts_count, + zero_ts_names={}, + include_details=bool(stats_request.get("include_details", True)), + ) + parser_stats_ms = 1000.0 * (time.perf_counter() - stats_start_s) + filtered_records = [] + output_queue.put({ + "kind": "finish_done", + "generation": generation, + "records": filtered_records, + "zero_ts_count": zero_ts_count, + "zero_ts_names": {}, + "raw_record_count": len(raw_records), + "stats": stats, + "stats_ready": stats_ready, + "parser_stats_ms": parser_stats_ms, + }) + except Exception as exc: # pragma: no cover - diagnostic worker path + output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) + else: + output_queue.put({"kind": "error", "generation": -1, "error": f"unknown parser message {kind!r}"}) + for shm in shared_blocks.values(): + shm.close() + + +class CuptiKernelTimer: + """Raw CUPTI Activity timer with out-of-process parsing for timed runs. + + CUPTI's callback gives us raw activity buffers. The callback only hands + shared-memory buffer metadata to a parser process, so the main process + avoids the cupti-python per-record object creation cost during the timed + path. A single local calibration replay may parse names in-process to + build an ordinal filter plan for a just-captured CUDA graph. + """ + + _instance = None + _import_error = None + + _request_callback_type = ctypes.CFUNCTYPE( + None, + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t), + ) + _complete_callback_type = ctypes.CFUNCTYPE( + None, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_size_t, + ) + + @classmethod + def get(cls) -> "CuptiKernelTimer": + if cls._instance is not None: + return cls._instance + if cls._import_error is not None: + raise cls._import_error + try: + cls._instance = cls() + return cls._instance + except ImportError as exc: # pragma: no cover - env-dependent + cls._import_error = exc + raise + + def __init__(self) -> None: + self._libcupti = _load_libcupti() + self._configure_functions() + self._lock = threading.Lock() + self._shared_buffers: dict[int, shared_memory.SharedMemory] = {} + self._buffer_id_by_ptr: dict[int, int] = {} + self._free_buffer_ids: list[int] = [] + self._local_completed: list[tuple[int, int]] = [] + self._mode = "drop" + self._generation = 0 + self._finish_results: dict[int, dict] = {} + self._parser_errors: list[str] = [] + self._filter_plan = () + self._last_start_timing: dict[str, float] = {} + self._last_stop_timing: dict[str, float] = {} + self._current_flush_period_ms = 0 + self._mp_ctx = mp.get_context(_MP_START_METHOD) + # Retry parser-process spawn: concurrent bench instances on the same + # node race on POSIX named semaphores in /dev/shm — child can die in + # pickle.load with FileNotFoundError in SemLock._rebuild before + # signalling ready_event. Detect early-dead child via is_alive() so + # we don't waste the full timeout, and retry up to 3x with jitter. + last_err = None + for _spawn_attempt in range(3): + self._parse_input_queue = self._mp_ctx.Queue() + self._parse_output_queue = self._mp_ctx.Queue() + ready_event = self._mp_ctx.Event() + self._parse_process = self._mp_ctx.Process( + target=_cupti_parser_worker, + args=(self._parse_input_queue, self._parse_output_queue, ready_event), + ) + self._parse_process.start() + deadline = time.time() + 30.0 + spawn_ok = False + while time.time() < deadline: + if ready_event.wait(timeout=0.5): + spawn_ok = True + break + if not self._parse_process.is_alive(): + break + if spawn_ok: + last_err = None + break + last_err = (f"attempt {_spawn_attempt + 1}: " + f"alive={self._parse_process.is_alive()}, " + f"exitcode={self._parse_process.exitcode}") + try: + if self._parse_process.is_alive(): + self._parse_process.terminate() + self._parse_process.join(timeout=2.0) + except Exception: + pass + time.sleep(0.5 + 0.5 * _spawn_attempt) + if last_err is not None: + raise RuntimeError( + f"CUPTI parser process did not initialize after 3 attempts: {last_err}" + ) + + self._set_zeroed_host_buffer_attr() + for _ in range(_CUPTI_HOST_BUFFER_COUNT): + self._free_buffer_ids.append(self._allocate_shared_buffer()) + + self._request_callback = self._request_callback_type(self._request_buffer) + self._complete_callback = self._complete_callback_type(self._complete_buffer) + self._check(self._libcupti.cuptiActivityRegisterCallbacks( + self._request_callback, + self._complete_callback, + )) + self._check(self._libcupti.cuptiActivityEnable(_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL)) + atexit.register(self.close) + + def _configure_functions(self) -> None: + self._libcupti.cuptiActivityRegisterCallbacks.argtypes = [ + self._request_callback_type, + self._complete_callback_type, + ] + self._libcupti.cuptiActivityRegisterCallbacks.restype = ctypes.c_int + self._libcupti.cuptiActivityEnable.argtypes = [ctypes.c_int] + self._libcupti.cuptiActivityEnable.restype = ctypes.c_int + self._libcupti.cuptiActivityFlushAll.argtypes = [ctypes.c_uint32] + self._libcupti.cuptiActivityFlushAll.restype = ctypes.c_int + self._libcupti.cuptiActivityFlushPeriod.argtypes = [ctypes.c_uint32] + self._libcupti.cuptiActivityFlushPeriod.restype = ctypes.c_int + self._libcupti.cuptiActivitySetAttribute.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_void_p, + ] + self._libcupti.cuptiActivitySetAttribute.restype = ctypes.c_int + _configure_cupti_get_next_record(self._libcupti) + + def _set_zeroed_host_buffer_attr(self) -> None: + value_obj = ctypes.c_uint8(1) + size_obj = ctypes.c_size_t(ctypes.sizeof(value_obj)) + result = self._libcupti.cuptiActivitySetAttribute( + _CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER, + ctypes.byref(size_obj), + ctypes.byref(value_obj), + ) + if result != _CUPTI_SUCCESS: + print( + "[WARN] CUPTI zeroed host-buffer attribute failed; " + f"continuing with default CUPTI buffer handling (CUptiResult={result}).", + file=sys.stderr, + ) + + def _check(self, result: int) -> None: + if result != _CUPTI_SUCCESS: + raise RuntimeError(f"CUPTI call failed with CUptiResult={result}") + + def _allocate_shared_buffer(self) -> int: + buffer_id = len(self._shared_buffers) + shm = shared_memory.SharedMemory(create=True, size=_CUPTI_HOST_BUFFER_BYTES) + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + ptr = ctypes.addressof(shared_char) + finally: + del shared_char + if ptr % 8 != 0: + shm.close() + shm.unlink() + raise RuntimeError("CUPTI shared-memory activity buffer was not 8-byte aligned") + self._shared_buffers[buffer_id] = shm + self._buffer_id_by_ptr[ptr] = buffer_id + return buffer_id + + def _buffer_ptr(self, buffer_id: int) -> int: + shm = self._shared_buffers[buffer_id] + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + return ctypes.addressof(shared_char) + finally: + del shared_char + + def _request_buffer(self, buffer, size, max_num_records) -> None: + with self._lock: + if self._free_buffer_ids: + buffer_id = self._free_buffer_ids.pop() + else: + buffer_id = self._allocate_shared_buffer() + ptr = self._buffer_ptr(buffer_id) + buffer[0] = ptr + size[0] = _CUPTI_HOST_BUFFER_BYTES + max_num_records[0] = 0 + + def _complete_buffer(self, context, stream_id, buffer, size, valid_size) -> None: + del context, stream_id, size + buffer_ptr = int(buffer) + valid_size_int = int(valid_size) + with self._lock: + mode = self._mode + generation = self._generation + buffer_id = self._buffer_id_by_ptr[buffer_ptr] + if valid_size_int == 0 or mode == "drop": + self._free_buffer_ids.append(buffer_id) + return + if mode == "local": + self._local_completed.append((buffer_id, valid_size_int)) + return + shm = self._shared_buffers[buffer_id] + self._parse_input_queue.put(("buffer", generation, buffer_id, shm.name, valid_size_int)) + + def _handle_parser_result(self, result: dict) -> None: + kind = result.get("kind") + if kind == "buffer_done": + with self._lock: + self._free_buffer_ids.append(int(result["buffer_id"])) + elif kind == "finish_done": + self._finish_results[int(result["generation"])] = result + elif kind == "error": + self._parser_errors.append(str(result.get("error"))) + + def _drain_parser_results(self) -> None: + while True: + try: + result = self._parse_output_queue.get_nowait() + except queue.Empty: + break + self._handle_parser_result(result) + + def is_generation_ready(self, generation: int) -> bool: + self._drain_parser_results() + return generation in self._finish_results or bool(self._parser_errors) + + def _flush(self, flag: int) -> None: + self._check(self._libcupti.cuptiActivityFlushAll(flag)) + + def _set_flush_period_ms(self, period_ms: int) -> None: + if period_ms == self._current_flush_period_ms: + return + self._check(self._libcupti.cuptiActivityFlushPeriod(period_ms)) + self._current_flush_period_ms = period_ms + + def _begin( + self, + mode: str, + filter_plan=(), + flush_period_ms: int = 0, + collect_timing: bool = False, + ) -> int: + start_timing: dict[str, float] = {} + with self._lock: + self._mode = "drop" + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._flush(1) + if collect_timing: + start_timing["forced_flush_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._drain_parser_results() + if collect_timing: + start_timing["drain_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) + with self._lock: + self._generation += 1 + generation = self._generation + self._mode = mode + self._local_completed = [] + self._filter_plan = filter_plan + if flush_period_ms > 0: + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._set_flush_period_ms(flush_period_ms) + if collect_timing: + start_timing["period_enable_ms"] = 1000.0 * ( + time.perf_counter() - phase_start_s + ) + self._last_start_timing = start_timing + return generation + + def capture_names(self, replay_fn) -> tuple[list[tuple], int, dict]: + """Run a small calibration replay and parse kernel names locally.""" + self._begin("local") + replay_fn() + torch.cuda.synchronize() + self._flush(0) + records: list[tuple] = [] + zero_ts_count = 0 + zero_ts_names: dict[str, int] = {} + with self._lock: + completed = list(self._local_completed) + self._local_completed = [] + self._mode = "drop" + for buffer_id, valid_size in completed: + ptr = self._buffer_ptr(buffer_id) + recs, zeros, zero_names = _parse_cupti_buffer_ptr( + self._libcupti, + ptr, + valid_size, + include_names=True, + ) + records.extend(recs) + zero_ts_count += zeros + for name, count in zero_names.items(): + zero_ts_names[name] = zero_ts_names.get(name, 0) + count + ctypes.memset(ptr, 0, _CUPTI_HOST_BUFFER_BYTES) + with self._lock: + self._free_buffer_ids.append(buffer_id) + records.sort(key=lambda r: r[1]) + return records, zero_ts_count, zero_ts_names + + def start( + self, + filter_plan=(), + flush_period_ms: int = 0, + collect_timing: bool = False, + ) -> None: + self._begin("parser", filter_plan, flush_period_ms, collect_timing) + + def stop_async( + self, + collect_timing: bool = False, + stats_request: dict | None = None, + ) -> tuple[int, dict[str, float]]: + stop_timing: dict[str, float] = {} + generation = self._generation + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._set_flush_period_ms(0) + if collect_timing: + stop_timing["period_disable_ms"] = 1000.0 * ( + time.perf_counter() - phase_start_s + ) + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._flush(0) + if collect_timing: + stop_timing["flush_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) + with self._lock: + self._mode = "drop" + filter_plan = self._filter_plan + self._parse_input_queue.put(("finish", generation, filter_plan, stats_request)) + self._last_stop_timing = stop_timing + return generation, stop_timing + + def wait_for_generation_result( + self, + generation: int, + stop_timing: dict[str, float] | None = None, + collect_timing: bool = False, + ) -> dict: + if stop_timing is None: + stop_timing = {} + phase_start_s = time.perf_counter() if collect_timing else 0.0 + deadline = time.perf_counter() + 10.0 + while time.perf_counter() < deadline: + result = self._finish_results.pop(generation, None) + if result is not None: + if collect_timing: + stop_timing["parser_wait_ms"] = 1000.0 * ( + time.perf_counter() - phase_start_s + ) + stop_timing["total_ms"] = ( + stop_timing.get("period_disable_ms", 0.0) + + stop_timing.get("flush_ms", 0.0) + + stop_timing["parser_wait_ms"] + ) + self._last_stop_timing = stop_timing + return result + timeout_s = max(0.0, min(0.01, deadline - time.perf_counter())) + try: + parser_result = self._parse_output_queue.get(timeout=timeout_s) + except queue.Empty: + continue + self._handle_parser_result(parser_result) + if self._parser_errors: + raise RuntimeError("CUPTI parser process failed: " + "; ".join(self._parser_errors)) + raise TimeoutError("Timed out waiting for CUPTI parser process") + + def wait_for_generation( + self, + generation: int, + stop_timing: dict[str, float] | None = None, + collect_timing: bool = False, + ) -> tuple[list[tuple], int, dict, int]: + result = self.wait_for_generation_result(generation, stop_timing, collect_timing) + return ( + list(result["records"]), + int(result["zero_ts_count"]), + dict(result["zero_ts_names"]), + int(result["raw_record_count"]), + ) + + def stop(self, collect_timing: bool = False) -> tuple[list[tuple], int, dict, int]: + generation, stop_timing = self.stop_async(collect_timing) + return self.wait_for_generation(generation, stop_timing, collect_timing) + + def last_start_timing(self) -> dict[str, float]: + return dict(self._last_start_timing) + + def last_stop_timing(self) -> dict[str, float]: + return dict(self._last_stop_timing) + + def close(self) -> None: + parse_process = getattr(self, "_parse_process", None) + if parse_process is not None and parse_process.is_alive(): + self._parse_input_queue.put(None) + parse_process.join(timeout=5.0) + if parse_process.is_alive(): + parse_process.terminate() + parse_process.join(timeout=1.0) + for shm in getattr(self, "_shared_buffers", {}).values(): + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + +# ============================================================================= +# Timing helpers +# ============================================================================= + + +def _stats_from_spans(spans_us: list[float]) -> dict: + """Compute median / p95 / p99 / n from a per-iter span list.""" + s = sorted(spans_us) + return { + "median": statistics.median(s), + "p95": s[int(0.95 * len(s))], + "p99": s[int(0.99 * len(s))], + "n": len(s), + } + + +def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, + zero_ts_count: int = 0, + zero_ts_names: dict | None = None, + include_details: bool = True): + """Bin a flat CUPTI kernel record stream into per-iter spans + per-kernel + relative timestamps. Used by both graph and eager CUPTI paths. + + `records` are tuples (name, start_ns, end_ns, ...) — see CuptiKernelTimer. + `expected_K` is the kernels-per-iter count the caller declares; we + validate the CUPTI total matches `expected_K * (warmup + iters)` exactly. + On mismatch we dump per-name record counts so missing or extra kernels + are obvious (most common cause: a new dispatch mode whose kernels lack + a matching entry in `_CUPTI_KEEP_KERNEL_SUBSTRINGS`, silently filtering + them out). + """ + records = [ + r for r in records + if r[0] is not None and any(s in r[0] for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS) + ] + records.sort(key=lambda r: r[1]) # by start_ns + + total = len(records) + expected_iters = warmup + iters + expected_total = expected_K * expected_iters + if total != expected_total: + from collections import Counter + name_counts = dict(Counter(r[0] for r in records)) + # Non-fatal: skip this cell instead of killing the whole sweep. + # Mismatch may be a CUPTI dropped-records issue (rare configs), + # not necessarily a K-table bug. Log so the user can investigate + # the specific cell post-hoc; return None so the caller can skip + # writing a JSON row. + zero_msg = "" + if zero_ts_count: + zero_msg = ( + f" + {zero_ts_count} records with start/end=0 " + f"(dropped by callback, breakdown {zero_ts_names}). " + f"Total observed kernel records (timed + zero-ts) = " + f"{total + zero_ts_count} / {expected_total}." + ) + print( + f"[WARN] CUPTI capture mismatch for {tag!r}: expected " + f"{expected_K} kernels/iter × {expected_iters} iters " + f"(warmup+iters) = {expected_total} records, got {total}. " + f"Kernel record counts: {name_counts}.{zero_msg} SKIPPING cell.", + file=sys.stderr, + flush=True, + ) + # Per-record dump: (name, start_ns_rel, end_ns_rel, corr_id, graph_id, stream_id). + # Times relative to first record so absolute ns isn't drowning output. + # Limit dump to first 30 records to avoid flooding logs at high K. + if records: + t0_ns = records[0][1] + for i, r in enumerate(records[:30]): + # r = (name, start_ns, end_ns, corr_id, graph_id, graph_node_id, stream_id) + rel_start = (r[1] - t0_ns) / 1000.0 # us + rel_end = (r[2] - t0_ns) / 1000.0 + print( + f" rec[{i:3d}] name={r[0]!r} start={rel_start:.2f}us " + f"end={rel_end:.2f}us corr={r[3]} graph={r[4]} stream={r[6]}", + file=sys.stderr, + flush=True, + ) + if len(records) > 30: + print(f" ... ({len(records) - 30} more records elided)", + file=sys.stderr, flush=True) + return None + K = expected_K + timed = records[warmup * K:] + + spans_us: list[float] = [] + per_kernel: dict[str, dict[str, list[float]]] = {} + for i in range(iters): + chunk = timed[i * K:(i + 1) * K] + iter_start_ns = min(r[1] for r in chunk) + iter_end_ns = max(r[2] for r in chunk) + spans_us.append((iter_end_ns - iter_start_ns) / 1000.0) + if include_details: + for r in chunk: + name = r[0] + slot = per_kernel.setdefault(name, {"start_us": [], "end_us": []}) + slot["start_us"].append((r[1] - iter_start_ns) / 1000.0) + slot["end_us"].append((r[2] - iter_start_ns) / 1000.0) + + out = _stats_from_spans(spans_us) + out["iters_us"] = spans_us + if include_details: + out["per_kernel"] = per_kernel + return out + + +_PRE_GRAPH_WARMUP_ITERS = 1 +_CUPTI_FILTER_PLAN_CACHE: dict[tuple, tuple[int, tuple[str | None, ...]]] = {} + + +class _HostTiming: + def __init__(self, enabled: bool) -> None: + self.enabled = enabled + self.values: dict[str, float | int | bool] = {} + self._total_start_s = time.perf_counter() if enabled else 0.0 + self._phase_start_s = 0.0 + + def start(self) -> None: + if self.enabled: + self._phase_start_s = time.perf_counter() + + def stop(self, key: str) -> None: + if self.enabled: + self.values[key] = 1000.0 * (time.perf_counter() - self._phase_start_s) + + def add(self, key: str, value: float | int | bool) -> None: + if self.enabled: + self.values[key] = value + + def stop_total(self) -> None: + if self.enabled: + self.values["total_ms"] = 1000.0 * (time.perf_counter() - self._total_start_s) + + def attach(self, stats: dict | None) -> None: + if self.enabled and stats is not None: + stats["host_timing"] = self.values + + +class _PendingCuptiStats: + + def __init__( + self, + timer: CuptiKernelTimer, + generation: int, + stop_timing: dict[str, float], + host_timing: _HostTiming, + *, + warmup: int, + iters: int, + tag: str, + expected_K: int, + expected_raw_record_count: int, + ) -> None: + self._timer = timer + self._generation = generation + self._stop_timing = stop_timing + self._host_timing = host_timing + self._warmup = warmup + self._iters = iters + self._tag = tag + self._expected_K = expected_K + self._expected_raw_record_count = expected_raw_record_count + + def is_ready(self) -> bool: + return self._timer.is_generation_ready(self._generation) + + def resolve(self) -> dict | None: + result = self._timer.wait_for_generation_result( + self._generation, + self._stop_timing, + collect_timing=self._host_timing.enabled, + ) + for key, value in self._timer.last_stop_timing().items(): + self._host_timing.add(f"cupti_stop_{key}", value) + raw_record_count = int(result["raw_record_count"]) + if raw_record_count != self._expected_raw_record_count: + print( + f"[WARN] CUPTI raw-record mismatch for {self._tag!r}: expected " + f"{self._expected_raw_record_count}, got {raw_record_count}. SKIPPING cell.", + file=sys.stderr, + ) + return None + + if result.get("stats_ready"): + stats = result.get("stats") + self._host_timing.add("stats_ms", 0.0) + self._host_timing.add("parser_stats_ms", float(result.get("parser_stats_ms", 0.0))) + else: + self._host_timing.start() + stats = _stats_from_cupti_records( + list(result["records"]), + self._warmup, + self._iters, + self._tag, + self._expected_K, + zero_ts_count=int(result["zero_ts_count"]), + zero_ts_names=dict(result["zero_ts_names"]), + ) + self._host_timing.stop("stats_ms") + self._host_timing.attach(stats) + return stats + + +def _target_name_or_none(name: str | None) -> str | None: + if name is None: + return None + if any(s in name for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS): + return name + return None + + +def _capture_group_graph( + args, + run_fn, + reset_fn, + group_iters: int, + graph_pre_iter_fn=None, +) -> torch.cuda.CUDAGraph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for j in range(group_iters): + if graph_pre_iter_fn is not None: + graph_pre_iter_fn(j) + reset_fn() + if args.l2_flush: + _l2_flush.fill_(0.0) + run_fn() + return graph + + +def _graph_group_iters(args, total_iters: int, pre_iter_fn, pre_iter_group_factory) -> int: + """Pick the graph-group size unconditionally; the caller is expected to + round total_iters up to a multiple of this so all iters fit in clean + replays. Sample arrays are pre-padded at allocation (see _sample_pnat + call site) so the per-replay window can index past the user-requested + iter count by up to group_iters-1 extra samples. + """ + if pre_iter_fn is not None and pre_iter_group_factory is None: + # Per-iter callback without a group-factory: can't batch. + return 1 + requested = getattr(args, "cuda_graph_group_iters", None) + if requested is None: + return ( + _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX + if pre_iter_group_factory is not None + else _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE + ) + return max(1, int(requested)) + + +def _get_cupti_filter_plan(timer: CuptiKernelTimer, graph, cache_key: tuple | None, + group_iters: int) -> tuple[int, tuple[str | None, ...]]: + full_cache_key = None if cache_key is None else (cache_key, group_iters) + if full_cache_key is not None: + cached = _CUPTI_FILTER_PLAN_CACHE.get(full_cache_key) + if cached is not None: + return cached + + records, zero_ts_count, zero_ts_names = timer.capture_names(graph.replay) + if zero_ts_count: + print( + f"[WARN] CUPTI calibration saw {zero_ts_count} zero-timestamp records " + f"(breakdown {zero_ts_names}); continuing with nonzero records.", + file=sys.stderr, + ) + ordinal_names = tuple(_target_name_or_none(r[0]) for r in records) + target_count = sum(name is not None for name in ordinal_names) + if target_count == 0: + raise RuntimeError("CUPTI calibration did not find any target kernel records") + plan = (len(records), ordinal_names) + if full_cache_key is not None: + _CUPTI_FILTER_PLAN_CACHE[full_cache_key] = plan + return plan + + +def _time_kernel_cuda_graph( + args, + run_fn, + reset_fn, + tag: str, + *, + expected_K: int, + pre_iter_fn=None, + pre_iter_group_factory=None, + iters_override: int | None = None, + cupti_plan_key: tuple | None = None, +) -> dict: + """CUDA-graph CUPTI timer (graph-per-iter design). + + Captures one CUDA graph holding a small group of logical iterations + (per-iter setup + reset + l2_flush + run_fn) and replays it enough + times to cover `warmup + iters`. + + Why graph-per-iter (vs the older "one giant graph holding all iters" + design): instantiating a CUDA graph is expensive — proportional to + graph size — so a single small graph instantiated once is much + cheaper than one big graph instantiated for each cell of a sweep. + Replays are cheap regardless. + + Mix cells use a per-replay device window: an outside-graph copy loads + the next group of PNAT/n_writes samples, then graph-captured per-iter + copies update kernel inputs before each reset + L2 flush + run. + + Pre-graph eager warmup: forces PyTorch's caching allocator + + Triton's autotune cache to settle before capture so the graph + doesn't bake in init-only allocations. + + ``iters_override`` (if not None) overrides ``args.iters`` for this + call. Used to give mix scenarios a higher iter count than pure + (more iters = more independent mix draws averaged in). + """ + host_timing = _HostTiming(bool(getattr(args, "host_timing", False))) + timer = CuptiKernelTimer.get() + warmup = args.warmup + iters = iters_override if iters_override is not None else args.iters + + # Pre-graph eager warmup: full per-iter chain once. This settles + # Triton/PyTorch setup and wrapper-side intermediate allocations; + # skipping it risks lazy work leaking into graph capture. + warmup_iters = _PRE_GRAPH_WARMUP_ITERS + host_timing.add("pre_graph_warmup_iters", warmup_iters) + host_timing.start() + for _ in range(warmup_iters): + reset_fn() + if pre_iter_fn is not None: + pre_iter_fn(0) + run_fn() + if warmup_iters > 0: + torch.cuda.synchronize() + host_timing.stop("pre_graph_warmup_ms") + + total_iters = warmup + iters + group_iters = _graph_group_iters(args, total_iters, pre_iter_fn, pre_iter_group_factory) + # Args are rounded at argparse-time so warmup+iters/mix_iters are already + # multiples of the relevant group_iters. Assert here to catch any caller + # bypassing argparse. + assert total_iters % group_iters == 0, ( + f"total_iters={total_iters} not a multiple of group_iters={group_iters}; " + f"args.warmup/iters/mix_iters should be rounded post-argparse." + ) + pre_replay_fn = None + graph_pre_iter_fn = None + if pre_iter_group_factory is not None and group_iters > 1: + pre_replay_fn, graph_pre_iter_fn = pre_iter_group_factory(group_iters) + + # Reset just before capture so warmup state changes don't bleed in. + host_timing.start() + reset_fn() + torch.cuda.synchronize() + host_timing.stop("pre_capture_reset_ms") + + # Capture a small group of identical logical iterations. Mix/pre_iter + # cells can group when they provide a graph-side pre-iter updater backed + # by a per-replay device window. + host_timing.start() + g = _capture_group_graph(args, run_fn, reset_fn, group_iters, graph_pre_iter_fn) + host_timing.stop("graph_capture_ms") + + if pre_replay_fn is not None: + host_timing.start() + pre_replay_fn(0) + torch.cuda.synchronize() + host_timing.stop("graph_preload_ms") + + plan_cache_key = None if cupti_plan_key is None else (cupti_plan_key, group_iters) + host_timing.add("cupti_plan_cached", ( + plan_cache_key is not None and plan_cache_key in _CUPTI_FILTER_PLAN_CACHE + )) + host_timing.start() + records_per_replay, ordinal_names = _get_cupti_filter_plan( + timer, + g, + cupti_plan_key, + group_iters, + ) + host_timing.stop("cupti_plan_ms") + target_count = sum(name is not None for name in ordinal_names) + expected_targets_per_replay = expected_K * group_iters + if target_count != expected_targets_per_replay: + print( + f"[WARN] CUPTI calibration mismatch for {tag!r}: expected " + f"{expected_targets_per_replay} target records in a {group_iters}-iter graph replay, " + f"got {target_count} target records out of {records_per_replay} total records.", + file=sys.stderr, + ) + + # Time: replay the grouped graph enough times to cover warmup+iters. + # Mix cells preload one device window per replay on the same stream. + # CUPTI records every kernel launch; _stats_from_cupti_records + # validates against expected_K and slices warmup off the front. + graph_replays = total_iters // group_iters + filter_plan = ((records_per_replay, ordinal_names),) * graph_replays + cupti_flush_period_ms = max(0, int(getattr(args, "cupti_flush_period_ms", 0))) + host_timing.start() + timer.start( + filter_plan, + flush_period_ms=cupti_flush_period_ms, + collect_timing=host_timing.enabled, + ) + host_timing.stop("cupti_start_ms") + for key, value in timer.last_start_timing().items(): + host_timing.add(f"cupti_start_{key}", value) + torch.cuda.nvtx.range_push(tag) + host_timing.start() + for i in range(graph_replays): + if pre_replay_fn is not None: + pre_replay_fn(i) + elif pre_iter_fn is not None: + pre_iter_fn(i) + g.replay() + host_timing.stop("graph_enqueue_ms") + host_timing.start() + torch.cuda.synchronize() + host_timing.stop("graph_sync_ms") + torch.cuda.nvtx.range_pop() + expected_raw_record_count = records_per_replay * graph_replays + host_timing.start() + if int(getattr(args, "cupti_defer_depth", 1)) > 1: + generation, stop_timing = timer.stop_async( + collect_timing=host_timing.enabled, + stats_request={ + "warmup": warmup, + "iters": iters, + "tag": tag, + "expected_K": expected_K, + "include_details": bool(getattr(args, "json_detailed", False)), + }, + ) + host_timing.stop("cupti_stop_ms") + for key, value in timer.last_stop_timing().items(): + host_timing.add(f"cupti_stop_{key}", value) + host_timing.stop_total() + host_timing.add("graph_group_iters", group_iters) + host_timing.add("graph_replays", graph_replays) + host_timing.add("cupti_records_per_replay", records_per_replay) + host_timing.add("cupti_target_records_per_replay", target_count) + host_timing.add("cupti_raw_records_expected", expected_raw_record_count) + host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) + return _PendingCuptiStats( + timer, + generation, + stop_timing, + host_timing, + warmup=warmup, + iters=iters, + tag=tag, + expected_K=expected_K, + expected_raw_record_count=expected_raw_record_count, + ) + + records, zero_ts_count, zero_ts_names, raw_record_count = timer.stop( + collect_timing=host_timing.enabled, + ) + host_timing.stop("cupti_stop_ms") + for key, value in timer.last_stop_timing().items(): + host_timing.add(f"cupti_stop_{key}", value) + if raw_record_count != expected_raw_record_count: + print( + f"[WARN] CUPTI raw-record mismatch for {tag!r}: expected " + f"{records_per_replay} total records/replay × {graph_replays} replays " + f"= {expected_raw_record_count}, got {raw_record_count}. SKIPPING cell.", + file=sys.stderr, + ) + return None + + host_timing.start() + stats = _stats_from_cupti_records( + records, + warmup, + iters, + tag, + expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names, + include_details=bool(getattr(args, "json_detailed", False)), + ) + host_timing.stop("stats_ms") + host_timing.stop_total() + host_timing.add("graph_group_iters", group_iters) + host_timing.add("graph_replays", graph_replays) + host_timing.add("cupti_records_per_replay", records_per_replay) + host_timing.add("cupti_target_records_per_replay", target_count) + host_timing.add("cupti_raw_records", raw_record_count) + host_timing.add("cupti_raw_records_expected", expected_raw_record_count) + host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) + host_timing.attach(stats) + return stats + + +def _time_kernel_eager( + args, + run_fn, + reset_fn, + tag: str, + *, + expected_K: int, + pre_iter_fn=None, + iters_override: int | None = None, + cupti_plan_key: tuple | None = None, +) -> dict: + """Non-graph CUPTI timer (for ncu wrapping, debugging, etc.). + + Each iter runs serially with sync between, but kernel start/end still + come from CUPTI — same accuracy as the graph path, just slower per-iter + (extra Python + sync overhead). + """ + host_timing = _HostTiming(bool(getattr(args, "host_timing", False))) + timer = CuptiKernelTimer.get() + warmup = args.warmup + iters = iters_override if iters_override is not None else args.iters + + del cupti_plan_key + + def _run_eager_loop(): + torch.cuda.nvtx.range_push(tag) + # Unified warmup+iters loop; CUPTI filters by warmup count internally. + for i in range(warmup + iters): + reset_fn() + if args.l2_flush: + _flush_l2() # includes synchronize + if pre_iter_fn is not None: + pre_iter_fn(i) + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + + host_timing.start() + records, zero_ts_count, zero_ts_names = timer.capture_names(_run_eager_loop) + host_timing.stop("timed_loop_and_cupti_parse_ms") + + host_timing.start() + stats = _stats_from_cupti_records( + records, + warmup, + iters, + tag, + expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names, + include_details=bool(getattr(args, "json_detailed", False)), + ) + host_timing.stop("stats_ms") + host_timing.stop_total() + host_timing.attach(stats) + return stats + + +def _run_kernel_untimed(args, run_fn, reset_fn, tag: str) -> dict: + """No in-bench timing: just run the kernels for an external profiler + (nsys / ncu) to time externally. Returns a stats dict full of zeros so + downstream code (table, JSON) doesn't break. + + Note: pre_iter_fn / iters_override aren't plumbed here yet — mix-mode + benchmarking relies on CUPTI. Add when a use-case lands. + """ + warmup = args.warmup + iters = args.iters + + if args.cuda_graph: + # Eager warmup before capture (Triton autotune) + reset_fn(); run_fn(); torch.cuda.synchronize() + reset_fn(); torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + for _ in range(warmup + iters): + reset_fn() + if args.l2_flush: + _l2_flush.fill_(0.0) + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_push(tag) + g.replay() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + else: + torch.cuda.nvtx.range_push(tag) + for _ in range(warmup + iters): + reset_fn() + if args.l2_flush: + _flush_l2() + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + + spans_us = [0.0] * iters + out = _stats_from_spans(spans_us) + out["iters_us"] = spans_us + out["per_kernel"] = {} + return out + + +def _time_kernel( + args, run_fn, reset_fn, tag: str, + *, + expected_K: int, + pre_iter_fn=None, + pre_iter_group_factory=None, + iters_override: int | None = None, + cupti_plan_key: tuple | None = None, +) -> dict: + """Dispatch to graph-CUPTI / eager-CUPTI / no-timer path. + + --cupti: in-process CUPTI Activity API timing (default). Use --no-cupti + when running under nsys (in-process CUPTI conflicts with nsys's own + subscriber); the bench then runs the kernels for nsys to time externally. + + `expected_K` is the kernels-per-iter count the caller declares + (computed via _kernels_per_iter_*). CUPTI paths validate against it + explicitly; the no-timer fallback ignores it (no records to validate). + """ + if not getattr(args, "cupti", True): + if pre_iter_fn is not None: + raise RuntimeError( + "_time_kernel: pre_iter_fn requires CUPTI (mix-mode); " + "got --no-cupti. Re-run with CUPTI on or plumb pre_iter_fn " + "through _run_kernel_untimed." + ) + return _run_kernel_untimed(args, run_fn, reset_fn, tag) + if args.cuda_graph: + return _time_kernel_cuda_graph( + args, run_fn, reset_fn, tag, + expected_K=expected_K, + pre_iter_fn=pre_iter_fn, + pre_iter_group_factory=pre_iter_group_factory, + iters_override=iters_override, + cupti_plan_key=cupti_plan_key, + ) + return _time_kernel_eager( + args, run_fn, reset_fn, tag, + expected_K=expected_K, + pre_iter_fn=pre_iter_fn, + iters_override=iters_override, + cupti_plan_key=cupti_plan_key, + ) + + +# Per-config benchmark (consolidated baseline + replay) + + +def _warm_one_config(args, cfg, baseline_fn) -> None: + """Module-level worker for the compile-warmup process pool. + + Module-level so ProcessPoolExecutor can pickle it (nested functions + aren't picklable). Each worker process holds its own GIL → no + serialization between concurrent compiles. + + ``cfg`` is a tuple of (outer_cfg, inner_overrides_or_list): + * outer_cfg = (batch, mtp_len, prev_ks, state_dtype, act_dtype, + sr_mode, rect, write_ckpt, mode, + sort_slots, reverse_nowrite, hardcode_sort) + * inner_overrides_or_list = dict of args attribute name -> value-string, + OR a list of such dicts. In the list form (CPS-grouped task) the + worker compiles each entry sequentially within the same process so + Triton's in-process kernel cache catches value-spec hits across + related entries (e.g. CPS={1,2} and {4,8} each form a `div_by_16` + spec bucket; the second compile in a bucket short-circuits). + + ``baseline_fn`` is optional — when ``None``, only the checkpointing + kernel is warmed (the baseline-selection kernel can be warmed once in + the parent if needed). This lets us avoid pickling C-extension + function references across processes. + """ + outer_cfg, inner_overrides_or_list = cfg + overrides_list = (inner_overrides_or_list + if isinstance(inner_overrides_or_list, list) + else [inner_overrides_or_list]) + (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, + rect, write_ckpt, mode, sort_slots, reverse_nowrite, hardcode_sort) = outer_cfg + import argparse as _ap + for inner_overrides in overrides_list: + # Fresh clone per entry: prevents knob-value leakage between + # consecutive cells in a CPS-grouped task (entries may set + # different non-CPS knobs in degenerate edge cases). + args_copy = _ap.Namespace(**vars(args)) + for k, v in inner_overrides.items(): + setattr(args_copy, k, v) + _bench_config( + args_copy, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, + sr_mode=sr_mode, rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, mode=mode, + sort_slots=sort_slots, reverse_nowrite=reverse_nowrite, + hardcode_sort=hardcode_sort, + warmup_only=True, + ) + + +def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, + baseline_fn, max_workers: int) -> None: + _cw_t0 = time.perf_counter() + def _cw(label: str) -> None: + dt = time.perf_counter() - _cw_t0 + print(f"[compile-warmup] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) + _cw("entered _compile_warmup_phase") + """Parallel compile-warmup using a ProcessPoolExecutor with `spawn` + start method. + + Each worker process holds its own GIL and its own CUDA context, so + Triton compiles (Python AST/codegen + LLVM/ptxas) run truly in + parallel. Previous ThreadPoolExecutor design hit GIL contention + in the Python codegen phase, capping throughput at ~1-2 cores even + with 28 threads (observed: 4 R threads vs 28 in pool). + + Compiled binaries land in Triton's on-disk cache (TRITON_CACHE_DIR + or default ~/.triton/cache). Workers share the cache via filesystem + — first to write any given (kernel_source × constexpr_set) hash + wins; concurrent writes to the SAME hash are wasteful but not + corrupting. + + spawn start method avoids inheriting parent CUDA state (which is + unsafe after fork on Linux with active CUDA contexts). Per-worker + import + CUDA init costs ~10s, amortized over each worker's many + compiles. baseline_fn is intentionally NOT passed to workers to + avoid pickling complications; the parent compiles the baseline + kernel itself before launching the pool when applicable. + """ + from concurrent.futures import ProcessPoolExecutor + import multiprocessing + + sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) + + rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) + write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) + modes_list = getattr(args, "modes_list", ["monolithic"]) + sort_list = getattr(args, "sort_slots_list", [False]) + rev_list = getattr(args, "reverse_nowrite_list", [False]) + hsort_list = getattr(args, "hardcode_sort_list", [False]) + + # Compile-warmup task enumeration: outer × inner cartesian. + # CRITICAL: only enumerate axes that change the kernel's COMPILE signature. + # Drop runtime axes (batch, prev_k) that produce identical kernel hashes — + # otherwise we'd pay ~50-100ms of bench setup per redundant cache-hit task. + # + # Batches collapsed to first only: batch is a runtime int passed to the + # kernel, not a constexpr; all batches share the same compiled kernel. + # prev_k is already a list passed into _bench_config (not enumerated here). + configs = [] + _compile_batches = batch_sizes[:1] # collapse runtime axis + for batch in _compile_batches: + for mtp_len in mtp_lengths: + prev_ks = _resolve_prev_ks(args, mtp_len) + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + for sr_mode in sr_modes_list: + for mode in modes_list: + effective_write_modes = ( + write_modes_list if mode == "monolithic" else [True] + ) + for write_ckpt in effective_write_modes: + if mode == "monolithic": + effective_rect_list = ( + [False] if write_ckpt else rect_list + ) + else: + effective_rect_list = rect_list + for rect in effective_rect_list: + # Note: "persistent_main" is included in + # the dl-family for sort/hsort sweep + # eligibility — it consumes the same + # slot_perm and benefits from the same + # write-first clustering. It additionally + # requires _n_writes (count of write + # slots) which the bench computes from + # the pure-scenario PNAT (mix scenarios + # not yet supported for persistent_main). + is_dl_family = mode in ( + "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", "persistent_main", + "persistent_dynamic", + ) + # Match the timed-run skip: sort=1 only + # makes sense when there's a mix scenario. + can_sort = is_dl_family and (args.mix_csv is not None) + effective_sort_list = ( + sort_list if can_sort else [False] + ) + effective_hsort_list = ( + hsort_list if can_sort else [False] + ) + for sort_slots in effective_sort_list: + effective_rev_list = ( + rev_list if sort_slots else [False] + ) + for reverse_nowrite in effective_rev_list: + for hardcode_sort in effective_hsort_list: + if sort_slots and hardcode_sort: + continue + configs.append(( + batch, mtp_len, prev_ks, + state_dtype, act_dtype, + sr_mode, rect, write_ckpt, mode, + sort_slots, reverse_nowrite, + hardcode_sort, + )) + + # Enumerate inner-knob signatures. Two paths: + # (1) --cell-list mode (preferred when set): pull exactly the cells + # that will be timed from args._cell_list_set. No synthetic + # cartesian — we only pre-compile what will run. + # (2) Sweep-args mode: cartesian over split-aware axes that read + # BOTH unsplit (args.X) and per-half (args.X_write/_nowrite) + # knob settings. Older code read only args.X and silently + # enumerated 1 inner combo when callers set only the per-half + # versions (all cell-list usage, plus any --block-size-m-write/ + # _nowrite CLI invocation), causing massive in-process JIT + # compile tax for persistent_main especially. + # + # In BOTH paths we GROUP tasks by non-CPS signature so each worker + # process compiles all CPS values for its group sequentially. + # NUM_PERSISTENT = CPS * num_sms is a runtime int but Triton auto- + # specializes on `div_by_16`, partitioning {CPS=1,2} (132,264) from + # {CPS=4,8} (528,1056) into two distinct compiled variants. By + # keeping all CPS variants for one (M,W,S,LS,TMA,...) signature in + # the same worker, the second compile in each spec bucket hits the + # in-process Triton cache (no disk-cache round trip). + + def _ps(val): + if val is None or (isinstance(val, str) and not val): + return [None] + if isinstance(val, str): + return [v.strip() for v in val.split(",") if v.strip()] + return [val] + + def _split_or_pair(shared_attr, w_attr, nw_attr): + """Return list of (write_val, nowrite_val) strings. + + Reads shared (args.X), write-side (args.X_write), and nowrite-side + (args.X_nowrite) values. If both per-half attrs are None, emits + tied pairs (v,v) over the shared values. If either per-half is + set, cartesian-iterates per-half values, falling back to shared + for whichever side is None. + """ + w = _ps(getattr(args, w_attr, None)) + nw = _ps(getattr(args, nw_attr, None)) + s = _ps(getattr(args, shared_attr, None)) + if w == [None] and nw == [None]: + return [(v, v) for v in s] + if w == [None]: + w = s + if nw == [None]: + nw = s + return [(a, b) for a in w for b in nw] + + _cw(f"built {len(configs)} outer configs") + cell_set = getattr(args, "_cell_list_set", set()) + cell_keys = getattr(args, "_cell_list_keys", ()) + # CPS keys are runtime ints (kernel value-specializes on `div_by_16`); + # cells differing only on CPS values can SHARE a worker so the second + # CPS value in a div_by_16 bucket hits the in-process Triton cache. + _cps_keys = ("cta_per_sm_write", "cta_per_sm_nowrite", "cta_per_sm") + + if cell_set: + # ============== CELL-LIST PATH ============== + # Build tasks DIRECTLY from cells. Each cell carries its OWN + # outer-axis values (RECT, MODE, SR, WC, SORT, REVN, HSORT) so we + # pair each cell with its specific outer config — NOT the union- + # cartesian of all cells' outer values. Previously the OUTER × + # CELL cartesian doubled task count when a cell-list spanned both + # RECT=0 and RECT=1 (or any other outer-axis split); half the + # tasks then failed the cell-list filter inside the worker and + # wasted dispatch overhead. This path is O(|unique cell groups|). + from collections import defaultdict as _dd + cell_groups: dict = _dd(list) + for tup in cell_set: + d = dict(zip(cell_keys, tup)) + cell_outer = ( + "SR" if d.get("SR", 0) else "RN", # sr_mode + bool(d.get("RECT", 0)), # rect + bool(d.get("WC", 1)), # write_ckpt + d.get("MODE", "monolithic"), # mode + bool(d.get("SORT", 0)), # sort_slots + bool(d.get("REVN", 0)), # reverse_nowrite + bool(d.get("HSORT", 0)), # hardcode_sort + ) + inner = {} + for k, v in d.items(): + if k in _CELL_LIST_KEY_TO_ARG: + inner[_CELL_LIST_KEY_TO_ARG[k]] = str(v) + non_cps_sig = tuple(sorted((k, v) for k, v in inner.items() if k not in _cps_keys)) + cell_groups[(cell_outer, non_cps_sig)].append(inner) + + # CLI-runtime axes (batch/mtp/dtype) are NOT in cell-list — they + # come from CLI args and cartesian here (typically just 1 combo). + cli_outers = [] + for _b in _compile_batches: + for _m in mtp_lengths: + _pk = _resolve_prev_ks(args, _m) + for _sd in state_dtypes: + for _ad in act_dtypes: + cli_outers.append((_b, _m, _pk, _sd, _ad)) + + tasks = [] + for cli_outer in cli_outers: + for (cell_outer, _sig), inner_list in cell_groups.items(): + outer_cfg = (*cli_outer, *cell_outer) + tasks.append((outer_cfg, inner_list)) + n_groups = len(cell_groups) + n_total_cells = sum(len(g) for g in cell_groups.values()) + n_outer_used = len(cli_outers) + else: + # ============== SWEEP-ARGS PATH ============== + # Build inner_dicts via cartesian over knob axes, then cross with + # the `configs` outer cartesian. Existing behavior. + m_pairs = _split_or_pair("block_size_m", "block_size_m_write", "block_size_m_nowrite") + w_pairs = _split_or_pair("num_warps", "num_warps_write", "num_warps_nowrite") + ns_pairs = _split_or_pair("num_stages", "num_stages_write", "num_stages_nowrite") + cps_pairs = _split_or_pair("cta_per_sm", "cta_per_sm_write", "cta_per_sm_nowrite") + ls_pairs = _split_or_pair("num_loop_stages", "num_loop_stages_write", "num_loop_stages_nowrite") + pw_vals = _ps(args.precompute_num_warps) + ps_vals = _ps(args.precompute_num_stages) + h_vals = _ps(args.heads_per_block) + mr_vals = _ps(args.maxnreg) + ct_vals = _ps(args.num_ctas) + fl_vals = _ps(args.flatten) + wsp_vals = _ps(args.warp_specialize) + trl_vals = _ps(args.use_tma_rect_load) + twl_vals = _ps(args.use_tma_replay_write_load) + tnl_vals = _ps(args.use_tma_replay_nowrite_load) + tws_vals = _ps(args.use_tma_replay_write_store) + import itertools as _it + inner_dicts = [] + for ((mw, mnw), (ww, wnw), (sw, snw), (cw, cnw), (lw, lnw), + pw, ps_, h, mr, ct, fl, wsp, + trl, twl, tnl, tws) in _it.product( + m_pairs, w_pairs, ns_pairs, cps_pairs, ls_pairs, + pw_vals, ps_vals, h_vals, mr_vals, ct_vals, + fl_vals, wsp_vals, + trl_vals, twl_vals, tnl_vals, tws_vals): + d = {} + for k, v in ( + ("block_size_m_write", mw), + ("block_size_m_nowrite", mnw), + ("num_warps_write", ww), + ("num_warps_nowrite", wnw), + ("num_stages_write", sw), + ("num_stages_nowrite", snw), + ("cta_per_sm_write", cw), + ("cta_per_sm_nowrite", cnw), + ("num_loop_stages_write", lw), + ("num_loop_stages_nowrite", lnw), + ("precompute_num_warps", pw), + ("precompute_num_stages", ps_), + ("heads_per_block", h), + ("maxnreg", mr), + ("num_ctas", ct), + ("flatten", fl), + ("warp_specialize", wsp), + ("use_tma_rect_load", trl), + ("use_tma_replay_write_load", twl), + ("use_tma_replay_nowrite_load", tnl), + ("use_tma_replay_write_store", tws), + ): + if v is not None: + d[k] = str(v) + inner_dicts.append(d) + + groups: dict = {} + for d in inner_dicts: + sig = tuple(sorted((k, v) for k, v in d.items() if k not in _cps_keys)) + groups.setdefault(sig, []).append(d) + tasks = [] + for outer in configs: + for sig, group in groups.items(): + tasks.append((outer, group)) + n_groups = len(groups) + n_total_cells = sum(len(g) for g in groups.values()) + n_outer_used = len(configs) + + # Shuffle ACROSS tasks (preserve within-group CPS sequence for in-process + # cache adjacency — within-group order is intentional, not shuffled). + import random as _r + _r.shuffle(tasks) + + _cw(f"built {len(tasks)} tasks covering {n_total_cells} cells in {n_groups} groups") + print(f"[compile-warmup] {len(tasks)} compile tasks " + f"({n_outer_used} outer × {n_groups} cell-groups " + f"covering {n_total_cells} cells, CPS-grouped" + + (", per-cell outer" if cell_set else "") + + f") across {max_workers} processes (ProcessPoolExecutor, {_MP_START_METHOD} start)") + t0 = time.perf_counter() + + ctx = multiprocessing.get_context(_MP_START_METHOD) + errors = [] + _cw("about to create ProcessPoolExecutor") + with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as ex: + _cw("ProcessPoolExecutor created, about to submit tasks") + # baseline_fn=None: workers compile only the checkpointing kernel. + # Baseline kernels (if any) get compiled lazily in the parent during + # the timing phase — usually just one extra compile, negligible. + futures = { + ex.submit(_warm_one_config, args, task, None): task + for task in tasks + } + _cw(f"submitted {len(futures)} tasks, waiting for results") + _n_done = 0 + for fut in futures: + try: + fut.result() + except Exception as e: + errors.append((futures[fut], e)) + _n_done += 1 + # Progress beacons at 10/25/50/75/100% to gauge effective parallelism. + if _n_done in (max(1, len(futures)//10), + max(1, len(futures)//4), + max(1, len(futures)//2), + max(1, (3*len(futures))//4), + len(futures)): + _cw(f"{_n_done}/{len(futures)} tasks complete") + + if errors: + for cfg, e in errors: + print(f"[compile-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", + file=sys.stderr) + raise errors[0][1] + + print(f"[compile-warmup] done in {time.perf_counter() - t0:.1f}s") + + +def _bench_config( + args, + batch: int, + mtp_len: int, + prev_ks: list[int], + state_dtype: torch.dtype, + act_dtype: torch.dtype, + baseline_fn, + sr_mode: str = "RN", + rectangle_for_nowrite: bool = False, + write_checkpoint: bool = True, + mode: str = "monolithic", + mix_samples_cpu=None, + mix_label: str = "", + sort_slots: bool = False, + reverse_nowrite: bool = False, + perm_samples_cpu=None, + hardcode_sort: bool = False, + mix_samples_sorted_cpu=None, + warmup_only: bool = False, +) -> None: + """ + Benchmark one (batch, mtp_len, dtype) configuration. + + Runs the baseline kernel (if baseline_fn is not None) followed by the + replay kernel for each prev_k value. Tensors are built once and + shared across all runs in this config. + + When ``warmup_only`` is True, calls each kernel exactly once instead of + timing it. Used by the parallel-warmup phase to populate Triton's + persistent compile cache across all configs concurrently. No timing + output is produced. + """ + state_dtype_name = str(state_dtype).split(".")[-1] + act_dtype_name = str(act_dtype).split(".")[-1] + + ( + state0, + state_scales0, + old_x0, + old_B0, + old_dt0, + old_dA_cumsum0, + cache_buf_idx0, + x, + dt, + B, + C, + A, + dt_bias, + D, + prev_tokens, + slot_perm_buf, + out_incr, + out_base, + intermediate_states_buffer, + xbc_input0, + conv_state0, + conv_weight, + conv_bias, + d_inner, + conv_dim, + ) = _build_tensors( + batch, + mtp_len, + state_dtype, + act_dtype, + args.tp_nheads, + args.head_dim, + args.d_state, + args.tp_ngroups, + max_window=getattr(args, "max_window", None) or None, + ) + + nheads = args.tp_nheads + ngroups = args.tp_ngroups + head_dim = args.head_dim + d_state = args.d_state + with_conv1d = getattr(args, "with_conv1d", False) + use_philox = (sr_mode == "SR") + variant_fn = _VARIANT_FNS[args.variant]() + + # SR rounding: allow fp16 and the quantized dtypes (int8/int16/fp8). + # bf16/fp32 SR is not supported (no PTX path for bf16; fp32 doesn't need + # rounding). When sweeping --sr-modes RN,SR over a mixed dtype set, + # silently skip the SR cell for unsupported dtypes — the RN cell still + # prints, and other dtypes still get their SR row. + rand_seed = None + _SR_SUPPORTED = ( + torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn, + ) + if use_philox: + if state_dtype not in _SR_SUPPORTED: + return + rand_seed = torch.randint(0, 2**62, (1,), device="cuda", dtype=torch.int64) + + is_quantized = state_dtype in (torch.int8, torch.int16, torch.float8_e4m3fn) + + state_work = state0.clone() + state_scales_work = state_scales0.clone() if state_scales0 is not None else None + old_x_work = old_x0.clone() + old_B_work = old_B0.clone() + old_dt_work = old_dt0.clone() + old_dA_cumsum_work = old_dA_cumsum0.clone() + cache_buf_idx_work = cache_buf_idx0.clone() + xbc_input_work = xbc_input0.clone() + conv_state_work = conv_state0.clone() + + def _reset(): + state_work.copy_(state0) + if state_scales_work is not None: + state_scales_work.copy_(state_scales0) + old_x_work.copy_(old_x0) + old_B_work.copy_(old_B0) + old_dt_work.copy_(old_dt0) + old_dA_cumsum_work.copy_(old_dA_cumsum0) + cache_buf_idx_work.copy_(cache_buf_idx0) + if with_conv1d: + conv_state_work.copy_(conv_state0) + + def _reset_conv1d_realistic(): + """Realistic reset: cold cache, L2 flush, then hot in_proj output.""" + # 1. Reset cold state (cache tensors, SSM state) + state_work.copy_(state0) + if state_scales_work is not None: + state_scales_work.copy_(state_scales0) + old_x_work.copy_(old_x0) + old_B_work.copy_(old_B0) + old_dt_work.copy_(old_dt0) + old_dA_cumsum_work.copy_(old_dA_cumsum0) + cache_buf_idx_work.copy_(cache_buf_idx0) + conv_state_work.copy_(conv_state0) + # 2. L2 flush (evicts cold state from cache) + if _l2_flush is not None: + _l2_flush.fill_(0.0) + # 3. Write hot tensors (simulates in_proj output landing in L2) + xbc_input_work.copy_(xbc_input0) + + # Silently skip the baseline row for any (baseline, state_dtype, SR) + # combo it can't run. Better than erroring on a partial sweep — our + # kernel rows still print. Compatibility: + # * Quantized states (int8 / int16 / fp8): no baseline supports them. + # * Triton baseline (selective_state_update): no rand_seed kwarg. + # * flashinfer baseline: rand_seed only on fp16 state. + def _baseline_supports() -> bool: + if baseline_fn is None: + return False + if is_quantized: + return False + if use_philox: + if args.baseline == "triton": + return False + if args.baseline == "flashinfer" and state_dtype != torch.float16: + return False + return True + + if baseline_fn is not None and not _baseline_supports(): + if not warmup_only: + sr_tag = " + SR" if use_philox else "" + print( + f"# Skipping {args.baseline} baseline for " + f"state_dtype={state_dtype_name}{sr_tag} (unsupported)." + ) + baseline_fn = None + + show_kernel_col = baseline_fn is not None + + def _conv1d_split(xbc_in, conv_st, launch_dependent_kernels=False): + """Run conv1d update and split output into (x, B, C) views. + + The input tensor's strides are preserved through conv1d and the + transpose+view chain. With the production-matching layout + (contiguous (batch*T, conv_dim) viewed as (batch, conv_dim, T)), + the output after transpose+view has stride(-1)==1 and + stride(1)==dim, satisfying both our kernel and flashinfer. + """ + xbc_result = causal_conv1d_update( + xbc_in, + conv_st, + conv_weight, + conv_bias, + activation="silu", + launch_dependent_kernels=launch_dependent_kernels, + ) + xbc_flat = xbc_result.transpose(1, 2).view(batch * mtp_len, conv_dim) + x_flat, B_flat, C_flat = torch.split( + xbc_flat, [d_inner, ngroups * d_state, ngroups * d_state], dim=-1 + ) + x_conv = x_flat.view(batch, mtp_len, nheads, head_dim) + B_conv = B_flat.view(batch, mtp_len, ngroups, d_state) + C_conv = C_flat.view(batch, mtp_len, ngroups, d_state) + return x_conv, B_conv, C_conv + + # --- Baseline --- + if baseline_fn is not None: + tag = f"base_b{batch}_mtp{mtp_len}_s{state_dtype_name}_a{act_dtype_name}" + + philox_kwargs = {} + if rand_seed is not None and args.baseline == "flashinfer": + philox_kwargs = {"rand_seed": rand_seed, "philox_rounds": args.philox_rounds} + + if with_conv1d: + + def _run_baseline(): + x_conv, B_conv, C_conv = _conv1d_split(xbc_input_work, conv_state_work) + baseline_fn( + state_work, + x=x_conv, + dt=dt, + A=A, + B=B_conv, + C=C_conv, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + out=out_base, + disable_state_update=True, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=mtp_len, + **philox_kwargs, + ) + else: + + def _run_baseline(): + baseline_fn( + state_work, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + out=out_base, + disable_state_update=True, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=mtp_len, + **philox_kwargs, + ) + + reset_fn = _reset_conv1d_realistic if with_conv1d else _reset + if warmup_only: + reset_fn() + _run_baseline() + torch.cuda.synchronize() + else: + stats = _time_kernel( + args, _run_baseline, reset_fn, tag, + expected_K=_kernels_per_iter_baseline(with_conv1d), + cupti_plan_key=( + "baseline", + args.baseline, + batch, + mtp_len, + state_dtype_name, + act_dtype_name, + with_conv1d, + bool(args.l2_flush), + bool(args.external_pdl), + bool(use_philox), + _kernels_per_iter_baseline(with_conv1d), + ), + ) + + _submit_result_job( + args, + stats, + show_kernel_col=show_kernel_col, + kernel_name=args.baseline, + batch=batch, + mtp_len=mtp_len, + prev_k="N/A", + state_dtype_name=state_dtype_name, + act_dtype_name=act_dtype_name, + skipped_tag=tag, + ) + + # --- Sweep parameter parsing (invariant across prev_k) --- + def _parse_sweep(val): + if val is None: + return [None] + return [int(v) for v in val.split(",")] + + block_size_m_values = _parse_sweep(args.block_size_m) + num_warps_values = _parse_sweep(args.num_warps) + num_stages_values = _parse_sweep(args.num_stages) + precompute_num_warps_values = _parse_sweep(args.precompute_num_warps) + precompute_num_stages_values = _parse_sweep(args.precompute_num_stages) + heads_per_block_values = _parse_sweep(args.heads_per_block) + maxnreg_values = _parse_sweep(args.maxnreg) + num_ctas_values = _parse_sweep(args.num_ctas) + # Persistent-only sweep dims; ignored when the cell's mode != persistent_main. + cta_per_sm_values = _parse_sweep(args.cta_per_sm) + num_loop_stages_values = _parse_sweep(args.num_loop_stages) + flatten_values = _parse_sweep(args.flatten) + warp_specialize_values = _parse_sweep(args.warp_specialize) + # Per-main split-knob sweeps. Default = same as the shared sweep (so each + # combo is tied). When set independently, the inner loop sweeps the + # cross-product (write × nowrite); --skip-diagonal drops the tied subset. + def _split_or_share(split_csv, shared_values): + return _parse_sweep(split_csv) if split_csv else shared_values + block_size_m_write_values = _split_or_share(args.block_size_m_write, block_size_m_values) + block_size_m_nowrite_values = _split_or_share(args.block_size_m_nowrite, block_size_m_values) + num_warps_write_values = _split_or_share(args.num_warps_write, num_warps_values) + num_warps_nowrite_values = _split_or_share(args.num_warps_nowrite, num_warps_values) + num_stages_write_values = _split_or_share(args.num_stages_write, num_stages_values) + num_stages_nowrite_values = _split_or_share(args.num_stages_nowrite, num_stages_values) + cta_per_sm_write_values = _split_or_share(args.cta_per_sm_write, cta_per_sm_values) + cta_per_sm_nowrite_values = _split_or_share(args.cta_per_sm_nowrite, cta_per_sm_values) + num_loop_stages_write_values = _split_or_share(args.num_loop_stages_write, num_loop_stages_values) + num_loop_stages_nowrite_values = _split_or_share(args.num_loop_stages_nowrite, num_loop_stages_values) + # Whether any *_write / *_nowrite knob was independently set — used by + # --skip-diagonal to know if the cross-product is non-trivial. Without + # any split, the per-main values == shared values and skip-diagonal is + # a no-op (which is correct). + _any_split = any(getattr(args, name) for name in ( + "block_size_m_write", "block_size_m_nowrite", + "num_warps_write", "num_warps_nowrite", + "num_stages_write", "num_stages_nowrite", + "cta_per_sm_write", "cta_per_sm_nowrite", + "num_loop_stages_write", "num_loop_stages_nowrite", + )) + # TMA toggles — independent 0/1 sweep per path. The skip-dupe at the + # top of the inner loop body collapses cells where a flag's path is + # unreachable, so e.g. monolithic + WC=True only runs the value=0 + # cells for nowrite-load and rect-load. + use_tma_rect_load_values = _parse_sweep(args.use_tma_rect_load) + use_tma_replay_write_load_values = _parse_sweep(args.use_tma_replay_write_load) + use_tma_replay_nowrite_load_values = _parse_sweep(args.use_tma_replay_nowrite_load) + use_tma_replay_write_store_values = _parse_sweep(args.use_tma_replay_write_store) + + # --- Replay kernel --- + # Cache T-axis capacity (for prev_k validity check on the nowrite path). + max_window = getattr(args, "max_window", 0) or mtp_len + + # Build the list of scenarios to time. A scenario is one cell in the + # output: pure-mode scenarios fill prev_tokens with one constant before + # the timing loop; mix-mode scenarios feed a pre-baked per-iter samples + # tensor, with the per-iter copy captured inside the CUDA graph. Pure + # and mix can coexist in one call so a single nsys trace covers both. + scenarios = [] + if not (getattr(args, "mix_only", False) and mix_samples_cpu is not None): + for prev_k in prev_ks: + # On the nowrite path, new tokens append at [prev_k, prev_k+T) of + # the active buffer, so prev_k+T must fit within max_window. + # mode != monolithic dispatches per-slot from PNAT, so any + # prev_k <= max_window is valid for those modes. + if mode == "monolithic" and not write_checkpoint and prev_k + mtp_len > max_window: + continue + scenarios.append({ + "label": f"k{prev_k}", + "print_label": prev_k, + "fill": prev_k, + "pre_iter": None, + "iters": None, # use args.iters + }) + # Mix scenario: skip on monolithic (mono on mixed PNAT corrupts the + # wrong-mode slots). Persistent_main + mix is now supported: bench + # pre-bakes both a per-iter PNAT samples tensor and a per-iter + # n_writes samples tensor; grouped graph capture copies window rows + # into kernel-input tensors (PNAT and n_writes_dev) before each + # in-graph L2 flush, so the timed kernels read PNAT cold. + if mix_samples_cpu is not None and mode != "monolithic": + device = state_work.device + # Hardcode-sort: per-iter prev_tokens are CPU-sorted write-first. + # Kernel runs USE_PERM=False but the EO gate sees clustered modes. + # Output is scrambled (we don't permute x/B/C/dt to match) but + # timing is meaningful — isolates clustering benefit from the + # per-program perm-load overhead in --sort-slots. + src = mix_samples_sorted_cpu if (hardcode_sort and mix_samples_sorted_cpu is not None) else mix_samples_cpu + samples_gpu = torch.from_numpy(src).to(device=device, dtype=torch.int32) + + # For persistent_main + mix: pre-compute the per-iter n_writes + # (count of slots needing the write path = PNAT+T > max_window) + # and the (1,) scratch the kernel reads from. Both halves of + # persistent_main always launch in mix scenarios (host can't + # cheaply read n_writes per iter without a sync), so the kernel's + # slot-range derivation must be correct from device n_writes. + # persistent_dynamic doesn't need n_writes (the kernel ignores + # n_writes_dev when IS_DYNAMIC=True via Triton DCE), but we still + # allocate a sentinel scratch so the wrapper API is uniform. + n_writes_samples_gpu = None + n_writes_dev_mix = None + # n_writes per iter = number of slots that overflow the window. + # Computed for ALL mix scenarios so the JSON output (--json-detailed) + # can pair each iter's span_us with its mix composition for downstream + # analysis (group iters by # writes → per-bucket median → analytic + # expectation under the steady-state PNAT distribution). + n_writes_per_iter_all = ((src + mtp_len) > max_window).sum(axis=1).astype(np.int32) + if mode in ("persistent_main", "persistent_dynamic"): + n_writes_samples_gpu = torch.from_numpy(n_writes_per_iter_all).to( + device=device, dtype=torch.int32 + ) + n_writes_dev_mix = torch.zeros(1, dtype=torch.int32, device=device) + + # Build _mix_pre_iter — the closure that runs OUTSIDE the captured + # graph between replays. Updates: prev_tokens (always), + # slot_perm_buf (when sort_slots), n_writes_dev_mix (persistent). + perm_samples_gpu = None + if sort_slots and perm_samples_cpu is not None: + perm_samples_gpu = torch.from_numpy(perm_samples_cpu).to( + device=device, dtype=torch.int32 + ) + if n_writes_samples_gpu is not None: + def _mix_pre_iter(i, _s=samples_gpu, _ps=perm_samples_gpu, + _ns=n_writes_samples_gpu, _pt=prev_tokens, + _pm=slot_perm_buf, _nw=n_writes_dev_mix): + _pt.copy_(_s[i]) + _pm.copy_(_ps[i]) + _nw.copy_(_ns[i:i+1]) + else: + def _mix_pre_iter(i, _s=samples_gpu, _ps=perm_samples_gpu, + _pt=prev_tokens, _pm=slot_perm_buf): + _pt.copy_(_s[i]) + _pm.copy_(_ps[i]) + else: + if n_writes_samples_gpu is not None: + def _mix_pre_iter(i, _s=samples_gpu, _ns=n_writes_samples_gpu, + _pt=prev_tokens, _nw=n_writes_dev_mix): + _pt.copy_(_s[i]) + _nw.copy_(_ns[i:i+1]) + else: + def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): + _pt.copy_(_s[i]) + + def _mix_pre_iter_group_factory( + group_iters, + _s=samples_gpu, + _ps=perm_samples_gpu, + _ns=n_writes_samples_gpu, + _pt=prev_tokens, + _pm=slot_perm_buf, + _nw=n_writes_dev_mix, + ): + sample_window = torch.empty( + (group_iters, _s.shape[1]), device=_s.device, dtype=_s.dtype, + ) + perm_window = ( + torch.empty((group_iters, _ps.shape[1]), device=_ps.device, dtype=_ps.dtype) + if _ps is not None else None + ) + nw_window = ( + torch.empty((group_iters,), device=_ns.device, dtype=_ns.dtype) + if _ns is not None else None + ) + + def _pre_replay(replay_idx): + start = replay_idx * group_iters + end = start + group_iters + sample_window.copy_(_s[start:end]) + if perm_window is not None: + perm_window.copy_(_ps[start:end]) + if nw_window is not None: + nw_window.copy_(_ns[start:end]) + + def _graph_pre_iter(j): + _pt.copy_(sample_window[j]) + if perm_window is not None: + _pm.copy_(perm_window[j]) + if nw_window is not None: + _nw.copy_(nw_window[j:j + 1]) + + return _pre_replay, _graph_pre_iter + + # Mix iters override: if --mix-iters set, use it; else use args.iters. + mix_iters = getattr(args, "mix_iters", None) + scenarios.append({ + "label": f"mix{mix_label}", + "print_label": "mix", + "fill": None, + "pre_iter": _mix_pre_iter, + "pre_iter_group_factory": _mix_pre_iter_group_factory, + "iters": mix_iters, # None => use args.iters + # Pass through to _run_incr so the wrapper receives _n_writes_dev + # (mix scenarios) instead of _n_writes (pure scenarios). + "n_writes_dev": n_writes_dev_mix, + # Full per-iter n_writes array (size = warmup + iters). Used by + # the JSON-detailed output to pair each iter's span with its + # mix composition for post-hoc bucketing analysis. + "n_writes_per_iter": n_writes_per_iter_all, + }) + + # Pure scenarios don't pre-allocate n_writes_dev; mix scenarios do. + # Default empty-halves skip: True for pure (host knows n_writes, + # production-equivalent host-skip), False for mix (host can't read + # device n_writes per iter without sync, must always launch both). + for scn in scenarios: + scenario_n_writes_dev = scn.get("n_writes_dev") # None for pure + scenario_skip_empty = scenario_n_writes_dev is None + if scn["fill"] is not None: + prev_tokens.fill_(scn["fill"]) + prev_k_for_print = scn["print_label"] + scenario_pre_iter = scn["pre_iter"] + scenario_pre_iter_group_factory = scn.get("pre_iter_group_factory") + scenario_iters = scn.get("iters") # None => use args.iters + tag = f"incr_b{batch}_mtp{mtp_len}_{scn['label']}_s{state_dtype_name}_a{act_dtype_name}" + + # Iteration over per-cell knob combos. + # When NO per-main split is requested (_any_split=False), each row in + # the cross-product gives the same value to both write_main and + # nowrite_main (current behavior — backward-compat). When ANY split + # IS requested, we iterate the write and nowrite axes independently + # (cross-product blowup is the user's responsibility — they typically + # pair this with --skip-diagonal to drop the tied subset). + if _any_split: + _iter_axes = ( + block_size_m_write_values, block_size_m_nowrite_values, + num_warps_write_values, num_warps_nowrite_values, + num_stages_write_values, num_stages_nowrite_values, + precompute_num_warps_values, + precompute_num_stages_values, + heads_per_block_values, + maxnreg_values, num_ctas_values, + cta_per_sm_write_values, cta_per_sm_nowrite_values, + num_loop_stages_write_values, num_loop_stages_nowrite_values, + flatten_values, warp_specialize_values, + use_tma_rect_load_values, + use_tma_replay_write_load_values, + use_tma_replay_nowrite_load_values, + use_tma_replay_write_store_values, + ) + else: + # Tied: one value per shared knob. Wrap in single-element list for + # uniform iteration; the body sets w/nw both to the shared value. + _iter_axes = ( + block_size_m_values, [None], + num_warps_values, [None], + num_stages_values, [None], + precompute_num_warps_values, + precompute_num_stages_values, + heads_per_block_values, + maxnreg_values, num_ctas_values, + cta_per_sm_values, [None], + num_loop_stages_values, [None], + flatten_values, warp_specialize_values, + use_tma_rect_load_values, + use_tma_replay_write_load_values, + use_tma_replay_nowrite_load_values, + use_tma_replay_write_store_values, + ) + # Iteration source: when --cell-list is active AND this is the main + # timing path (not a compile-warmup worker), iterate the cell set + # DIRECTLY (one yield per cell). The earlier design iterated the + # full inner cartesian and filtered each iteration via membership in + # args._cell_list_set — that's O(cartesian) which blows up to + # billions of iterations when the cell-list spans wide split-knob + # values (CPS, LS, M, W, S each contributing a Wx*Wnw factor on top + # of TMA flags), producing 50+ min of CPU spin per bench call before + # any actual timing. Direct iteration is O(|cell_list|). + # + # IMPORTANT exception for workers (warmup_only=True): _warm_one_config + # clamps args.*_write/_nowrite via inner_overrides to single values, + # making the cartesian 1×1×...×1 = 1 iter, which is exactly the one + # cell that worker was given. If we used cell-list-direct iteration + # here, every worker would iterate ALL 2884 cells instead of just + # its assigned one — turning compile-warmup into 28-way duplication. + # (Observed: 256 tasks in 233s under that bug vs ~18s correct.) + if getattr(args, "_cell_list_set", None) and not warmup_only: + def _gen_from_cell_list(): + keys = args._cell_list_keys + for tup in args._cell_list_set: + d = dict(zip(keys, tup)) + yield ( + d.get("Mw"), d.get("Mnw"), + d.get("Ww"), d.get("Wnw"), + d.get("Sw"), d.get("Snw"), + d.get("pW"), d.get("pS"), + d.get("H"), + d.get("R"), d.get("CT"), + d.get("CPSw"), d.get("CPSnw"), + d.get("LSw"), d.get("LSnw"), + d.get("FL"), d.get("WS"), + d.get("TMARL"), d.get("TMAWL"), + d.get("TMANL"), d.get("TMAWS"), + ) + _iter_source = _gen_from_cell_list() + else: + _iter_source = itertools.product(*_iter_axes) + + for ( + block_size_m_w, + block_size_m_nw, + num_warps_w, + num_warps_nw, + num_stages_w, + num_stages_nw, + precompute_num_warps, + precompute_num_stages, + heads_per_block, + maxnreg, + num_ctas, + cta_per_sm_w, + cta_per_sm_nw, + num_loop_stages_w, + num_loop_stages_nw, + flatten, + warp_specialize, + use_tma_rect_load, + use_tma_replay_write_load, + use_tma_replay_nowrite_load, + use_tma_replay_write_store, + ) in _iter_source: + # When tied, _nw values were placeholder None; fill from _w (the + # shared value). When split, _w and _nw came from independent lists. + if not _any_split: + block_size_m_nw = block_size_m_w + num_warps_nw = num_warps_w + num_stages_nw = num_stages_w + cta_per_sm_nw = cta_per_sm_w + num_loop_stages_nw = num_loop_stages_w + # Skip-diagonal: when split is on, drop the tied subset (same as a + # prior shared-knob sweep would cover). + if _any_split and args.skip_diagonal and ( + block_size_m_w == block_size_m_nw and + num_warps_w == num_warps_nw and + num_stages_w == num_stages_nw and + cta_per_sm_w == cta_per_sm_nw and + num_loop_stages_w == num_loop_stages_nw + ): + continue + # Backward-compat aliases used by the existing body below. When + # tied, these are simply the shared value. When split, the + # _write copy is used for sweep_tag and grouping (a stable choice + # so the tag is unique per (write, nowrite) combo). + block_size_m = block_size_m_w + num_warps = num_warps_w + num_stages = num_stages_w + cta_per_sm = cta_per_sm_w + num_loop_stages = num_loop_stages_w + # Skip-dupe for TMA flag sweeps: a flag whose code path isn't + # reachable in this cell produces identical timing for value=0 + # and value=1. We canonicalize by skipping value=1 cells when + # the flag's path is unreachable. Path reachability rules: + # * write path (replay write-load + write-store): mono+WC=True, + # OR any non-monolithic mode. + # * rect path (rect-load): rectangle_for_nowrite=True AND a + # nowrite path exists in this mode (mono+WC=False, OR any + # non-monolithic mode). + # * replay-nowrite path (nowrite-load): nowrite path exists + # AND rect isn't taking it: mono+WC=False+rect=False, OR + # any non-monolithic mode with rect=False. + _is_mono = (mode == "monolithic") + _write_path = (_is_mono and write_checkpoint) or (not _is_mono) + _rect_path = rectangle_for_nowrite and ( + (not _is_mono) or (_is_mono and not write_checkpoint) + ) + _replay_nowrite_path = ( + (_is_mono and not write_checkpoint and not rectangle_for_nowrite) + or ((not _is_mono) and not rectangle_for_nowrite) + ) + def _set(v): # flag set to a non-zero sweep value + return v is not None and v != 0 + if (_set(use_tma_rect_load) and not _rect_path + or _set(use_tma_replay_write_load) and not _write_path + or _set(use_tma_replay_nowrite_load) and not _replay_nowrite_path + or _set(use_tma_replay_write_store) and not _write_path): + continue + + # Pre-allocate n_writes_dev tensor OUTSIDE the captured graph for + # persistent modes in pure scenarios. Mix scenarios already have + # `scenario_n_writes_dev` pre-allocated. The wrapper's fallback + # `torch.tensor([...], device=...)` allocation would invalidate + # the CUDA-graph capture stream — must allocate here, before the + # `_run_incr` lambda (which is what gets captured) is defined. + # For persistent_dynamic the kernel ignores the value (IS_DYNAMIC + # DCE's the load); we still need a valid pointer. For + # persistent_main pure, the value is constant per cell so we set + # it once here. + _n_writes_dev_pure: torch.Tensor | None = None + _host_n_writes_pure: int | None = None + if mode in ("persistent_main", "persistent_dynamic") and scenario_n_writes_dev is None: + _n_writes_dev_pure = torch.zeros(1, dtype=torch.int32, device=state_work.device) + if mode == "persistent_main": + scn_fill = scn["fill"] + is_write_scenario_local = (scn_fill + mtp_len) > max_window + _host_n_writes_pure = batch if is_write_scenario_local else 0 + _n_writes_dev_pure.fill_(_host_n_writes_pure) + + def _run_incr( + block_size_m=block_size_m, + num_warps=num_warps, + num_stages=num_stages, + precompute_num_warps=precompute_num_warps, + precompute_num_stages=precompute_num_stages, + heads_per_block=heads_per_block, + maxnreg=maxnreg, + num_ctas=num_ctas, + cta_per_sm=cta_per_sm, + num_loop_stages=num_loop_stages, + flatten=flatten, + warp_specialize=warp_specialize, + use_tma_rect_load=use_tma_rect_load, + use_tma_replay_write_load=use_tma_replay_write_load, + use_tma_replay_nowrite_load=use_tma_replay_nowrite_load, + use_tma_replay_write_store=use_tma_replay_write_store, + ): + if with_conv1d: + x_call, B_call, C_call = _conv1d_split( + xbc_input_work, conv_state_work, launch_dependent_kernels=args.external_pdl + ) + extra_kwargs = {"launch_with_pdl": args.external_pdl} + else: + x_call, B_call, C_call = x, B, C + extra_kwargs = {} + # write_checkpoint is only meaningful for the checkpointing + # variant; replay variant ignores the kwarg. state_scales + # is also checkpointing-only (replay kernel doesn't quantize). + if args.variant == "checkpointing": + extra_kwargs["write_checkpoint"] = write_checkpoint + extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite + extra_kwargs["mode"] = mode + if sort_slots: + extra_kwargs["slot_perm"] = slot_perm_buf + # reverse_nowrite is meaningful in two ways: + # - with slot_perm: walk the perm tail-first + # - without slot_perm (hardcode-sort): walk pid_b + # itself tail-first via the REVERSE_PERM constexpr + if sort_slots or (hardcode_sort and reverse_nowrite): + extra_kwargs["reverse_nowrite"] = reverse_nowrite + if state_scales_work is not None: + extra_kwargs["state_scales"] = state_scales_work + if use_tma_rect_load: # 1 → True, 0/None → False + extra_kwargs["_use_tma_rect_load"] = True + if use_tma_replay_write_load: + extra_kwargs["_use_tma_replay_write_load"] = True + if use_tma_replay_nowrite_load: + extra_kwargs["_use_tma_replay_nowrite_load"] = True + if use_tma_replay_write_store: + extra_kwargs["_use_tma_replay_write_store"] = True + # persistent_main needs n_writes (count of write-mode + # slots in the pre-sorted batch) as a host-side int. + # Pure scenarios: every slot has the same PNAT, so + # n_writes is either 0 (all nowrite) or batch (all + # write) depending on whether PNAT+T overflows the + # window. Mix scenarios are skipped earlier. + if mode in ("persistent_main", "persistent_dynamic"): + # Per-cell sweep values for persistent-only knobs. + # Apply to both persistent variants. _parse_sweep + # returns [None] when the user didn't pass the flag, + # in which case we leave the wrapper's defaults. + if cta_per_sm is not None: + extra_kwargs["_cta_per_sm"] = cta_per_sm + if num_loop_stages is not None: + extra_kwargs["_num_loop_stages"] = num_loop_stages + if flatten is not None: + extra_kwargs["_flatten"] = bool(flatten) + if warp_specialize is not None: + extra_kwargs["_warp_specialize"] = bool(warp_specialize) + if mode in ("persistent_main", "persistent_dynamic"): + # persistent_main + mix REQUIRES sort: the kernel + # partitions slots [0, n_writes) = write half, + # [n_writes, batch) = nowrite half. This only holds + # if PNAT is monotone (writes first), which sort + # provides via either: + # sort_slots=1 → USE_PERM reads slot_perm to remap + # hardcode_sort=1 → PNAT itself is CPU-pre-sorted + # persistent_dynamic doesn't need sort (per-slot + # runtime dispatch); persistent_main pure scenarios + # are trivially sorted (homogeneous PNAT). + if (mode == "persistent_main" + and scenario_n_writes_dev is not None + and not (sort_slots or hardcode_sort)): + raise AssertionError( + "persistent_main + mix requires sort_slots=1 " + "or hardcode_sort=1 — kernel partitions slots " + "by index, which is only valid when PNAT is " + "monotone (writes first). Without sort, the " + "partition silently mismatches actual slot " + "modes. Re-run with --sort-slots 1 or " + "--hardcode-sort 1." + ) + # n_writes plumbing: pure scenarios pass an int + # (host knows the value, can host-skip empty halves); + # mix scenarios pass a (1,) device tensor updated + # per iter by the benchmark pre-iter path. + # _persistent_skip_empty_halves=False on mix so both + # halves always launch (kernel uses device n_writes + # to derive its slot range). + if scenario_n_writes_dev is not None: + # Mix path: caller-allocated tensor, updated per + # iter by scenario_pre_iter outside capture. + extra_kwargs["_n_writes_dev"] = scenario_n_writes_dev + extra_kwargs["_persistent_skip_empty_halves"] = False + elif mode == "persistent_main": + # Pure: caller pre-allocated `_n_writes_dev_pure` + # outside this lambda (so the alloc doesn't land + # inside the captured graph). Pass both the + # tensor and the host int so the wrapper can use + # host-skip when `_persistent_skip_empty_halves`. + extra_kwargs["_n_writes"] = _host_n_writes_pure + extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure + extra_kwargs["_persistent_skip_empty_halves"] = scenario_skip_empty + elif mode == "persistent_dynamic": + # persistent_dynamic pure: kernel ignores n_writes + # via IS_DYNAMIC DCE, but the wrapper needs a + # valid (1,) tensor pointer. Pass the pre-allocated + # zero tensor to avoid any in-capture alloc. + extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure + variant_fn( + state_work, + old_x_work, + old_B_work, + old_dt_work, + old_dA_cumsum_work, + cache_buf_idx_work, + prev_tokens, + x=x_call, + dt=dt, + A=A, + B=B_call, + C=C_call, + out=out_incr, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + rand_seed=rand_seed, + philox_rounds=args.philox_rounds, + use_internal_pdl=args.internal_pdl, + _block_size_m=block_size_m, + _num_warps=num_warps, + _num_stages=num_stages, + _precompute_num_warps=precompute_num_warps, + _precompute_num_stages=precompute_num_stages, + _heads_per_block=heads_per_block, + _maxnreg=maxnreg, + _num_ctas=num_ctas, + # Per-main overrides (None = tied to shared above; explicit + # only when the inner loop is iterating split axes). + _block_size_m_write=block_size_m_w if _any_split else None, + _block_size_m_nowrite=block_size_m_nw if _any_split else None, + _num_warps_write=num_warps_w if _any_split else None, + _num_warps_nowrite=num_warps_nw if _any_split else None, + _num_stages_write=num_stages_w if _any_split else None, + _num_stages_nowrite=num_stages_nw if _any_split else None, + _cta_per_sm_write=cta_per_sm_w if _any_split else None, + _cta_per_sm_nowrite=cta_per_sm_nw if _any_split else None, + _num_loop_stages_write=num_loop_stages_w if _any_split else None, + _num_loop_stages_nowrite=num_loop_stages_nw if _any_split else None, + **extra_kwargs, + ) + + parts = [] + # When tied (not _any_split), emit the shared single-value tag + # (M=8 etc). When split, emit explicit Mw / Mnw tags so cells + # with the same shared value but different per-main values get + # unique JSON keys. + def _emit_split(name_w, name_nw, val_w, val_nw): + if val_w is None and val_nw is None: + return + if not _any_split or val_w == val_nw: + parts.append(f"{name_w[:-1]}={val_w}") # strip the 'w' suffix + else: + parts.append(f"{name_w}={val_w}") + parts.append(f"{name_nw}={val_nw}") + _emit_split("Mw", "Mnw", block_size_m_w, block_size_m_nw) + _emit_split("Ww", "Wnw", num_warps_w, num_warps_nw) + _emit_split("Sw", "Snw", num_stages_w, num_stages_nw) + if precompute_num_warps is not None: + parts.append(f"pW={precompute_num_warps}") + if precompute_num_stages is not None: + parts.append(f"pS={precompute_num_stages}") + if heads_per_block is not None: + parts.append(f"H={heads_per_block}") + if maxnreg is not None: + parts.append(f"R={maxnreg}") + if num_ctas is not None: + parts.append(f"CT={num_ctas}") + # Persistent-only knobs (only meaningful when MODE=persistent_main; + # printed unconditionally so output rows are uniformly comparable + # across modes when the user passed these sweeps). + _emit_split("CPSw", "CPSnw", cta_per_sm_w, cta_per_sm_nw) + _emit_split("LSw", "LSnw", num_loop_stages_w, num_loop_stages_nw) + if flatten is not None: + parts.append(f"FL={flatten}") + if warp_specialize is not None: + parts.append(f"WS={warp_specialize}") + # TMA sweep tags. Four wrapper-level flags map to three + # kernel-level constexprs (rect-load and replay-nowrite-load + # share `USE_TMA_LOAD_NOWRITE`, picked by the wrapper based on + # RECTANGLE). TMARL specifically gates the rectangle path's + # state load; TMANL specifically gates the replay-style + # nowrite path's state load. Distinct because their measured + # perf profiles differ (see CHECKPOINTING_DESIGN.md item #17: + # rect TMA is "not a win" while replay-nowrite TMA is the + # biggest measured win at int8 b>=64). + if use_tma_rect_load is not None: + parts.append(f"TMARL={use_tma_rect_load}") # rect path load + if use_tma_replay_write_load is not None: + parts.append(f"TMAWL={use_tma_replay_write_load}") # replay-write load + if use_tma_replay_nowrite_load is not None: + parts.append(f"TMANL={use_tma_replay_nowrite_load}") # replay-NOWRITE load (NOT rect) + if use_tma_replay_write_store is not None: + parts.append(f"TMAWS={use_tma_replay_write_store}") # replay-write store + parts.append(f"SR={1 if use_philox else 0}") + parts.append(f"RECT={1 if rectangle_for_nowrite else 0}") + parts.append(f"WC={1 if write_checkpoint else 0}") + parts.append(f"MODE={mode}") + parts.append(f"SORT={1 if sort_slots else 0}") + parts.append(f"REVN={1 if reverse_nowrite else 0}") + parts.append(f"HSORT={1 if hardcode_sort else 0}") + sweep_suffix = (" " + ",".join(parts)) if parts else "" + sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") + + reset_fn = _reset_conv1d_realistic if with_conv1d else _reset + # --cell-list filter: only time cells whose (canonical-knob-values) + # tuple is in the loaded set. Robust to bench gaining new knobs + # (old cell-list files keep working: any keys they don't list + # become wildcards that retain CLI defaults). + if args._cell_list_keys: + _tup = _current_cell_tuple(args, locals()) + if _tup is None or _tup not in args._cell_list_set: + continue + # Resume from JSONL: skip cells already recorded. Built the same + # way _print_row builds JSON keys; must stay in sync. + done_keys = getattr(args, "_done_keys", None) + if done_keys: + # One key per scenario (k=6, k=11, mix) — skip the whole cell + # only if ALL of its scenarios are already done. We don't + # know which scenarios will be emitted here without + # re-evaluating the inner scenario loop; conservatively skip + # only when the prev_k_for_print's specific key is done. + _resume_key = _build_json_key( + args.variant, batch, mtp_len, prev_k_for_print, + state_dtype_name, sweep_suffix, args.tp_size, + ) + if _resume_key in done_keys: + continue + if warmup_only: + reset_fn() + if scenario_pre_iter is not None: + scenario_pre_iter(0) + _run_incr() + torch.cuda.synchronize() + else: + # Inline retry: CUPTI sometimes loses records under PDL + + # high cell count; retrying the SAME cell often catches it + # because the failure is transient at the kernel-launch level. + # Per --cupti-retry budget. On final failure, append tag to + # the skipped list for an external rerun in a fresh process. + defer_results = ( + args.cuda_graph + and getattr(args, "cupti", True) + and int(getattr(args, "cupti_defer_depth", 1)) > 1 + ) + retry_budget = 0 if defer_results else max(0, getattr(args, "cupti_retry", 1)) + stats = None + expected_K = _kernels_per_iter_incremental( + mode, with_conv1d=with_conv1d, + persistent_skip_empty=scenario_skip_empty, + ) + plan_key = ( + "incremental", + args.variant, + mode, + batch, + mtp_len, + state_dtype_name, + act_dtype_name, + with_conv1d, + bool(args.l2_flush), + bool(args.external_pdl), + bool(args.internal_pdl), + bool(use_philox), + bool(rectangle_for_nowrite), + bool(write_checkpoint), + bool(sort_slots), + bool(reverse_nowrite), + bool(hardcode_sort), + scenario_pre_iter is not None, + expected_K, + ) + for attempt in range(retry_budget + 1): + stats = _time_kernel( + args, _run_incr, reset_fn, sweep_tag, + expected_K=expected_K, + pre_iter_fn=scenario_pre_iter, + pre_iter_group_factory=scenario_pre_iter_group_factory, + iters_override=scenario_iters, + cupti_plan_key=plan_key, + ) + if stats is not None: + break + if attempt < retry_budget: + print( + f"[retry] CUPTI mismatch on {sweep_tag!r}; " + f"retrying ({attempt + 1}/{retry_budget})", + file=sys.stderr, + flush=True, + ) + if stats is None: + args._skipped_cells.append(sweep_tag) + + # Attach n_writes_per_iter when it is needed for scoring. + # For pure scenarios, n_writes is constant: 0 (nowrite) or + # batch (write), determined by scn["fill"] + mtp_len > max_window. + # For mix, scn carries the precomputed per-iter array. + per_iter_nw = None + if stats is not None and ( + getattr(args, "json_detailed", False) or scn["fill"] is None + ): + eff_iters = scenario_iters if scenario_iters is not None else args.iters + if scn["fill"] is not None: + if getattr(args, "json_detailed", False): + # Pure scenario: constant n_writes for every iter. + is_write = (scn["fill"] + mtp_len > max_window) + per_iter_nw = [batch if is_write else 0] * eff_iters + else: + # Mix scenario: slice off warmup, keep timed iters. + nw_full = scn.get("n_writes_per_iter") + if nw_full is not None: + per_iter_nw = nw_full[args.warmup:args.warmup + eff_iters].tolist() + else: + per_iter_nw = None + + if stats is not None: + _submit_result_job( + args, + stats, + show_kernel_col=show_kernel_col, + kernel_name=args.variant, + batch=batch, + mtp_len=mtp_len, + prev_k=prev_k_for_print, + state_dtype_name=state_dtype_name, + act_dtype_name=act_dtype_name, + sweep_suffix=sweep_suffix, + per_iter_nw=per_iter_nw, + skipped_tag=sweep_tag, + ) + + +# Map full torch dtype name → short tag used in JSON keys (matches collect.py). +_DTYPE_SHORT = { + "float32": "fp32", "bfloat16": "bf16", "float16": "fp16", + "int8": "int8", "int16": "int16", "float8_e4m3fn": "fp8", +} + + +def _build_json_key( + kernel_name, batch, mtp_len, prev_k, state_dtype_name, sweep_suffix, tp_size +): + """Build a key matching collect.py's kernel_data.json convention: + + incremental/{batch}/{mtp}/{sd}/k{k}/{sweep_parts}/tp{tp} + triton/{batch}/{mtp}/{sd}/tp{tp} + flashinfer/{batch}/{mtp}/{sd}/tp{tp} + + `kernel_name` is what _print_row receives: variant name for the timed + kernel (replay/checkpointing) or baseline name for the baseline row. + Variant rows collapse to "incremental" — the variant choice is captured + by the sweep flags collect.py would otherwise apply via --variant. + """ + if kernel_name in ("replay", "checkpointing"): + kind = "incremental" + else: + kind = kernel_name # "triton" / "flashinfer" + + sd = _DTYPE_SHORT.get(state_dtype_name, state_dtype_name) + parts = [kind, str(batch), str(mtp_len), sd] + if prev_k != "N/A": + parts.append(f"k{prev_k}") + if sweep_suffix: + # sweep_suffix format: " M=4,W=1,S=1,SR=0,RECT=0,WC=1" + # collect.py format: "M4_W1_S1_SR0_RECT0_WC0" + # Strip leading/trailing whitespace, drop '=', commas → underscores. + parts.append( + sweep_suffix.strip().replace("=", "").replace(",", "_") + ) + parts.append(f"tp{tp_size}") + return "/".join(parts) + + +def _print_row( + show_kernel_col, + kernel_name, + batch, + mtp_len, + prev_k, + state_dtype_name, + act_dtype_name, + stats, + sweep_suffix="", + tp_size=None, + json_detailed=False, + jsonl_path=None, + jsonl_host=None, + jsonl_gpu=None, +): + """Print one summary row and append the result to the JSONL sidecar. + + `stats` is a dict from _time_kernel: {median, p95, p99, n, iters_us, + [n_writes_per_iter], [per_kernel]}. The summary table only shows the + headline percentiles. JSONL captures the compact per-iter spans + + n_writes_per_iter by default; with json_detailed=True it also captures + per-kernel data. + + When `jsonl_path` is provided, appends one JSON line per row to the + JSONL sidecar (crash-safe incremental persistence; lets a killed sweep + resume from the last completed cell on rerun, even across hosts). Open + per-write because `args` is pickled to ProcessPoolExecutor workers and + file handles aren't picklable. JSONL is the canonical artifact — the + bench no longer writes a final `.json` summary; use `jsonl_to_json.py` + if a one-shot `.json` snapshot is needed. + """ + kernel_col = f"{kernel_name:>11} | " if show_kernel_col else "" + print( + f"| {kernel_col}{batch:>5} | {mtp_len:>7} | {str(prev_k):>6} | " + f"{state_dtype_name:>11} | {act_dtype_name:>9} | " + f"{stats['median']:>9.2f} | {stats['p95']:>7.2f} | {stats['p99']:>7.2f} |" + f"{sweep_suffix}" + ) + if jsonl_path is not None: + key = _build_json_key( + kernel_name, batch, mtp_len, prev_k, state_dtype_name, + sweep_suffix, tp_size, + ) + if json_detailed: + row_stats = stats + else: + row_stats = { + k: stats[k] + for k in ("median", "p95", "p99", "n", "iters_us", "n_writes_per_iter") + if k in stats + } + if "host_timing" in stats: + row_stats["host_timing"] = stats["host_timing"] + # Append to JSONL sidecar if a path is set (incremental persistence). + # Open per-write because args is pickled to ProcessPoolExecutor + # workers, and file handles aren't picklable. A clean SIGTERM or + # Python exception will leave the file consistent up to the last + # newline; catastrophic kills can leave a partial last line, which + # the resume reader tolerates via json.JSONDecodeError pass. + if jsonl_path is not None: + # Wall-clock timestamp (float seconds since UNIX epoch) at write + # time. Lets post-hoc analysis diff consecutive rows to derive + # per-cell wall budget and identify startup-bound vs steady-state + # segments (cells/sec, downtime between bench invocations) without + # needing to instrument the bench's outer loops separately. + import time as _time + rec = {"key": key, "stats": row_stats, "t": _time.time()} + if jsonl_host is not None: + rec["host"] = jsonl_host + if jsonl_gpu is not None: + rec["gpu"] = jsonl_gpu + with open(jsonl_path, "a") as f: + f.write(json.dumps(rec) + "\n") + + +def _finish_result_job(args, job: dict) -> None: + result = job["result"] + if isinstance(result, _PendingCuptiStats): + stats = result.resolve() + else: + stats = result + + if stats is None: + skipped_tag = job.get("skipped_tag") + if skipped_tag is not None: + args._skipped_cells.append(skipped_tag) + return + + per_iter_nw = job.get("per_iter_nw") + if per_iter_nw is not None: + stats["n_writes_per_iter"] = per_iter_nw + + _print_row( + job["show_kernel_col"], + job["kernel_name"], + job["batch"], + job["mtp_len"], + job["prev_k"], + job["state_dtype_name"], + job["act_dtype_name"], + stats, + job.get("sweep_suffix", ""), + tp_size=args.tp_size, + json_detailed=getattr(args, "json_detailed", False), + jsonl_path=getattr(args, "_jsonl_path", None), + jsonl_host=getattr(args, "_jsonl_host", None), + jsonl_gpu=getattr(args, "_jsonl_gpu", None), + ) + + +def _drain_pending_results(args, *, force: bool = False) -> None: + pending_results = getattr(args, "_pending_results", None) + if not pending_results: + return + + max_pending = max(1, int(getattr(args, "cupti_defer_depth", 1))) + while pending_results: + first_result = pending_results[0]["result"] + should_block = force or len(pending_results) >= max_pending + if ( + not should_block + and isinstance(first_result, _PendingCuptiStats) + and not first_result.is_ready() + ): + break + job = pending_results.pop(0) + _finish_result_job(args, job) + + +def _submit_result_job( + args, + result, + *, + show_kernel_col, + kernel_name, + batch, + mtp_len, + prev_k, + state_dtype_name, + act_dtype_name, + sweep_suffix="", + per_iter_nw=None, + skipped_tag=None, +) -> None: + job = { + "result": result, + "show_kernel_col": show_kernel_col, + "kernel_name": kernel_name, + "batch": batch, + "mtp_len": mtp_len, + "prev_k": prev_k, + "state_dtype_name": state_dtype_name, + "act_dtype_name": act_dtype_name, + "sweep_suffix": sweep_suffix, + "per_iter_nw": per_iter_nw, + "skipped_tag": skipped_tag, + } + if isinstance(result, _PendingCuptiStats): + args._pending_results.append(job) + _drain_pending_results(args) + else: + _finish_result_job(args, job) + + +# Cell-list mode — canonical knob-key mapping to argparse args + local +# loop variable. See _load_cell_list_into_args / inner-loop filter. +# +# Each entry: cell-key → (args attribute name, comma-separated string flag) +# For split (write/nowrite) knobs, we use Xw / Xnw keys. Tied forms (M, W, +# S, CPS, LS) accepted on load and expanded to their w/nw variants. +_CELL_LIST_KEY_TO_ARG = { + "Mw": "block_size_m_write", + "Mnw": "block_size_m_nowrite", + "Ww": "num_warps_write", + "Wnw": "num_warps_nowrite", + "Sw": "num_stages_write", + "Snw": "num_stages_nowrite", + "CPSw": "cta_per_sm_write", + "CPSnw": "cta_per_sm_nowrite", + "LSw": "num_loop_stages_write", + "LSnw": "num_loop_stages_nowrite", + "pW": "precompute_num_warps", + "pS": "precompute_num_stages", + "H": "heads_per_block", + "R": "maxnreg", + "CT": "num_ctas", + "FL": "flatten", + "WS": "warp_specialize", + "TMARL": "use_tma_rect_load", + "TMAWL": "use_tma_replay_write_load", + "TMANL": "use_tma_replay_nowrite_load", + "TMAWS": "use_tma_replay_write_store", + "RECT": "rectangle_for_nowrite", + "WC": "write_modes", + "SORT": "sort_slots", + "REVN": "reverse_nowrite", + "HSORT": "hardcode_sort", + # MODE and SR get special handling (string values): + # MODE → args.modes (single mode name) + # SR → args.sr_modes ("RN" if 0, "SR" if 1) +} + +# Split-knob tied form: "M" expands to both "Mw" and "Mnw". +_CELL_LIST_TIED_EXPANSIONS = { + "M": ("Mw", "Mnw"), + "W": ("Ww", "Wnw"), + "S": ("Sw", "Snw"), + "CPS": ("CPSw", "CPSnw"), + "LS": ("LSw", "LSnw"), +} + + +def _normalize_cell(cell: dict) -> dict: + """Expand tied-form keys (M, W, S, CPS, LS) to their w/nw variants. + Returns a new dict with only canonical split-or-plain keys. + """ + out = dict(cell) + for tied, (w_key, nw_key) in _CELL_LIST_TIED_EXPANSIONS.items(): + if tied in out: + v = out.pop(tied) + out.setdefault(w_key, v) + out.setdefault(nw_key, v) + return out + + +def _load_cell_list_into_args(args) -> None: + """Read --cell-list JSON, normalize, override args.* knob ranges, and + populate args._cell_list_keys + args._cell_list_set for the inner-loop + filter. Errors out if cells aren't uniform (different key sets). + """ + with open(args.cell_list) as f: + raw = json.load(f) + if not isinstance(raw, list): + sys.exit(f"--cell-list: expected JSON list, got {type(raw).__name__}") + cells = [_normalize_cell(c) for c in raw] + if not cells: + print("[cell-list] empty list — nothing to time", file=sys.stderr) + return + # All cells must share the same key set (uniform schema) + keys0 = frozenset(cells[0].keys()) + for i, c in enumerate(cells[1:], start=1): + if frozenset(c.keys()) != keys0: + sys.exit( + f"--cell-list: cells must have uniform key sets; cell[0] " + f"has {sorted(keys0)} but cell[{i}] has {sorted(c.keys())}" + ) + + # Auto-cover: collect per-knob value set across all cells + cover: dict = {} + for c in cells: + for k, v in c.items(): + cover.setdefault(k, set()).add(v) + # Apply overrides + for key, vals in cover.items(): + if key in _CELL_LIST_KEY_TO_ARG: + arg_name = _CELL_LIST_KEY_TO_ARG[key] + vals_str = ",".join(str(v) for v in sorted(vals)) + setattr(args, arg_name, vals_str) + elif key == "MODE": + args.modes = ",".join(sorted({str(v) for v in vals})) + elif key == "SR": + args.sr_modes = ",".join(sorted({"SR" if v else "RN" for v in vals})) + else: + print(f"[cell-list] WARNING: unknown key {key!r} in cells; " + f"will not override any args.* attribute (the value will " + f"still be matched in the filter if a matching local var " + f"is in scope)", file=sys.stderr) + + # Canonical key order (sorted) for tuple matching in the inner loop + args._cell_list_keys = tuple(sorted(keys0)) + args._cell_list_set = { + tuple(c[k] for k in args._cell_list_keys) for c in cells + } + print(f"[cell-list] loaded {len(cells)} cells with keys " + f"{list(args._cell_list_keys)}; overrode args.* to auto-cover", + file=sys.stderr) + + +# Maps cell-list key → name of the local variable in _bench_config's inner +# loop. Used to extract the "current cell" tuple for the filter check. +# Keep in sync with the loop-variable names; the filter is lenient about +# missing names (it picks them up from the inner scope at runtime). +_CELL_LIST_KEY_TO_LOCAL = { + "Mw": "block_size_m_w", + "Mnw": "block_size_m_nw", + "Ww": "num_warps_w", + "Wnw": "num_warps_nw", + "Sw": "num_stages_w", + "Snw": "num_stages_nw", + "CPSw": "cta_per_sm_w", + "CPSnw": "cta_per_sm_nw", + "LSw": "num_loop_stages_w", + "LSnw": "num_loop_stages_nw", + "pW": "precompute_num_warps", + "pS": "precompute_num_stages", + "H": "heads_per_block", + "R": "maxnreg", + "CT": "num_ctas", + "FL": "flatten", + "WS": "warp_specialize", + "TMARL": "use_tma_rect_load", + "TMAWL": "use_tma_replay_write_load", + "TMANL": "use_tma_replay_nowrite_load", + "TMAWS": "use_tma_replay_write_store", + "RECT": "rectangle_for_nowrite", + "WC": "write_checkpoint", + "MODE": "mode", + "SORT": "sort_slots", + "REVN": "reverse_nowrite", + "HSORT": "hardcode_sort", + "SR": "use_philox", +} + + +def _current_cell_tuple(args, locals_dict: dict) -> tuple | None: + """Build the (key1=val1, key2=val2, ...) tuple for the current inner-loop + iteration, matching args._cell_list_keys' order. Used by the inner-loop + filter to check membership in args._cell_list_set. Returns None if any + expected local is missing (the bench evolved a knob name — caller skips). + """ + if not args._cell_list_keys: + return None + vals = [] + for k in args._cell_list_keys: + local_name = _CELL_LIST_KEY_TO_LOCAL.get(k, k) + if local_name not in locals_dict: + return None + v = locals_dict[local_name] + # Coerce bools to ints to match cell-list JSON (1/0) + if isinstance(v, bool): + v = int(v) + vals.append(v) + return tuple(vals) + + +# Main benchmark loop + + +def _run_benchmark(args) -> None: + # Phase-timing markers — emit timestamped checkpoints so a captured-stdout + # run can later attribute wall time to setup vs compile-warmup vs prewarm + # vs timing. Single-line format makes log-grepping trivial. + _phase_t0 = time.perf_counter() + def _phase(label: str) -> None: + dt = time.perf_counter() - _phase_t0 + print(f"[phase] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) + _phase("enter _run_benchmark") + + # Pending-results FIFO for srxl's deferred CUPTI parsing pipeline. Each + # entry holds a _PendingCuptiStats handle; _drain_pending_results pulls + # ready entries and routes them to _print_row (which appends to JSONL). + args._pending_results = [] + + # JSONL incremental sidecar. Path = `.jsonl`. Each completed + # cell appends one line `{"key": , "stats": {...}, "host": }` + # to this file as it finishes timing. On startup we read this sidecar (if + # present) and populate _done_keys so a killed bench can resume without + # redoing already-timed cells. Crash-safe by construction: append-only + # writes survive SIGTERM/SIGKILL/reboot mid-sweep. + # + # Resume is host-blind: _done_keys includes records from any host, so a + # bench restarted on a different node fills in the missing cells without + # redoing cells already covered elsewhere. Cross-host *timings* aren't + # directly comparable, but each JSONL record carries its `host` stamp so + # the analyzer can group/compare per host. This bench no longer writes a + # final `.json` summary — the JSONL is the canonical artifact; use the + # `jsonl_to_json.py` helper if a one-shot `.json` snapshot is needed. + # + # Note: we store only paths/strings on `args` because args is pickled to + # ProcessPoolExecutor workers during compile-warmup, and file handles + # (TextIOWrapper) aren't picklable. _print_row open-appends per cell. + args._jsonl_path = None + args._done_keys: set[str] = set() + args._jsonl_host = None # hostname stamp for the current run + args._jsonl_gpu = None # GPU device id stamp (current process visibility) + if getattr(args, "json_output", None): + import socket + args._jsonl_host = socket.gethostname() + # Capture GPU id once at startup. Used by the oracle-cache layer in + # search_driver to attribute timings to a specific (host, gpu) pair + # for cross-process pruning. os.environ['CUDA_VISIBLE_DEVICES'] + # is the right source pre-torch-init (it's what the harness sets); + # post-init we could use torch.cuda.current_device() but we keep it + # to env to avoid forcing a CUDA init at this point in startup. + args._jsonl_gpu = os.environ.get("CUDA_VISIBLE_DEVICES", "") + args._jsonl_path = args.json_output + ".jsonl" + # Read existing JSONL if present: load every record's key into the + # skip set regardless of host (gap-fill on a new node). + if os.path.exists(args._jsonl_path): + n_loaded = 0 + host_counts: dict[str, int] = {} + with open(args._jsonl_path) as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except json.JSONDecodeError: + # Tolerate partial last line from a crash mid-write. + continue + k = rec.get("key") + if k is None: + continue + args._done_keys.add(k) + n_loaded += 1 + rec_host = rec.get("host") + if rec_host: + host_counts[rec_host] = host_counts.get(rec_host, 0) + 1 + if n_loaded: + host_summary = ", ".join( + f"{h}={n}" for h, n in sorted(host_counts.items()) + ) if host_counts else "(no host stamps)" + print( + f"[resume] {args._jsonl_path}: loaded {n_loaded} prior " + f"cell results across hosts [{host_summary}]; sweep will " + f"skip them. New cells stamp host={args._jsonl_host}.", + file=sys.stderr, + ) + + # Sidecar metadata: cmd, host, tp_size, variant, cupti, etc. Written + # once at startup; helps later analysis identify how this JSONL was + # produced even though there's no top-level .json wrapper anymore. + meta_path = args.json_output + ".meta.json" + meta_payload = { + "timestamp": datetime.now().isoformat(), + "host": args._jsonl_host, + "cmd": " ".join(sys.argv), + "tp_size": getattr(args, "tp_size", None), + "warmup": getattr(args, "warmup", None), + "iters": getattr(args, "iters", None), + "variant": getattr(args, "variant", None), + "cupti": getattr(args, "cupti", False), + } + # Append to a list so successive runs (gap-fill, retry) keep history. + existing_meta = [] + if os.path.exists(meta_path): + try: + with open(meta_path) as f: + existing_meta = json.load(f) + if not isinstance(existing_meta, list): + existing_meta = [existing_meta] + except (OSError, json.JSONDecodeError): + existing_meta = [] + existing_meta.append(meta_payload) + # Bench is sometimes invoked with --json-output pointing into a dir + # the caller hasn't created (subprocess driver, search loop, etc.). + # Ensure the dir exists before writing the meta sidecar OR the JSONL. + os.makedirs(os.path.dirname(os.path.abspath(meta_path)), exist_ok=True) + tmp = meta_path + ".tmp" + with open(tmp, "w") as f: + json.dump(existing_meta, f, indent=2) + os.replace(tmp, meta_path) + + # Skipped cells accumulator — populated by _bench_config when CUPTI capture + # mismatch causes a cell to be skipped. Written to args.skipped_output + # (or derived from json_output) at end of run. + args._skipped_cells = [] + + # Cell-list filter (replaces the old --retry-cells tag-string filter). + # When set, the sweep iterates ONLY the cells described in the list. + # + # Each entry in the JSON file is a dict of canonical knob keys → values, + # using the same names that appear in the sweep_tag (Mw/Mnw, Ww/Wnw, + # Sw/Snw, pW, pS, H, R, CT, CPSw/CPSnw, LSw/LSnw, FL, WS, TMARL, + # TMAWL, TMANL, TMAWS, SR, RECT, WC, MODE, SORT, REVN, HSORT). Each + # cell may also use the tied forms M / W / S / CPS / LS (single value + # applied to both write and nowrite halves). + # + # On load we: + # - Override the bench's CLI knob args (`args.block_size_m_write`, + # etc.) with the union of values present across all cells per knob, + # so the cartesian iteration auto-covers the list. + # - Build `args._cell_list_keys` (the canonical key order used by + # every cell — must be uniform across the list) and + # `args._cell_list_set` (frozen tuples for O(1) membership check + # inside the inner loop). + # + # In the inner loop, we build the current iteration's tuple and skip + # cells not in the set. Dict-matching is robust to bench gaining new + # knobs (old cell-list files keep working — newly-added knobs simply + # aren't matched on, so they retain CLI defaults). + # Cell-list state may already have been populated by main() (so that + # the args.*_list derivations downstream see the override). Default to + # empty if not. + _phase(f"done loading _done_keys ({len(args._done_keys)} entries)") + + if not hasattr(args, "_cell_list_keys"): + args._cell_list_keys: tuple = () + args._cell_list_set: set = set() + if getattr(args, "cell_list", None): + _load_cell_list_into_args(args) + _phase(f"done loading cell-list ({len(args._cell_list_set)} cells)") + + assert args.nheads % args.tp_size == 0, ( + f"nheads ({args.nheads}) must be divisible by tp_size ({args.tp_size})" + ) + assert args.ngroups % args.tp_size == 0, ( + f"ngroups ({args.ngroups}) must be divisible by tp_size ({args.tp_size})" + ) + args.tp_nheads = args.nheads // args.tp_size + args.tp_ngroups = args.ngroups // args.tp_size + + batch_sizes = [int(x) for x in args.batch_sizes.split(",")] + mtp_lengths = [int(x) for x in args.mtp_lengths.split(",")] + + dtype_map = { + "bf16": torch.bfloat16, + "fp32": torch.float32, + "fp16": torch.float16, + "int8": torch.int8, + "int16": torch.int16, + "fp8": torch.float8_e4m3fn, + } + state_dtypes = [dtype_map[s] for s in args.state_dtypes.split(",")] + act_dtypes = [dtype_map[s] for s in args.act_dtypes.split(",")] + + # Resolve baseline function + if args.baseline == "flashinfer": + from flashinfer.mamba import selective_state_update as baseline_fn + elif args.baseline == "triton": + baseline_fn = selective_state_update + else: + baseline_fn = None + + # --with-conv1d uses its own realistic L2 flush (cold cache flush then + # hot in_proj write). Override the generic l2_flush to avoid double-flushing. + if args.with_conv1d: + args.l2_flush = False + _init_l2_flush() # still needed for the realistic reset's flush step + elif args.l2_flush: + _init_l2_flush() + + _phase("about to enter compile-warmup") + if args.compile_threads > 0: + _compile_warmup_phase( + args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, + baseline_fn, max_workers=args.compile_threads, + ) + _phase("returned from compile-warmup") + + # Pre-warm the per-(state_dtype, act_dtype, mtp_len, ...) tensor cache at + # the largest requested batch size. Without this, the timing loop would + # progressively grow the cache as it encounters larger batches (e.g., + # iterate 1 -> 16 -> 64 -> 128 -> 512 = 5 separate growth allocations, + # each freeing the previous buffers). Pre-warming at max-batch up front + # makes every subsequent timing cell a view-slice (zero alloc cost). + _max_batch = max(batch_sizes) + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + for mtp_len in mtp_lengths: + _build_tensors( + _max_batch, mtp_len, state_dtype, act_dtype, + args.tp_nheads, args.head_dim, args.d_state, args.tp_ngroups, + max_window=getattr(args, "max_window", None) or None, + ) + _phase("done tensor prewarm — entering timing") + + if args.profile: + torch.cuda.cudart().cudaProfilerStart() + + # Print header + if baseline_fn is not None: + print( + f"| {'kernel':>11} | {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " + f"{'state_dtype':>11} | {'act_dtype':>9} | " + f"{'median_us':>9} | {'p95_us':>7} | {'p99_us':>7} |" + ) + print( + f"|{'-' * 13}|{'-' * 7}|{'-' * 9}|{'-' * 8}|" + f"{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" + ) + else: + print( + f"| {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " + f"{'state_dtype':>11} | {'act_dtype':>9} | " + f"{'median_us':>9} | {'p95_us':>7} | {'p99_us':>7} |" + ) + print( + f"|{'-' * 7}|{'-' * 9}|{'-' * 8}|{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" + ) + + sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) + rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) + write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) + modes_list = getattr(args, "modes_list", ["monolithic"]) + sort_list = getattr(args, "sort_slots_list", [False]) + rev_list = getattr(args, "reverse_nowrite_list", [False]) + hsort_list = getattr(args, "hardcode_sort_list", [False]) + + # Pre-load AL distribution for mix mode (if --mix-csv set). + mix_al = None + mix_label = "" + if args.mix_csv is not None: + from pathlib import Path as _Path + from checkpoint_mix_sim import load_al_distribution as _load_al + mix_label = _Path(args.mix_csv).stem + # T (= mtp_len) varies per cell; load once with the LARGEST mtp so + # we have enough columns; the loader normalizes the dist anyway. + mix_al = _load_al(_Path(args.mix_csv), T=max(mtp_lengths), column=args.mix_csv_column) + + for batch in batch_sizes: + for mtp_len in mtp_lengths: + # Resolve prev_k fractions → clamped integers in [0, mtp_len] + prev_ks = _resolve_prev_ks(args, mtp_len) + + # Pre-generate mix samples once per (batch, mtp_len) cell so all + # tuning configs see the same per-iter prev_tokens vectors — + # tuning differences become signal, mix-noise is shared. + # Size the sample buffer for the LARGER of args.iters and + # args.mix_iters since mix scenarios use mix_iters. + mix_samples_cpu = None + perm_samples_cpu = None # per-iter slot perm sorted write-first + mix_samples_sorted_cpu = None # per-iter prev_tokens, write-first + if mix_al is not None: + from checkpoint_mix_sim import sample_steady_state_pnat as _sample_pnat + _max_window = getattr(args, "max_window", 0) or mtp_len + _max_iters = max(args.iters, getattr(args, "mix_iters", None) or args.iters) + mix_samples_cpu = _sample_pnat( + mix_al, T=mtp_len, window=_max_window, batch=batch, + K=args.warmup + _max_iters, seed=args.mix_seed, + ) + if any(sort_list) or any(hsort_list): + # write-first stable argsort: kind='stable' preserves + # original-slot order within each mode group. + write_mask = ( + mix_samples_cpu + mtp_len > _max_window + ).astype(np.int8) # 1 = write, 0 = nowrite + perm_idx = np.argsort( + -write_mask, kind="stable", axis=-1 + ).astype(np.int32) + if any(sort_list): + perm_samples_cpu = perm_idx + if any(hsort_list): + # Apply the perm to the prev_tokens samples themselves. + # Result row i = mix_samples_cpu[i] reordered such + # that write-mode entries come first. + mix_samples_sorted_cpu = np.take_along_axis( + mix_samples_cpu, perm_idx, axis=-1 + ).astype(mix_samples_cpu.dtype) + + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + for sr_mode in sr_modes_list: + for mode in modes_list: + # Non-monolithic modes ignore write_checkpoint + # (per-slot from PNAT) — collapse the sweep so we + # don't duplicate identical cells. + effective_write_modes = ( + write_modes_list if mode == "monolithic" else [True] + ) + for write_ckpt in effective_write_modes: + # Rectangle is meaningful for: nowrite cells in + # monolithic; always for dynamic / doublelaunch + # (constexpr knob). + if mode == "monolithic": + effective_rect_list = ( + [False] if write_ckpt else rect_list + ) + else: + effective_rect_list = rect_list + for rect in effective_rect_list: + # Sort/reverse only meaningful for the + # dl-family early-out kernels AND only + # against the mix scenario (the actual + # sort experiment). Pure k= scenarios + # under sort=1 would just run a + # USE_PERM=True kernel against an + # identity perm — same data point as + # sort=0 + extra compile. Skip sort=1 + # when no mix is configured; mono / + # dynamic also skip sort=1; reverse=1 + # with sort=0 is a no-op (skip). + # Note: "persistent_main" is included in + # the dl-family for sort/hsort sweep + # eligibility — it consumes the same + # slot_perm and benefits from the same + # write-first clustering. It additionally + # requires _n_writes (count of write + # slots) which the bench computes from + # the pure-scenario PNAT (mix scenarios + # not yet supported for persistent_main). + is_dl_family = mode in ( + "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", "persistent_main", + "persistent_dynamic", + ) + can_sort = ( + is_dl_family and mix_samples_cpu is not None + ) + effective_sort_list = ( + sort_list if can_sort else [False] + ) + effective_hsort_list = ( + hsort_list if can_sort else [False] + ) + for sort_slots in effective_sort_list: + for hardcode_sort in effective_hsort_list: + # sort_slots and hardcode_sort + # are alternative experiments + # for the same idea — skip the + # combined cell to avoid double + # interpretation. + if sort_slots and hardcode_sort: + continue + # rev=1 is meaningful with EITHER + # sort_slots=1 (perm-based) or + # hardcode_sort=1 (raw pid_b + # subtraction in unsorted-perm + # path). rev=1 with both 0 is + # a no-op. + effective_rev_list = ( + rev_list if (sort_slots or hardcode_sort) else [False] + ) + for reverse_nowrite in effective_rev_list: + _bench_config( + args, batch, mtp_len, + prev_ks, state_dtype, + act_dtype, baseline_fn, + sr_mode=sr_mode, + rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, + mode=mode, + mix_samples_cpu=mix_samples_cpu, + mix_label=mix_label, + sort_slots=sort_slots, + reverse_nowrite=reverse_nowrite, + perm_samples_cpu=perm_samples_cpu, + hardcode_sort=hardcode_sort, + mix_samples_sorted_cpu=mix_samples_sorted_cpu, + ) + + _drain_pending_results(args, force=True) + + if args.profile: + torch.cuda.cudart().cudaProfilerStop() + + # JSONL is the canonical artifact (written incrementally per cell with + # host stamps). No clean-exit `.json` write — use `jsonl_to_json.py` to + # materialize a snapshot when an analyzer wants one. + if args.json_output and args._jsonl_path is not None: + print(f"\nJSONL results: {args._jsonl_path} " + f"(meta sidecar: {args.json_output}.meta.json)") + + # Write the skipped-cells sidecar. Caller can convert this list to a + # --cell-list JSON (one dict per skipped cell) to drive a retry pass in + # a fresh process. + skipped_path = getattr(args, "skipped_output", None) + if skipped_path is None and args.json_output: + # Derive default: foo.json -> foo.skipped.json + skipped_path = args.json_output.rsplit(".", 1)[0] + ".skipped.json" + if skipped_path is not None and args._skipped_cells: + payload = { + "metadata": { + "timestamp": datetime.now().isoformat(), + "cmd": " ".join(sys.argv), + "skipped_count": len(args._skipped_cells), + }, + "skipped": args._skipped_cells, + } + tmp = skipped_path + ".tmp" + with open(tmp, "w") as f: + json.dump(payload, f, indent=2) + os.replace(tmp, skipped_path) + print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " + f"tags written to: {skipped_path}", file=sys.stderr) + elif args._skipped_cells: + # No output path but there are skipped cells — emit a stderr summary. + print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " + f"first 5: {args._skipped_cells[:5]}", file=sys.stderr) + + +# CLI + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Benchmark replay_selective_state_update Triton kernel", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--nheads", + type=int, + default=NHEADS, + help="Full-model nheads (divided by --tp-size for per-GPU slice)", + ) + parser.add_argument( + "--ngroups", + type=int, + default=NGROUPS, + help="Full-model ngroups (divided by --tp-size for per-GPU slice)", + ) + parser.add_argument( + "--head-dim", type=int, default=HEAD_DIM, help="Head dimension (not TP-split)" + ) + parser.add_argument( + "--d-state", type=int, default=D_STATE, help="SSM state dimension (not TP-split)" + ) + parser.add_argument( + "--tp-size", + type=int, + default=TP_SIZE, + help="Tensor parallel size; divides nheads and ngroups", + ) + parser.add_argument( + "--batch-sizes", default="1,2,4,8", help="Comma-separated decode batch sizes" + ) + parser.add_argument( + "--mtp-lengths", + default="1,2,4,8", + help="Comma-separated per-request sequence lengths (num_draft_tokens + 1 target)", + ) + parser.add_argument( + "--state-dtypes", + default="fp32", + help="Comma-separated state dtypes: fp16,bf16,fp32,int8,int16,fp8. " + "Quantized dtypes (int8/int16/fp8) require the checkpointing variant " + "and skip baselines (selective_state_update doesn't accept them).", + ) + parser.add_argument( + "--act-dtypes", + default="bf16", + help="Comma-separated activation dtypes for x/B/C/dt: fp32,bf16", + ) + parser.add_argument("--warmup", type=int, default=4, + help="Number of warmup iterations. Default aligns with " + "the graph group-iters (default 4 for mix scenarios) so " + "warmup + iters / mix-iters lands on a clean multiple " + "without per-args rounding overhead. Earlier default of " + "20 was overkill for steady-state warming.") + parser.add_argument("--iters", type=int, default=100, help="Number of timed iterations") + parser.add_argument( + "--compile-threads", + type=int, + default=64, + help="Number of THREADS used in the compile-warmup phase (one call " + "per (batch, mtp_len, prev_k, dtype, sweep) cell, parallelized over " + "N threads). Triton compile releases the GIL, so threads compile " + "in parallel and populate the persistent cache for free hits during " + "the sequential timed phase. 0 disables the phase. Default 64.", + ) + parser.add_argument( + "--mp-start-method", + choices=("spawn", "forkserver"), + default="spawn", + help="multiprocessing start method for compile-warmup workers AND " + "the CUPTI parser child process. 'spawn' (default) is robust but " + "each child re-imports the bench module (~15s torch+triton import " + "cost). 'forkserver' starts a server once, preloads the bench " + "module ONCE, then forks children cheaply (~1s each). When 4 " + "benches run concurrently with --compile-threads 26 each, spawn " + "still incurs 4*26=104 imports per round; forkserver cuts this to " + "4 (one per server).", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Wrap timed region in cudaProfilerStart/Stop (for ncu --target-processes all)", + ) + parser.add_argument( + "--l2-flush", + action=argparse.BooleanOptionalAction, + default=True, + help="L2 eviction between iterations", + ) + parser.add_argument( + "--cuda-graph", + action=argparse.BooleanOptionalAction, + default=True, + help="Capture all warmup + timed iterations in a " + "single CUDA graph with per-iteration events " + "inside the graph, eliminating all host overhead.", + ) + parser.add_argument( + "--cuda-graph-group-iters", + type=int, + default=None, + help="Capture this many logical benchmark iterations per graph " + "replay when warmup + iters is divisible by this value. Default " + f"auto-selects {_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE} for pure " + f"cells and {_DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX} for mix cells. " + "Mix cells use a per-replay device window so they can group " + "iterations too.", + ) + parser.add_argument( + "--cupti", + action=argparse.BooleanOptionalAction, + default=True, + help="Time kernels via CUPTI Activity API (1 ns from the GPU " + "profiling fabric); per-iter span = max(kernel_end) - " + "min(kernel_start). Default ON. --no-cupti disables in-bench " + "timing entirely (kernels still run, but median/p95/p99 are zero) " + "— use when wrapping the bench in nsys/ncu, where the external " + "profiler provides timings and our CUPTI subscriber would conflict.", + ) + parser.add_argument( + "--cupti-flush-period-ms", + type=int, + default=0, + help="If >0, ask CUPTI to periodically flush activity buffers during " + "the timed CUDA-graph region. This can overlap raw-buffer parsing with " + "long timed cells; 0 leaves flushing explicit at the end of each cell.", + ) + parser.add_argument( + "--cupti-defer-depth", + type=int, + default=4, + help="Maximum number of CUDA-graph CUPTI timing results that may be " + "left for the parser process while the main process starts later cells. " + "1 preserves synchronous per-cell parsing and inline retry behavior.", + ) + parser.add_argument( + "--json-output", + default=None, + help="If set, write per-cell results to this JSON file in the " + "shape consumed by collect.py / report.py. See the 'JSON output " + "schema' section at the top of this file.", + ) + parser.add_argument( + "--json-detailed", + action=argparse.BooleanOptionalAction, + default=False, + help="When --json-output is set, also include per_kernel " + "(per-iter relative start/end timestamps for each kernel). Compact " + "JSON always includes iters_us, and mix rows include " + "n_writes_per_iter. Default off keeps records compact.", + ) + parser.add_argument( + "--host-timing", + action=argparse.BooleanOptionalAction, + default=False, + help="Attach benchmark host-side phase timings to JSON/JSONL results. " + "Useful for diagnosing benchmark overhead, but it adds roughly 1 KB " + "per compact JSONL row and several perf_counter calls per cell.", + ) + parser.add_argument( + "--cupti-retry", + type=int, + default=1, + help="On CUPTI capture mismatch (kernel record count != expected), " + "retry the cell this many times in-process before giving up. CUPTI " + "gets racy after thousands of cells in one process (PDL + small " + "kernels occasionally lose records); a single retry usually catches " + "transient cases. Set 0 to disable and skip on first mismatch.", + ) + parser.add_argument( + "--skipped-output", + default=None, + help="Path to write the list of cells that failed CUPTI capture even " + "after --cupti-retry retries (JSON list of sweep_tag strings). " + "Default: derived from --json-output by replacing .json with " + ".skipped.json.", + ) + parser.add_argument( + "--cell-list", + default=None, + help="Path to a JSON list of cell dicts (one per cell to time). " + "Each dict has canonical knob keys → values: Mw, Mnw, Ww, Wnw, Sw, " + "Snw, pW, pS, H, R, CT, CPSw, CPSnw, LSw, LSnw, FL, WS, TMARL, " + "TMAWL, TMANL, TMAWS, SR, RECT, WC, MODE, SORT, REVN, HSORT (tied " + "forms M / W / S / CPS / LS are also accepted and auto-expanded). " + "When set, bench's CLI knob ranges are auto-overridden to the " + "per-knob union across all cells, and the inner-loop filter skips " + "any iteration whose knob-value tuple isn't in the list. All cells " + "must share the same key set (uniform schema).", + ) + parser.add_argument( + "--prev-tokens-fracs", + default="0,0.5,1.0", + type=lambda s: [float(x) for x in s.split(",")], + help="Fractions of mtp_len to use as prev_num_accepted_tokens " + "for the replay kernel sweep. Values are rounded " + "and clamped to [0, mtp_len].", + ) + parser.add_argument( + "--baseline", + default=None, + nargs="?", + const="triton", + choices=[None, "triton", "flashinfer"], + help="Baseline to benchmark alongside the replay kernel. " + "'triton': native Triton selective_state_update. " + "'flashinfer': flashinfer selective_state_update (same signature). " + "Pass --baseline alone for 'triton'. Default: no baseline.", + ) + parser.add_argument( + "--output", + default=None, + help="Path to save results (file or directory). " + "If a directory, writes benchmark_replay_.txt inside it.", + ) + parser.add_argument( + "--block-size-m", + type=str, + default=None, + help="Override BLOCK_SIZE_M: single value or comma-separated sweep (e.g. '4,8,16,32').", + ) + parser.add_argument( + "--num-warps", + type=str, + default=None, + help="Override num_warps: single value or comma-separated sweep (e.g. '1,2,4').", + ) + parser.add_argument( + "--internal-pdl", + action=argparse.BooleanOptionalAction, + default=True, + help="Internal PDL between precompute and main kernels (default: on).", + ) + parser.add_argument( + "--num-stages", + type=str, + default=None, + help="Override num_stages for the main kernel (comma-separated sweep).", + ) + parser.add_argument( + "--block-size-m-write", type=str, default=None, + help="Sweep BLOCK_SIZE_M for the WRITE main only (overrides --block-size-m " + "for the write half). Tied to --block-size-m if unset.", + ) + parser.add_argument( + "--block-size-m-nowrite", type=str, default=None, + help="Sweep BLOCK_SIZE_M for the NOWRITE main only. Tied to --block-size-m if unset.", + ) + parser.add_argument( + "--num-warps-write", type=str, default=None, + help="Sweep num_warps for the WRITE main only. Tied to --num-warps if unset.", + ) + parser.add_argument( + "--num-warps-nowrite", type=str, default=None, + help="Sweep num_warps for the NOWRITE main only. Tied to --num-warps if unset.", + ) + parser.add_argument( + "--num-stages-write", type=str, default=None, + help="Sweep num_stages for the WRITE main only. Tied to --num-stages if unset.", + ) + parser.add_argument( + "--num-stages-nowrite", type=str, default=None, + help="Sweep num_stages for the NOWRITE main only. Tied to --num-stages if unset.", + ) + parser.add_argument( + "--cta-per-sm-write", type=str, default=None, + help="Sweep cta_per_sm for the WRITE persistent_main only. Tied to --cta-per-sm if unset.", + ) + parser.add_argument( + "--cta-per-sm-nowrite", type=str, default=None, + help="Sweep cta_per_sm for the NOWRITE persistent_main only. Tied to --cta-per-sm if unset.", + ) + parser.add_argument( + "--num-loop-stages-write", type=str, default=None, + help="Sweep num_loop_stages for the WRITE persistent_main only. Tied to --num-loop-stages if unset.", + ) + parser.add_argument( + "--num-loop-stages-nowrite", type=str, default=None, + help="Sweep num_loop_stages for the NOWRITE persistent_main only. Tied to --num-loop-stages if unset.", + ) + parser.add_argument( + "--skip-diagonal", action=argparse.BooleanOptionalAction, default=False, + help="When sweeping any per-main *_write / *_nowrite knobs, skip cells " + "where ALL splittable knobs satisfy write_value == nowrite_value (i.e. " + "the 'diagonal' that's already covered by a prior shared-knob sweep). " + "Useful for incremental sweeps that extend earlier results without redoing " + "the tied-knob cells.", + ) + parser.add_argument( + "--precompute-num-warps", + type=str, + default=None, + help="Override num_warps for precompute kernel (comma-separated sweep).", + ) + parser.add_argument( + "--precompute-num-stages", + type=str, + default=None, + help="Override num_stages for precompute kernel (comma-separated sweep).", + ) + parser.add_argument( + "--max-window", + type=int, + default=16, + help="Cache T-axis capacity (max replay buffer length). Default 16 " + "matches Nemotron-3-Super-120B production. Pass 0 to fall back to " + "mtp_len (degenerate every-step-checkpoint case, mostly unused).", + ) + parser.add_argument( + "--prev-tokens-int", + type=lambda s: [int(x) for x in s.split(",")] if s else None, + default=None, + help="Absolute prev_num_accepted_tokens values to test, comma-separated " + "(e.g. '0,10,11,16'). Clamped to [0, max_window]. When set, " + "overrides --prev-tokens-fracs.", + ) + parser.add_argument( + "--write-checkpoint", + action=argparse.BooleanOptionalAction, + default=True, + help="Whether the checkpointing kernel should write the post-replay " + "state to HBM. True = checkpoint step (default). False = " + "non-checkpoint step (skip state HBM write + Philox). No effect on " + "the replay variant. Ignored if --write-modes is set.", + ) + parser.add_argument( + "--write-modes", + type=str, + default=None, + help="Comma-separated 0/1 values to sweep both write modes in a " + "single nsys process — for apples-to-apples comparison of write " + "vs nowrite (replay) vs nowrite (rectangle) within one timeline. " + "Skips silently for (write=False, prev_k+T>max_window) combos. " + "When set, overrides --write-checkpoint.", + ) + parser.add_argument( + "--with-conv1d", + action="store_true", + help="Include conv1d kernel before replay SSM. " + "Uses realistic L2 flush: cold caches flushed, hot in_proj output " + "kept warm. Measures conv1d → precompute → main span.", + ) + parser.add_argument( + "--external-pdl", + action=argparse.BooleanOptionalAction, + default=True, + help="External PDL: conv1d launches dependents, precompute waits. " + "Only relevant with --with-conv1d. --no-external-pdl disables.", + ) + parser.add_argument( + "--heads-per-block", + type=str, + default=None, + help="Override HEADS_PER_BLOCK for precompute kernel (comma-separated sweep).", + ) + parser.add_argument( + "--maxnreg", + type=str, + default=None, + help="Override maxnreg for the main kernel (comma-separated sweep).", + ) + parser.add_argument( + "--num-ctas", + type=str, + default=None, + help="Override num_ctas for the main kernel (comma-separated sweep).", + ) + parser.add_argument( + "--cta-per-sm", + type=str, + default=None, + help="CTAs per SM in the 1D persistent grid for mode=persistent_main " + "(comma-separated sweep). num_persistent = cta_per_sm × NUM_SMS. " + "Default = 1 (one CTA per SM). Replaces the old --num-persistent. " + "Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--num-loop-stages", + type=str, + default=None, + help="num_stages on the inner tl.range(...) persistent loop for " + "mode=persistent_main (comma-separated sweep). Default = 2. Note: " + "this is loop-level, NOT the kernel-arg num_stages (which only " + "pipelines dot-feeding loads). Watch Triton issue #8259 — " + "num_stages>1 + flatten=True can corrupt stores in non-dot kernels. " + "Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--flatten", + type=str, + default=None, + help="`flatten` arg on tl.range(...) for mode=persistent_main " + "(comma-separated 0/1 sweep). Default = 1. Ignored for " + "non-persistent_main modes.", + ) + parser.add_argument( + "--warp-specialize", + type=str, + default=None, + help="`warp_specialize` arg on tl.range(...) for mode=persistent_main " + "(comma-separated 0/1 sweep). Default = 0. Triton 3.6 only " + "supports it on simple matmul loops; our scan loop probably won't " + "pattern-match — exposed as a knob for sweep experiments. Requires " + "num_warps >= 4 if 1. Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--sr-modes", + type=str, + default="RN", + help="Comma-separated rounding modes to sweep: any combination of " + "{RN, SR}. SR (stochastic rounding) is silently skipped for state " + "dtypes that don't support it (bf16, fp32). Default 'RN' matches " + "legacy --philox-rounding=False behavior.", + ) + parser.add_argument( + "--rectangle-for-nowrite", + type=str, + default="0", + help="Comma-separated 0/1 values: 0 = replay-style nowrite kernel, " + "1 = dedicated rectangle nowrite kernel. Sweep both with '0,1' to " + "compare in one invocation. Silently no-op for write cells (the " + "write path always uses replay-style). Only applies to the " + "checkpointing variant.", + ) + parser.add_argument( + "--use-tma-rect-load", + type=str, + default=None, + help="Comma-separated 0/1 sweep. Use TMA (host-built tensor " + "descriptor) for state load in the rectangle nowrite path. " + "Cells where the rect path isn't reachable (e.g. mode=monolithic " + "+ WC=True) skip the value=1 case as a dupe.", + ) + parser.add_argument( + "--use-tma-replay-write-load", + type=str, + default=None, + help="Comma-separated 0/1 sweep. TMA state LOAD in replay main " + "when WC=True. Independent from nowrite-load and rect TMA — see " + "CHECKPOINTING_DESIGN.md item #17 for measured perf.", + ) + parser.add_argument( + "--use-tma-replay-nowrite-load", + type=str, + default=None, + help="Comma-separated 0/1 sweep. TMA state LOAD in replay main " + "when WC=False. Design doc reports the largest win on this path " + "(int8 b>=64: -8 to -12%%).", + ) + parser.add_argument( + "--use-tma-replay-write-store", + type=str, + default=None, + help="Comma-separated 0/1 sweep. TMA state STORE in replay main " + "(WC=True path only — no-op for WC=False). Independent from all " + "load TMA flags.", + ) + parser.add_argument( + "--modes", + type=str, + default="monolithic", + help="Comma-separated dispatch modes to sweep, any of " + "{monolithic,dynamic,doublelaunch}. monolithic = today's behavior " + "(one kernel pair, write_checkpoint applied to whole batch); " + "dynamic = single kernel pair that dispatches per-slot at runtime " + "based on PNAT (rectangle_for_nowrite picks RECTANGLE constexpr); " + "doublelaunch = two kernel pairs launched in sequence with " + "EARLY_OUT=True, each handling slots whose mode matches it. " + "Only applies to the checkpointing variant; non-monolithic modes " + "ignore --write-modes (per-slot from PNAT).", + ) + parser.add_argument( + "--mix-csv", + type=str, + default=None, + help="Path to AL histogram CSV (cols: AL, count). When set, an " + "additional 'mix' cell is emitted per (batch, mtp, dtype, sr, " + "mode, RECT, M, W, ...) combo where prev_tokens varies per iter, " + "drawn from the steady-state PNAT distribution induced by the " + "AL histogram. Mix cells run only on dynamic and doublelaunch " + "modes (mono on a mixed batch corrupts wrong-mode slots). " + "Each iteration of the captured CUDA graph has a different " + "pre-baked prev_tokens vector; warmup iters use distinct samples " + "from the timed iters so nsys-included warmup leaks don't bias.", + ) + parser.add_argument( + "--mix-csv-column", + type=int, + default=1, + help="Column index (0-based) in the AL histogram CSV for the " + "count/probability column. Default 1 (second column).", + ) + parser.add_argument( + "--mix-seed", + type=int, + default=42, + help="RNG seed for the steady-state PNAT sampler. Same seed " + "across runs => same per-slot samples for reproducible " + "comparisons.", + ) + parser.add_argument( + "--sort-slots", + type=str, + default="0", + help="Comma-separated 0/1. When 1, mix scenarios pre-sort slots " + "write-first (write slots at the head of slot_perm, nowrite at the " + "tail) and the dl-family kernels read pid_b through that perm — " + "clusters early-outs at one end of the grid. Only meaningful for " + "doublelaunch/dlgrouped/maindl with mix scenarios; mono/dynamic " + "and pure-batch cells skip sort=1.", + ) + parser.add_argument( + "--reverse-nowrite", + type=str, + default="0", + help="Comma-separated 0/1. When 1 (and --sort-slots 1), the " + "nowrite-side kernels in dlgrouped/doublelaunch/maindl walk the " + "perm in reverse so both halves of the dl chain front-load real " + "work. reverse=1 with sort=0 is skipped (no perm to reverse).", + ) + parser.add_argument( + "--hardcode-sort", + type=str, + default="0", + help="Comma-separated 0/1. When 1, the per-iter prev_tokens " + "samples are pre-sorted write-first OFFLINE (CPU-side) before " + "the timed region — kernel runs unchanged (USE_PERM=False) but " + "the EO gate sees sorted PNAT so early-outs cluster naturally. " + "Zero per-program load cost vs --sort-slots; output is " + "scrambled (we don't permute x/B/C/dt) but timing is meaningful. " + "Used to isolate whether clustering helps independent of the " + "perm-load overhead in the sort-slots path.", + ) + parser.add_argument( + "--mix-iters", + type=int, + default=None, + help="Iteration count override for mix scenarios (each iter is a " + "different per-slot prev_tokens draw). Default (None) uses " + "--iters. Mix scenarios benefit from more iters since each " + "iter samples a different mix; pure scenarios don't.", + ) + parser.add_argument( + "--mix-only", + action=argparse.BooleanOptionalAction, + default=False, + help="When --mix-csv is set, emit only mix scenarios and skip the " + "pure prev_k sibling scenarios. Default: false.", + ) + parser.add_argument( + "--philox-rounding", + action="store_true", + help="DEPRECATED — equivalent to --sr-modes SR. Retained for " + "backward compatibility; use --sr-modes for new scripts. fp16 SR " + "and fp8 SR require sm_100a (Blackwell B200+).", + ) + parser.add_argument( + "--philox-rounds", + type=int, + default=5, + help="Number of Philox PRNG rounds. Default 5 matches the " + "Nemotron-3-Super-120B production config (mamba_ssm_philox_rounds=5 " + "in examples/configs and tests/integration/perf configs). The " + "wrapper's generic fallback default is 10; callers without explicit " + "config see 10. Only consulted when --philox-rounding is enabled.", + ) + parser.add_argument( + "--variant", + choices=["replay", "checkpointing"], + default="replay", + help="Which kernel to time as the 'replay' row. 'replay' = today's " + "kernel (selective_state_update.py:replay). 'checkpointing' = " + "checkpointing_state_update.py. Both share the same wrapper signature.", + ) + parser.add_argument( + "--full-import", + action="store_true", + help="Use standard tensorrt_llm import path instead of fast direct " + "module loading. Slower (~40s startup) but guaranteed correct " + "if the fast path breaks due to package changes.", + ) + args = parser.parse_args() + if args.mix_only and args.mix_csv is None: + parser.error("--mix-only requires --mix-csv") + + # Round iter counts up so warmup + iters (and warmup + mix_iters) are clean + # multiples of the graph group-iters used downstream. Default mix group is + # 4, default pure group is 2. An explicit --cuda-graph-group-iters can + # request a larger group. We round to the max of the two so all scenarios + # in a single run (pure + mix) share a clean total_iters. The cost is at + # most (group-1) extra iters per scenario — negligible — and the win is + # that graph_group_iters never falls back to 1 (which caused ~5x slowdown + # in observed benchmark walls). + _group_for_rounding = max( + _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX, + _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE, + getattr(args, "cuda_graph_group_iters", None) or 0, + ) + def _round_iters_to_group(name, val): + total = args.warmup + val + if total % _group_for_rounding == 0: + return val + new_total = ((total + _group_for_rounding - 1) // _group_for_rounding) * _group_for_rounding + new_val = new_total - args.warmup + print(f"[bench] rounding --{name} {val} → {new_val} so warmup+{name} " + f"({new_total}) is a multiple of graph group_iters={_group_for_rounding}", + file=sys.stderr) + return new_val + args.iters = _round_iters_to_group("iters", args.iters) + if getattr(args, "mix_iters", None): + args.mix_iters = _round_iters_to_group("mix-iters", args.mix_iters) + + # Cell-list (if any) must be applied BEFORE the post-argparse string→list + # derivations below — those build args.*_list from args.* strings, so a + # cell-list override of e.g. args.modes='maindl' needs to land before + # args.modes_list is computed. The function populates args._cell_list_keys + # and args._cell_list_set, plus overrides args.* knob strings to the + # per-knob union of values across the listed cells. + if getattr(args, "cell_list", None): + _load_cell_list_into_args(args) + + # Backward-compat: --philox-rounding implies --sr-modes SR if --sr-modes + # was left at the default. If both are set explicitly, error. + sr_modes_default = (args.sr_modes == "RN") + if args.philox_rounding: + if not sr_modes_default and args.sr_modes != "SR": + parser.error( + "--philox-rounding (deprecated) is incompatible with explicit " + f"--sr-modes={args.sr_modes!r}. Use --sr-modes SR (or " + "RN,SR) instead and drop --philox-rounding." + ) + args.sr_modes = "SR" + + sr_modes = [m.strip() for m in args.sr_modes.split(",") if m.strip()] + for m in sr_modes: + if m not in ("RN", "SR"): + parser.error(f"--sr-modes value must be RN or SR, got {m!r}") + args.sr_modes_list = sr_modes + + rect_modes = [v.strip() for v in args.rectangle_for_nowrite.split(",") if v.strip()] + rect_list = [] + for v in rect_modes: + if v not in ("0", "1"): + parser.error(f"--rectangle-for-nowrite value must be 0 or 1, got {v!r}") + rect_list.append(v == "1") + args.rectangle_for_nowrite_list = rect_list + + sort_modes = [v.strip() for v in (args.sort_slots or "0").split(",") if v.strip()] + sort_list = [] + for v in sort_modes: + if v not in ("0", "1"): + parser.error(f"--sort-slots value must be 0 or 1, got {v!r}") + sort_list.append(v == "1") + args.sort_slots_list = sort_list + + rev_modes = [v.strip() for v in (args.reverse_nowrite or "0").split(",") if v.strip()] + rev_list = [] + for v in rev_modes: + if v not in ("0", "1"): + parser.error(f"--reverse-nowrite value must be 0 or 1, got {v!r}") + rev_list.append(v == "1") + args.reverse_nowrite_list = rev_list + + hsort_modes = [v.strip() for v in (args.hardcode_sort or "0").split(",") if v.strip()] + hsort_list = [] + for v in hsort_modes: + if v not in ("0", "1"): + parser.error(f"--hardcode-sort value must be 0 or 1, got {v!r}") + hsort_list.append(v == "1") + args.hardcode_sort_list = hsort_list + + if args.write_modes is not None: + wm = [v.strip() for v in args.write_modes.split(",") if v.strip()] + write_list = [] + for v in wm: + if v not in ("0", "1"): + parser.error(f"--write-modes value must be 0 or 1, got {v!r}") + write_list.append(v == "1") + args.write_modes_list = write_list + else: + args.write_modes_list = [args.write_checkpoint] + + modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] + valid_modes = { + "monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", + "dl_write_only", "persistent_main", "persistent_dynamic", + } + for m in modes_raw: + if m not in valid_modes: + parser.error( + f"--modes value must be one of {sorted(valid_modes)}, got {m!r}" + ) + args.modes_list = modes_raw or ["monolithic"] + return args + + +class _Tee: + """Write to both stdout and a file simultaneously.""" + + def __init__(self, path: str): + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + self._file = open(path, "w") # noqa: SIM115 + self._stdout = sys.stdout + + def write(self, data): + self._stdout.write(data) + self._file.write(data) + + def flush(self): + self._stdout.flush() + self._file.flush() + + def close(self): + self._file.close() + + +if __name__ == "__main__": + _args = _parse_args() + + # Configure multiprocessing start method early — must be before any + # mp.get_context() that uses the chosen method. For forkserver, also + # add this file's dir to sys.path so the forkserver can import this + # module by basename for preload (otherwise it tries to import + # __main__, which is a different beast across processes). + if _args.mp_start_method == "forkserver": + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + mp.set_start_method("forkserver", force=True) + try: + mp.set_forkserver_preload([ + "benchmark_replay_selective_state_update", + ]) + except Exception as _e: + print(f"[warn] set_forkserver_preload failed: {_e!r}; " + f"forks will still work but pay full import cost", + file=sys.stderr) + _MP_START_METHOD = _args.mp_start_method + + _out_path = None + if _args.output != "-": + _ts = datetime.now().strftime("%Y%m%d_%H%M%S") + _fname = f"benchmark_replay_{_ts}.txt" + if _args.output is None: + _out_path = os.path.expanduser(f"~/nemo_logs/{_fname}") + elif os.path.isdir(_args.output) or _args.output.endswith("/"): + _out_path = os.path.join(_args.output, _fname) + else: + _out_path = _args.output + + if _out_path: + _tee = _Tee(_out_path) + sys.stdout = _tee + print(f"# benchmark_replay_selective_state_update {datetime.now().isoformat()}") + print(f"# cmd: {' '.join(sys.argv)}") + + try: + _run_benchmark(_args) + finally: + if _out_path: + sys.stdout = _tee._stdout + _tee.close() + print(f"\nResults saved to: {_out_path}") diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py index 88e08c380870..d0a80cc844f4 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math + import pytest import torch import torch.nn.functional as F @@ -21,6 +23,8 @@ from einops import repeat from tensorrt_llm._torch.modules.mamba.checkpointing_state_update import ( + _stochastic_round_int8_packed, + _stochastic_round_int16_packed, checkpointing_state_update, ) from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update @@ -2003,6 +2007,92 @@ def test_checkpointing_heads_per_block_multistep( # the test exercises the exact PTX form independent of wrapper changes. +@triton.jit +def _packed_int8_sr_kernel(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + (offs // 4)) + y = _stochastic_round_int8_packed(x, rand, offs) + tl.store(out_ptr + offs, y.to(tl.int8)) + + +@triton.jit +def _packed_int16_sr_kernel(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + (offs // 2)) + y = _stochastic_round_int16_packed(x, rand, offs) + tl.store(out_ptr + offs, y.to(tl.int16)) + + +def _bitrev_int(x: int, bits: int) -> int: + out = 0 + for _ in range(bits): + out = (out << 1) | (x & 1) + x >>= 1 + return out + + +def _rand_words(rand: torch.Tensor) -> list[int]: + return [int(v) & 0xFFFFFFFF for v in rand.cpu().tolist()] + + +def test_packed_int_sr_matches_reference(): + device = "cuda" + n = 1024 + offs = torch.arange(n, device=device, dtype=torch.float32) + x = ((offs % 37) - 18.0) + (((offs * 13.0) % 97.0) + 0.3) / 128.0 + + torch.manual_seed(42) + rand_i8 = torch.randint(-(2**31), 2**31, (n // 4,), device=device, dtype=torch.int32) + out_i8 = torch.empty(n, device=device, dtype=torch.int8) + _packed_int8_sr_kernel[(1,)](x, rand_i8, out_i8, BLOCK=n) + + x_cpu = x.cpu().tolist() + rand_i8_words = _rand_words(rand_i8) + ref_i8 = [] + for i, value in enumerate(x_cpu): + word = rand_i8_words[i // 4] + low = word & 0x0000FFFF + high = (word >> 16) & 0x0000FFFF + pos = i & 3 + if pos == 0: + rand16 = low + elif pos == 1: + rand16 = _bitrev_int(low, 16) + elif pos == 2: + rand16 = high + else: + rand16 = _bitrev_int(high, 16) + ref_i8.append(math.floor(value + rand16 / float(1 << 16))) + + torch.testing.assert_close( + out_i8.cpu().to(torch.int16), + torch.tensor(ref_i8, dtype=torch.int16), + rtol=0, + atol=0, + ) + + rand_i16 = torch.randint(-(2**31), 2**31, (n // 2,), device=device, dtype=torch.int32) + out_i16 = torch.empty(n, device=device, dtype=torch.int16) + _packed_int16_sr_kernel[(1,)](x, rand_i16, out_i16, BLOCK=n) + + rand_i16_words = _rand_words(rand_i16) + ref_i16 = [] + for i, value in enumerate(x_cpu): + word = rand_i16_words[i // 2] + rand_bits = word if (i & 1) == 0 else _bitrev_int(word, 32) + rand24 = rand_bits & 0x00FFFFFF + ref_i16.append(math.floor(value + rand24 / float(1 << 24))) + + torch.testing.assert_close( + out_i16.cpu(), + torch.tensor(ref_i16, dtype=torch.int16), + rtol=0, + atol=0, + ) + + @triton.jit def _bracket_kernel_fp8(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): offs = tl.arange(0, BLOCK) @@ -2138,4 +2228,3 @@ def test_sr_grid_bracket(state_dtype): "bug (cvt.rs source-register order)." ) - diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_refactored.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_refactored.py new file mode 100644 index 000000000000..2e87b9963715 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_refactored.py @@ -0,0 +1,2230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import math + +import pytest +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import repeat + +from tensorrt_llm._torch.modules.mamba.checkpointing_state_update_refactored import ( + _stochastic_round_int8_packed, + _stochastic_round_int16_packed, + checkpointing_state_update, +) +from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update +from tensorrt_llm._utils import get_sm_version + +# Philox stochastic rounding uses PTX cvt.rs.f16x2.f32 which requires sm >= 100. +_skip_pre_sm100 = pytest.mark.skipif( + get_sm_version() < 100, reason="Philox stochastic rounding needs sm >= 100" +) + +# Configs derived from NVIDIA-Nemotron-3-Super-120B-A12B Mamba2 parameters +# (nheads=128, headdim=64, d_state=128, ngroups=8) with TP split applied: +# TP=8: nheads=16, ngroups=1 — primary production config +# TP=4: nheads=32, ngroups=2 — exercises ngroups>1 (grouped B/C path) +_CONFIGS = [ + # (nheads, head_dim, d_state, ngroups) + (16, 64, 128, 1), # TP=8 production config + (32, 64, 128, 2), # TP=4, ngroups>1 (more heads than B/C groups) +] + +# Quantized state dtypes and their representable-magnitude limits (== QUANT_MAX +# in the kernel). fp8_e4m3fn cells require SM 89+ for the fp32↔fp8 cvt PTX +# instructions; SR variants of fp16/fp8 additionally need SM 100+. +_QUANT_MAX_BY_DTYPE = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, +} + + +def _quantize_state(state_fp32: torch.Tensor, state_dtype: torch.dtype, quant_max: float): + """Quantize fp32 state to (state_quant, decode_scale) using the same + per-(head, dim) channel scheme the kernel does on store. decode_scale = + max_abs_per_channel / quant_max (= 1/encode_scale). + """ + amax = state_fp32.abs().amax(dim=-1) # (cache, nheads, head_dim) + encode_scale = quant_max / amax.clamp(min=1e-30) + decode_scale = 1.0 / encode_scale + scaled = state_fp32 * encode_scale.unsqueeze(-1) + if state_dtype == torch.float8_e4m3fn: + # Native cast does RN at the fp8 grid; explicit round() would destroy + # sub-integer precision (matches the kernel's fp8 RN path). + state_quant = scaled.clamp(-quant_max, quant_max).to(state_dtype) + else: + state_quant = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) + return state_quant, decode_scale + + +def _dequantize_state(state_quant: torch.Tensor, decode_scale: torch.Tensor): + return state_quant.to(torch.float32) * decode_scale.unsqueeze(-1) + + +def _maybe_skip_dtype(state_dtype, use_sr): + """Skip on insufficient SM. fp8 e4m3fn (any) needs SM 89+; fp16/fp8 SR + needs SM 100+; int8/int16 (RN or SR) runs anywhere.""" + if state_dtype == torch.float8_e4m3fn and get_sm_version() < 89: + pytest.skip("fp8_e4m3fn requires SM 89+ (Ada Lovelace / Hopper / Blackwell)") + if use_sr and state_dtype in (torch.float16, torch.float8_e4m3fn) and get_sm_version() < 100: + pytest.skip(f"{state_dtype} stochastic rounding requires SM 100+ (Blackwell B200+)") + + +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize( + "state_dtype", + [ + torch.float16, + torch.bfloat16, + torch.float32, + torch.int8, + torch.int16, + torch.float8_e4m3fn, + ], + ids=["fp16", "bf16", "fp32", "int8", "int16", "fp8"], +) +@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) +@pytest.mark.parametrize( + "T", [6, 10, 16, 27, 32, 55], ids=["T6", "T10", "T16", "T27", "T32", "T55"] +) +@pytest.mark.parametrize( + "write_checkpoint,rectangle_for_nowrite", + [ + (True, False), # write path (rectangle_for_nowrite is ignored) + (False, False), # nowrite path via replay-style kernels + (False, True), # nowrite path via dedicated rectangle kernels + ], + ids=["write", "no_write_replay", "no_write_rectangle"], +) +@pytest.mark.parametrize( + "mode", + ["monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl"], + ids=["monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl"], +) +def test_checkpointing_state_update( + nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T, + write_checkpoint, rectangle_for_nowrite, mode, +): + """ + Verify that: + checkpointing_state_update(state0, old_caches, k, new_x, ...) + produces the same output as: + selective_state_update(state_after_k_old_tokens, new_x, ...) + and writes state_after_k_old_tokens back to the state tensor. + + Quantized state dtypes (int8/int16/fp8) follow the same flow with + a per-(head, dim) channel decode-scale tensor; comparison is done + via dequant(state, scales) against the fp32 reference. + """ + _maybe_skip_dtype(state_dtype, use_sr=False) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + + batch = 2 + device = "cuda" + dtype = torch.bfloat16 # input activations are bf16 + assert nheads % ngroups == 0 + + # Cache T-axis size (max_window). Use the kernel's BLOCK_SIZE_T as the + # ceiling — this is what the wrapper allows and enables PNAT-aware writes + # at [PNAT, PNAT+T) for no-checkpoint mode. For T=6 that's 16 (production + # max_window); for larger T it scales with np2(T). + max_window = max(triton.next_power_of_2(T), 16) + + if paged_cache: + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + else: + cache_size = batch + state_batch_indices = None + + torch.manual_seed(42) + + # A: (nheads, head_dim, d_state) with stride(-2)=0, stride(-1)=0 [tie_hdim] + A_base = -torch.rand(nheads, device=device) - 0.5 # float32, negative + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + + # dt_bias: (nheads, head_dim) with stride(-1)=0 [tie_hdim] + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + + # D: (nheads, head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + # Initial SSM state (cache_size slots). Quantized dtypes need a separate + # init: derive scales from a fp32 source so the quantized state isn't + # garbage on dequant. ref_input_state is what the fp32 reference run + # sees — for non-quant it's state0 (cast to fp32 inside reference); for + # quant it's the lossy dequant of state0 (matches what the kernel sees + # internally on load). + if is_quantized: + state0_fp32 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + state0, state0_scales = _quantize_state(state0_fp32, state_dtype, quant_max) + ref_input_state = _dequantize_state(state0, state0_scales) + else: + state0 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state0_scales = None + ref_input_state = state0.float() + + # Old inputs: up to `max_window` tokens per batch request, so the test + # loop can probe PNAT > T-1 (which the prior T-token setup couldn't + # reach). step1_T = max_window covers the full PNAT range we sweep. + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) # stride(-1)=0 + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + # Capture intermediate SSM states using selective_state_update across + # all step1_T positions — gives us reference states for k ∈ [0, step1_T]. + states_buffer_f32 = torch.zeros( + cache_size, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = ( + state_batch_indices + if paged_cache + else torch.arange(batch, device=device, dtype=torch.int32) + ) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, + dt1, + A, + B1, + C1, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + # Build cache tensors for the replay kernel. + # old_x: (cache, max_window, nheads, dim) bf16 — single-buffered + # old_B: (cache, 2, max_window, ngroups, dstate) bf16 — double-buffered + # old_dt: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous + # old_dA_cumsum: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous + # cache_buf_idx: random 0s and 1s to verify indexing correctness + old_x = torch.zeros(cache_size, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) + + # Fill each slot's active buffer (= cache_buf_idx) with step 1's data at + # positions [0:step1_T) = [0:max_window). Whole buffer covered so PNAT + # values up to max_window are exercised. Inactive buffer has random + # garbage to catch indexing bugs. + slots = state_batch_indices if paged_cache else slice(None) + old_x[slots, :step1_T] = x1 + + # Compute processed dt and dA_cumsum for step 1 + dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) + + # Write to each slot's active buffer based on its cache_buf_idx + slot_indices = state_batch_indices.tolist() if paged_cache else list(range(cache_size)) + for i, slot in enumerate(slot_indices): + buf = cache_buf_idx[slot].item() + batch_idx = i # maps slot back to the batch index + old_B[slot, buf, :step1_T] = B1[batch_idx] + old_dt[slot, buf, :, :step1_T] = dt1[batch_idx].T # (step1_T, nheads) → (nheads, step1_T) + old_dA_cumsum[slot, buf, :, :step1_T] = dA_cumsum1[batch_idx].T + + # Main loop: test each k (number of old tokens replayed). + # write_checkpoint=False (nowrite): k ∈ [0, max_window-T] — new tokens + # append at [k, k+T) of the active buffer; need k+T ≤ max_window. + # write_checkpoint=True (write): k ∈ [max_window-T+1, max_window] — + # new tokens land in the staging buffer at [0, T); k > max_window-T + # captures the overflow case that triggers a checkpoint in production. + # Combined sweep covers the full k ∈ [0, max_window] with the + # appropriate boundary handling per mode. + if write_checkpoint: + k_lo = max(0, max_window - T + 1) + k_hi = max_window + 1 # exclusive + else: + k_lo = 0 + k_hi = max_window - T + 1 # exclusive + for k in range(k_lo, k_hi): + torch.manual_seed(k + 100) + + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + # Reference (fp32, starting from the same lossy-or-not state the + # kernel sees). + ref_state_f32 = ref_input_state.clone() + if k > 0: + ref_state_f32[slots] = states_buffer_f32[slots, k - 1] + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=(state_batch_indices if paged_cache else None), + out=ref_out, + ) + + # Replay kernel — clone caches into mutable working copies that we + # can inspect AFTER the call to verify cache postconditions. + test_state = state0.clone() + test_scales = state0_scales.clone() if is_quantized else None + prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_x_w = old_x.clone() + old_B_w = old_B.clone() + old_dt_w = old_dt.clone() + old_dA_cumsum_w = old_dA_cumsum.clone() + # cache_buf_idx stays at its random values — each slot reads from its own buffer + + checkpointing_state_update( + test_state, + old_x_w, + old_B_w, + old_dt_w, + old_dA_cumsum_w, + cache_buf_idx.clone(), + prev_tokens, + x=x2, + dt=dt2, + A=A, + B=B2, + C=C2, + out=test_out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + state_scales=test_scales, + write_checkpoint=write_checkpoint, + rectangle_for_nowrite=rectangle_for_nowrite, + mode=mode, + ) + + # Tolerance rationale: the replay kernel uses bf16 tl.dot for four + # matmuls (dB_scaled @ old_x, C @ state, CB_scaled @ x, and C @ B in + # precompute). The reference (selective_state_update) and flashinfer + # baseline use fp32 element-wise MACs. The bf16 input casts lose the + # dt_bias/A-derived bits that the baselines keep — per-element rounding, + # not accumulating. Prefill (ssd_chunk_scan) does identical bf16 tl.dot + # casts, so we match prefill precision exactly. Empirical: max ~1.0 at + # T<=16, ~2.0 at T=32-55; mean ~0.014; <0.02% of elements exceed 0.5. + # State dtype (fp16/bf16/fp32) doesn't shift the error — bf16 dot + # inputs dominate, not state storage. + # + # Quantized states add a per-element state quant error eps that + # propagates through C @ state in the output dot. With dstate=128 + # and C ~ N(0,1), the output channel std from this noise is roughly + # eps * sqrt(128/3) ≈ 6.5 * eps. Stack with the bf16 baseline: + # out_atol = bf16_atol + 6.5 * eps_max + # where eps_max is the worst-case per-element error at the + # post-replay state magnitude (T=55 → amax ≈ 23). + # + # Per-element error (eps_max for T=55): + # int8 (uniform grid): amax/(2*127) ≈ 0.091 + # int16 (uniform grid): amax/(2*32767) ≈ 3.5e-4 + # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 (worst-case + # cell at top of channel; smaller for + # smaller-magnitude elements) + out_atol = ( + {torch.int8: 1.6, torch.int16: 1.05, torch.float8_e4m3fn: 4.0}[state_dtype] + if is_quantized else 1.0 + ) + out_rtol = ( + {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] + if is_quantized else 2e-2 + ) + out_diff = (test_out.float() - ref_out.float()).abs() + out_max = out_diff.max().item() + out_mean = out_diff.mean().item() + try: + torch.testing.assert_close( + test_out, ref_out, rtol=out_rtol, atol=out_atol, + msg=f"Output mismatch at k={k}", + ) + except AssertionError: + print( + f"k={k} out: max={out_max:.4f} mean={out_mean:.4f} " + f"nan={torch.isnan(test_out).any().item()} " + f"inf={torch.isinf(test_out).any().item()}" + ) + raise + + # State expectation depends on write_checkpoint: + # True → kernel writes the post-replay state; expect the + # selective_state_update reference's state at step k-1. + # False → kernel skips the HBM store; state must be UNCHANGED + # from the input (state0; for quant, scales also unchanged). + if is_quantized: + if write_checkpoint: + # Compare via dequant against the fp32 reference state. + expected_fp32 = ( + ref_input_state[slots] if k == 0 else states_buffer_f32[slots, k - 1] + ) + actual_fp32 = _dequantize_state(test_state[slots], test_scales[slots]) + # State diff = bf16_replay_error + quant_error (per element). + # The bf16 component is the SAME error source the non-quant + # test absorbs in its atol=1.0 baseline (replay's tl.dot is + # bf16-input fp32-accum; per-element error ~ 2^-7 * amax, + # empirically ≤ ~0.2 at T=55 amax≈23). Quant adds: + # int8: amax/(2*127) ≈ 0.091 worst-case + # int16: amax/(2*32767) ≈ 3.5e-4 (negligible vs bf16) + # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 worst-case + # Atol = bf16_baseline (1.0) + quant_eps_max. + state_atol = { + torch.int8: 1.1, torch.int16: 1.0, torch.float8_e4m3fn: 2.5, + }[state_dtype] + state_rtol = { + torch.int8: 5e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 1e-1, + }[state_dtype] + try: + torch.testing.assert_close( + actual_fp32, expected_fp32, + rtol=state_rtol, atol=state_atol, + msg=f"State mismatch at k={k} dtype={state_dtype}", + ) + except AssertionError: + diff = (actual_fp32 - expected_fp32).abs() + print( + f"k={k} state(dequant): max={diff.max().item():.4f} " + f"mean={diff.mean().item():.4f}" + ) + raise + # Scales sanity (fp32, finite, positive). + assert test_scales.dtype == torch.float32 + assert torch.isfinite(test_scales[slots]).all(), ( + f"state_scales has non-finite values at k={k}" + ) + assert (test_scales[slots] > 0).all(), ( + f"state_scales has non-positive values at k={k}" + ) + else: + # No write: raw quant state and scales unchanged. Use + # torch.equal for byte-level equality (dtype-agnostic; works + # for int8 / int16 / fp8 alike). + assert torch.equal(test_state[slots], state0[slots]), ( + f"Quant state changed at k={k} write_checkpoint=False" + ) + assert torch.equal(test_scales[slots], state0_scales[slots]), ( + f"State scales changed at k={k} write_checkpoint=False" + ) + else: + if write_checkpoint: + expected_state = ( + state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) + ) + else: + expected_state = state0[slots] + state_diff = (test_state[slots].float() - expected_state.float()).abs() + state_max = state_diff.max().item() + state_mean = state_diff.mean().item() + try: + torch.testing.assert_close( + test_state[slots], + expected_state, + rtol=2e-2, + atol=1.0 if write_checkpoint else 0.0, + msg=f"State mismatch at k={k} (write_checkpoint={write_checkpoint})", + ) + except AssertionError: + print( + f"k={k} state: max={state_max:.4f} mean={state_mean:.4f} " + f"nan={torch.isnan(test_state).any().item()} " + f"inf={torch.isinf(test_state).any().item()}" + ) + raise + + # --- Cache postconditions --- + # Compute step 2's processed values (what the kernel should have + # stored at [write_offset : write_offset+T) of write_buf): + # write_buf = (1 - active_buf) if write_checkpoint else active_buf + # write_offset = 0 if write_checkpoint else k + # Untouched cache regions must equal their pre-call snapshots + # (old_x / old_B / old_dt / old_dA_cumsum captured before the call). + dt2_proc = F.softplus(dt2_base.float() + dt_bias_base.float()[None, None, :]) # (B,T,H) + dA_cumsum2 = torch.cumsum(A_base.float()[None, None, :] * dt2_proc, dim=1) + write_offset = 0 if write_checkpoint else k + + for batch_idx, slot in enumerate(slot_indices): + active = cache_buf_idx[slot].item() + wb = (1 - active) if write_checkpoint else active + + # --- old_x (single-buffered): write at [write_offset : +T) of slot --- + written_x = old_x_w[slot, write_offset : write_offset + T] + torch.testing.assert_close( + written_x, x2[batch_idx], rtol=0, atol=0, + msg=f"old_x written region wrong at k={k} write={write_checkpoint}", + ) + # Untouched ranges of old_x[slot] + if write_offset > 0: + torch.testing.assert_close( + old_x_w[slot, :write_offset], old_x[slot, :write_offset], + rtol=0, atol=0, + msg=f"old_x [0:{write_offset}) modified at k={k} write={write_checkpoint}", + ) + if write_offset + T < max_window: + torch.testing.assert_close( + old_x_w[slot, write_offset + T:], old_x[slot, write_offset + T:], + rtol=0, atol=0, + msg=f"old_x [{write_offset+T}:) modified at k={k} write={write_checkpoint}", + ) + + # --- old_B (double-buffered): write at write_buf, [write_offset:+T) --- + torch.testing.assert_close( + old_B_w[slot, wb, write_offset : write_offset + T], + B2[batch_idx], rtol=0, atol=0, + msg=f"old_B written region wrong at k={k} write={write_checkpoint}", + ) + # Other-buffer (= 1-wb) untouched + torch.testing.assert_close( + old_B_w[slot, 1 - wb], old_B[slot, 1 - wb], + rtol=0, atol=0, + msg=f"old_B inactive buffer modified at k={k} write={write_checkpoint}", + ) + + # --- old_dt (double-buffered, fp32, layout (heads, T)): --- + torch.testing.assert_close( + old_dt_w[slot, wb, :, write_offset : write_offset + T], + dt2_proc[batch_idx].T, + rtol=1e-4, atol=1e-4, + msg=f"old_dt written region wrong at k={k} write={write_checkpoint}", + ) + torch.testing.assert_close( + old_dt_w[slot, 1 - wb], old_dt[slot, 1 - wb], + rtol=0, atol=0, + msg=f"old_dt inactive buffer modified at k={k} write={write_checkpoint}", + ) + + # --- old_dA_cumsum (double-buffered, fp32, layout (heads, T)): --- + torch.testing.assert_close( + old_dA_cumsum_w[slot, wb, :, write_offset : write_offset + T], + dA_cumsum2[batch_idx].T, + rtol=1e-4, atol=1e-4, + msg=f"old_dA_cumsum written region wrong at k={k} write={write_checkpoint}", + ) + torch.testing.assert_close( + old_dA_cumsum_w[slot, 1 - wb], old_dA_cumsum[slot, 1 - wb], + rtol=0, atol=0, + msg=f"old_dA_cumsum inactive buf modified at k={k} write={write_checkpoint}", + ) + + +@pytest.mark.parametrize( + "mode,rectangle_for_nowrite", + [ + ("dynamic", False), + ("dynamic", True), + ("doublelaunch", False), + ("doublelaunch", True), + ("dlgrouped", False), + ("dlgrouped", True), + ("maindl", False), + ("maindl", True), + ], + ids=[ + "dynamic_replay", + "dynamic_rectangle", + "doublelaunch_replay", + "doublelaunch_rectangle", + "dlgrouped_replay", + "dlgrouped_rectangle", + "maindl_replay", + "maindl_rectangle", + ], +) +def test_checkpointing_state_update_mixed_mode(mode, rectangle_for_nowrite): + """ + Mixed-mode dispatch: a batch where some slots have PNAT triggering + write and others triggering nowrite, exercising the per-slot dispatch + of mode={dynamic, doublelaunch}. + + Setup: 4 slots, max_window=16, T=6. + pnat_per_slot = [3, 10, 12, 16] + slots 0, 1: nowrite (pnat + T <= max_window) + slots 2, 3: write (pnat + T > max_window) + + Reference: per-slot post-replay state computed from the captured + state evolution, then selective_state_update for the new step. + Output and post-replay state are verified per-slot. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + cache_size = batch + device = "cuda" + dtype = torch.bfloat16 + state_dtype = torch.bfloat16 + + # PNAT mix: write threshold is pnat + T > max_window → pnat >= 11. + pnat_per_slot = torch.tensor([3, 10, 12, 16], device=device, dtype=torch.int32) + # Per-slot dispatch destinations under each mode (for the postcondition checks). + pnat_means_write = (pnat_per_slot + T > max_window).tolist() # [F, F, T, T] + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + ref_input_state = state0.float() + + # Old inputs spanning the full window (so any pnat 0..max_window is exercised). + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + # Capture per-step intermediate SSM states across the window. + states_buffer_f32 = torch.zeros( + cache_size, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + # Build the cache tensors with old data on each slot's active buffer. + old_x = torch.zeros(cache_size, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + old_dA_cumsum = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(cache_size): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + # New-step inputs. + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + # Reference: per-slot post-replay state, then selective_state_update. + ref_state_f32 = ref_input_state.clone() + for i in range(cache_size): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, + x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + out=ref_out, + ) + + # Kernel call. + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_x_w = old_x.clone() + old_B_w = old_B.clone() + old_dt_w = old_dt.clone() + old_dA_cumsum_w = old_dA_cumsum.clone() + cache_buf_idx_w = cache_buf_idx.clone() + + checkpointing_state_update( + test_state, + old_x_w, old_B_w, old_dt_w, old_dA_cumsum_w, + cache_buf_idx_w, + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + rectangle_for_nowrite=rectangle_for_nowrite, + mode=mode, + # write_checkpoint is ignored in dynamic / doublelaunch. + ) + + # Output: every slot must match its reference (bf16 atol consistent + # with existing tests). + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (mode={mode}, rect={rectangle_for_nowrite})", + ) + + # State postconditions per slot: + # write slots (pnat + T > max_window): state in HBM is the post-replay + # fp32 reference (cast back to state_dtype with bf16 atol). + # nowrite slots: HBM state is unchanged (still state0 bitwise). + for i in range(cache_size): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), + ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (mode={mode})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], + rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM was modified (mode={mode})", + ) + + +@pytest.mark.parametrize( + "mode,rectangle_for_nowrite,reverse_nowrite", + [ + ("doublelaunch", False, False), + ("doublelaunch", False, True), + ("doublelaunch", True, False), + ("doublelaunch", True, True), + ("dlgrouped", False, False), + ("dlgrouped", False, True), + ("dlgrouped", True, False), + ("dlgrouped", True, True), + ("maindl", False, False), + ("maindl", False, True), + ("maindl", True, False), + ("maindl", True, True), + ], + ids=lambda v: str(v), +) +def test_checkpointing_state_update_sorted_dispatch(mode, rectangle_for_nowrite, reverse_nowrite): + """ + Sort-driven dispatch: caller pre-sorts slots write-first via slot_perm. + Verifies the perm-aware kernels remap pid_b correctly so each slot's + work lands at the right grid program — i.e. slot S's output and HBM + state still match the reference under any permutation. + + Setup mirrors test_checkpointing_state_update_mixed_mode but with + slot_perm = [2, 3, 0, 1] (write slots 2, 3 first; nowrite 0, 1 after). + reverse_nowrite=True walks the perm tail-first on the nowrite-side, + so e.g. rectangle main reads perm[B-1-pid_grid] = [1, 0, 3, 2]. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + + pnat_per_slot = torch.tensor([3, 10, 12, 16], device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() # [F, F, T, T] + # Write-first perm: indices 2, 3 (write) then 0, 1 (nowrite). + slot_perm = torch.tensor([2, 3, 0, 1], device=device, dtype=torch.int32) + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=dtype) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + test_state, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + rectangle_for_nowrite=rectangle_for_nowrite, + mode=mode, + slot_perm=slot_perm, + reverse_nowrite=reverse_nowrite, + ) + + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (mode={mode}, rect={rectangle_for_nowrite}, rev={reverse_nowrite})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (mode={mode}, rev={reverse_nowrite})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM modified (mode={mode}, rev={reverse_nowrite})", + ) + + +@pytest.mark.parametrize( + "scenario,pnat_per_slot_list,n_writes_expected,slot_perm_list", + [ + # All-write: every slot has PNAT triggering write + # (PNAT + T > max_window). No permutation needed. + ("all_write", [12, 13, 14, 15], 4, [0, 1, 2, 3]), + # All-nowrite: every slot fits in the window. n_writes = 0. + ("all_nowrite", [3, 4, 5, 6], 0, [0, 1, 2, 3]), + # Mixed (write-first sorted via slot_perm): physical slots 2, 3 + # are writes; physical slots 0, 1 are nowrites. slot_perm + # remaps grid pid_b 0..3 to physical slots 2, 3, 0, 1 — so the + # first n_writes=2 grid programs hit write slots and the rest + # hit nowrite slots. + ("mixed_sorted", [3, 10, 12, 16], 2, [2, 3, 0, 1]), + ], + ids=["all_write", "all_nowrite", "mixed_sorted"], +) +def test_checkpointing_state_update_persistent_main( + scenario, pnat_per_slot_list, n_writes_expected, slot_perm_list, +): + """ + Persistent-CTA main kernel: 1D-grid kernel that loops over + (slot, M-tile, head) work units via tl.range. Caller pre-sorts + slots write-first and passes _n_writes (count of write slots) so + the kernel can split the persistent loop into write and nowrite + halves with the right WRITE_CHECKPOINT constexpr each time. + + Setup mirrors test_checkpointing_state_update_sorted_dispatch + (same fixed seeds, same input shapes) so the reference state + evolution is identical and we can compare per-slot output and + HBM-state postconditions to the same reference. + + Cases: + - all_write (n_writes=B): every slot exercises the + WRITE_CHECKPOINT=True branch of the persistent loop. + - all_nowrite (n_writes=0): every slot exercises the + WRITE_CHECKPOINT=False branch. Verifies the kernel handles + the "write half is empty" launch (n_slots=0 → early return). + - mixed_sorted: slots [2, 3] are writes, slots [0, 1] are + nowrites. slot_perm = [2, 3, 0, 1]. Persistent kernel + should call its impl with pid_b ∈ {2, 3} for the write half + and pid_b ∈ {0, 1} for the nowrite half, even though the + grid pid_b_grid is 0..n_slots-1 in each. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + + pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() + slot_perm = torch.tensor(slot_perm_list, device=device, dtype=torch.int32) + # Sanity: caller-supplied n_writes must match the actual count of + # write slots in the post-perm order. + write_count = sum(pnat_means_write) + assert write_count == n_writes_expected, ( + f"test setup error: expected {n_writes_expected} writes, got {write_count}" + ) + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=dtype + ) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + batch, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + test_state, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + mode="persistent_main", + slot_perm=slot_perm, + _n_writes=n_writes_expected, + ) + + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (scenario={scenario})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (scenario={scenario})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", + ) + + +@pytest.mark.parametrize( + "scenario,pnat_per_slot_list", + [ + # All-write: every slot has PNAT triggering write (PNAT + T > max_window). + ("all_write", [12, 13, 14, 15]), + # All-nowrite: every slot fits in the window. + ("all_nowrite", [3, 4, 5, 6]), + # Mixed: some slots write, some nowrite. No pre-sort needed; the + # dynamic kernel dispatches per-slot at runtime via PNAT load. + ("mixed_unsorted", [3, 12, 10, 15]), + ], + ids=["all_write", "all_nowrite", "mixed_unsorted"], +) +def test_checkpointing_state_update_persistent_dynamic( + scenario, pnat_per_slot_list, +): + """ + Persistent-dynamic kernel: 1D persistent-CTA grid covering the full + batch, with runtime per-slot WRITE_CHECKPOINT branch derived from + each slot's PNAT. Single launch, no half-split, no n_writes needed, + no slot_perm needed (handles unsorted batches natively). + + Same setup as test_checkpointing_state_update_persistent_main; we + verify all three scenarios — including a mixed-unsorted batch the + persistent_main kernel can't handle without pre-sorting. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + + pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=dtype + ) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + batch, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + test_state, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + mode="persistent_dynamic", + ) + + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (scenario={scenario})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (scenario={scenario})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", + ) + + +@pytest.mark.parametrize( + "scenario,pnat_per_slot_list,n_writes_expected,slot_perm_list", + [ + # Same input shape as the mixed_sorted case in + # test_checkpointing_state_update_persistent_main, but exercises + # the device-tensor n_writes plumbing + skip_empty_halves=False + # path that the bench's mix-mode benchmarking depends on. + ("mixed_sorted", [3, 10, 12, 16], 2, [2, 3, 0, 1]), + # Boundary: n_writes=batch — when skip_empty_halves=False, the + # nowrite half launches with an empty slot range and the kernel + # must do nothing useful (tl.range covers 0 iterations). + ("all_write_noskip", [12, 13, 14, 15], 4, [0, 1, 2, 3]), + # Boundary: n_writes=0 — write half launches with empty range. + ("all_nowrite_noskip", [3, 4, 5, 6], 0, [0, 1, 2, 3]), + ], + ids=["mixed_sorted", "all_write_noskip", "all_nowrite_noskip"], +) +def test_checkpointing_state_update_persistent_main_device_n_writes( + scenario, pnat_per_slot_list, n_writes_expected, slot_perm_list, +): + """ + Persistent_main with the device-tensor n_writes plumbing. + + The bench's mix-mode benchmarking captures a single CUDA graph and + replays it many times with varying per-iter n_writes. To do that the + wrapper takes `_n_writes_dev` (a (1,) int32 device tensor) instead of + `_n_writes` (host int), and the kernel reads `n_writes` from device + memory at entry — same captured pointer across replays, value can + change between replays via an outside-graph copy. Also exercises + `_persistent_skip_empty_halves=False`: both halves of persistent_main + always launch, even when one half has no work (mix scenarios can't + cheaply know n_writes host-side per iter to skip). + + Verifies: + 1. Kernel reads device n_writes correctly (output matches reference). + 2. Empty-half launches don't corrupt state (skip_empty_halves=False + + n_writes=0 / =batch boundary cases). + 3. slot_perm + USE_PERM still works through the device-n_writes path. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + + pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() + slot_perm = torch.tensor(slot_perm_list, device=device, dtype=torch.int32) + write_count = sum(pnat_means_write) + assert write_count == n_writes_expected, ( + f"test setup error: expected {n_writes_expected} writes, got {write_count}" + ) + + # Device-tensor n_writes (the new path). Caller mutates between iters + # in mix-mode benchmarking; we only run one iter here so a single fill. + n_writes_dev = torch.tensor([n_writes_expected], device=device, dtype=torch.int32) + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=dtype + ) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + batch, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + test_state, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + mode="persistent_main", + slot_perm=slot_perm, + # NEW PATHS: + _n_writes_dev=n_writes_dev, # device tensor (not host int) + _persistent_skip_empty_halves=False, # both halves always launch + ) + + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (scenario={scenario})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (scenario={scenario})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", + ) + + +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize( + "state_dtype", + [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], + ids=["fp16", "int8", "int16", "fp8"], +) +@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) +@pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) +def test_checkpointing_state_update_philox(state_dtype, nheads, head_dim, d_state, ngroups, paged_cache, T): + """ + Verify that Philox stochastic rounding produces correct results across + all SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). + + Runs our kernel twice with identical inputs — once without rand_seed + (deterministic RN), once with rand_seed (Philox SR) — and confirms: + - Outputs are within bf16-dot tolerance (state perturbation ≤ 1 ULP). + - State dtype is preserved. + - State difference is bounded by ~1 ULP of the chosen grid. + """ + _maybe_skip_dtype(state_dtype, use_sr=True) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + + batch = 2 + device = "cuda" + dtype = torch.bfloat16 + assert nheads % ngroups == 0 + + if paged_cache: + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + else: + cache_size = batch + state_batch_indices = None + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + if is_quantized: + state0_fp32 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + state0, state0_scales = _quantize_state(state0_fp32, state_dtype, quant_max) + else: + state0 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state0_scales = None + + # Cache tensors + old_x = torch.randn(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) + + # New token inputs + x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) + B = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + prev_tokens = torch.full((cache_size,), T // 2, device=device, dtype=torch.int32) + + common_kwargs = dict( + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + ) + + # --- Run without rounding (deterministic RN store) --- + state_no_round = state0.clone() + scales_no_round = state0_scales.clone() if is_quantized else None + out_no_round = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_no_round, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + out=out_no_round, + state_scales=scales_no_round, + **common_kwargs, + ) + + # --- Run with Philox rounding --- + rand_seed = torch.tensor([12345], device=device, dtype=torch.int64) + state_rounded = state0.clone() + scales_rounded = state0_scales.clone() if is_quantized else None + out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_rounded, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + out=out_rounded, + rand_seed=rand_seed, + philox_rounds=10, + state_scales=scales_rounded, + **common_kwargs, + ) + + # Outputs should be nearly identical — rounding only perturbs the + # post-replay state by ±1 ULP before the output phase reads it. + # Out_atol = bf16_baseline + 6.5 * per_elem_ULP_after_dequant: + # non-quant fp16: fp16 ULP at typical magnitude is tiny → 1.0 + # int8: amax/127 ≈ 23/127 → 6.5*0.18 ≈ 1.2 + bf16_baseline + # int16: amax/32767 ≈ 7e-4 → ~bf16_baseline only + # fp8: amax/14 ≈ 23/14 → 6.5*1.6 ≈ 10.7 + bf16_baseline + out_atol = ( + {torch.int8: 1.5, torch.int16: 1.0, torch.float8_e4m3fn: 6.0}[state_dtype] + if is_quantized else 1.0 + ) + out_rtol = ( + {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] + if is_quantized else 2e-2 + ) + torch.testing.assert_close( + out_rounded, out_no_round, rtol=out_rtol, atol=out_atol, + msg=f"Output diverged with Philox rounding ({state_dtype})", + ) + + # State dtype preserved. + assert state_rounded.dtype == state_dtype + + # State diff between RN and SR is bounded by 1 quant cell per element. + # Per-channel decode_scale varies by 10x+ across channels (amax depends + # on randn extremes), so a single flat atol can't bound it accurately — + # use per-channel ULP-aware comparison. + slots = state_batch_indices if paged_cache else slice(None) + if is_quantized: + rounded_fp32 = _dequantize_state(state_rounded[slots], scales_rounded[slots]) + no_round_fp32 = _dequantize_state(state_no_round[slots], scales_no_round[slots]) + diff = (rounded_fp32 - no_round_fp32).abs() + # Per-element bound = max(decode_scale_no_round, decode_scale_rounded). + # decode_scale is shape (cache, nheads, dim); broadcast over dstate. + scale_bound = torch.maximum( + scales_no_round[slots], scales_rounded[slots] + ).unsqueeze(-1) + # int8 / int16: 1 cell after dequant = decode_scale exactly. + # fp8_e4m3: variable grid; the largest cell within a channel scaled + # to fit ±448 is at the channel's max-magnitude element, where the + # cell is ~32x larger than the average. Bound = decode_scale * 32. + # Apply a 1.5x slack pad for floating-point compare quirks at the + # exact-cell boundary. + cell_pad = ( + 32.0 if state_dtype == torch.float8_e4m3fn else 1.0 + ) + bound = scale_bound * (cell_pad * 1.5) + if not (diff <= bound).all(): + offenders = (diff > bound).sum().item() + n_total = diff.numel() + pytest.fail( + f"State RN-SR diff exceeds 1 cell per element for " + f"{offenders}/{n_total} elements ({state_dtype}). " + f"max_diff={diff.max().item():.4g}, " + f"max_bound={bound.max().item():.4g}." + ) + else: + # fp16 ULP depends on magnitude — rtol absorbs that. + torch.testing.assert_close( + state_rounded[slots], + state_no_round[slots], + rtol=2e-3, + atol=0.2, + msg=f"State diverged with Philox rounding ({state_dtype})", + ) + + +@pytest.mark.parametrize( + "state_dtype", + [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], + ids=["fp16", "int8", "int16", "fp8"], +) +def test_philox_rounding_unbiased(state_dtype): + """ + Verify that Philox stochastic rounding is unbiased across all + SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). + + Captures the true fp32 post-replay state by running with fp32 storage, + then runs the kernel with the target dtype + Philox SR. Compares the + SR rounding residual against the deterministic-RN residual: SR should + have mean residual closer to zero than RN, since RN has a systematic + round-to-nearest-even bias and SR is unbiased by construction. + + Uses a large batch (16) for ~2M state elements — plenty of statistics. + """ + _maybe_skip_dtype(state_dtype, use_sr=True) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + batch, T = 16, 6 + device = "cuda" + dtype = torch.bfloat16 + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + # fp32 reference state — replay produces values that don't fit cleanly + # in the target dtype's grid, exposing the rounding bias. + state0_fp32 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + + old_x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) + + x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt_val = repeat(dt_base, "b t h -> b t h p", p=head_dim) + B = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + prev_tokens = torch.full((batch,), T, device=device, dtype=torch.int32) + + common_kwargs = dict( + x=x, dt=dt_val, A=A, B=B, C=C, D=D, dt_bias=dt_bias, dt_softplus=True, + ) + + # 1. fp32 state — captures true post-replay fp32 state. + state_fp32 = state0_fp32.clone() + out_fp32 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_fp32, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), prev_tokens, out=out_fp32, **common_kwargs, + ) + + # 2. Target dtype + Philox SR. For quant we also need scales (derived + # from the same per-channel amax used by the kernel on store). + rand_seed = torch.tensor([99999], device=device, dtype=torch.int64) + if is_quantized: + state_rounded, scales_rounded = _quantize_state(state0_fp32, state_dtype, quant_max) + else: + state_rounded = state0_fp32.to(state_dtype) + scales_rounded = None + out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_rounded, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), prev_tokens, out=out_rounded, + rand_seed=rand_seed, philox_rounds=10, + state_scales=scales_rounded, + **common_kwargs, + ) + + # Compute residuals. For non-quant: stochastic_residual = SR(fp32) - + # fp32, deterministic_residual = RN(fp32) - fp32. For quant: dequant + # both, comparing in fp32. + if is_quantized: + fp32_vals = state_fp32.flatten() + stochastic_residual = ( + _dequantize_state(state_rounded, scales_rounded).flatten() - fp32_vals + ) + # Deterministic reference: do the same per-channel quant on the + # captured fp32 state, then dequant. This is what the kernel would + # have produced with rand_seed=None. + det_quant, det_scales = _quantize_state(state_fp32, state_dtype, quant_max) + deterministic_residual = ( + _dequantize_state(det_quant, det_scales).flatten() - fp32_vals + ) + else: + fp32_vals = state_fp32.flatten() + stochastic_residual = state_rounded.float().flatten() - fp32_vals + deterministic_residual = fp32_vals.to(state_dtype).float() - fp32_vals + + # Only consider elements where rounding matters (non-zero residual possible). + nonzero_mask = deterministic_residual.abs() > 0 + num_nonzero = nonzero_mask.sum().item() + assert num_nonzero > 1000, f"Too few roundable elements: {num_nonzero}" + + stochastic_mean = stochastic_residual[nonzero_mask].mean().item() + stochastic_std = stochastic_residual[nonzero_mask].std().item() + deterministic_mean = deterministic_residual[nonzero_mask].mean().item() + + # SE-based bias check. An unbiased estimator's sample mean has standard + # error SE = std / sqrt(n). We require |sr_mean| < K*SE (K=4 ≈ ~3.2e-5 + # one-sided false-positive rate). This auto-calibrates per dtype: + # * int16: residual std ~1e-4 → SE ~9e-8 (very tight bound) + # * int8: residual std ~3e-2 → SE ~2e-5 + # * fp8: residual std ~1e-1 → SE ~9e-5 (loosest, magnitude-driven) + # The previous fixed-1e-5 threshold was below SE for int8/fp8 and would + # always fail by chance. Note the |sr|<|det| fallback was also dropped: + # on Gaussian (symmetric) inputs RN's bias is ~0 by symmetry, so SR vs RN + # is just two unbiased estimators racing — unreliable as a unbias test. + se_sr = stochastic_std / (num_nonzero ** 0.5) + K = 4 + assert abs(stochastic_mean) < K * se_sr, ( + f"SR mean exceeds {K}*SE (likely biased) ({state_dtype}): " + f"stochastic_mean={stochastic_mean:.3e}, " + f"SE={se_sr:.3e} (K*SE={K * se_sr:.3e}), " + f"deterministic_mean={deterministic_mean:.3e} (for reference), " + f"n_elements={num_nonzero}" + ) + + +# HEADS_PER_BLOCK > 1 test. The default heuristic only picks HPB > 1 at large +# total_heads (>= 256-512), which the main test with batch=2 never reaches. +# This test overrides _heads_per_block to exercise the two-loop structure in +# the precompute kernel (store-then-reload of per-head dt/dA_cumsum). +# Configs: (nheads=16, ngroups=1) and (nheads=32, ngroups=2) both have +# heads_per_group=16. The heuristic caps HPB at min(2|4, hpg), so HPB=2, 4. +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) +@pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) +def test_checkpointing_heads_per_block( + nheads, + head_dim, + d_state, + ngroups, + state_dtype, + T, + heads_per_block, +): + # PDL flags use wrapper defaults; trimming the parametrize keeps this + # suite fast. Coverage of {launch_with_pdl, use_internal_pdl} variations + # lives in the dedicated correctness tests above (test_checkpointing_state_update). + batch = 8 + """ + Verify checkpointing_state_update produces correct results when + _heads_per_block > 1, exercising the precompute kernel's two-loop + structure (store per-head dt/dA_cumsum in loop 1, reload in loop 2). + """ + device = "cuda" + dtype = torch.bfloat16 + + if nheads % heads_per_block != 0: + pytest.skip(f"nheads ({nheads}) not divisible by heads_per_block ({heads_per_block})") + if heads_per_block > nheads // ngroups: + pytest.skip( + f"heads_per_block ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})" + ) + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + cache_size = batch + state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + + x1 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + cache_size, T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + state0.clone(), + x1, + dt1, + A, + B1, + C1, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) + + old_x[:] = x1 + dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) + + for slot in range(cache_size): + buf = cache_buf_idx[slot].item() + old_B[slot, buf] = B1[slot] + old_dt[slot, buf] = dt1[slot].T + old_dA_cumsum[slot, buf] = dA_cumsum1[slot].T + + k = T + torch.manual_seed(123) + + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = state0.float().clone() + ref_state_f32[:] = states_buffer_f32[:, k - 1] + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + out=ref_out, + ) + + test_state = state0.clone() + prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + + checkpointing_state_update( + test_state, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + x=x2, + dt=dt2, + A=A, + B=B2, + C=C2, + out=test_out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + _heads_per_block=heads_per_block, + ) + + torch.testing.assert_close( + test_out, + ref_out, + rtol=2e-2, + atol=1.0, + msg=f"Output mismatch with HPB={heads_per_block}, T={T}, " + f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", + ) + + expected_state = states_buffer_f32[:, k - 1].to(state_dtype) + torch.testing.assert_close( + test_state, + expected_state, + rtol=2e-2, + atol=1.0, + msg=f"State mismatch with HPB={heads_per_block}, T={T}, " + f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", + ) + + +# HPB > 1 multi-step test. Production chains decode steps; bugs in +# buffer ordering or stale cache values accumulate across steps and can +# be invisible in a single-step test. +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("T", [6, 16], ids=["T6", "T16"]) +@pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) +@pytest.mark.parametrize("paged_cache", [False, True], ids=["contig", "paged"]) +def test_checkpointing_heads_per_block_multistep( + nheads, head_dim, d_state, ngroups, state_dtype, T, heads_per_block, paged_cache +): + """ + Chain N decode steps with HPB > 1 and verify each step's output matches + a fresh reference. A bug that mixes up WRITE/READ buffers, writes wrong + data to cache, or races in the two-loop structure would accumulate + across steps. + """ + batch = 2 + device = "cuda" + dtype = torch.bfloat16 + n_steps = 8 + + if nheads % heads_per_block != 0: + pytest.skip(f"nheads ({nheads}) not divisible by HPB ({heads_per_block})") + if heads_per_block > nheads // ngroups: + pytest.skip(f"HPB ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})") + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + if paged_cache: + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + slots = state_batch_indices + else: + cache_size = batch + state_batch_indices = None + slots = slice(None) + + all_x = [] + all_dt = [] + all_B = [] + all_C = [] + for step in range(n_steps): + torch.manual_seed(1000 + step) + all_x.append(torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype)) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + all_dt.append(repeat(dt_base, "b t h -> b t h p", p=head_dim)) + all_B.append(torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype)) + all_C.append(torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype)) + + torch.manual_seed(999) + state_init = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + + ref_state = state_init.float().clone() + ref_outs = [] + ref_slots = ( + state_batch_indices + if paged_cache + else torch.arange(batch, device=device, dtype=torch.int32) + ) + for step in range(n_steps): + out_step = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state, + all_x[step], + all_dt[step], + A, + all_B[step], + all_C[step], + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=ref_slots, + out=out_step, + ) + ref_outs.append(out_step) + + test_state = state_init.clone() + old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.zeros(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) + + for step in range(n_steps): + k = T if step > 0 else 0 + prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + + checkpointing_state_update( + test_state, + old_x, + old_B, + old_dt, + old_dA_cumsum, + cache_buf_idx, + prev_tokens, + x=all_x[step], + dt=all_dt[step], + A=A, + B=all_B[step], + C=all_C[step], + out=test_out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + _heads_per_block=heads_per_block, + ) + + if paged_cache: + cache_buf_idx[slots] = 1 - cache_buf_idx[slots] + else: + cache_buf_idx[:] = 1 - cache_buf_idx + + torch.testing.assert_close( + test_out, + ref_outs[step], + rtol=2e-2, + atol=2.0, + msg=f"Output mismatch at step {step} with HPB={heads_per_block}, " + f"T={T}, nheads={nheads}, ngroups={ngroups}, " + f"state_dtype={state_dtype}, paged_cache={paged_cache}", + ) + + +# ----- SR grid-bracket tests (fp8 and fp16) ----- +# +# Verify that each PTX SR output lands on the destination dtype's grid as +# a bracket neighbour of the fp32 input. Catches byte-order traps in the +# inline-asm source-register specifier: +# * fp8: cvt.rs.satfinite.e4m3x4.f32 with pack=4, asm "{$4,$3,$2,$1}" +# * fp16: cvt.rs.f16x2.f32 with pack=2, asm "$0, $2, $1, $3" +# The unbiased test (test_philox_rounding_unbiased) wouldn't catch a +# shuffle: outputs that are still on-grid but swapped within a pack still +# average correctly. Only the per-element bracket check exposes it. +# +# Both kernels are inline copies of the production helpers — kept here so +# the test exercises the exact PTX form independent of wrapper changes. + + +@triton.jit +def _packed_int8_sr_kernel(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + (offs // 4)) + y = _stochastic_round_int8_packed(x, rand, offs) + tl.store(out_ptr + offs, y.to(tl.int8)) + + +@triton.jit +def _packed_int16_sr_kernel(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + (offs // 2)) + y = _stochastic_round_int16_packed(x, rand, offs) + tl.store(out_ptr + offs, y.to(tl.int16)) + + +def _bitrev_int(x: int, bits: int) -> int: + out = 0 + for _ in range(bits): + out = (out << 1) | (x & 1) + x >>= 1 + return out + + +def _rand_words(rand: torch.Tensor) -> list[int]: + return [int(v) & 0xFFFFFFFF for v in rand.cpu().tolist()] + + +def test_packed_int_sr_matches_reference(): + device = "cuda" + n = 1024 + offs = torch.arange(n, device=device, dtype=torch.float32) + x = ((offs % 37) - 18.0) + (((offs * 13.0) % 97.0) + 0.3) / 128.0 + + torch.manual_seed(42) + rand_i8 = torch.randint(-(2**31), 2**31, (n // 4,), device=device, dtype=torch.int32) + out_i8 = torch.empty(n, device=device, dtype=torch.int8) + _packed_int8_sr_kernel[(1,)](x, rand_i8, out_i8, BLOCK=n) + + x_cpu = x.cpu().tolist() + rand_i8_words = _rand_words(rand_i8) + ref_i8 = [] + for i, value in enumerate(x_cpu): + word = rand_i8_words[i // 4] + low = word & 0x0000FFFF + high = (word >> 16) & 0x0000FFFF + pos = i & 3 + if pos == 0: + rand16 = low + elif pos == 1: + rand16 = _bitrev_int(low, 16) + elif pos == 2: + rand16 = high + else: + rand16 = _bitrev_int(high, 16) + ref_i8.append(math.floor(value + rand16 / float(1 << 16))) + + torch.testing.assert_close( + out_i8.cpu().to(torch.int16), + torch.tensor(ref_i8, dtype=torch.int16), + rtol=0, + atol=0, + ) + + rand_i16 = torch.randint(-(2**31), 2**31, (n // 2,), device=device, dtype=torch.int32) + out_i16 = torch.empty(n, device=device, dtype=torch.int16) + _packed_int16_sr_kernel[(1,)](x, rand_i16, out_i16, BLOCK=n) + + rand_i16_words = _rand_words(rand_i16) + ref_i16 = [] + for i, value in enumerate(x_cpu): + word = rand_i16_words[i // 2] + rand_bits = word if (i & 1) == 0 else _bitrev_int(word, 32) + rand24 = rand_bits & 0x00FFFFFF + ref_i16.append(math.floor(value + rand24 / float(1 << 24))) + + torch.testing.assert_close( + out_i16.cpu(), + torch.tensor(ref_i16, dtype=torch.int16), + rtol=0, + atol=0, + ) + + +@triton.jit +def _bracket_kernel_fp8(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + offs) + y = tl.inline_asm_elementwise( + asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", + constraints="=r,r,r,r,r,r,r,r,r", + args=(x, rand), + dtype=tl.float8e4nv, + is_pure=True, + pack=4, + ) + tl.store(out_ptr + offs, y) + + +@triton.jit +def _bracket_kernel_fp16(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + offs) + y = tl.inline_asm_elementwise( + asm="""{ + cvt.rs.f16x2.f32 $0, $2, $1, $3; + }""", + constraints=("=r,r,r,r,r"), + args=(x, rand), + dtype=tl.float16, + is_pure=True, + pack=2, + ) + tl.store(out_ptr + offs, y) + + +_BRACKET_KERNEL = { + torch.float8_e4m3fn: _bracket_kernel_fp8, + torch.float16: _bracket_kernel_fp16, +} + + +def _build_finite_grid(dtype: torch.dtype, device: str) -> torch.Tensor: + """Reinterpret all bit patterns of ``dtype`` as floats; return sorted + unique finite values (drops ±inf, NaNs).""" + if dtype == torch.float8_e4m3fn: + ints = torch.arange(256, dtype=torch.uint8, device=device) + full = ints.view(torch.float8_e4m3fn).to(torch.float32) + elif dtype == torch.float16: + # int16 view of all 65536 patterns (covers fp16 normals + subnormals + # + ±inf + NaN; we filter to finite below). + ints = torch.arange(65536, dtype=torch.int32, device=device).to(torch.int16) + full = ints.view(torch.float16).to(torch.float32) + else: + raise ValueError(f"Unsupported bracket-test dtype: {dtype}") + return full[torch.isfinite(full)].sort()[0].unique() + + +def _build_bracket_inputs(dtype: torch.dtype, n: int, device: str) -> torch.Tensor: + """Test inputs spanning the dtype's grid range. Includes on-grid points + so we exercise the no-rounding case; for fp8 also includes overflow to + test saturation (PTX `cvt.rs.satfinite.e4m3x4.f32` clamps in-op). + + fp16 inputs are kept inside the finite range — `cvt.rs.f16x2.f32` does + NOT have a `satfinite` modifier and produces ±inf for OOR inputs (not + a saturate-to-±max). The kernel only ever sees in-range fp32 state in + practice (state_amax is always ≪ fp16_max), so the test mirrors that. + """ + grid = _build_finite_grid(dtype, device) + g_min, g_max = grid[0].item(), grid[-1].item() + x = torch.empty(n, device=device, dtype=torch.float32) + if dtype == torch.float8_e4m3fn: + # 1.5x range exercises saturation; satfinite handles it in-op. + x.uniform_(g_min * 1.5, g_max * 1.5) + else: # fp16: four magnitude bands, all within finite range. + x[: n // 4].uniform_(-1.0, 1.0) + x[n // 4 : n // 2].uniform_(-100, 100) + x[n // 2 : 3 * n // 4].uniform_(-1000, 1000) + x[3 * n // 4 :].uniform_(g_min * 0.99, g_max * 0.99) + return x, grid + + +@_skip_pre_sm100 +@pytest.mark.parametrize( + "state_dtype", + [torch.float8_e4m3fn, torch.float16], + ids=["fp8", "fp16"], +) +def test_sr_grid_bracket(state_dtype): + """Verify SR PTX outputs each lie on the destination grid as a bracket + neighbour of the fp32 input.""" + device = "cuda" + n = 1024 # multiple of both pack=4 (fp8) and pack=2 (fp16) + + torch.manual_seed(42) + x, grid_finite = _build_bracket_inputs(state_dtype, n, device) + g_min, g_max = grid_finite[0].item(), grid_finite[-1].item() + + # Bracket [lo, hi] in the destination grid for each input. For + # out-of-range inputs the bracket is the saturating endpoint pair. + x_clamped = x.clamp(g_min, g_max) + idx = torch.searchsorted(grid_finite, x_clamped, right=False).clamp( + min=1, max=len(grid_finite) - 1 + ) + lo = grid_finite[idx - 1] + hi = grid_finite[idx] + # For x exactly on grid, idx points at it; lo = grid[i-1], hi = x — the + # bracket allows out==hi (=x) which is what RN-on-grid produces. + + kernel = _BRACKET_KERNEL[state_dtype] + + for seed in range(4): + torch.manual_seed(seed) + # int32 for raw random bits — PTX takes the bit pattern, sign + # interpretation doesn't matter. + rand = torch.randint(-(2**31), 2**31, (n,), device=device, dtype=torch.int32) + out = torch.empty(n, device=device, dtype=state_dtype) + kernel[(1,)](x, rand, out, BLOCK=n) + out_fp32 = out.to(torch.float32) + + on_grid = (out_fp32 == lo) | (out_fp32 == hi) + if not on_grid.all(): + offenders = ~on_grid + n_off = offenders.sum().item() + sample = ( + x[offenders][:5].tolist(), + lo[offenders][:5].tolist(), + hi[offenders][:5].tolist(), + out_fp32[offenders][:5].tolist(), + ) + pytest.fail( + f"{state_dtype} SR output not on grid bracket for {n_off}/{n} " + f"elements (seed={seed}). x={sample[0]} lo={sample[1]} " + f"hi={sample[2]} out={sample[3]}. Likely the PTX byte-order " + "bug (cvt.rs source-register order)." + ) + From 42b138248a07b8fae788f85a99d7d3469e8413f2 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Sun, 17 May 2026 23:21:37 -0700 Subject: [PATCH 52/89] mamba_checkpointing: promote pd WC_IS_CONSTEXPR refactor to live kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a WC_IS_CONSTEXPR constexpr param to _persistent_main_impl that, when True, forces is_write = WRITE_CHECKPOINT (constexpr-fold) regardless of IS_DYNAMIC. _persistent_main_kernel uses this only in the RECT=1 is_w=True arm: passes WRITE_CHECKPOINT=True literal + WC_IS_CONSTEXPR=True, so the inner DCEs the nowrite codepath under IS_DYNAMIC=True too — same codegen quality as persistent_main mode. RECT=0 path passes WC_IS_CONSTEXPR=False: inner keeps the original runtime is_write branch under IS_DYNAMIC=True, no binary-doubling regression. Measured at b=1024 fp16 SR persistent_dynamic dyn-shape (M=32 W=2 pW=1 H=8 TMA=1001 RECT=1) sweep over S × CPS × LS: uniform -3.7% improvement across all 160 cells, zero regressions. Neutral at pd RECT=0 shape (b=1024 and b=1), neutral at pers_main RECT=0 and RECT=1 (b=512). 9/9 dedicated test_..._persistent_main and test_..._persistent_dynamic pytests pass. Drops the _refactored.py staging files (kernel/bench/test). Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update.py | 95 +- .../checkpointing_state_update_refactored.py | 5510 ----------------- ...eplay_selective_state_update_refactored.py | 4935 --------------- ...t_checkpointing_state_update_refactored.py | 2230 ------- 4 files changed, 66 insertions(+), 12704 deletions(-) delete mode 100644 tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_refactored.py delete mode 100644 tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_refactored.py delete mode 100644 tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_refactored.py diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py index b6ebe6afbda9..96116649b7ff 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update.py @@ -3308,21 +3308,39 @@ def _persistent_main_impl( PHILOX_ROUNDS: tl.constexpr, QUANT_MAX: tl.constexpr, WRITE_CHECKPOINT: tl.constexpr, + # IS_DYNAMIC: kept in the signature for caller-side bookkeeping (the + # outer _persistent_main_kernel still inspects it to decide the slot- + # IS_DYNAMIC: when True (persistent_dynamic), is_write is per-slot from + # PNAT. When False (persistent_main), is_write is constexpr from + # WRITE_CHECKPOINT. See also WC_IS_CONSTEXPR below. IS_DYNAMIC: tl.constexpr, - # TMA flags — picked inside body based on is_write (which is constexpr - # from WRITE_CHECKPOINT when IS_DYNAMIC=False, or runtime from PNAT - # when IS_DYNAMIC=True). use_tma_load = USE_TMA_LOAD_WRITE if is_write - # else USE_TMA_LOAD_NOWRITE — constexpr-folds in non-dynamic mode, - # runtime ternary in dynamic. + # WC_IS_CONSTEXPR: when True, force is_write = WRITE_CHECKPOINT (constexpr) + # regardless of IS_DYNAMIC. Callers in RECT=1 use this in the is_w=True + # arm of _persistent_main_kernel (we know all slots that reach this call + # need is_write=True because is_w was the PNAT-derived runtime check, and + # this arm only fires when is_w is True). Passing WRITE_CHECKPOINT=True + # as a literal at the call site + WC_IS_CONSTEXPR=True here lets the inner + # body DCE the nowrite path under IS_DYNAMIC=True too — same codegen + # quality as persistent_main mode (-3.7% measured at b=1024 dyn-shape). + # When False (RECT=0 callers, where both write and nowrite slots are + # dispatched to ONE call), use the original runtime is_write under + # IS_DYNAMIC=True; avoids the binary-doubling regression that two + # specialized calls would cause. + WC_IS_CONSTEXPR: tl.constexpr = False, + # TMA flags — picked inside body based on is_write. When is_write is + # constexpr (either IS_DYNAMIC=False or WC_IS_CONSTEXPR=True), the + # use_tma_load = USE_TMA_LOAD_WRITE if is_write else USE_TMA_LOAD_NOWRITE + # ternary constexpr-folds and only one TMA load form survives. USE_TMA_LOAD_WRITE: tl.constexpr = False, USE_TMA_LOAD_NOWRITE: tl.constexpr = False, USE_TMA_STORE: tl.constexpr = False, ): - # IS_DYNAMIC: when False, WRITE_CHECKPOINT is the constexpr write/nowrite - # selector (caller pre-sorts and splits halves). When True, WRITE_CHECKPOINT - # is ignored; the impl computes is_write at runtime per work-item from - # the loaded PNAT. Used by mode="persistent_dynamic" — single kernel, - # one launch, no half-split, runtime per-slot dispatch. + # IS_DYNAMIC: kernel-mode label, used by the OUTER _persistent_main_kernel + # to decide slot-range derivation and outer is_w dispatch strategy + # (constexpr WC for persistent_main; runtime is_w split -> 2 specialized + # impl calls for persistent_dynamic). Inside this impl, IS_DYNAMIC is + # NOT consulted at runtime -- WRITE_CHECKPOINT is the only constexpr that + # gates the write/nowrite codegen, in BOTH modes. # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized # state dtype (int8 / int16 / float8e4nv) and only those. @@ -3345,11 +3363,19 @@ def _persistent_main_impl( active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - # Resolve is_write: constexpr from caller (persistent_main) or runtime - # from PNAT (persistent_dynamic). When IS_DYNAMIC=False, is_write - # collapses to a constexpr 0/1 and the downstream `if is_write:` - # blocks DCE the dead path at compile time. - if IS_DYNAMIC: + # Resolve is_write: see WC_IS_CONSTEXPR / IS_DYNAMIC docs in the param + # list above. Three cases: + # - WC_IS_CONSTEXPR=True (RECT=1 is_w=True arm callers): use WC + # constexpr. Caller knows the slot needs write; inner DCEs nowrite + # paths. Avoids the binary-doubling overhead that calling the impl + # twice would cause, while still constexpr-DCEing the nowrite half. + # - IS_DYNAMIC=True (RECT=0 caller, persistent_dynamic): runtime + # branch on PNAT. Both write and nowrite codegen live in one body + # (no bloat) — same as the pre-refactor behavior. + # - IS_DYNAMIC=False (persistent_main): WC constexpr from caller. + if WC_IS_CONSTEXPR: + is_write: tl.constexpr = WRITE_CHECKPOINT + elif IS_DYNAMIC: is_write = (prev_num_accepted_tokens + T) > MAX_REPLAY_BUFFER_LENGTH else: is_write = WRITE_CHECKPOINT @@ -3385,18 +3411,12 @@ def _persistent_main_impl( state_ptrs = ( state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate ) - # Load state. Branch on is_write (constexpr in non-dynamic mode, runtime - # in dynamic), then constexpr-pick TMA-vs-tl.load per side. - # - # Non-dynamic (is_write is constexpr = WRITE_CHECKPOINT): outer `if` - # DCE's, only the matching side's constexpr-gated load survives. - # - # Dynamic (is_write is runtime per-CTA): both write and nowrite blocks - # emit; each contains exactly one of (TMA load, tl.load) after the - # constexpr USE_TMA_LOAD_* gate resolves. The descriptor is real iff - # any of the 4 TMA flags is on at the wrapper (line 4712-4713 of this - # file); the constexpr gating guarantees we never call .load() on the - # plain-tensor fallback path, so this stays compilation-safe. + # Load state. Branch on is_write (constexpr = WRITE_CHECKPOINT in BOTH + # modes after the outer-dispatch refactor), then constexpr-pick TMA-vs- + # tl.load per side. Outer `if` DCE's, only the matching side's + # constexpr-gated load survives -- same compile-time picking for both + # persistent_main and persistent_dynamic (the latter dispatches at the + # outer kernel level so each impl instance sees a constexpr WC). if is_write: if USE_TMA_LOAD_WRITE: state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) @@ -4158,6 +4178,18 @@ def _persistent_main_kernel( else: is_w = WRITE_CHECKPOINT if is_w: + # Pass WRITE_CHECKPOINT=True constexpr to specialize this + # impl call for the write path. Under IS_DYNAMIC=True, the + # kernel-level WRITE_CHECKPOINT is False (launcher default), + # but the OUTER is_w branch we are inside narrows the + # runtime path to writes-only, so we override to True here + # so the impl's constexpr-gated `if is_write:` blocks DCE + # to the write-only codegen. Under IS_DYNAMIC=False + # (persistent_main), the kernel-level WRITE_CHECKPOINT is + # itself True for this half (write half launches with + # WC=True), and the outer is_w = WRITE_CHECKPOINT = True + # constexpr-folds; passing literal True here is consistent + # and constexpr-equivalent. _persistent_main_impl( pid_m, pid_b, pid_h, state_ptr, state_tma_descriptor, state_scales_ptr, @@ -4185,10 +4217,11 @@ def _persistent_main_kernel( BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - WRITE_CHECKPOINT, IS_DYNAMIC, + True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) + True, # WC_IS_CONSTEXPR — force inner to use WC constexpr # 3 TMA flags: write-load fires here (we're in the # is_write branch), nowrite-load is dead (no slot - # reaches it), store fires when WC=True at runtime. + # reaches it), store fires (write path). USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, ) else: @@ -4228,6 +4261,9 @@ def _persistent_main_kernel( # impl picks USE_TMA_LOAD_WRITE vs USE_TMA_LOAD_NOWRITE based on # its computed is_write — constexpr-folds when is_write is # constexpr (non-dyn), runtime branch when IS_DYNAMIC=True. + # (Reverted from outer two-call dispatch: that doubled the + # compiled body size under IS_DYNAMIC=True and regressed RECT=0 + # perf by ~+24%.) _persistent_main_impl( pid_m, pid_b, pid_h, state_ptr, state_tma_descriptor, state_scales_ptr, @@ -4256,6 +4292,7 @@ def _persistent_main_kernel( BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, WRITE_CHECKPOINT, IS_DYNAMIC, + False, # WC_IS_CONSTEXPR=False — RECT=0 has both write/nowrite slots in one call USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, ) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_refactored.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_refactored.py deleted file mode 100644 index 96116649b7ff..000000000000 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_refactored.py +++ /dev/null @@ -1,5510 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. -# -# Adapted from: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py -# SPDX-FileCopyrightText: Copyright contributors to the sglang project -# -# Copyright (c) 2024, Tri Dao, Albert Gu. -# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py - -import torch -import triton -import triton.language as tl - -from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID -from tensorrt_llm._utils import get_sm_version - -from .softplus import softplus - - -# Lazy global allocator for Triton TMA tensor descriptors. Required by any -# host- or device-built tensor_descriptor; without it Triton raises at first -# launch. See TMA backlog item #17 / scratch experiment notes. -_TMA_ALLOCATOR_SET = False - - -def _ensure_tma_allocator() -> None: - global _TMA_ALLOCATOR_SET - if _TMA_ALLOCATOR_SET: - return - - def _alloc_fn(size, alignment, stream): - # Triton expects an int8 buffer of `size` bytes; alignment is enforced - # by the allocator returning a buffer satisfying it (PyTorch's - # cudaMalloc-backed tensors are 256B-aligned, so we're fine). - return torch.empty(size, device="cuda", dtype=torch.int8) - - triton.set_allocator(_alloc_fn) - _TMA_ALLOCATOR_SET = True - - -@triton.jit -def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: - """Stochastic rounding: fp32 pair → fp16x2 using Philox random bits. - - Uses PTX cvt.rs.f16x2.f32 which rounds each fp32 value to fp16 using - the random bits to break ties, avoiding systematic rounding bias that - accumulates over many decode steps with fp16 state. - - Adapted from flashinfer (Apache-2.0, vLLM/mamba lineage). - """ - return tl.inline_asm_elementwise( - asm="""{ - cvt.rs.f16x2.f32 $0, $2, $1, $3; - }""", - constraints=("=r,r,r,r,r"), - args=(x, rand), - dtype=tl.float16, - is_pure=True, - pack=2, - ) - - -@triton.jit -def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: - """Stochastic rounding: fp32 quad → fp8 e4m3 using Philox random bits. - - Uses PTX cvt.rs.satfinite.e4m3x4.f32 which combines stochastic rounding - and saturating cast in a single op (output is final fp8, no separate - clamp needed). The reversed source-register order {$4,$3,$2,$1} is - load-bearing — PTX packs leftmost source into the high byte but Triton's - pack=4 is little-endian, so the natural {$1,$2,$3,$4} order would - silently shuffle every group of 4 contiguous outputs. - - Requires SM_100a+ (Blackwell B200). Caller must gate at the wrapper - level — this kernel does not check. - - Adapted from vLLM PR #40012 (Apache-2.0). - """ - return tl.inline_asm_elementwise( - asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", - constraints="=r,r,r,r,r,r,r,r,r", - args=(x, rand), - dtype=tl.float8e4nv, - is_pure=True, - pack=4, - ) - - -@triton.jit -def _bitrev32(x: tl.tensor) -> tl.tensor: - return tl.inline_asm_elementwise( - asm="brev.b32 $0, $1;", - constraints="=r,r", - args=(x,), - dtype=tl.uint32, - is_pure=True, - pack=1, - ) - - -@triton.jit -def _stochastic_round_int8_packed( - x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor -) -> tl.tensor: - """Stochastic rounding for int8 using one random uint32 per 4 values.""" - low = rand & 0x0000FFFF - high = (rand >> 16) & 0x0000FFFF - low_rev = _bitrev32(low) >> 16 - high_rev = _bitrev32(high) >> 16 - rand_pos = offs_n & 3 - rand16 = tl.where( - rand_pos == 0, - low, - tl.where(rand_pos == 1, low_rev, tl.where(rand_pos == 2, high, high_rev)), - ) - rand01 = rand16.to(tl.float32) * (1.0 / float(1 << 16)) - return tl.extra.cuda.libdevice.floor(x + rand01) - - -@triton.jit -def _stochastic_round_int16_packed( - x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor -) -> tl.tensor: - """Stochastic rounding for int16 using one random uint32 per 2 values.""" - rand_bits = tl.where((offs_n & 1) == 0, rand, _bitrev32(rand)) - rand01 = (rand_bits & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) - return tl.extra.cuda.libdevice.floor(x + rand01) - - -# Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, -# old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. -# Grid: (batch, nheads // HEADS_PER_BLOCK). - - -@triton.jit() -def _replay_precompute_impl( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, - decay_vec_ptr, - # Cache pointers (both buffers reachable via stride_*_dbuf). This - # kernel writes to either the active (= cache_buf_idx) or inactive - # (= 1 - cache_buf_idx) buffer depending on WRITE_CHECKPOINT — see - # comment block at top of kernel body. - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - # Double-buffer index (per cache slot) — selects this step's "active" - # buffer (= where the historical inputs for this step live). - cache_buf_idx_ptr, - # Per-request accepted-tokens count (already-cached old tokens at - # [0, PNAT) of the active buffer; new tokens this step go after them - # on no-checkpoint steps). - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - # Slot permutation: maps grid program_id -> original slot index. - # When USE_PERM=False, pid_b = tl.program_id(0) (today's behavior) and - # this ptr is unused. When USE_PERM=True, pid_b = perm[pid_grid] (or - # perm[B-1-pid_grid] if REVERSE_PERM=True). - slot_perm_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides: (cache, 2, T, ngroups, dstate) - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, - # Slot permutation flags. USE_PERM=True gates a slot_perm_ptr load that - # remaps grid program_id -> original slot index. REVERSE_PERM=True walks - # the perm from the tail (B-1-pid_grid). Used by sorted-dispatch - # variants of dl/dlgrouped/maindl to cluster early-outs at one end of - # the grid; ignored by monolithic / dynamic. - USE_PERM: tl.constexpr, - REVERSE_PERM: tl.constexpr, - # Checkpointing flag — selects target buffer + offset for new-token - # cache writes. See "Cache write semantics" block below. - # Runtime (not constexpr): the only WRITE_CHECKPOINT-dependent code in - # this body is the write_buf/write_offset selection, which is plain - # arithmetic — no constexpr-shaped tile or whole-block gate. Letting - # it be runtime lets the dynamic dispatch kernel call us once with the - # per-slot needs_write flag instead of inlining two specializations. - write_checkpoint, -): - pid_grid = tl.program_id(axis=0) - # REVERSE_PERM walks the grid tail-first regardless of USE_PERM — - # combined with hardcode-sorted prev_tokens (kernel-side USE_PERM=False - # but PNAT pre-sorted write-first), reverse traversal makes the nowrite - # half front-load real work. - pid_grid_eff = (tl.num_programs(axis=0) - 1 - pid_grid) if REVERSE_PERM else pid_grid - if USE_PERM: - pid_b = tl.load(slot_perm_ptr + pid_grid_eff) - else: - pid_b = pid_grid_eff - pid_hg = tl.program_id(axis=1) # head-group index - first_head = pid_hg * HEADS_PER_BLOCK - - # Resolve cache index for writes - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - - # --- Cache write semantics --- - # cache_buf_idx names this step's "active" buffer — the one with the - # historical inputs at [0, PNAT). The other buffer is "staging". - # - # Where do we write new tokens this step? - # WRITE_CHECKPOINT=False (no overflow): append to ACTIVE buffer at - # offset [PNAT : PNAT+T). Caller does NOT flip cache_buf_idx - # afterward; PNAT_next = PNAT + accepted. [0, PNAT) preserved. - # WRITE_CHECKPOINT=True (would overflow): write to STAGING buffer at - # [0, T). Caller flips cache_buf_idx afterward; next step's - # active = the one we just wrote. PNAT_next = accepted. Old - # data in the previous active buffer is folded into state via - # the replay update and discarded. This matches today's replay - # kernel behavior exactly. - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - if write_checkpoint: - write_buf = 1 - buf_active - write_offset = 0 - else: - write_buf = buf_active - write_offset = prev_num_accepted_tokens - - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - t_mask = offs_t < T - n_mask = offs_n < dstate - - # Causal mask is shared across all heads (depends only on offs_t) - causal_mask = offs_t[:, None] >= offs_t[None, :] - valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] - - # --- Vectorized pre-wait phase across HEADS_PER_BLOCK heads --- - # Compute dt, dA_cumsum, decay_vec as (H, T) tiles. Pre-compute - # scale_combo = decay_matrix * dt[:, None, :] as an (H, T, T) tile that - # stays in registers across gdc_wait — eliminates the post-wait reload - # of dt + dA_cumsum and the per-head loop. - offs_h = tl.arange(0, HEADS_PER_BLOCK) - heads_block = first_head + offs_h # (H,) - - # Load dt (H, T) - dt_addrs = ( - dt_ptr + pid_b * stride_dt_batch - + heads_block[:, None] * stride_dt_head - + offs_t[None, :] * stride_dt_T - ) - dt = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) - if HAS_DT_BIAS: - dt_bias = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) - dt = dt + dt_bias[:, None] - if DT_SOFTPLUS: - dt = softplus(dt) - - A = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) # (H,) - dA_cumsum = tl.cumsum(A[:, None] * dt, axis=1) # (H, T) - decay_vec = tl.exp(dA_cumsum) # (H, T) - - # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of write_buf. - old_dt_addrs = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + heads_block[:, None] * stride_old_dt_head - + (write_offset + offs_t)[None, :] * stride_old_dt_T - ) - tl.store(old_dt_addrs, dt, mask=t_mask[None, :]) - - old_dA_cumsum_addrs = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + heads_block[:, None] * stride_old_dA_cumsum_head - + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T - ) - tl.store(old_dA_cumsum_addrs, dA_cumsum, mask=t_mask[None, :]) - - # decay_vec scratch — always at offs_t. - decay_vec_addrs = ( - decay_vec_ptr + pid_b * stride_dv_batch - + heads_block[:, None] * stride_dv_head - + offs_t[None, :] * stride_dv_t - ) - tl.store(decay_vec_addrs, decay_vec, mask=t_mask[None, :]) - - # scale_combo (H, T, T) = exp(dA_cumsum[h, t1] - dA_cumsum[h, t2]) * dt[h, t2] - # Stays live across gdc_wait — used post-wait to compute CB_scaled. - decay_matrix = tl.exp(dA_cumsum[:, :, None] - dA_cumsum[:, None, :]) # (H, T, T) - scale_combo = decay_matrix * dt[:, None, :] # (H, T, T) - - # --- Wait for upstream kernel (external PDL) before loading B and C --- - # All dt processing above is independent of conv1d outputs. - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- - group_idx = first_head // nheads_ngroups_ratio - C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group - B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group - - C_all = tl.load( - C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - B_all = tl.load( - B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - - # Compute raw CB once — shared across all heads in this block - raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) - - # Store B to cache at [write_offset : write_offset+T) of write_buf. - if first_head % nheads_ngroups_ratio == 0: - old_B_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + write_buf * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - tl.store( - old_B_base - + (write_offset + offs_t)[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - B_all, - mask=t_mask[:, None] & n_mask[None, :], - ) - - # --- Vectorized post-wait phase: scale_combo (H, T, T) is still live in - # registers from pre-wait; multiply by raw_CB (T, T), apply causal mask, - # store as one (H, T, T) tile. --- - CB_scaled_block = tl.where( - valid_mask[None, :, :], - raw_CB[None, :, :] * scale_combo, - 0.0, - ) # (H, T, T) - cb_scaled_addrs = ( - cb_scaled_ptr + pid_b * stride_cb_batch - + heads_block[:, None, None] * stride_cb_head - + offs_t[None, :, None] * stride_cb_t - + offs_t[None, None, :] * stride_cb_j - ) # (H, T, T) - cb_store_mask = ( - (offs_t[None, :, None] < BLOCK_SIZE_T) - & (offs_t[None, None, :] < BLOCK_SIZE_T) - ) - tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) - - -# Replay-style precompute kernel. Thin wrapper around _replay_precompute_impl -# that carries the @triton.heuristics for constexpr derivation; called from -# the Python wrapper on the replay-style path (write or replay-nowrite). -@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.jit() -def _checkpointing_precompute_kernel( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, - decay_vec_ptr, - # Cache pointers - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, - WRITE_CHECKPOINT: tl.constexpr, - EARLY_OUT: tl.constexpr, - USE_PERM: tl.constexpr, - REVERSE_PERM: tl.constexpr, -): - # Hoisted PDL signal: fire as the first thing every program does, so - # main can start its setup regardless of how this program ends (pad, - # early-out, or full body). PDL signals are idempotent; main's - # gdc_wait still gates on prerequisite-kernel completion for - # correctness. - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - # Per-program early-out gate (option-2 double-launch). When EARLY_OUT - # is False the entire block is constexpr-folded out and the wrapper is - # just an impl call. When True, this kernel only runs for slots whose - # (PNAT + T > MAX) status matches WRITE_CHECKPOINT. - if EARLY_OUT: - pid_grid_eo = tl.program_id(axis=0) - pid_grid_eo_eff = (tl.num_programs(axis=0) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo - if USE_PERM: - pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) - else: - pid_b_eo = pid_grid_eo_eff - if HAS_CACHE_BATCH_INDICES: - cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) - if cbi_eo == pad_slot_id: - return - else: - cbi_eo = pid_b_eo.to(tl.int64) - pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) - if (pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH) != WRITE_CHECKPOINT: - return - _replay_precompute_impl( - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - cb_scaled_ptr, - decay_vec_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - pad_slot_id, - T, - dstate, - nheads_ngroups_ratio, - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - DT_SOFTPLUS, - HAS_DT_BIAS, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - LAUNCH_WITH_PDL, - HEADS_PER_BLOCK, - USE_PERM, - REVERSE_PERM, - WRITE_CHECKPOINT, - ) - - -# Rectangle precompute kernel: produces a (T, K) CB rectangle that combines -# old-token (B from cache, k ∈ [0, PNAT)) and new-token (B from input, k ∈ -# [MAX-T, MAX) at compile-time-static shift) contributions in a single matmul. -# Used only on no-checkpoint steps (nowrite path); pairs with -# `_rectangle_main_kernel`. K-axis size = max(np2(MAX_REPLAY_BUFFER_LENGTH), -# 16); the static layout is sound because nowrite implies PNAT + T <= -# MAX_REPLAY_BUFFER_LENGTH, so old [0, PNAT) and new [MAX-T, MAX) never -# overlap. Also folds total_decay into decay_vec at precomp time so main -# can skip materializing a state_prev_decayed (M, dstate) tile. - - -@triton.jit() -def _rectangle_precompute_impl( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle - decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) - # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite - # path: read from buf_active at [0, PNAT), write new tokens at - # [PNAT, PNAT+T) of buf_active (same buffer). - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - # Slot permutation: see _replay_precompute_impl for semantics. - slot_perm_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # rectangle K-axis bound - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides (rectangle: (batch, nheads, T, K)) - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides: (cache, 2, T_max, ngroups, dstate) - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides: (cache, 2, nheads, T_max) - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides: (cache, 2, nheads, T_max) - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, - # Slot permutation flags — see _replay_precompute_impl. - USE_PERM: tl.constexpr, - REVERSE_PERM: tl.constexpr, -): - pid_grid = tl.program_id(axis=0) - # REVERSE_PERM walks the grid tail-first regardless of USE_PERM — - # combined with hardcode-sorted prev_tokens (kernel-side USE_PERM=False - # but PNAT pre-sorted write-first), reverse traversal makes the nowrite - # half front-load real work. - pid_grid_eff = (tl.num_programs(axis=0) - 1 - pid_grid) if REVERSE_PERM else pid_grid - if USE_PERM: - pid_b = tl.load(slot_perm_ptr + pid_grid_eff) - else: - pid_b = pid_grid_eff - pid_hg = tl.program_id(axis=1) - first_head = pid_hg * HEADS_PER_BLOCK - - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - write_buf = buf_active - write_offset = prev_num_accepted_tokens - - # Static rectangle K-axis layout: old at [0, PNAT) (mask is_old_k); - # new at [MAX-T, MAX) at compile-time shift K_NEW_SHIFT = MAX - T. - # PNAT + T <= MAX is guaranteed on the nowrite path → no overlap. - K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T - - offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis (output rows) - offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis (rectangle input cols) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - t_mask = offs_t < T - n_mask = offs_n < dstate - - # K-axis masks (approach C: runtime PNAT-offset instead of K_NEW_SHIFT) - # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. - is_old_k = offs_k < prev_num_accepted_tokens - safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - prev_num_accepted_tokens - is_new_k = (k_new_idx >= 0) & (k_new_idx < T) - safe_k_new = tl.where(is_new_k, k_new_idx, 0) - - # Loop 1: per-head dt processing. dt → dt_processed → dA_cumsum → - # decay_vec_new (= exp(cumAdt_new)). Stored to write_buf for next step. - # decay_vec_full (= total_decay * decay_vec_new) is finalized in loop 2 - # once total_decay is loaded; loop 1 stores raw decay_vec_new to scratch. - for h_local in range(HEADS_PER_BLOCK): - head_idx = first_head + h_local - - dt_base = dt_ptr + pid_b * stride_dt_batch + head_idx * stride_dt_head - dt = tl.load(dt_base + offs_t * stride_dt_T, mask=t_mask, other=0.0).to(tl.float32) - if HAS_DT_BIAS: - dt_bias = tl.load(dt_bias_ptr + head_idx * stride_dt_bias_head).to(tl.float32) - dt = dt + dt_bias - if DT_SOFTPLUS: - dt = softplus(dt) - - A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) - dA_cumsum = tl.cumsum(A * dt, axis=0) - - # Store dt and dA_cumsum to write_buf at [write_offset, write_offset+T) - # for next step's replay/rectangle use. - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - tl.store( - old_dt_base + (write_offset + offs_t) * stride_old_dt_T, - dt, - mask=t_mask, - ) - - old_dA_cumsum_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - tl.store( - old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, - dA_cumsum, - mask=t_mask, - ) - - # ---- Hoisted: cache-only loads independent of conv1d ---- - # old_B (group-level, BLOCK_K × BLOCK_DSTATE = ~8KB tile) and the - # decay_vec_full per-head pre-compute (which writes to DRAM and doesn't - # need cross-gdc_wait variables) are issued BEFORE gdc_wait so their - # HBM latency overlaps with conv1d. Per-head factor_dt/exp_diff stay - # below gdc_wait — they need cross-iteration spans, which Triton can't - # express without a DRAM round-trip; the per-head LOADS in the post- - # wait loop are small and cheap, so leave them. - group_idx = first_head // nheads_ngroups_ratio - - # Group-level: old B from active buffer at [0, PNAT) of the K-axis. - old_B_read_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + buf_active * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - old_B_load = tl.load( - old_B_read_base - + safe_old_k[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - mask=is_old_k[:, None] & n_mask[None, :], - other=0.0, - ) - - # Pre-wait: vectorized across HEADS_PER_BLOCK heads. Compute decay_vec_full - # (H, T) and combo = factor_dt * exp_diff (H, T, K). Store decay_vec_full; - # combo_block stays in registers across gdc_wait — used directly post-wait - # to compute rect_CB_scaled without a global memory roundtrip. - prev_k_idx = tl.minimum( - tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 - ) - offs_h = tl.arange(0, HEADS_PER_BLOCK) - heads_block = first_head + offs_h # (H,) - - # Per-head bases (H,) — broadcast with offs_k or offs_t for 2D loads. - old_dt_read_h = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + buf_active * stride_old_dt_dbuf - + heads_block * stride_old_dt_head - ) - old_dA_cumsum_read_h = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_active * stride_old_dA_cumsum_dbuf - + heads_block * stride_old_dA_cumsum_head - ) - old_dt_write_h = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + heads_block * stride_old_dt_head - ) - old_dA_cumsum_write_h = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + heads_block * stride_old_dA_cumsum_head - ) - - # (H, K) loads at [0, PNAT) — old data from previous step. - hk_mask = is_old_k[None, :] # (1, K) - old_dt_all = tl.load( - old_dt_read_h[:, None] + safe_old_k[None, :] * stride_old_dt_T, - mask=hk_mask, other=0.0, - ).to(tl.float32) - old_dA_cumsum_all = tl.load( - old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, - mask=hk_mask, other=0.0, - ).to(tl.float32) - # (H,) scalar-per-head: total_dA_cumsum at prev_k_idx. - total_dA_cumsum = tl.load( - old_dA_cumsum_read_h + prev_k_idx * stride_old_dA_cumsum_T - ).to(tl.float32) - # (H, T) loads at [PNAT, PNAT+T) — this step's dA_cumsum_new from loop 1. - ht_mask = t_mask[None, :] # (1, T) - dA_cumsum_new = tl.load( - old_dA_cumsum_write_h[:, None] - + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T, - mask=ht_mask, other=0.0, - ).to(tl.float32) - # (H, K) loads at K_NEW_SHIFT-shifted positions for new tokens. - hkn_mask = is_new_k[None, :] - dt_at_kn = tl.load( - old_dt_write_h[:, None] - + (write_offset + safe_k_new)[None, :] * stride_old_dt_T, - mask=hkn_mask, other=0.0, - ).to(tl.float32) - dA_cumsum_at_kn = tl.load( - old_dA_cumsum_write_h[:, None] - + (write_offset + safe_k_new)[None, :] * stride_old_dA_cumsum_T, - mask=hkn_mask, other=0.0, - ).to(tl.float32) - - # decay_vec_full = total_decay * exp(cumAdt_new). (H, T). - total_decay = tl.where( - prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0 - ) # (H,) - decay_vec_full_block = total_decay[:, None] * tl.exp(dA_cumsum_new) # (H, T) - decay_vec_addrs = ( - decay_vec_ptr - + pid_b * stride_dv_batch - + heads_block[:, None] * stride_dv_head - + offs_t[None, :] * stride_dv_t - ) # (H, T) - tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) - - # combo_block = factor_dt * exp_diff — (H, T, K). Stays in registers - # across gdc_wait. - factor_dt = tl.where(is_old_k[None, :], old_dt_all, dt_at_kn) # (H, K) - s_k = tl.where( - is_old_k[None, :], - total_dA_cumsum[:, None] - old_dA_cumsum_all, - -dA_cumsum_at_kn, - ) # (H, K) - # exp_diff (H, T, K) = exp(s_k (H, 1, K) + dA_cumsum_new (H, T, 1)). - exp_diff = tl.exp(s_k[:, None, :] + dA_cumsum_new[:, :, None]) - combo_block = factor_dt[:, None, :] * exp_diff # (H, T, K) - - # ---- gdc_wait: from here on we depend on conv1d's outputs ---- - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - # Conv1d outputs: B and C - C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group - B_new_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group - - C_all = tl.load( - C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - B_new_orig = tl.load( - B_new_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - B_new_shifted = tl.load( - B_new_base + safe_k_new[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, - mask=is_new_k[:, None] & n_mask[None, :], - other=0.0, - ) - # Disjoint masks: old at [0, PNAT), new at [K_NEW_SHIFT, K_NEW_SHIFT+T). - B_combined = old_B_load + B_new_shifted - raw_rect_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_combined).to(tl.bfloat16)) - - # Append new B to cache at [PNAT, PNAT+T) of write_buf (once per group). - if first_head % nheads_ngroups_ratio == 0: - old_B_write_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + write_buf * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - tl.store( - old_B_write_base - + (write_offset + offs_t)[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - B_new_orig, - mask=t_mask[:, None] & n_mask[None, :], - ) - - # Causal mask (BLOCK_SIZE_T × BLOCK_SIZE_K, shared across heads). - # Approach C: new tokens at runtime [PNAT, PNAT+T) instead of K_NEW_SHIFT. - t_idx_2d = offs_t[:, None] - k_idx_2d = offs_k[None, :] - is_old_k_2d = k_idx_2d < prev_num_accepted_tokens - k_new_idx_2d = k_idx_2d - prev_num_accepted_tokens - is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) - causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] - - # Post-wait vectorized: combo_block (H, T, K) is still live in registers. - # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store as - # one (H, T, K) tile. - rect_CB_scaled_block = tl.where( - causal_combined[None, :, :], - raw_rect_CB[None, :, :] * combo_block, - 0.0, - ) # (H, T, K) - cb_scaled_addrs = ( - cb_scaled_ptr - + pid_b * stride_cb_batch - + heads_block[:, None, None] * stride_cb_head - + offs_t[None, :, None] * stride_cb_t - + offs_k[None, None, :] * stride_cb_j - ) # (H, T, K) - cb_store_mask_3d = ( - (offs_t[None, :, None] < BLOCK_SIZE_T) - & (offs_k[None, None, :] < BLOCK_SIZE_K) - ) # (1, T, K) → broadcasts to (H, T, K) - tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) - - -# Rectangle precompute kernel. Thin wrapper around _rectangle_precompute_impl -# that carries the @triton.heuristics for constexpr derivation; called from -# the Python wrapper on the rectangle nowrite path. -@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.jit() -def _rectangle_precompute_kernel( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, - decay_vec_ptr, - # Cache pointers - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, - EARLY_OUT: tl.constexpr, - USE_PERM: tl.constexpr, - REVERSE_PERM: tl.constexpr, -): - # Hoisted PDL signal: fire as the first thing every program does, so - # main can start its setup regardless of how this program ends. - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - # Per-program early-out gate. Rectangle is nowrite-only, so EARLY_OUT - # skips slots whose PNAT + T > MAX (slots that would need write). - if EARLY_OUT: - pid_grid_eo = tl.program_id(axis=0) - pid_grid_eo_eff = (tl.num_programs(axis=0) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo - if USE_PERM: - pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) - else: - pid_b_eo = pid_grid_eo_eff - if HAS_CACHE_BATCH_INDICES: - cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) - if cbi_eo == pad_slot_id: - return - else: - cbi_eo = pid_b_eo.to(tl.int64) - pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) - if pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH: - return - _rectangle_precompute_impl( - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - cb_scaled_ptr, - decay_vec_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - pad_slot_id, - T, - MAX_REPLAY_BUFFER_LENGTH, - dstate, - nheads_ngroups_ratio, - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - DT_SOFTPLUS, - HAS_DT_BIAS, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - BLOCK_SIZE_K, - LAUNCH_WITH_PDL, - HEADS_PER_BLOCK, - USE_PERM, - REVERSE_PERM, - ) - - -# Dynamic precompute kernel. Single launchable kernel that, per program, -# reads PNAT and dispatches to one of the existing impls: -# -# if pnat + T > MAX: replay_precompute_impl(WRITE_CHECKPOINT=True) -# else if RECTANGLE: rectangle_precompute_impl -# else: replay_precompute_impl(WRITE_CHECKPOINT=False) -# -# RECTANGLE is constexpr (compile-time tuning param); the inner branch -# is folded so only one of the two nowrite paths is emitted per -# specialization. Reg envelope = max(replay_write, X) where X depends -# on RECTANGLE. cb_scaled is allocated (T, K) by the wrapper regardless; -# replay paths write to the first T columns, rectangle writes the full K. -@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.jit() -def _dynamic_precompute_kernel( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, - decay_vec_ptr, - # Cache pointers - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, - # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. - RECTANGLE: tl.constexpr, -): - # Hoisted PDL signal: fire as the first thing every program does. - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - - pid_b = tl.program_id(axis=0) - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH - # write_checkpoint is now runtime in replay precompute, so a single - # call site handles both write and nowrite for the replay branch. - # Take rectangle only when RECTANGLE is True AND this slot doesn't - # need write; everything else funnels into replay. - if needs_write_runtime or not RECTANGLE: - _replay_precompute_impl( - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - cb_scaled_ptr, - decay_vec_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) - pad_slot_id, - T, - dstate, - nheads_ngroups_ratio, - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - DT_SOFTPLUS, - HAS_DT_BIAS, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - LAUNCH_WITH_PDL, - HEADS_PER_BLOCK, - False, # USE_PERM - False, # REVERSE_PERM - needs_write_runtime, - ) - else: - _rectangle_precompute_impl( - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - cb_scaled_ptr, - decay_vec_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) - pad_slot_id, - T, - MAX_REPLAY_BUFFER_LENGTH, - dstate, - nheads_ngroups_ratio, - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - DT_SOFTPLUS, - HAS_DT_BIAS, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - BLOCK_SIZE_K, - LAUNCH_WITH_PDL, - HEADS_PER_BLOCK, - False, # USE_PERM - False, # REVERSE_PERM - ) - - -# Main kernel: tl.dot replay + precomputed CB output. -# Grid: (cdiv(dim, M), batch, nheads). - - -@triton.jit() -def _replay_main_impl( - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or - # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor - # USE_TMA_STORE is enabled (kernel ignores it via constexpr). - state_tma_descriptor, - # Per-(cache, head, dim) decode scale, fp32, only consulted when QUANT_MAX>0. - # Layout (cache, nheads, dim) — broadcast over dstate at load/store. - state_scales_ptr, - # Cache READ pointers (read-buffer from previous step) - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - # Cache WRITE pointer (write-buffer for old_x only; B/dt/dA_cumsum written by precompute) - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - # New input pointers - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - # Precomputed pointers - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - # Slot permutation: see _replay_precompute_impl for semantics. - slot_perm_ptr, - # Stochastic rounding - rand_seed_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # cache T-axis capacity (= max_window) - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides: (cache, nheads, dim) — only used when QUANT_MAX>0 - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides: (cache, T, nheads, dim) — single-buffered - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides: (cache, 2, T, ngroups, dstate) - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_WINDOW: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, - # State quantization: 0.0 means non-quantized (fp16/bf16/fp32); >0 means - # quantized (int8=127, int16=32767, fp8_e4m3fn=448). Single in-kernel - # switch for the dequant-on-load and encode-on-store paths. Wrapper sets - # this from state.dtype; kernel-entry static_assert below pins the - # invariant that it must coincide with int8/int16/float8e4nv state dtype. - QUANT_MAX: tl.constexpr, - # Checkpointing flag - WRITE_CHECKPOINT: tl.constexpr, # When True: quantize+write post-replay state to HBM (checkpoint step). - # When False: skip state write entirely (non-checkpoint step). - # The rectangle non-checkpoint path is implemented in - # _rectangle_main_kernel (separate kernel pair, picked by - # the wrapper via rectangle_for_nowrite=True). - # When True: signal PDL dependents at the very top of every program - # (including pad/early-out programs). Used by doublelaunch and maindl - # so the next kernel (the second main, or the second precompute) can - # start its setup while this main is still computing. Default False - # for monolithic / dynamic / the LAST main in dl/maindl chains. - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - # Slot permutation flags — see _replay_precompute_impl. - USE_PERM: tl.constexpr, - REVERSE_PERM: tl.constexpr, - # TMA toggles — gated per-path inside the body (write-load picks - # USE_TMA_LOAD_WRITE; nowrite-load picks USE_TMA_LOAD_NOWRITE; store - # only fires on the write path and uses USE_TMA_STORE). The wrapper - # passes write_load_value when WC=True and nowrite_load_value when - # WC=False; the unused flag is dummy False. Both are constexpr; - # is_write here is constexpr (= WRITE_CHECKPOINT), so use_tma_load - # constexpr-folds. - USE_TMA_LOAD_WRITE: tl.constexpr = False, - USE_TMA_LOAD_NOWRITE: tl.constexpr = False, - USE_TMA_STORE: tl.constexpr = False, -): - # Hoisted PDL signal: fire as the first thing every program does, so - # downstream kernels can start setup regardless of how this program - # ends (pad / early-out / full body). - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - - # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized - # state dtype (int8 / int16 / float8e4nv) and only those. Cheap - # insurance against a wrapper bug that desynchronizes the two. - tl.static_assert( - (QUANT_MAX > 0.0) - == ( - (state_ptr.dtype.element_ty == tl.int8) - or (state_ptr.dtype.element_ty == tl.int16) - or (state_ptr.dtype.element_ty == tl.float8e4nv) - ), - "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", - ) - - pid_m = tl.program_id(axis=0) - pid_grid_b = tl.program_id(axis=1) - pid_grid_b_eff = (tl.num_programs(axis=1) - 1 - pid_grid_b) if REVERSE_PERM else pid_grid_b - if USE_PERM: - pid_b = tl.load(slot_perm_ptr + pid_grid_b_eff) - else: - pid_b = pid_grid_b_eff - pid_h = tl.program_id(axis=2) - - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - # Active buffer (= cache_buf_idx) holds the historical inputs for this - # step at [0, PNAT). The replay phase reads from there. The new-tokens - # write target depends on WRITE_CHECKPOINT — see Cache write semantics - # block in the precompute kernel for the full rationale. - active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - if WRITE_CHECKPOINT: - write_buf = 1 - active_buf # noqa: F841 — old_x is single-buffered (no use here) - write_offset = 0 - else: - write_buf = active_buf # noqa: F841 — old_x is single-buffered (no use here) - write_offset = prev_num_accepted_tokens - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) - # Replay axis: separate from offs_t. Spans [0, BLOCK_SIZE_WINDOW) ⊇ - # [0, MAX_REPLAY_BUFFER_LENGTH); used for old-token loads (mask: offs_window < PNAT). - offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) - m_mask = offs_m < dim - n_mask = offs_n < dstate - t_mask = offs_t < T - - # Load state. state_tma_descriptor is a host-built tensor_descriptor - # over a flat (cache*nheads*dim, dstate) view of state when any TMA - # path is enabled; raw `state_ptr` is the underlying tensor and is - # always passed. state_ptrs / state_ptr_raw are the raw-pointer view - # used for !TMA load and store paths. offs_y is the flat row index - # for TMA load/store; computed unconditionally (cheap int math; DCE'd - # when no TMA path is reachable). - state_mask = m_mask[:, None] & n_mask[None, :] - offs_y = ( - cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) - + pid_h * dim - + pid_m * BLOCK_SIZE_M - ) - state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - ) - # Pick LOAD flag based on WRITE_CHECKPOINT (constexpr). Use an if/else - # with tl.constexpr annotations on each branch — a plain ternary binds - # the result to a Python (non-constexpr) name and the `if use_tma_load` - # gate below becomes a runtime branch, which forces BOTH the - # `state_tma_descriptor.load(...)` and `tl.load(state_ptrs, ...)` paths - # to compile. When TMA is off, the wrapper passes the plain state - # tensor as state_tma_descriptor (no `.load()` method) → compile fails. - # Wrapper passes USE_TMA_LOAD_WRITE = write-load value when WC=True - # (NOWRITE flag is dummy False then), and the converse when WC=False. - if WRITE_CHECKPOINT: - use_tma_load: tl.constexpr = USE_TMA_LOAD_WRITE - else: - use_tma_load: tl.constexpr = USE_TMA_LOAD_NOWRITE - if use_tma_load: - state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) - else: - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) - # Dequantize on load (per-(head, dim) decode scale, broadcast over dstate). - # Only consulted when QUANT_MAX>0 — non-quantized paths skip entirely. - if QUANT_MAX > 0.0: - state_scales_base = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - ) - decode_scale = tl.load( - state_scales_base + offs_m * stride_state_scales_dim, - mask=m_mask, - other=1.0, - ).to(tl.float32) - state = state * decode_scale[:, None] - - # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) - group_idx = pid_h // nheads_ngroups_ratio - - # Old-token mask along the WINDOW (replay) axis. PNAT ≤ MAX ≤ - # BLOCK_SIZE_WINDOW, so this enables all valid old-cache positions. - # (Distinct from t_mask = offs_t < T which gates output T-rows only.) - old_window_mask = offs_window < prev_num_accepted_tokens - - # Load precomputed dt and dA_cumsum from READ buffer at [0, PNAT). - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + active_buf * stride_old_dt_dbuf - + pid_h * stride_old_dt_head - ) - old_dt_all = tl.load( - old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 - ).to(tl.float32) - - old_dA_cumsum_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + active_buf * stride_old_dA_cumsum_dbuf - + pid_h * stride_old_dA_cumsum_head - ) - old_dA_cumsum_all = tl.load( - old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, - mask=old_window_mask, other=0.0, - ).to(tl.float32) - - # Load dA_cumsum at prev_k-1 directly via pointer math (avoids masked reduction). - # Clamp to [0, MAX-1] defensively — caller contract gives PNAT ≤ MAX. - prev_k_idx = tl.minimum( - tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 - ) - total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( - tl.float32 - ) - - # Step 0 invariant: PNAT=0 means `state` is already last step's state (not - # two back). coeff is all-zero (old_window_mask all-false), total_decay - # is 1.0, so the replay leaves `state` unchanged — cache contents don't matter. - coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all - - # Load old_x at [0, PNAT) of the WINDOW axis (single-buffered cache). - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head - old_x_all = tl.load( - old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - mask=old_window_mask[:, None] & m_mask[None, :], - other=0.0, - ) - - # Load old_B from READ buffer at [0, PNAT) of the WINDOW axis. - old_B_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + active_buf * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - old_B_all = tl.load( - old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - mask=old_window_mask[:, None] & n_mask[None, :], - other=0.0, - ).to(tl.float32) - - # Scale B by coefficients - dB_scaled = coeff[:, None] * old_B_all - - # Apply total decay to initial state FIRST, then add contributions - total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) - state *= total_decay - - # tl.dot fast-forward: old_x^T @ dB_scaled → (M, dstate) - state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) - - # Write post-replay state — only on checkpoint steps. When - # WRITE_CHECKPOINT is False, the replay computed `state` is local-only and - # discarded; skipping the HBM store + Philox path is the main performance - # win of replay-style checkpointing on the common (non-checkpoint) step. - if WRITE_CHECKPOINT: - if USE_RS_ROUNDING: - # Generate random tensor for stochastic rounding. The amount of - # randomness needed depends on the SR codegen path: - # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) - # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) - # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs - # int16 SR (24b + bitrev32): 1 b32 per 2 outputs - # The PTX cvt.rs.* instructions consume a single 32-bit random - # and split the bits internally for 2 or 4 conversions. The - # tl.inline_asm_elementwise wrapper has uniform `pack` across all - # args, so it provides 2 (fp16) or 4 (fp8) rand inputs per asm - # call but only the first is read; the others are dead. Generate - # only what's actually consumed and broadcast to fill the unused - # slots — saves Philox rounds proportionally. - if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: - RAND_DIVISOR: tl.constexpr = 4 # fp8 SR - elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: - RAND_DIVISOR: tl.constexpr = 4 # int8 SR - elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: - RAND_DIVISOR: tl.constexpr = 2 # int16 SR - elif QUANT_MAX == 0.0: - RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) - else: - RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized - - rand_seed = tl.load(rand_seed_ptr) - base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head - # Number of unique randoms per row = dstate / RAND_DIVISOR. - # randint4x emits 4 randoms per offset, so use that / 4 offsets. - offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) - rand_offsets_q = ( - base_rand - + offs_m[:, None] * stride_state_dim - + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) - ) # (M, dstate / (4*RAND_DIVISOR)) - if PHILOX_ROUNDS > 0: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) - else: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) - r01 = tl.join(r0, r1) - r23 = tl.join(r2, r3) - r0123 = tl.join(r01, r23) - rand_compact = tl.reshape( - r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) - ) - # Broadcast each unique rand to RAND_DIVISOR adjacent positions - # in the dstate axis. Pack-group (pack=2 fp16 / pack=4 fp8) - # consumes adjacent positions; the unique rand lands at the - # asm's read slot ($3 fp16 / $5 fp8); duplicates feed the dead - # slots ($4 fp16; $6/$7/$8 fp8). Triton's broadcast_to is - # stride-0 in IR. - # - # Tested zero-fill alternative (tl.join with zeros): essentially - # equivalent register count (95 vs 96 at one config) and same - # timing. ptxas does not use RZ for the dead asm input slots - # in either case; the extra ~15 regs vs pre-fix come from - # rand_compact's lifetime across the asm call, not from the - # fill pattern. Broadcast wins on simplicity. - if RAND_DIVISOR > 1: - rand_3d = rand_compact[:, :, None] - rand_3d = tl.broadcast_to( - rand_3d, - (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), - ) - rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) - else: - rand = rand_compact - - if QUANT_MAX > 0.0: - # Quantized state path: int8 / int16 / fp8_e4m3fn (RN or SR). - # 1) Per-(head, dim) channel scale via amax over dstate. - amax = tl.max(tl.abs(state), axis=1) # (M,) - encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) # (M,) - decode_scale = 1.0 / encode_scale # (M,) - # 2) Store decode_scale (1/encode) so reads do a single multiply. - state_scales_ptrs = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - + offs_m * stride_state_scales_dim - ) - tl.store(state_scales_ptrs, decode_scale, mask=m_mask) - # 3) Scale state into quant range — into a NEW variable so the - # downstream output phase still sees the dequantized fp32 state. - state_q = state * encode_scale[:, None] - # 4) Round per dtype. Order matters: handle fp8 SR first (PTX - # combines round + saturating cast in one op, output is final fp8 - # so we store and finish on that branch). Other branches share - # the clamp + cast tail below. - if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): - # fp8_e4m3fn + SR — PTX cvt.rs.satfinite.e4m3x4.f32. Output - # is final fp8 (saturate included); store directly. - _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_q_fp8sr) - else: - tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) - else: - if USE_RS_ROUNDING: - # int8 / int16 + SR — uniform-noise + floor. - # int8 packs 4 values per random u32 using 16-bit chunks - # and bitrev16; int16 packs 2 values per random u32 using - # 24-bit uniforms from the direct/reversed u32. - # (fp8 SR was handled by the early branch above.) - tl.static_assert( - (state_ptrs.dtype.element_ty == tl.int8) - or (state_ptrs.dtype.element_ty == tl.int16), - "Quantized SR fall-through expects int8 or int16; " - "fp8 SR is handled by the prior branch.", - ) - if state_ptrs.dtype.element_ty == tl.int8: - state_q = _stochastic_round_int8_packed( - state_q, rand, offs_n[None, :] - ) - else: - state_q = _stochastic_round_int16_packed( - state_q, rand, offs_n[None, :] - ) - elif state_ptrs.dtype.element_ty != tl.float8e4nv: - # int8 / int16 + RN — explicit round before clamp. - # fp8 + RN deliberately skips this — explicit round() would - # destroy fp8 sub-integer precision; native cast at store - # does RN at the fp8 grid resolution. - tl.static_assert( - (state_ptrs.dtype.element_ty == tl.int8) - or (state_ptrs.dtype.element_ty == tl.int16), - "Quantized RN with explicit round() expects int8 or int16.", - ) - state_q = tl.extra.cuda.libdevice.round(state_q) - # Clamp + cast tail: int8/int16 (RN+SR) and fp8 RN. - # fp8 RN reaches here without prior round() — .to(float8e4nv) - # does native RN at the fp8 grid. - state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) - _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_q_cast) - else: - tl.store(state_ptrs, _state_q_cast, mask=state_mask) - elif USE_RS_ROUNDING: - # Non-quantized + SR: only fp16 (bf16 has no PTX SR cast; fp32 - # doesn't need rounding). - tl.static_assert( - state_ptrs.dtype.element_ty == tl.float16, - "Non-quantized SR only supports fp16 state.", - ) - _state_sr = _stochastic_round_fp16x2(state, rand) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_sr) - else: - tl.store(state_ptrs, _state_sr, mask=state_mask) - else: - # Non-quantized + RN: fp16 / bf16 / fp32 native cast. - _state_cast = state.to(state_ptrs.dtype.element_ty) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_cast) - else: - tl.store(state_ptrs, _state_cast, mask=state_mask) - - # Phase 2: Output using precomputed CB_scaled and decay_vec - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head - C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group - if HAS_Z: - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - - if HAS_D: - D = tl.load( - D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - - # Wait for precompute kernel (PDL) before reading its outputs. - # With chained PDL (conv1d → precompute → main), gdc_wait() ensures - # precompute has completed — which transitively ensures conv1d has - # completed (precompute waited on conv1d via its own gdc_wait). - # All loads below (x, C from conv1d; CB_scaled, decay_vec from precompute) - # are safe after this point. - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - # Load conv1d outputs: C_all and x_all - C_all = tl.load( - C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - - x_all = tl.load( - x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=t_mask[:, None] & m_mask[None, :], - other=0.0, - ) - # Store new x to old_x cache at [write_offset : write_offset+T). - # old_x is single-buffered: write goes to the active buffer regardless; - # replay already read positions [0, PNAT) so write_offset = PNAT (no - # overlap) on no-checkpoint steps. On checkpoint steps write_offset = 0 - # (cache reset; old data folded into state via replay update). - tl.store( - old_x_base - + (write_offset + offs_t)[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - x_all, - mask=t_mask[:, None] & m_mask[None, :], - ) - x_all = x_all.to(tl.float32) - - # Load precomputed CB_scaled and decay_vec - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head - CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), - other=0.0, - ).to(tl.float32) - - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( - tl.float32 - ) - - # init_out = C_all @ state^T * decay_vec - init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] - - # cb_out = CB_scaled @ x_all - cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) - - out_all = init_out + cb_out - - if HAS_D: - out_all = out_all + x_all * D[None, :] - - if HAS_Z: - z_all = tl.load( - z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, - mask=t_mask[:, None] & m_mask[None, :], other=0.0, - ).to(tl.float32) - out_all_z = out_all * z_all * tl.sigmoid(z_all) - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) - else: - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) - - -# Replay-style main kernel. Thin wrapper around _replay_main_impl that carries -# the @triton.heuristics for constexpr derivation; called from the Python -# wrapper on the replay-style path (write or replay-nowrite). -@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) -@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics( - {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} -) -@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_WINDOW": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.jit() -def _checkpointing_main_kernel( - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or - # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor - # USE_TMA_STORE is enabled (kernel ignores it via constexpr). - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - rand_seed_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_WINDOW: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, - QUANT_MAX: tl.constexpr, - WRITE_CHECKPOINT: tl.constexpr, - EARLY_OUT: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - USE_PERM: tl.constexpr, - REVERSE_PERM: tl.constexpr, - # 3 TMA flags passed through to _replay_main_impl (which picks - # USE_TMA_LOAD_WRITE vs NOWRITE based on WRITE_CHECKPOINT). Wrapper - # passes write_load_value when WC=True (NOWRITE flag dummy False), - # and the converse when WC=False. - USE_TMA_LOAD_WRITE: tl.constexpr = False, - USE_TMA_LOAD_NOWRITE: tl.constexpr = False, - USE_TMA_STORE: tl.constexpr = False, -): - # Hoisted PDL signal: fire as the first thing every program does. - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - # Per-program early-out gate. Signal-then-skip lets the next kernel - # in dl/maindl chains start regardless of early-out outcome. - if EARLY_OUT: - pid_grid_eo = tl.program_id(axis=1) - pid_grid_eo_eff = (tl.num_programs(axis=1) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo - if USE_PERM: - pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) - else: - pid_b_eo = pid_grid_eo_eff - if HAS_CACHE_BATCH_INDICES: - cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) - if cbi_eo == pad_slot_id: - return - else: - cbi_eo = pid_b_eo.to(tl.int64) - pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) - if (pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH) != WRITE_CHECKPOINT: - return - _replay_main_impl( - state_ptr, - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - rand_seed_ptr, - pad_slot_id, - T, - MAX_REPLAY_BUFFER_LENGTH, - dim, - dstate, - nheads_ngroups_ratio, - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_D_head, - stride_D_dim, - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - BLOCK_SIZE_M, - HAS_D, - HAS_Z, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, - USE_RS_ROUNDING, - PHILOX_ROUNDS, - QUANT_MAX, - WRITE_CHECKPOINT, - LAUNCH_DEPENDENT_KERNELS, - USE_PERM, - REVERSE_PERM, - USE_TMA_LOAD_WRITE, - USE_TMA_LOAD_NOWRITE, - USE_TMA_STORE, - ) - - -# Rectangle main kernel (nowrite-only): no replay step, no state HBM write, -# no SR codegen. state_out is computed from state_prev directly using the -# precomp-folded decay_vec_full; token_out is a single rectangle matmul over -# the (T, K) CB rectangle and (K, M) x_combined. Pairs with -# `_rectangle_precompute_kernel`. - - -@triton.jit() -def _rectangle_main_impl( - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as - # replay path). Used when USE_TMA_LOAD; ignored otherwise (kernel - # branches via constexpr). Wrapper passes the same descriptor as - # for replay paths — single underlying memory, consumed by per-path - # constexpr gates. - state_tma_descriptor, - state_scales_ptr, # only consulted when QUANT_MAX > 0 - old_x_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, # rectangle (batch, nheads, T, K) - decay_vec_ptr, # folded (batch, nheads, T) — total_decay * exp(cumAdt_new[t]) - state_batch_indices_ptr, - # Slot permutation: see _replay_precompute_impl for semantics. - slot_perm_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides (no quant-store path; state read-only for state_out) - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides: (cache, nheads, dim) — fp32, broadcast over dstate - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides: (cache, T_max, nheads, dim) — single-buffered - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides (rectangle (batch, nheads, T, K)) - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - QUANT_MAX: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - # Slot permutation flags — see _replay_precompute_impl. - USE_PERM: tl.constexpr, - REVERSE_PERM: tl.constexpr, - USE_TMA_LOAD: tl.constexpr = False, -): - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - - pid_m = tl.program_id(axis=0) - pid_grid_b = tl.program_id(axis=1) - pid_grid_b_eff = (tl.num_programs(axis=1) - 1 - pid_grid_b) if REVERSE_PERM else pid_grid_b - if USE_PERM: - pid_b = tl.load(slot_perm_ptr + pid_grid_b_eff) - else: - pid_b = pid_grid_b_eff - pid_h = tl.program_id(axis=2) - - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - write_offset = prev_num_accepted_tokens - - # Static rectangle K-axis layout (matches precompute). - K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_k = tl.arange(0, BLOCK_SIZE_K) - m_mask = offs_m < dim - n_mask = offs_n < dstate - t_mask = offs_t < T - - # K-axis masks (approach C: PNAT-runtime offset, matches precompute). - # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. - is_old_k = offs_k < prev_num_accepted_tokens - safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - prev_num_accepted_tokens - is_new_k = (k_new_idx >= 0) & (k_new_idx < T) - safe_k_new = tl.where(is_new_k, k_new_idx, 0) - - # Load state. Read-only — no HBM write on the nowrite path. - # Quant scale hoist (backlog #16): for QUANT_MAX > 0 paths, defer the - # `* decode_scale` to AFTER the C @ state dot — applied to the (T, M) - # dot output instead of broadcast-multiplied into the (M, dstate) state - # tile. Algebra-equivalent (decode_scale is per-M, commutes with the - # matmul over dstate). Saves M·dstate fp32 muls (replaced by T·M), - # but the bigger potential win is shorter register lifetime for state - # (kept as native int8/int16/fp8 until just before the dot, where Triton - # casts to bf16 — vs current fp32 tile across the whole kernel). Only - # applies in rectangle main (no `state += dot` here). - if USE_TMA_LOAD: - # TMA descriptor (host-built) over the flat 2D view of state: - # shape=[cache_size * nheads * dim, dstate], strides=[dstate, 1]. - # Convert (cache, head, m) → flat row index using existing strides: - # rows-per-cache-slot = stride_state_batch / stride_state_dim - # rows-per-head = stride_state_head / stride_state_dim = dim (constexpr) - # rows-per-m = 1 - # Diagnostic: prior in-kernel descriptor attempts emitted - # ttng.tensormap_create setup (divergent shared-mem write) which - # blows up branch count; host-built descriptors avoid that. - offs_y = ( - cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) - + pid_h * dim - + pid_m * BLOCK_SIZE_M - ) - state = state_tma_descriptor.load([offs_y, 0]) - else: - state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - ) - state_mask = m_mask[:, None] & n_mask[None, :] - state = tl.load(state_ptrs, mask=state_mask, other=0.0) - if QUANT_MAX > 0.0: - state_scales_base = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - ) - decode_scale = tl.load( - state_scales_base + offs_m * stride_state_scales_dim, - mask=m_mask, other=1.0, - ).to(tl.float32) - # state stays in native quant dtype — cast happens inside the dot below. - else: - state = state.to(tl.float32) - - # Group / pointer offset setup - group_idx = pid_h // nheads_ngroups_ratio - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head - C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group - if HAS_Z: - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head - - if HAS_D: - D = tl.load( - D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - - # Hoist: old_x from cache doesn't depend on conv1d/precompute, so issue - # the load BEFORE gdc_wait so its HBM latency overlaps with conv1d. - old_x_load = tl.load( - old_x_base - + safe_old_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - mask=is_old_k[:, None] & m_mask[None, :], - other=0.0, - ).to(tl.float32) - - # PDL gate: precompute outputs (cb_scaled, decay_vec_full) become safe - # after gdc_wait. conv1d outputs (x, C) also gated by the chained PDL. - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - # Load C and x (conv1d outputs after PDL wait) - C_all = tl.load( - C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - # Single (BLOCK_K, M) load at PNAT-offset positions (approach C): - # K-axis [PNAT, PNAT+T) gets new tokens directly from x[0:T, :] via - # safe_k_new = offs_k - PNAT. Cache layout matches K-axis layout, so - # one load serves both matmul and cache write. - x_K = tl.load( - x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=is_new_k[:, None] & m_mask[None, :], - other=0.0, - ) - # Cache write: store at offs_k directly (is_new_k mask makes offs_k land - # at [PNAT, PNAT+T) in the cache, which is exactly write_offset+0..T-1). - tl.store( - old_x_base - + offs_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - x_K, - mask=is_new_k[:, None] & m_mask[None, :], - ) - - x_K_f32 = x_K.to(tl.float32) - # Matmul side: K-axis aligned; sum with old_x_load. - x_combined = old_x_load + x_K_f32 - - # T-axis view for D feedthrough / Z-gating: extract via (T, K) selection. - # Only materialized when needed. - if HAS_D or HAS_Z: - sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) - x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) - else: - x_all = x_K_f32 # placeholder; unused - - # Load precomputed rectangle CB and folded decay_vec. - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head - CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), - other=0.0, - ).to(tl.float32) - - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec_full = tl.load( - decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 - ).to(tl.float32) - - # state_out: state_prev contribution to output, with decay folded post-matmul. - # No state_prev_decayed (M, dstate) materialization — state is consumed - # directly by the matmul, then decay_vec_full multiplies the (T, M) result. - # For QUANT_MAX > 0 (#16 hoist): decode_scale also applies post-matmul - # at (T, M) granularity instead of pre-multiplied into the (M, dstate) - # state tile. Triton's tl.dot(a.to(bf16), b.to(bf16)) handles the - # int8/fp8 → bf16 cast inside the dot's input prep. - state_out = ( - tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) - * decay_vec_full[:, None] - ) - if QUANT_MAX > 0.0: - state_out = state_out * decode_scale[None, :] - - # token_out: combined old + new tokens contribution via the rectangle. - token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) - - out_all = state_out + token_out - - if HAS_D: - out_all = out_all + x_all * D[None, :] - - if HAS_Z: - z_all = tl.load( - z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, - mask=t_mask[:, None] & m_mask[None, :], other=0.0, - ).to(tl.float32) - out_all_z = out_all * z_all * tl.sigmoid(z_all) - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) - else: - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) - - -# Rectangle main kernel. Thin wrapper around _rectangle_main_impl that carries -# the @triton.heuristics for constexpr derivation; called from the Python -# wrapper on the rectangle nowrite path. -@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) -@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics( - {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} -) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.jit() -def _rectangle_main_kernel( - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor for state (used when - # USE_TMA_LOAD). Same descriptor as replay paths use; gate via - # constexpr. Wrapper passes the unified state_tma_descriptor. - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - QUANT_MAX: tl.constexpr, - EARLY_OUT: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - USE_PERM: tl.constexpr, - REVERSE_PERM: tl.constexpr, - USE_TMA_LOAD: tl.constexpr = False, -): - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - # Per-program early-out gate. Rectangle is nowrite-only. - if EARLY_OUT: - pid_grid_eo = tl.program_id(axis=1) - pid_grid_eo_eff = (tl.num_programs(axis=1) - 1 - pid_grid_eo) if REVERSE_PERM else pid_grid_eo - if USE_PERM: - pid_b_eo = tl.load(slot_perm_ptr + pid_grid_eo_eff) - else: - pid_b_eo = pid_grid_eo_eff - if HAS_CACHE_BATCH_INDICES: - cbi_eo = tl.load(state_batch_indices_ptr + pid_b_eo).to(tl.int64) - if cbi_eo == pad_slot_id: - return - else: - cbi_eo = pid_b_eo.to(tl.int64) - pnat_eo = tl.load(prev_num_accepted_tokens_ptr + cbi_eo) - if pnat_eo + T > MAX_REPLAY_BUFFER_LENGTH: - return - _rectangle_main_impl( - state_ptr, - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - pad_slot_id, - T, - MAX_REPLAY_BUFFER_LENGTH, - dim, - dstate, - nheads_ngroups_ratio, - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_D_head, - stride_D_dim, - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - BLOCK_SIZE_M, - HAS_D, - HAS_Z, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - BLOCK_SIZE_K, - LAUNCH_WITH_PDL, - QUANT_MAX, - LAUNCH_DEPENDENT_KERNELS, - USE_PERM, - REVERSE_PERM, - USE_TMA_LOAD, - ) - - -# Dynamic main kernel. Single launchable kernel that, per program, reads -# PNAT and dispatches to one of the existing impls: -# -# if pnat + T > MAX: replay_main_impl(WRITE_CHECKPOINT=True) -# elif RECTANGLE (constexpr): rectangle_main_impl -# else: replay_main_impl(WRITE_CHECKPOINT=False) -# -# Unlike precompute, WRITE_CHECKPOINT stays constexpr in the main impl — -# the body has constexpr-gated state-write code (quant + Philox + HBM -# store) where folding meaningfully shrinks the codegen. So this kernel -# has TWO replay call sites (one per WRITE_CHECKPOINT specialization) -# both inlined, with a runtime branch picking which runs. Reg envelope = -# max(replay_write, X) where X = rectangle_nowrite (RECTANGLE=True) or -# replay_nowrite (RECTANGLE=False). cb_scaled is allocated (T, K) by the -# wrapper regardless. -@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) -@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics( - {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} -) -@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_WINDOW": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.jit() -def _dynamic_main_kernel( - # Pointers — union of replay-main and rectangle-main pointer args. - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as - # other main kernels). Currently always passed as dummy `state_ptr` - # by launch_dynamic_main since dynamic doesn't expose TMA toggles - # yet — kept in the signature for uniformity with replay/rect/persistent. - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - rand_seed_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides (replay only; passed but unused on rectangle path) - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides (replay only) - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides (replay only) - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_WINDOW: tl.constexpr, # for replay path - BLOCK_SIZE_K: tl.constexpr, # for rectangle path - LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, - QUANT_MAX: tl.constexpr, - # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. - RECTANGLE: tl.constexpr, - # Default False — dynamic main is normally terminal in its chain. - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - # 3 TMA toggles (matching dual-path kernel scheme): - # USE_TMA_LOAD_WRITE — write path's state load - # USE_TMA_LOAD_NOWRITE — nowrite path's state load (rect when - # RECTANGLE, else replay-nowrite) - # USE_TMA_STORE — write path's state store (no-op for nowrite) - USE_TMA_LOAD_WRITE: tl.constexpr = False, - USE_TMA_LOAD_NOWRITE: tl.constexpr = False, - USE_TMA_STORE: tl.constexpr = False, -): - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - pid_b = tl.program_id(axis=1) - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - if pnat_local + T > MAX_REPLAY_BUFFER_LENGTH: - # Write slot — replay-style write (WRITE_CHECKPOINT=True constexpr). - _replay_main_impl( - state_ptr, - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) - rand_seed_ptr, - pad_slot_id, - T, - MAX_REPLAY_BUFFER_LENGTH, - dim, - dstate, - nheads_ngroups_ratio, - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_D_head, - stride_D_dim, - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - BLOCK_SIZE_M, - HAS_D, - HAS_Z, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, - USE_RS_ROUNDING, - PHILOX_ROUNDS, - QUANT_MAX, - True, # WRITE_CHECKPOINT (constexpr) - False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top - False, # USE_PERM - False, # REVERSE_PERM - USE_TMA_LOAD_WRITE, # write-load fires here - USE_TMA_LOAD_NOWRITE, # nowrite-load: dummy at this site - USE_TMA_STORE, # store fires (WC=True) - ) - else: - if RECTANGLE: - _rectangle_main_impl( - state_ptr, - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) - pad_slot_id, - T, - MAX_REPLAY_BUFFER_LENGTH, - dim, - dstate, - nheads_ngroups_ratio, - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_D_head, - stride_D_dim, - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - BLOCK_SIZE_M, - HAS_D, - HAS_Z, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - BLOCK_SIZE_K, - LAUNCH_WITH_PDL, - QUANT_MAX, - False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top - False, # USE_PERM - False, # REVERSE_PERM - USE_TMA_LOAD_NOWRITE, # rect-load TMA flag - ) - else: - # Replay-style nowrite (WRITE_CHECKPOINT=False constexpr). - _replay_main_impl( - state_ptr, - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - state_batch_indices_ptr, # slot_perm_ptr unused (USE_PERM=False) - rand_seed_ptr, - pad_slot_id, - T, - MAX_REPLAY_BUFFER_LENGTH, - dim, - dstate, - nheads_ngroups_ratio, - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_D_head, - stride_D_dim, - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - BLOCK_SIZE_M, - HAS_D, - HAS_Z, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, - USE_RS_ROUNDING, - PHILOX_ROUNDS, - QUANT_MAX, - False, # WRITE_CHECKPOINT (constexpr) - False, # LAUNCH_DEPENDENT_KERNELS — already signaled at top - False, # USE_PERM - False, # REVERSE_PERM - USE_TMA_LOAD_WRITE, # write-load: dummy at this site - USE_TMA_LOAD_NOWRITE, # nowrite-load fires here - USE_TMA_STORE, # store: dummy (WC=False) - ) - - -# Python wrapper - - -_QUANT_MAX_BY_DTYPE = { - torch.int8: 127.0, - torch.int16: 32767.0, - torch.float8_e4m3fn: 448.0, -} - - -# ============================================================================ -# Persistent main kernel — 1D grid, persistent CTA loop with tl.range -# ============================================================================ -# -# Design (see ~/dev/scripts/mamba_replay/kernel_microbenchmarks/PERSISTENT_KERNELS.md -# for the full strawman): -# -# * Outer 1D grid of `NUM_PERSISTENT` CTAs (start at NUM_SMS, sweep upward). -# * Inside the kernel, a `tl.range(pid, total_work, NUM_PERSISTENT, flatten=True, -# num_stages=NUM_STAGES)` loop iterates over (slot, M_tile, head) work units. -# * Hard-sort PNAT host-side and pass `n_writes` as a runtime int32 scalar: -# the launcher invokes the kernel twice — once with slot_offset=0, -# n_slots=n_writes, WRITE_CHECKPOINT=True, and once with -# slot_offset=n_writes, n_slots=B-n_writes, WRITE_CHECKPOINT=False. -# * `_persistent_main_impl` is a copy of `_replay_main_impl`'s body with the -# program_id reads replaced by parameters and the slot_perm logic moved into -# the persistent loop wrapper. No code shared with the existing kernels; -# easy to delete if the experiment is abandoned. -# -# Notes: -# * `flatten=True` is canonical for Triton 3.6 persistent kernels (matches the -# upstream `_p_matmul_ogs.py` and tutorial 09). Combined with `num_stages=2` -# it pipelines the loop body — but watch open issue triton-lang/triton#8259 -# which reports this combo can corrupt stores in non-dot loops. First run -# correctness check is critical. -# * Warp specialization (`warp_specialize=True`) is NOT enabled — Triton 3.6 -# only supports it for simple matmul loops and our scan won't pattern-match. -# * No 2CTA cluster mode — that's dot-only per the kernel-tileir-optimization -# skill classification. - - -@triton.jit() -def _persistent_main_impl( - # Per-work-unit indices (computed by the persistent wrapper). - # `pid_b` is the post-perm slot index (caller has already applied any - # slot permutation and slot_offset). - pid_m, - pid_b, - pid_h, - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or - # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor - # USE_TMA_STORE is enabled (kernel ignores it via constexpr). - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - rand_seed_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_WINDOW: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, - QUANT_MAX: tl.constexpr, - WRITE_CHECKPOINT: tl.constexpr, - # IS_DYNAMIC: kept in the signature for caller-side bookkeeping (the - # outer _persistent_main_kernel still inspects it to decide the slot- - # IS_DYNAMIC: when True (persistent_dynamic), is_write is per-slot from - # PNAT. When False (persistent_main), is_write is constexpr from - # WRITE_CHECKPOINT. See also WC_IS_CONSTEXPR below. - IS_DYNAMIC: tl.constexpr, - # WC_IS_CONSTEXPR: when True, force is_write = WRITE_CHECKPOINT (constexpr) - # regardless of IS_DYNAMIC. Callers in RECT=1 use this in the is_w=True - # arm of _persistent_main_kernel (we know all slots that reach this call - # need is_write=True because is_w was the PNAT-derived runtime check, and - # this arm only fires when is_w is True). Passing WRITE_CHECKPOINT=True - # as a literal at the call site + WC_IS_CONSTEXPR=True here lets the inner - # body DCE the nowrite path under IS_DYNAMIC=True too — same codegen - # quality as persistent_main mode (-3.7% measured at b=1024 dyn-shape). - # When False (RECT=0 callers, where both write and nowrite slots are - # dispatched to ONE call), use the original runtime is_write under - # IS_DYNAMIC=True; avoids the binary-doubling regression that two - # specialized calls would cause. - WC_IS_CONSTEXPR: tl.constexpr = False, - # TMA flags — picked inside body based on is_write. When is_write is - # constexpr (either IS_DYNAMIC=False or WC_IS_CONSTEXPR=True), the - # use_tma_load = USE_TMA_LOAD_WRITE if is_write else USE_TMA_LOAD_NOWRITE - # ternary constexpr-folds and only one TMA load form survives. - USE_TMA_LOAD_WRITE: tl.constexpr = False, - USE_TMA_LOAD_NOWRITE: tl.constexpr = False, - USE_TMA_STORE: tl.constexpr = False, -): - # IS_DYNAMIC: kernel-mode label, used by the OUTER _persistent_main_kernel - # to decide slot-range derivation and outer is_w dispatch strategy - # (constexpr WC for persistent_main; runtime is_w split -> 2 specialized - # impl calls for persistent_dynamic). Inside this impl, IS_DYNAMIC is - # NOT consulted at runtime -- WRITE_CHECKPOINT is the only constexpr that - # gates the write/nowrite codegen, in BOTH modes. - - # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized - # state dtype (int8 / int16 / float8e4nv) and only those. - tl.static_assert( - (QUANT_MAX > 0.0) - == ( - (state_ptr.dtype.element_ty == tl.int8) - or (state_ptr.dtype.element_ty == tl.int16) - or (state_ptr.dtype.element_ty == tl.float8e4nv) - ), - "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", - ) - - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - # Resolve is_write: see WC_IS_CONSTEXPR / IS_DYNAMIC docs in the param - # list above. Three cases: - # - WC_IS_CONSTEXPR=True (RECT=1 is_w=True arm callers): use WC - # constexpr. Caller knows the slot needs write; inner DCEs nowrite - # paths. Avoids the binary-doubling overhead that calling the impl - # twice would cause, while still constexpr-DCEing the nowrite half. - # - IS_DYNAMIC=True (RECT=0 caller, persistent_dynamic): runtime - # branch on PNAT. Both write and nowrite codegen live in one body - # (no bloat) — same as the pre-refactor behavior. - # - IS_DYNAMIC=False (persistent_main): WC constexpr from caller. - if WC_IS_CONSTEXPR: - is_write: tl.constexpr = WRITE_CHECKPOINT - elif IS_DYNAMIC: - is_write = (prev_num_accepted_tokens + T) > MAX_REPLAY_BUFFER_LENGTH - else: - is_write = WRITE_CHECKPOINT - if is_write: - write_buf = 1 - active_buf # noqa: F841 - write_offset = 0 - else: - write_buf = active_buf # noqa: F841 - write_offset = prev_num_accepted_tokens - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) - m_mask = offs_m < dim - n_mask = offs_n < dstate - t_mask = offs_t < T - - # Load state. state_tma_descriptor is a host-built tensor_descriptor - # over a flat (cache*nheads*dim, dstate) view of state when any TMA - # path is enabled; raw `state_ptr` is the underlying tensor and is - # always passed. state_ptrs / state_ptr_raw are the raw-pointer view - # used for !TMA load and store paths. offs_y is the flat row index - # for TMA load/store; computed unconditionally (cheap int math; DCE'd - # when no TMA path is reachable). - state_mask = m_mask[:, None] & n_mask[None, :] - offs_y = ( - cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) - + pid_h * dim - + pid_m * BLOCK_SIZE_M - ) - state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - ) - # Load state. Branch on is_write (constexpr = WRITE_CHECKPOINT in BOTH - # modes after the outer-dispatch refactor), then constexpr-pick TMA-vs- - # tl.load per side. Outer `if` DCE's, only the matching side's - # constexpr-gated load survives -- same compile-time picking for both - # persistent_main and persistent_dynamic (the latter dispatches at the - # outer kernel level so each impl instance sees a constexpr WC). - if is_write: - if USE_TMA_LOAD_WRITE: - state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) - else: - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) - else: - if USE_TMA_LOAD_NOWRITE: - state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) - else: - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) - if QUANT_MAX > 0.0: - state_scales_base = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - ) - decode_scale = tl.load( - state_scales_base + offs_m * stride_state_scales_dim, - mask=m_mask, - other=1.0, - ).to(tl.float32) - state = state * decode_scale[:, None] - - # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) - group_idx = pid_h // nheads_ngroups_ratio - - old_window_mask = offs_window < prev_num_accepted_tokens - - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + active_buf * stride_old_dt_dbuf - + pid_h * stride_old_dt_head - ) - old_dt_all = tl.load( - old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 - ).to(tl.float32) - - old_dA_cumsum_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + active_buf * stride_old_dA_cumsum_dbuf - + pid_h * stride_old_dA_cumsum_head - ) - old_dA_cumsum_all = tl.load( - old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, - mask=old_window_mask, other=0.0, - ).to(tl.float32) - - prev_k_idx = tl.minimum( - tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 - ) - total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( - tl.float32 - ) - - coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all - - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head - old_x_all = tl.load( - old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - mask=old_window_mask[:, None] & m_mask[None, :], - other=0.0, - ) - - old_B_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + active_buf * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - old_B_all = tl.load( - old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - mask=old_window_mask[:, None] & n_mask[None, :], - other=0.0, - ).to(tl.float32) - - dB_scaled = coeff[:, None] * old_B_all - - total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) - state *= total_decay - - state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) - - if is_write: - if USE_RS_ROUNDING: - # Generate random tensor for stochastic rounding. The amount of - # randomness needed depends on the SR codegen path: - # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) - # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) - # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs - # int16 SR (24b + bitrev32): 1 b32 per 2 outputs - # The PTX cvt.rs.* instructions consume a single 32-bit random - # and split the bits internally for 2 or 4 conversions. Generate - # only what's actually consumed and broadcast to fill the unused - # slots — saves Philox rounds proportionally. - if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: - RAND_DIVISOR: tl.constexpr = 4 # fp8 SR - elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: - RAND_DIVISOR: tl.constexpr = 4 # int8 SR - elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: - RAND_DIVISOR: tl.constexpr = 2 # int16 SR - elif QUANT_MAX == 0.0: - RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) - else: - RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized - - rand_seed = tl.load(rand_seed_ptr) - base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head - # Number of unique randoms per row = dstate / RAND_DIVISOR. - # randint4x emits 4 randoms per offset, so use that / 4 offsets. - offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) - rand_offsets_q = ( - base_rand - + offs_m[:, None] * stride_state_dim - + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) - ) # (M, dstate / (4*RAND_DIVISOR)) - if PHILOX_ROUNDS > 0: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) - else: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) - r01 = tl.join(r0, r1) - r23 = tl.join(r2, r3) - r0123 = tl.join(r01, r23) - rand_compact = tl.reshape( - r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) - ) - # Broadcast each unique rand to RAND_DIVISOR adjacent positions. - # Pack-group (pack=2 fp16 / pack=4 fp8) consumes adjacent positions; - # the unique rand lands at the asm's read slot; duplicates feed - # the dead slots. Triton's broadcast_to is stride-0 in IR. - if RAND_DIVISOR > 1: - rand_3d = rand_compact[:, :, None] - rand_3d = tl.broadcast_to( - rand_3d, - (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), - ) - rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) - else: - rand = rand_compact - - if QUANT_MAX > 0.0: - amax = tl.max(tl.abs(state), axis=1) - encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) - decode_scale = 1.0 / encode_scale - state_scales_ptrs = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - + offs_m * stride_state_scales_dim - ) - tl.store(state_scales_ptrs, decode_scale, mask=m_mask) - state_q = state * encode_scale[:, None] - if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): - _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_q_fp8sr) - else: - tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) - else: - if USE_RS_ROUNDING: - tl.static_assert( - (state_ptrs.dtype.element_ty == tl.int8) - or (state_ptrs.dtype.element_ty == tl.int16), - "Quantized SR fall-through expects int8 or int16; " - "fp8 SR is handled by the prior branch.", - ) - if state_ptrs.dtype.element_ty == tl.int8: - state_q = _stochastic_round_int8_packed( - state_q, rand, offs_n[None, :] - ) - else: - state_q = _stochastic_round_int16_packed( - state_q, rand, offs_n[None, :] - ) - elif state_ptrs.dtype.element_ty != tl.float8e4nv: - tl.static_assert( - (state_ptrs.dtype.element_ty == tl.int8) - or (state_ptrs.dtype.element_ty == tl.int16), - "Quantized RN with explicit round() expects int8 or int16.", - ) - state_q = tl.extra.cuda.libdevice.round(state_q) - state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) - _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_q_cast) - else: - tl.store(state_ptrs, _state_q_cast, mask=state_mask) - elif USE_RS_ROUNDING: - tl.static_assert( - state_ptrs.dtype.element_ty == tl.float16, - "Non-quantized SR only supports fp16 state.", - ) - _state_sr = _stochastic_round_fp16x2(state, rand) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_sr) - else: - tl.store(state_ptrs, _state_sr, mask=state_mask) - else: - _state_cast = state.to(state_ptrs.dtype.element_ty) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_cast) - else: - tl.store(state_ptrs, _state_cast, mask=state_mask) - - # Phase 2: Output using precomputed CB_scaled and decay_vec - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head - C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group - if HAS_Z: - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - - if HAS_D: - D = tl.load( - D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - C_all = tl.load( - C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - - x_all = tl.load( - x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=t_mask[:, None] & m_mask[None, :], - other=0.0, - ) - tl.store( - old_x_base - + (write_offset + offs_t)[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - x_all, - mask=t_mask[:, None] & m_mask[None, :], - ) - x_all = x_all.to(tl.float32) - - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head - CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), - other=0.0, - ).to(tl.float32) - - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( - tl.float32 - ) - - init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] - cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) - out_all = init_out + cb_out - - if HAS_D: - out_all = out_all + x_all * D[None, :] - - if HAS_Z: - z_all = tl.load( - z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, - mask=t_mask[:, None] & m_mask[None, :], other=0.0, - ).to(tl.float32) - out_all_z = out_all * z_all * tl.sigmoid(z_all) - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) - else: - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) - - -# `_persistent_rectangle_impl`: rectangle nowrite path for the persistent -# kernel. Body is a copy of `_rectangle_main_impl` with `pid_m`/`pid_b`/`pid_h` -# lifted to args (same pattern as `_persistent_main_impl` vs `_replay_main_impl`). -# Called only for nowrite slots when the kernel runs with RECTANGLE=True. -# Dropped from the rect impl: LAUNCH_DEPENDENT_KERNELS / USE_PERM / REVERSE_PERM -# (kernel-level, signalled once at top); slot_perm_ptr (kernel resolves perm). -@triton.jit() -def _persistent_rectangle_impl( - # Per-work-unit indices (computed by the persistent wrapper). - pid_m, - pid_b, - pid_h, - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as - # replay path). Used when USE_TMA_LOAD; ignored otherwise. - state_tma_descriptor, - state_scales_ptr, # only consulted when QUANT_MAX > 0 - old_x_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides (rectangle (batch, nheads, T, K)) - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - QUANT_MAX: tl.constexpr, - USE_TMA_LOAD: tl.constexpr = False, -): - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - write_offset = prev_num_accepted_tokens - - # Static rectangle K-axis layout (matches precompute). - K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_k = tl.arange(0, BLOCK_SIZE_K) - m_mask = offs_m < dim - n_mask = offs_n < dstate - t_mask = offs_t < T - - # K-axis masks (approach C: PNAT-runtime offset, matches precompute). - is_old_k = offs_k < prev_num_accepted_tokens - safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - prev_num_accepted_tokens - is_new_k = (k_new_idx >= 0) & (k_new_idx < T) - safe_k_new = tl.where(is_new_k, k_new_idx, 0) - - # Load state. Quant scale hoist: defer `* decode_scale` post-matmul. - if USE_TMA_LOAD: - offs_y = ( - cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) - + pid_h * dim - + pid_m * BLOCK_SIZE_M - ) - state = state_tma_descriptor.load([offs_y, 0]) - else: - state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - ) - state_mask = m_mask[:, None] & n_mask[None, :] - state = tl.load(state_ptrs, mask=state_mask, other=0.0) - if QUANT_MAX > 0.0: - state_scales_base = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - ) - decode_scale = tl.load( - state_scales_base + offs_m * stride_state_scales_dim, - mask=m_mask, other=1.0, - ).to(tl.float32) - else: - state = state.to(tl.float32) - - # Group / pointer offset setup - group_idx = pid_h // nheads_ngroups_ratio - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head - C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group - if HAS_Z: - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head - - if HAS_D: - D = tl.load( - D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - - # Hoist: old_x doesn't depend on conv1d/precompute; load before gdc_wait. - old_x_load = tl.load( - old_x_base - + safe_old_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - mask=is_old_k[:, None] & m_mask[None, :], - other=0.0, - ).to(tl.float32) - - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - C_all = tl.load( - C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - x_K = tl.load( - x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=is_new_k[:, None] & m_mask[None, :], - other=0.0, - ) - tl.store( - old_x_base - + offs_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - x_K, - mask=is_new_k[:, None] & m_mask[None, :], - ) - - x_K_f32 = x_K.to(tl.float32) - x_combined = old_x_load + x_K_f32 - - if HAS_D or HAS_Z: - sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) - x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) - else: - x_all = x_K_f32 # placeholder; unused - - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head - CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), - other=0.0, - ).to(tl.float32) - - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec_full = tl.load( - decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 - ).to(tl.float32) - - state_out = ( - tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) - * decay_vec_full[:, None] - ) - if QUANT_MAX > 0.0: - state_out = state_out * decode_scale[None, :] - - token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) - - out_all = state_out + token_out - - if HAS_D: - out_all = out_all + x_all * D[None, :] - - if HAS_Z: - z_all = tl.load( - z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, - mask=t_mask[:, None] & m_mask[None, :], other=0.0, - ).to(tl.float32) - out_all_z = out_all * z_all * tl.sigmoid(z_all) - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) - else: - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) - - -# Persistent main kernel: 1D grid, persistent CTA loop. -# Heuristics mirror those of `_checkpointing_main_kernel`. -@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) -@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics( - {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} -) -@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_WINDOW": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.heuristics( - {"NUM_PID_M_BLOCKS": lambda args: triton.cdiv(args["dim"], args["BLOCK_SIZE_M"])} -) -@triton.jit() -def _persistent_main_kernel( - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view. - # Shared across BOTH the replay path (consumed by _persistent_main_impl - # when USE_TMA_LOAD_*/STORE) AND the rectangle path (consumed by - # _persistent_rectangle_impl when USE_TMA_LOAD) — same descriptor, same - # block_shape, just gated by separate constexprs per impl. Wrapper sets - # this to a TensorDescriptor when ANY of the three TMA flags is on, else - # to `state_ptr` (raw); each impl ignores it via its own constexpr when - # not consuming it. - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - rand_seed_ptr, - pad_slot_id, - # Persistent-loop work-distribution scalars. Caller pre-sorts the batch - # write-first; the kernel uses (n_writes, batch_total, WRITE_CHECKPOINT) - # to derive its own slot range. Write half processes [0, n_writes), - # nowrite half processes [n_writes, batch_total). - # - # n_writes_ptr is a device pointer to a (1,) int32 tensor. Reading - # from device memory (rather than taking a Python int kernel arg) is - # required so mix-mode benchmarking can vary n_writes per iter inside - # a captured CUDA graph — the source tensor's contents change, the - # pointer doesn't. Cost: one int load per kernel launch (~negligible). - # When IS_DYNAMIC=True the value is unused (Triton DCEs the load). - n_writes_ptr, # int32 *: device-side count of write-mode slots - batch_total, # int32: total slot count - nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_WINDOW: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, - QUANT_MAX: tl.constexpr, - WRITE_CHECKPOINT: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - USE_PERM: tl.constexpr, - # NUM_PERSISTENT: runtime int (not constexpr). Used ONLY as the loop - # stride in `tl.range(pid, total_work, NUM_PERSISTENT, ...)`. Making it - # runtime collapses the cta_per_sm tuning dim from the kernel's compile - # signature: 8 CPS values used to mean 8x recompiles; now they share one - # compiled kernel. Work decomposition (pid_m, pid_b_local, pid_h) does - # NOT depend on NUM_PERSISTENT — it uses constexpr NUM_PID_M_BLOCKS and - # runtime n_slots_local — so loop unrolling and flatten=/num_stages=/ - # warp_specialize= optimizations on `tl.range` operate independently of - # the stride value. - NUM_PERSISTENT, - NUM_LOOP_STAGES: tl.constexpr, - NUM_PID_M_BLOCKS: tl.constexpr, - FLATTEN: tl.constexpr, - WARP_SPECIALIZE: tl.constexpr, - IS_DYNAMIC: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr = 16, # rectangle K-axis (heuristic-derived) - RECTANGLE: tl.constexpr = False, # when True, dispatch nowrite slots to _persistent_rectangle_impl - # 3 TMA toggles per the 3 live paths per-compilation: - # USE_TMA_LOAD_WRITE — replay-style state load when is_write - # USE_TMA_LOAD_NOWRITE — nowrite-path state load (rect when RECTANGLE, - # else replay-nowrite) - # USE_TMA_STORE — replay-style state store (only fires on write - # path; no-op when not is_write) - # Wrapper picks USE_TMA_LOAD_NOWRITE = _use_tma_rect_load (if rectangle) - # or _use_tma_replay_nowrite_load (if not). - USE_TMA_LOAD_WRITE: tl.constexpr = False, - USE_TMA_LOAD_NOWRITE: tl.constexpr = False, - USE_TMA_STORE: tl.constexpr = False, -): - # PDL signal: fire once at kernel entry (not per work unit). - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - - # Load runtime n_writes from device memory. Read once at kernel entry; - # used only by the !IS_DYNAMIC slot-range derivation below. Triton - # DCEs the load when IS_DYNAMIC=True (n_writes is dead there). - n_writes = tl.load(n_writes_ptr) - - # Derive this kernel's slot range. Two modes: - # IS_DYNAMIC=False (persistent_main): caller pre-sorts and splits halves; - # slot range is [0, n_writes) when WRITE_CHECKPOINT else [n_writes, batch_total) - # IS_DYNAMIC=True (persistent_dynamic): single launch covers full batch; - # each work-item dispatches via runtime PNAT check inside the impl. - if IS_DYNAMIC: - slot_lo = 0 - slot_hi = batch_total - else: - if WRITE_CHECKPOINT: - slot_lo = 0 - slot_hi = n_writes - else: - slot_lo = n_writes - slot_hi = batch_total - n_slots_local = slot_hi - slot_lo - - pid = tl.program_id(axis=0) - total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads - - # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) - # with pid_m varying fastest (M-tile cache locality on state load), then - # slot, then head — mirrors the existing 3D grid's axis ordering - # (axis=0 fastest = pid_m). - for tile_id in tl.range( - pid, total_work, NUM_PERSISTENT, - flatten=FLATTEN, num_stages=NUM_LOOP_STAGES, warp_specialize=WARP_SPECIALIZE, - ): - pid_m = tile_id % NUM_PID_M_BLOCKS - pid_b_local = (tile_id // NUM_PID_M_BLOCKS) % n_slots_local - pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) - # Translate local slot index → global slot index. When USE_PERM is - # set, the caller-provided slot_perm gives the original slot index - # for the post-sort position. - pid_b_grid = pid_b_local + slot_lo - if USE_PERM: - pid_b = tl.load(slot_perm_ptr + pid_b_grid) - else: - pid_b = pid_b_grid - - # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle - # impl. Pre-load cache_batch_idx + PNAT here only for the RECTANGLE - # path's branch decision. Both impls re-load and handle pad_slot_id - # internally (Triton's L1 cache makes the duplicate loads ~free). - if RECTANGLE: - if HAS_CACHE_BATCH_INDICES: - cbi_pre = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - is_pad = cbi_pre == pad_slot_id - else: - cbi_pre = pid_b.to(tl.int64) - is_pad = False - if not is_pad: - pnat_pre = tl.load(prev_num_accepted_tokens_ptr + cbi_pre) - if IS_DYNAMIC: - is_w = (pnat_pre + T) > MAX_REPLAY_BUFFER_LENGTH - else: - is_w = WRITE_CHECKPOINT - if is_w: - # Pass WRITE_CHECKPOINT=True constexpr to specialize this - # impl call for the write path. Under IS_DYNAMIC=True, the - # kernel-level WRITE_CHECKPOINT is False (launcher default), - # but the OUTER is_w branch we are inside narrows the - # runtime path to writes-only, so we override to True here - # so the impl's constexpr-gated `if is_write:` blocks DCE - # to the write-only codegen. Under IS_DYNAMIC=False - # (persistent_main), the kernel-level WRITE_CHECKPOINT is - # itself True for this half (write half launches with - # WC=True), and the outer is_w = WRITE_CHECKPOINT = True - # constexpr-folds; passing literal True here is consistent - # and constexpr-equivalent. - _persistent_main_impl( - pid_m, pid_b, pid_h, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, - stride_old_B_group, stride_old_B_dstate, - stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, - stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) - True, # WC_IS_CONSTEXPR — force inner to use WC constexpr - # 3 TMA flags: write-load fires here (we're in the - # is_write branch), nowrite-load is dead (no slot - # reaches it), store fires (write path). - USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, - ) - else: - # Rectangle nowrite: pass state_ptr (raw, always) + - # state_tma_descriptor (the single unified descriptor — - # same memory replay paths use). Rect impl gates use - # of the descriptor via its USE_TMA_LOAD constexpr. - _persistent_rectangle_impl( - pid_m, pid_b, pid_h, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - state_batch_indices_ptr, pad_slot_id, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, - LAUNCH_WITH_PDL, QUANT_MAX, - USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle - ) - # else: pad slot — skip both impls (both would early-return anyway) - else: - # No rectangle path — single _persistent_main_impl call covers - # both write and nowrite slots via WC constexpr (non-dynamic) or - # runtime is_write (IS_DYNAMIC=True). Pass all 3 TMA flags; - # impl picks USE_TMA_LOAD_WRITE vs USE_TMA_LOAD_NOWRITE based on - # its computed is_write — constexpr-folds when is_write is - # constexpr (non-dyn), runtime branch when IS_DYNAMIC=True. - # (Reverted from outer two-call dispatch: that doubled the - # compiled body size under IS_DYNAMIC=True and regressed RECT=0 - # perf by ~+24%.) - _persistent_main_impl( - pid_m, pid_b, pid_h, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, - stride_old_B_group, stride_old_B_dstate, - stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, - stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - WRITE_CHECKPOINT, IS_DYNAMIC, - False, # WC_IS_CONSTEXPR=False — RECT=0 has both write/nowrite slots in one call - USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, - ) - - -# ============================================================================ -# Python wrapper -# ============================================================================ - - -def checkpointing_state_update( - state: torch.Tensor, - old_x: torch.Tensor, - old_B: torch.Tensor, - old_dt: torch.Tensor, - old_dA_cumsum: torch.Tensor, - cache_buf_idx: torch.Tensor, - prev_num_accepted_tokens: torch.Tensor, - x: torch.Tensor, - dt: torch.Tensor, - A: torch.Tensor, - B: torch.Tensor, - C: torch.Tensor, - out: torch.Tensor, - D: torch.Tensor | None = None, - z: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - dt_softplus: bool = False, - state_batch_indices: torch.Tensor | None = None, - pad_slot_id: int = PAD_SLOT_ID, - rand_seed: torch.Tensor | None = None, - philox_rounds: int = 10, - state_scales: torch.Tensor | None = None, - launch_with_pdl=False, - use_internal_pdl=True, - write_checkpoint: bool = True, - rectangle_for_nowrite: bool = False, - mode: str = "monolithic", - # Slot permutation: int32 (batch,) tensor mapping grid program_id -> - # original slot index. When provided, dl-family kernels (doublelaunch / - # dlgrouped / maindl) read pid_b through this perm so callers can pre-sort - # slots (e.g. write-first) to cluster early-outs at one end of the grid. - # Ignored by monolithic / dynamic. None => identity (today's behavior). - slot_perm: torch.Tensor | None = None, - # When True and slot_perm is provided, the nowrite-side kernels in - # dlgrouped/doublelaunch traverse the perm in reverse (B-1-pid_grid). - # Combined with a write-first sort, this front-loads real work in BOTH - # halves of the dl chain (writes from the head, nowrites from the tail). - reverse_nowrite: bool = False, - _block_size_m: int | None = None, - _num_warps: int | None = None, - _num_stages: int | None = None, - _precompute_num_warps: int | None = None, - _precompute_num_stages: int | None = None, - _heads_per_block: int | None = None, - _maxnreg: int | None = None, - _num_ctas: int | None = None, - # Per-main knobs (override shared values for one half of the dl-family / - # persistent_main launches). Default None = tied to the shared value - # (backward compat). The two main kernels (write vs nowrite) have - # different per-slot work — write does a state shift + store, nowrite - # just appends — so the optimum (M, W, S, H) can differ. Precompute - # knobs are intentionally NOT split: shared precompute wins (cheaper - # launch, hotter precompute outputs in L2). Persistent CPS / LS knobs - # are also split per-main since the two persistent_main launches have - # different grid sizes. - _block_size_m_write: int | None = None, - _block_size_m_nowrite: int | None = None, - _num_warps_write: int | None = None, - _num_warps_nowrite: int | None = None, - _num_stages_write: int | None = None, - _num_stages_nowrite: int | None = None, - # Note: heads_per_block / precompute_num_warps are NOT split — they only - # affect the precompute kernel, which is shared across write/nowrite. - # TMA state-tensor toggles — 4 independent paths (see CHECKPOINTING_DESIGN.md - # item #17 for measured perf profiles). Each is False=raw load/store, True= - # use a host-built TMA tensor_descriptor for that path. - _use_tma_rect_load: bool = False, # rect kernel's state load (nowrite-only) - _use_tma_replay_write_load: bool = False, # replay-style state load when WC=True - _use_tma_replay_write_store: bool = False, # replay-style state store when WC=True - _use_tma_replay_nowrite_load: bool = False, # replay-style state load when WC=False - # Persistent-mode bench kwargs (only consulted when mode == "persistent_main"): - # _n_writes : int — count of write-mode slots in the (pre-sorted) batch. - # Required when mode == "persistent_main"; the persistent kernel uses - # it as a runtime int32 to compute total_work for write/nowrite halves. - # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally - # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. Default = 1. - # _num_loop_stages : int — `num_stages` arg on the inner `tl.range(...)` - # persistent loop. Note: this is loop-level, NOT the kernel-arg - # `num_stages` (which only pipelines dot-feeding loads). Default 2. - # _flatten : bool — `flatten` arg on `tl.range(...)`. Default True - # (the canonical Triton 3.6 persistent idiom). - # _warp_specialize : bool — `warp_specialize` arg on `tl.range(...)`. - # Default False. Triton 3.6 only supports it on simple matmul loops; - # our scan loop probably won't pattern-match — but exposed as a knob - # for sweep experiments. Requires num_warps >= 4 if True. - _n_writes: int | None = None, - # Optional pre-allocated (1,) int32 device tensor for the persistent - # kernel's n_writes input. Bench passes this in mix scenarios so the - # captured CUDA graph can read varying n_writes per iter without - # re-capture. When None and `_n_writes` is provided, we allocate a - # scratch tensor and fill from `_n_writes` (pure scenarios). - _n_writes_dev: torch.Tensor | None = None, - # When True, persistent_main host-skips empty-half launches (n_writes=0 - # or =batch in pure scenarios). Default True preserves today's behavior. - # Set False to always launch both halves — used by mix scenarios (where - # host can't cheaply read n_writes per iter) and for fair K-consistent - # comparisons. - _persistent_skip_empty_halves: bool = True, - _cta_per_sm: int | None = None, - _num_loop_stages: int | None = None, - _flatten: bool | None = None, - _warp_specialize: bool | None = None, - # Per-main persistent-specific knobs. Same rationale as the BLOCK_SIZE_M - # split above: the two persistent_main launches (write half vs nowrite - # half) have different grid sizes and per-work-item costs, so they may - # want different cta_per_sm / num_loop_stages. - _cta_per_sm_write: int | None = None, - _cta_per_sm_nowrite: int | None = None, - _num_loop_stages_write: int | None = None, - _num_loop_stages_nowrite: int | None = None, -): - """ - Replay SSM state update with precomputed CB and tl.dot fast-forward. - - Two-kernel architecture: - 1. Precompute kernel: computes CB_scaled and decay_vec from B, C, dt, A. - Writes processed dt/dA_cumsum/B to double-buffered cache for next step. - 2. Main kernel: replays old tokens via tl.dot fast-forward on cached data, - then computes output using precomputed CB_scaled and new x/C inputs. - - PDL (Programmatic Dependent Launch) chain: - conv1d → (external PDL) → precompute → (internal PDL) → main - External PDL: precompute starts while conv1d is running; gdc_wait() - in precompute blocks until conv1d completes before loading B/C. - Internal PDL: main starts while precompute is running; main's replay - phase uses only cached data from the previous step. gdc_wait() in - main blocks until precompute completes before loading conv1d outputs - (x, C) and precompute outputs (CB_scaled, decay_vec). - - Uses double-buffered cache tensors. cache_buf_idx[slot] indicates which - buffer (0 or 1) to READ from for replay. The WRITE buffer is 1 - read. - Caller must flip cache_buf_idx[slot] after each call. - - Arguments: - state: (cache, nheads, dim, dstate) in-place. After the call, contains - the state after replaying prev_num_accepted_tokens old tokens. - old_x: (cache, T, nheads, dim) bf16 — old x cache (single-buffered). - old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. - old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. - old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. - cache_buf_idx: (cache,) int32 — which buffer to read (0 or 1). - prev_num_accepted_tokens: (cache,) int32. - x: (batch, T, nheads, dim) new token inputs. - dt: (batch, T, nheads, dim) with stride(-1)==0 (tie_hdim). - A: (nheads, dim, dstate) with stride(-1)==0, stride(-2)==0 (tie_hdim). - B: (batch, T, ngroups, dstate). - C: (batch, T, ngroups, dstate). - out: (batch, T, nheads, dim) preallocated output. - D: (nheads, dim) optional feed-through parameter. - z: (batch, T, nheads, dim) optional silu gate. - dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). - state_batch_indices: (batch,) optional cache slot mapping. - rand_seed: optional single-element int64 CUDA tensor for Philox PRNG seed. - When provided, state is stochastically rounded on store. Supported - for state.dtype in (fp16, int8, int16, fp8_e4m3fn); other dtypes - silently use deterministic rounding. fp16+SR and fp8+SR both - require sm_100a (Blackwell B200+) — wrapper asserts this loudly. - philox_rounds: number of Philox PRNG rounds (default 10). - state_scales: required when state.dtype in (int8, int16, fp8_e4m3fn). - Shape (cache_size, nheads, dim), fp32. Per-(head, dim) channel - decode scale (= 1 / encode_scale). The kernel writes scales on - checkpoint steps and reads them on load (broadcast over dstate). - Ignored for non-quantized state dtypes. - launch_with_pdl: enable external PDL (conv1d → precompute chain). - Defaults False; caller opts in when the upstream chain is PDL-safe. - Ignored on hardware that doesn't support PDL (sm < 90). - use_internal_pdl: enable internal PDL (precompute → main overlap). - Defaults True; override for testing only. - Ignored on hardware that doesn't support PDL (sm < 90). - - _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, - _precompute_num_warps, _precompute_num_stages, _heads_per_block, - _maxnreg, _num_ctas) are benchmark-only overrides; production callers - should leave them None to use the heuristic-tuned defaults. - """ - # PDL needs sm >= 90. - if get_sm_version() < 90: - launch_with_pdl = False - use_internal_pdl = False - - # Mode selection: - # mode="monolithic" (default): today's behavior. write_checkpoint and - # rectangle_for_nowrite together pick a single kernel pair for the - # whole batch. Calls the corresponding kernel pair with EARLY_OUT=False. - # mode="dynamic": single kernel pair (_dynamic_*_kernel) that dispatches - # per-slot at runtime based on PNAT. RECTANGLE constexpr (= - # rectangle_for_nowrite) picks whether the nowrite path is rectangle - # or replay-nowrite. write_checkpoint is ignored (per-slot from PNAT). - # mode="doublelaunch": two kernel pairs launched in sequence, each with - # EARLY_OUT=True, partitioning the batch by PNAT-derived mode. - # Write half: replay-write. Nowrite half: rectangle if - # rectangle_for_nowrite else replay-nowrite. write_checkpoint ignored. - # mode="dlgrouped": same 4 kernels as doublelaunch, but reordered to - # launch both precomputes first, then both mains. Lets the GPU - # run precomp1 || precomp2 in parallel before the mains start. - # write_checkpoint ignored. - # mode="maindl": shared (dynamic) precompute + doublelaunched main. - # One precompute call (_dynamic_precompute_kernel) handles per-slot - # dispatch, then two main kernels with EARLY_OUT=True for the write - # and nowrite halves. Strictly fewer kernel launches than - # doublelaunch (3 vs 4) at the cost of dispatch precompute's wider - # reg envelope. write_checkpoint ignored. - assert mode in ( - "monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", - "dl_write_only", "persistent_main", "persistent_dynamic", - ), ( - f"unknown mode {mode!r}; expected one of " - "'monolithic', 'dynamic', 'doublelaunch', 'dlgrouped', 'maindl', " - "'dl_write_only', 'persistent_main', or 'persistent_dynamic'" - ) - use_rectangle = rectangle_for_nowrite and not write_checkpoint - - # --- Hardware support gates --- - # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX - # instructions (Ada Lovelace introduced them; Hopper/Blackwell carry them). - if state.dtype == torch.float8_e4m3fn: - assert get_sm_version() >= 89, ( - "fp8_e4m3fn state requires SM 89+ (Ada Lovelace / Hopper / Blackwell) " - f"for fp32↔fp8 cvt PTX instructions; current SM is {get_sm_version()}." - ) - - # PTX cvt.rs.* (stochastic rounding) family lands on Blackwell only. - # Wrapper fails loud; framework decides fall-back (e.g. drop SR, use RN). - # int8 / int16 SR uses pure-Triton libdevice.floor + uniform noise — no - # PTX SR instruction needed, runs anywhere. - if rand_seed is not None: - if state.dtype == torch.float16: - assert get_sm_version() >= 100, ( - "fp16 stochastic rounding (PTX cvt.rs.f16x2.f32) requires " - f"sm_100a (Blackwell B200+); current SM is {get_sm_version()}." - ) - elif state.dtype == torch.float8_e4m3fn: - assert get_sm_version() >= 100, ( - "fp8 stochastic rounding (PTX cvt.rs.satfinite.e4m3x4.f32) " - f"requires sm_100a (Blackwell B200+); current SM is {get_sm_version()}." - ) - - # --- Unsqueeze inputs to canonical shapes --- - if state.dim() == 3: - state = state.unsqueeze(1) - if x.dim() == 2: - x = x.unsqueeze(1) - if x.dim() == 3: - x = x.unsqueeze(1) - if dt.dim() == 2: - dt = dt.unsqueeze(1) - if dt.dim() == 3: - dt = dt.unsqueeze(1) - if A.dim() == 2: - A = A.unsqueeze(0) - if B.dim() == 2: - B = B.unsqueeze(1) - if B.dim() == 3: - B = B.unsqueeze(1) - if C.dim() == 2: - C = C.unsqueeze(1) - if C.dim() == 3: - C = C.unsqueeze(1) - if D is not None and D.dim() == 1: - D = D.unsqueeze(0) - if z is not None: - if z.dim() == 2: - z = z.unsqueeze(1) - if z.dim() == 3: - z = z.unsqueeze(1) - if dt_bias is not None and dt_bias.dim() == 1: - dt_bias = dt_bias.unsqueeze(0) - if out.dim() == 2: - out = out.unsqueeze(1) - if out.dim() == 3: - out = out.unsqueeze(1) - - cache_size, nheads, dim, dstate = state.shape - batch, T, _, _ = x.shape - ngroups = B.shape[2] - assert nheads % ngroups == 0 - - # --- Quantization plumbing --- - # QUANT_MAX > 0 ⇔ state is int8 / int16 / fp8_e4m3fn. Kernel-entry - # static_assert on the Triton side mirrors this invariant. - quant_max = _QUANT_MAX_BY_DTYPE.get(state.dtype, 0.0) - is_quantized = quant_max > 0.0 - if is_quantized: - assert state_scales is not None, ( - f"state.dtype={state.dtype} requires state_scales tensor " - "(shape (cache_size, nheads, dim), fp32)." - ) - assert state_scales.shape == (cache_size, nheads, dim), ( - f"state_scales shape mismatch: expected {(cache_size, nheads, dim)}, " - f"got {state_scales.shape}." - ) - assert state_scales.dtype == torch.float32, ( - f"state_scales must be fp32, got {state_scales.dtype}." - ) - assert state_scales.device == state.device - - # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the - # placeholder degenerate case max_window = T (every step is a checkpoint - # step). For real replay-style checkpointing, max_window > T and - # `prev_num_accepted_tokens` can be 0..max_window. - max_window = old_x.shape[1] - assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" - # Replay-style code path uses BLOCK_SIZE_T = max(np2(T), 16) for the - # combined T-axis (T_new tile size) and reuses it for window loads. Until - # the heuristic is generalized to track max_window separately, require - # max_window to fit within that tile. - block_size_t = max(triton.next_power_of_2(T), 16) - assert max_window <= block_size_t, ( - f"max_window={max_window} exceeds BLOCK_SIZE_T={block_size_t} " - f"derived from T={T}; extend the heuristic to include max_window." - ) - - assert x.shape == (batch, T, nheads, dim) - assert dt.shape == x.shape - assert A.shape == (nheads, dim, dstate) - assert B.shape == (batch, T, ngroups, dstate) - assert C.shape == B.shape - assert old_x.shape == (cache_size, max_window, nheads, dim) - assert old_B.shape == (cache_size, 2, max_window, ngroups, dstate) - assert old_dt.shape == (cache_size, 2, nheads, max_window) - assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_window) - assert cache_buf_idx.shape == (cache_size,) - assert prev_num_accepted_tokens.shape == (cache_size,) - - tie_hdim = ( - A.stride(-1) == 0 - and A.stride(-2) == 0 - and dt.stride(-1) == 0 - and (dt_bias is None or dt_bias.stride(-1) == 0) - ) - assert tie_hdim - - device = x.device - BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) - # Rectangle K-axis bound = window (max_window). Computed unconditionally - # so the launch sites can refer to it; only used on the rectangle path. - BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) - - # Allocate precomputed intermediates (per-call, not cached). Always - # allocate (T, K) — the largest layout that any path uses. Replay-style - # paths only touch the first T columns; rectangle/dynamic use the full K. - # The few extra unused columns per row are negligible (~6KB per layer at - # production sizes) and let the dispatch helpers share one buffer. - cb_scaled = torch.empty( - batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 - ) - decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) - - z_strides = ( - (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) if z is not None else (0, 0, 0, 0) - ) - - # Kernel tuning: BLOCK_SIZE_M, num_warps, HEADS_PER_BLOCK, precompute_num_warps. - # Dtype-aware heuristic from B200 sweeps (batch 1-512, T=6/32, TP=8, conv1d + - # chained PDL). Keyed on total_heads, BLOCK_SIZE_T, and state dtype; 16-bit - # states prefer different tiles from fp32 due to lower bandwidth. Philox - # gets its own branch — stochastic rounding shifts compute toward CUDA cores, - # so small-batch configs want more warps to hide the extra work. - total_heads = batch * nheads - heads_per_group = nheads // ngroups - state_is_16bit = state.dtype in (torch.float16, torch.bfloat16) - use_philox = rand_seed is not None - if BLOCK_SIZE_T <= 16: - if use_philox and state_is_16bit: - # Philox: more warps at small batch to hide CUDA core work. - # At large batch, converges to non-Philox fp16 config. - if total_heads <= 16: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 4, 4, 4, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - elif state_is_16bit: - if total_heads <= 16: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 4, 1 - elif total_heads <= 64: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 2, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 1, - min(2, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - else: # fp32 state (no Philox — fp32 doesn't need stochastic rounding) - if total_heads <= 32: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 4, 1 - elif total_heads <= 64: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - elif total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 2, 2, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 1, 2, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 2, - min(2, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - else: # T > 16 - if state_is_16bit: - if total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 16, - 1, - 4, - min(2, heads_per_group), - ) - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 1, - min(4, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 4, - min(2, heads_per_group), - ) - else: # fp32 state - if total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 2, - 4, - min(2, heads_per_group), - ) - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 2, - min(4, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 4, - min(2, heads_per_group), - ) - if _block_size_m is not None: - BLOCK_SIZE_M = _block_size_m - if _num_warps is not None: - num_warps = _num_warps - if _heads_per_block is not None: - heads_per_block = _heads_per_block - if _precompute_num_warps is not None: - precompute_num_warps = _precompute_num_warps - - # Per-main knob resolution: each _*_{write,nowrite} arg, if not None, - # overrides the corresponding shared value for ONE main launch only. - # Default (None) = tied to shared value (current behavior). - BLOCK_SIZE_M_WRITE = _block_size_m_write if _block_size_m_write is not None else BLOCK_SIZE_M - BLOCK_SIZE_M_NOWRITE = _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M - NUM_WARPS_WRITE = _num_warps_write if _num_warps_write is not None else num_warps - NUM_WARPS_NOWRITE = _num_warps_nowrite if _num_warps_nowrite is not None else num_warps - NUM_STAGES_WRITE = _num_stages_write if _num_stages_write is not None else _num_stages - NUM_STAGES_NOWRITE = _num_stages_nowrite if _num_stages_nowrite is not None else _num_stages - # Persistent-only per-main: - CTA_PER_SM_WRITE = _cta_per_sm_write if _cta_per_sm_write is not None else _cta_per_sm - CTA_PER_SM_NOWRITE = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _cta_per_sm - NUM_LOOP_STAGES_WRITE = _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages - NUM_LOOP_STAGES_NOWRITE = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages - - HAS_CACHE_BATCH_INDICES = state_batch_indices is not None - - assert nheads % heads_per_block == 0, ( - f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" - ) - assert heads_per_block <= heads_per_group, ( - f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" - ) - - # state_scales pointer + strides: real tensor when quantized, otherwise - # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). - if is_quantized: - state_scales_arg = state_scales - state_scales_strides = ( - state_scales.stride(0), - state_scales.stride(1), - state_scales.stride(2), - ) - else: - state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 - state_scales_strides = (0, 0, 0) - - # Per-path TMA descriptors for state — write-side and nowrite-side. Each - # kernel launch consumes the descriptor whose block_shape[0] matches its - # BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need - # distinct descriptors; otherwise the descriptor's block_shape[0] would - # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot / arithmetic - # on the loaded tile fails shape inference at compile time - # ("Cannot make_shape_compatible: incompatible dimensions"). When Mw == - # Mnw (tied, the common case) the two descriptors are the same object. - # Same memory (state's flat 2D view, shape (cache*nheads*dim, dstate)) - # and same dstate block_shape — only block_shape[0] differs. - # When no TMA flag is on, both variables hold the raw `state` tensor as a - # dummy; kernels never reference it because their constexprs are all - # False (Triton DCEs the dead branches). - # `triton.set_allocator()` must run before any descriptor-using launch. - if (_use_tma_rect_load or _use_tma_replay_write_load - or _use_tma_replay_write_store or _use_tma_replay_nowrite_load): - from triton.tools.tensor_descriptor import TensorDescriptor - _ensure_tma_allocator() - assert state.is_contiguous(), "TMA state requires contiguous state" - assert state.stride(-1) == 1, "TMA state requires inner stride 1" - _state_flat = state.view(-1, state.shape[-1]) - _dstate_pow2 = triton.next_power_of_2(dstate) - state_tma_descriptor_write = TensorDescriptor.from_tensor( - _state_flat, block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], - ) - if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: - state_tma_descriptor_nowrite = state_tma_descriptor_write - else: - state_tma_descriptor_nowrite = TensorDescriptor.from_tensor( - _state_flat, block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], - ) - else: - state_tma_descriptor_write = state # dummy; all consuming constexprs False - state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False - - # Slot permutation — pointer + USE_PERM gate. When the caller provides - # a perm tensor the dl-family launches read pid_b through it; otherwise - # we pass any valid pointer (state_batch_indices) and USE_PERM=False so - # the kernel falls back to pid_grid. Sort-driven dispatch (write-first - # clustering) is opt-in per call; monolithic / dynamic ignore the flag. - if slot_perm is not None: - assert slot_perm.dtype in (torch.int32, torch.int64), ( - f"slot_perm must be int32/int64, got {slot_perm.dtype}" - ) - assert slot_perm.numel() >= batch, ( - f"slot_perm has {slot_perm.numel()} entries; need >= batch ({batch})" - ) - slot_perm_arg = slot_perm - use_perm = True - else: - # Any valid ptr — gated by USE_PERM=False at compile time. - slot_perm_arg = state_batch_indices if state_batch_indices is not None else state - use_perm = False - - # Grid for main kernels (M tiling × batch × nheads). - def main_grid(META): - return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) - - precomp_grid = (batch, nheads // heads_per_block) - d_strides = (D.stride(0), D.stride(1)) if D is not None else (0, 0) - - # ---- Launch helpers (close over locals) ------------------------------- - # Each helper is a thin closure that calls one Triton kernel with the - # full positional + kwarg argument list. Mode-dependent constexprs - # (write_checkpoint, early_out, rectangle) are passed in. - - def launch_replay_precompute(write_checkpoint: bool, early_out: bool, - reverse_perm: bool = False): - _checkpointing_precompute_kernel[precomp_grid]( - dt, dt_bias, A, B, C, - cb_scaled, decay_vec, - old_B, old_dt, old_dA_cumsum, - cache_buf_idx, prev_num_accepted_tokens, - state_batch_indices, slot_perm_arg, pad_slot_id, - T, max_window, dstate, nheads // ngroups, - dt.stride(0), dt.stride(1), dt.stride(2), - dt_bias.stride(0) if dt_bias is not None else 0, - A.stride(0), - B.stride(0), B.stride(1), B.stride(2), B.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - dt_softplus, - HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, - LAUNCH_WITH_PDL=launch_with_pdl, - LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, - HEADS_PER_BLOCK=heads_per_block, - WRITE_CHECKPOINT=write_checkpoint, - EARLY_OUT=early_out, - USE_PERM=use_perm, - REVERSE_PERM=reverse_perm, - num_warps=precompute_num_warps, - **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), - launch_pdl=launch_with_pdl, - ) - - def launch_rectangle_precompute(early_out: bool, reverse_perm: bool = False): - _rectangle_precompute_kernel[precomp_grid]( - dt, dt_bias, A, B, C, - cb_scaled, decay_vec, - old_B, old_dt, old_dA_cumsum, - cache_buf_idx, prev_num_accepted_tokens, - state_batch_indices, slot_perm_arg, pad_slot_id, - T, max_window, dstate, nheads // ngroups, - dt.stride(0), dt.stride(1), dt.stride(2), - dt_bias.stride(0) if dt_bias is not None else 0, - A.stride(0), - B.stride(0), B.stride(1), B.stride(2), B.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - dt_softplus, - HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, - LAUNCH_WITH_PDL=launch_with_pdl, - LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, - HEADS_PER_BLOCK=heads_per_block, - EARLY_OUT=early_out, - USE_PERM=use_perm, - REVERSE_PERM=reverse_perm, - num_warps=precompute_num_warps, - **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), - launch_pdl=launch_with_pdl, - ) - - def launch_dynamic_precompute(rectangle: bool): - _dynamic_precompute_kernel[precomp_grid]( - dt, dt_bias, A, B, C, - cb_scaled, decay_vec, - old_B, old_dt, old_dA_cumsum, - cache_buf_idx, prev_num_accepted_tokens, - state_batch_indices, pad_slot_id, - T, max_window, dstate, nheads // ngroups, - dt.stride(0), dt.stride(1), dt.stride(2), - dt_bias.stride(0) if dt_bias is not None else 0, - A.stride(0), - B.stride(0), B.stride(1), B.stride(2), B.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - dt_softplus, - HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, - LAUNCH_WITH_PDL=launch_with_pdl, - LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, - HEADS_PER_BLOCK=heads_per_block, - RECTANGLE=rectangle, - num_warps=precompute_num_warps, - **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), - launch_pdl=launch_with_pdl, - ) - - def launch_replay_main(write_checkpoint: bool, early_out: bool, - launch_dependent_kernels: bool = False, - reverse_perm: bool = False): - # Per-main knob selection: write vs nowrite branches use independent - # M / num_warps / num_stages / heads_per_block values. Grid is - # M-dependent so it must be a closure over the selected M. - _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE - _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE - _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE - # Per-path TMA descriptor: block_shape[0] must match the kernel's - # BLOCK_SIZE_M (`_bsm`); see the descriptor build block above. - _desc = (state_tma_descriptor_write if write_checkpoint - else state_tma_descriptor_nowrite) - def _main_grid_local(META, _bsm=_bsm): - return (triton.cdiv(dim, _bsm), batch, nheads) - _checkpointing_main_kernel[_main_grid_local]( - state, _desc, state_scales_arg, old_x, - old_B, old_dt, old_dA_cumsum, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - _bsm, - LAUNCH_WITH_PDL=use_internal_pdl, - PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, - QUANT_MAX=quant_max, - WRITE_CHECKPOINT=write_checkpoint, - EARLY_OUT=early_out, - LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, - USE_PERM=use_perm, - REVERSE_PERM=reverse_perm, - # Per-launch WC fixes which LOAD flag is "live"; pass write-load - # value when WC=True (NOWRITE flag dummy False), else converse. - USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), - USE_TMA_LOAD_NOWRITE=bool(_use_tma_replay_nowrite_load and not write_checkpoint), - USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), - num_warps=_nw, - **({"num_stages": _ns} if _ns else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), - launch_pdl=use_internal_pdl, - ) - - def launch_rectangle_main(early_out: bool, - launch_dependent_kernels: bool = False, - reverse_perm: bool = False): - # Rectangle is the nowrite-side path; use the nowrite-main knobs. - _bsm = BLOCK_SIZE_M_NOWRITE - _nw = NUM_WARPS_NOWRITE - _ns = NUM_STAGES_NOWRITE - def _main_grid_local(META, _bsm=_bsm): - return (triton.cdiv(dim, _bsm), batch, nheads) - # Rectangle is always the nowrite-side path; descriptor block_shape[0] - # must match BLOCK_SIZE_M_NOWRITE (= _bsm here). - _rectangle_main_kernel[_main_grid_local]( - state, state_tma_descriptor_nowrite, state_scales_arg, old_x, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, slot_perm_arg, pad_slot_id, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - _bsm, - LAUNCH_WITH_PDL=use_internal_pdl, - QUANT_MAX=quant_max, - EARLY_OUT=early_out, - LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, - USE_PERM=use_perm, - REVERSE_PERM=reverse_perm, - USE_TMA_LOAD=bool(_use_tma_rect_load), - num_warps=_nw, - **({"num_stages": _ns} if _ns else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), - launch_pdl=use_internal_pdl, - ) - - def launch_dynamic_main(rectangle: bool, - launch_dependent_kernels: bool = False): - # Dynamic mode uses a single BLOCK_SIZE_M (no M-split inside this - # kernel); BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE by the wrapper's tied - # convention, so the write-side descriptor matches. - _dynamic_main_kernel[main_grid]( - state, state_tma_descriptor_write, state_scales_arg, old_x, - old_B, old_dt, old_dA_cumsum, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, rand_seed, pad_slot_id, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - BLOCK_SIZE_M, - LAUNCH_WITH_PDL=use_internal_pdl, - PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, - QUANT_MAX=quant_max, - RECTANGLE=rectangle, - LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, - # 3 TMA flags. NOWRITE_LOAD picks rect-load vs replay-nowrite-load - # based on RECTANGLE constexpr (only one is reachable per compile). - USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load), - USE_TMA_LOAD_NOWRITE=bool(_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load), - USE_TMA_STORE=bool(_use_tma_replay_write_store), - num_warps=num_warps, - **({"num_stages": _num_stages} if _num_stages else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), - launch_pdl=use_internal_pdl, - ) - - # ---- launch_persistent_main ------------------------------------------ - # Persistent-CTA main kernel. Single launch covers `n_slots` slots - # starting at `slot_offset`. Caller invokes twice: once for the write - # half (slot_offset=0, n_slots=n_writes, write_checkpoint=True) and - # once for the nowrite half (slot_offset=n_writes, - # n_slots=batch-n_writes, write_checkpoint=False). Hard-sort - # contract: caller has pre-sorted slots so [0, n_writes) are writes - # and [n_writes, batch) are nowrites. - - # Resolve persistent-mode bench knobs. Defaults: cta_per_sm = 1 - # (one CTA per SM, matches upstream `_p_matmul_ogs.py`); num_loop_stages - # = 2 (matches in-tree `swiglu` precedent for non-dot persistent loops); - # flatten = True (canonical Triton 3.6 idiom); warp_specialize = False. - _num_sms = torch.cuda.get_device_properties(device).multi_processor_count - cta_per_sm_arg = _cta_per_sm if _cta_per_sm else 1 - num_persistent_arg = cta_per_sm_arg * _num_sms - num_loop_stages_arg = _num_loop_stages if _num_loop_stages else 2 - flatten_arg = True if _flatten is None else bool(_flatten) - warp_specialize_arg = False if _warp_specialize is None else bool(_warp_specialize) - # Per-launch work-item count. At small batch, total_work may be < the - # full persistent grid; capping `grid` at `min(NUM_PERSISTENT, total_work)` - # avoids launching empty CTAs that pay setup cost for no work. Correctness: - # the kernel's `tl.range(pid, total_work, NUM_PERSISTENT)` ensures each - # tile_id is covered exactly once across all live pids in [0, grid) when - # grid <= NUM_PERSISTENT (each CTA does 1 tile; loop step >= total_work - # exits immediately) AND when grid == NUM_PERSISTENT (each CTA loops over - # multiple tiles). NUM_PERSISTENT is now a runtime int (see kernel def - # docstring at _persistent_main_kernel) so changing cta_per_sm does NOT - # trigger a new Triton compile — same kernel binary, different loop step. - # (Named UPPERCASE for historical Triton-style consistency only; not - # constexpr.) - _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M - - def launch_persistent_main(write_checkpoint: bool, - n_writes_dev: torch.Tensor, - *, - host_n_writes: int | None = None, - skip_empty_halves: bool = True, - launch_dependent_kernels: bool = False, - rectangle: bool = False): - # `n_writes_dev` is a (1,) int32 device tensor; the kernel reads - # the count from device memory. `host_n_writes` is the same value - # known host-side (when available — pure scenarios) and lets us - # skip the launch entirely if its half is empty. In mix scenarios - # the host doesn't know n_writes per iter without a sync, so - # `host_n_writes is None` and `skip_empty_halves` is forced False - # — both halves always launch and the kernel processes whatever - # range device-n_writes implies. - if skip_empty_halves and host_n_writes is not None: - n_slots_for_kernel = host_n_writes if write_checkpoint else (batch - host_n_writes) - if n_slots_for_kernel <= 0: - return - # Per-main knob selection. The two persistent_main launches (write - # half vs nowrite half) get independent BLOCK_SIZE_M / num_warps / - # num_stages / cta_per_sm / num_loop_stages. See the per-main args - # block in the wrapper signature. - _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE - _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE - _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE - _cps = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE - _cps = _cps if _cps else 1 - _nls = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE - _nls = _nls if _nls else 2 - _num_persistent = _cps * _num_sms - _num_pid_m_local = (dim + _bsm - 1) // _bsm - # Grid sizing: cap at min(full persistent grid, actual total_work). - # `n_slots` for this launch is `host_n_writes` (write half) / `batch - - # host_n_writes` (nowrite half) when host knows it (pure); else upper - # bound `batch` for mix scenarios where host can't read n_writes_dev - # without a sync. Upper-bound is fine — the kernel's runtime check - # only iterates actual work; the only cost of overcounting is a few - # extra CTAs. - if host_n_writes is not None: - _n_slots_for_launch = host_n_writes if write_checkpoint else (batch - host_n_writes) - else: - _n_slots_for_launch = batch - _total_work_launch = max(1, _n_slots_for_launch * _num_pid_m_local * nheads) - grid = (min(_num_persistent, _total_work_launch),) - # Per-path TMA descriptor — block_shape[0] must match _bsm. - _desc = (state_tma_descriptor_write if write_checkpoint - else state_tma_descriptor_nowrite) - _persistent_main_kernel[grid]( - state, _desc, state_scales_arg, old_x, - old_B, old_dt, old_dA_cumsum, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, - n_writes_dev, batch, nheads, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - _bsm, - LAUNCH_WITH_PDL=use_internal_pdl, - PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, - QUANT_MAX=quant_max, - WRITE_CHECKPOINT=write_checkpoint, - LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, - USE_PERM=use_perm, - NUM_PERSISTENT=_num_persistent, - NUM_LOOP_STAGES=_nls, - FLATTEN=flatten_arg, - WARP_SPECIALIZE=warp_specialize_arg, - IS_DYNAMIC=False, - RECTANGLE=rectangle, - # 3 TMA flags. IS_DYNAMIC=False: WC fixed per launch; impl - # constexpr-folds the LOAD pick. When WC=True (write half), - # NOWRITE_LOAD is dummy False; when WC=False, WRITE_LOAD/STORE - # dummy False. NOWRITE_LOAD picks rect-load (RECTANGLE) or - # replay-nowrite-load. - USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), - USE_TMA_LOAD_NOWRITE=bool( - (_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load) - and not write_checkpoint - ), - USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), - num_warps=_nw, - **({"num_stages": _ns} if _ns else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), - launch_pdl=use_internal_pdl, - ) - - def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, - launch_dependent_kernels: bool = False, - rectangle: bool = False): - # Single-launch persistent kernel covering the whole batch with - # runtime per-slot WRITE_CHECKPOINT branch. No half-split, no - # n_writes needed (the kernel ignores n_writes_dev when - # IS_DYNAMIC=True; Triton DCEs the load). is_write is computed - # at runtime per work-item from the loaded PNAT. - # We still pass `n_writes_dev` (the same tensor the persistent_main - # path uses) so the kernel signature is uniform; the value is - # immaterial. - # Grid sizing: cap at total_work (= batch * num_pid_m * nheads) for - # the dynamic case (full-batch coverage); see launch_persistent_main - # comment for correctness rationale. - _total_work_launch = max(1, batch * _num_pid_m * nheads) - grid = (min(num_persistent_arg, _total_work_launch),) - # Persistent-dynamic kernel uses a single BLOCK_SIZE_M (same as the - # wrapper's BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE tied convention), so - # the write-side descriptor matches. Both write and nowrite slots - # in this kernel share that BSM. - _persistent_main_kernel[grid]( - state, state_tma_descriptor_write, state_scales_arg, old_x, - old_B, old_dt, old_dA_cumsum, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, - n_writes_dev, batch, nheads, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - BLOCK_SIZE_M, - LAUNCH_WITH_PDL=use_internal_pdl, - PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, - QUANT_MAX=quant_max, - WRITE_CHECKPOINT=False, - LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, - USE_PERM=use_perm, - NUM_PERSISTENT=num_persistent_arg, - NUM_LOOP_STAGES=num_loop_stages_arg, - FLATTEN=flatten_arg, - WARP_SPECIALIZE=warp_specialize_arg, - IS_DYNAMIC=True, - RECTANGLE=rectangle, - # 3 TMA flags. IS_DYNAMIC=True: is_write is runtime per slot; - # impl's load TMA picks per-slot (constexpr ternary becomes a - # runtime branch — both load forms emitted, ~negligible cost). - # NOWRITE_LOAD picks rect-load when RECTANGLE, else - # replay-nowrite-load. STORE only fires on runtime is_write. - USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load), - USE_TMA_LOAD_NOWRITE=bool( - _use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load - ), - USE_TMA_STORE=bool(_use_tma_replay_write_store), - num_warps=num_warps, - **({"num_stages": _num_stages} if _num_stages else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), - launch_pdl=use_internal_pdl, - ) - - # ---- Mode dispatch ---------------------------------------------------- - with torch.cuda.device(device.index): - if mode == "monolithic": - if use_rectangle: - launch_rectangle_precompute(early_out=False) - launch_rectangle_main(early_out=False) - else: - launch_replay_precompute(write_checkpoint=write_checkpoint, early_out=False) - launch_replay_main(write_checkpoint=write_checkpoint, early_out=False) - elif mode == "dynamic": - launch_dynamic_precompute(rectangle=rectangle_for_nowrite) - launch_dynamic_main(rectangle=rectangle_for_nowrite) - elif mode == "maindl": - # Shared dispatch precompute, doublelaunched main. Precompute - # runs once with per-slot dispatch (saves the second precomp - # empty-grid tax of doublelaunch). Mains stay split with - # EARLY_OUT so each retains its constexpr-specialized reg - # envelope. First main signals PDL dependents so the second - # main can start its setup while the first is still computing. - # Dynamic precompute doesn't support sort (no early-out to - # cluster), so the perm only flows into the two EARLY_OUT mains. - launch_dynamic_precompute(rectangle=rectangle_for_nowrite) - launch_replay_main(write_checkpoint=True, early_out=True, - launch_dependent_kernels=True) - if rectangle_for_nowrite: - launch_rectangle_main(early_out=True, reverse_perm=reverse_nowrite) - else: - launch_replay_main(write_checkpoint=False, early_out=True, - reverse_perm=reverse_nowrite) - elif mode == "dlgrouped": - # Same 4 kernels as doublelaunch but reordered: both precomputes - # first, then both mains. Lets the GPU run precomp1 || precomp2 - # in parallel (they're tiny grids) before the mains start, vs - # doublelaunch's interleaved precomp1→main1→precomp2→main2. - # First main signals PDL so the second main's setup overlaps. - # When slot_perm + reverse_nowrite are set, the nowrite-side - # walks the perm in reverse so both kernels front-load real work. - launch_replay_precompute(write_checkpoint=True, early_out=True) - if rectangle_for_nowrite: - launch_rectangle_precompute(early_out=True, reverse_perm=reverse_nowrite) - else: - launch_replay_precompute(write_checkpoint=False, early_out=True, - reverse_perm=reverse_nowrite) - launch_replay_main(write_checkpoint=True, early_out=True, - launch_dependent_kernels=True) - if rectangle_for_nowrite: - launch_rectangle_main(early_out=True, reverse_perm=reverse_nowrite) - else: - launch_replay_main(write_checkpoint=False, early_out=True, - reverse_perm=reverse_nowrite) - elif mode == "dl_write_only": - # Debug-only: just the write half of doublelaunch. EARLY_OUT=True - # means nowrite slots still pay the EO-gate tax (PNAT load + branch), - # but no nowrite-side kernels run. Used to isolate "is the sort - # regression in the write-side kernels?". - launch_replay_precompute(write_checkpoint=True, early_out=True) - launch_replay_main(write_checkpoint=True, early_out=True, - launch_dependent_kernels=False) - elif mode == "persistent_dynamic": - # Single-launch persistent kernel covering the full batch. - # Each work-item dispatches via runtime PNAT check (is_write = - # (pnat + T) > MAX). No n_writes/half-split — kernel ignores - # n_writes_dev when IS_DYNAMIC=True (Triton DCEs the load). - # We still need a valid pointer to satisfy the kernel arg - # signature; allocate or reuse `_n_writes_dev`. - n_writes_dev_local = ( - _n_writes_dev if _n_writes_dev is not None - else torch.zeros(1, dtype=torch.int32, device=device) - ) - launch_dynamic_precompute(rectangle=rectangle_for_nowrite) - launch_persistent_dynamic_main( - n_writes_dev_local, - launch_dependent_kernels=False, - rectangle=rectangle_for_nowrite, - ) - elif mode == "persistent_main": - # Persistent-CTA main kernel. Reuses maindl's precompute - # structure (one shared dynamic_precompute that dispatches - # per-slot at runtime based on PNAT) followed by two - # persistent_main launches (write half + nowrite half). - # - # Hard-sort contract: caller has pre-sorted slots host-side so - # PNAT is monotone (writes first). Pass the perm via - # slot_perm + USE_PERM. - # - # n_writes is read by the kernel from a (1,) int32 device - # tensor. The caller can provide: - # * _n_writes_dev only (mix): a pre-filled (1,) int32 tensor - # it updates per iter via pre_iter_fn outside the captured - # graph. host can't cheaply read it without a sync, so both - # halves always launch. - # * _n_writes only (non-graph callers, e.g. unit tests): host - # int. We allocate the scratch tensor on the fly. CANNOT - # be used inside CUDA-graph capture — alloc inside capture - # invalidates the stream. - # * Both (pure under graph capture): caller pre-allocates the - # tensor outside capture and tells us the host value too. - # We skip the internal allocation and apply host-skip when - # _persistent_skip_empty_halves=True. This is the - # production-equivalent path the bench's pure cells take. - if _n_writes_dev is not None: - n_writes_dev_local = _n_writes_dev # no allocation - if _n_writes is not None: - # Caller provided both: pure scenario with pre-allocated - # tensor. Use host_n_writes for the skip-empty fast path. - assert 0 <= _n_writes <= batch, ( - f"_n_writes={_n_writes} must be in [0, batch={batch}]" - ) - host_n_writes_local = _n_writes - skip_empty_local = _persistent_skip_empty_halves - else: - # Mix: host doesn't know n_writes without a sync. - host_n_writes_local = None - skip_empty_local = False - else: - # No pre-allocated tensor. Fall back to on-the-fly alloc - # from _n_writes (host int). NOT graph-capture-safe. - assert _n_writes is not None, ( - "mode='persistent_main' requires either _n_writes " - "(host int, non-graph callers) or _n_writes_dev (device " - "tensor, recommended for graph-capture callers)." - ) - assert 0 <= _n_writes <= batch, ( - f"_n_writes={_n_writes} must be in [0, batch={batch}]" - ) - n_writes_dev_local = torch.tensor( - [_n_writes], dtype=torch.int32, device=device, - ) - host_n_writes_local = _n_writes - skip_empty_local = _persistent_skip_empty_halves - # rectangle_for_nowrite=True: precompute populates cb_scaled - # for the rect path; nowrite half uses the rectangle impl; - # write half always replay-style (rect doesn't apply). - launch_dynamic_precompute(rectangle=rectangle_for_nowrite) - launch_persistent_main( - write_checkpoint=True, - n_writes_dev=n_writes_dev_local, - host_n_writes=host_n_writes_local, - skip_empty_halves=skip_empty_local, - launch_dependent_kernels=True, - rectangle=False, # write always replay-style - ) - launch_persistent_main( - write_checkpoint=False, - n_writes_dev=n_writes_dev_local, - host_n_writes=host_n_writes_local, - skip_empty_halves=skip_empty_local, - launch_dependent_kernels=False, - rectangle=rectangle_for_nowrite, - ) - else: # mode == "doublelaunch" - # Write half: always replay-style write. First main signals - # PDL dependents so the second precompute can start its setup - # while the first main is still computing. - launch_replay_precompute(write_checkpoint=True, early_out=True) - launch_replay_main(write_checkpoint=True, early_out=True, - launch_dependent_kernels=True) - # Nowrite half: rectangle if asked, else replay-nowrite. - if rectangle_for_nowrite: - launch_rectangle_precompute(early_out=True, reverse_perm=reverse_nowrite) - launch_rectangle_main(early_out=True, reverse_perm=reverse_nowrite) - else: - launch_replay_precompute(write_checkpoint=False, early_out=True, - reverse_perm=reverse_nowrite) - launch_replay_main(write_checkpoint=False, early_out=True, - reverse_perm=reverse_nowrite) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_refactored.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_refactored.py deleted file mode 100644 index 298de246193f..000000000000 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_refactored.py +++ /dev/null @@ -1,4935 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. -"""Standalone benchmark for replay_selective_state_update (Triton kernel). - -Suitable for nsight-compute (ncu) and nsight-systems (nsys) capture. - -Fixed model config: NVIDIA-Nemotron-3-Super-120B-A12B at TP=8 - nheads=16, head_dim=64, d_state=128, ngroups=1 - -mtp_len is the per-request sequence length processed by replay: in MTP it -equals num_draft_tokens + 1 target token, so --mtp-lengths 6 models 5 drafts -+ 1 target. - -Baseline kernel (--baseline [triton|flashinfer]): - Calls selective_state_update with T=mtp_len tokens and disable_state_update=True, - matching the MTP scoring pass in mamba2_mixer.py exactly. - -Timing methodology -================== - -All in-bench timing comes from CUPTI's Activity API (1 ns kernel -timestamps from the GPU profiling fabric). cudaEvent.elapsed_time() was -removed — its ~0.5 us resolution overshoots CUPTI by ~50% on short kernels -in graphs, and we have no other use for it here. See the CUPTI block -lower in this file for the timer source. - -Three modes: - - --cupti --cuda-graph (default) - Capture a small CUDA graph for the cell, replay it for warmup + timed - iterations, and read kernel start/end from CUPTI. Raw CUPTI buffers are - parsed out-of-process on the timed path, with a cached ordinal plan used - to keep only the kernels we care about. - - --cupti --no-cuda-graph - Eager loop with CUPTI. Per-kernel timestamps are still accurate, but - the per-iter SPAN (max(end) - min(start)) now includes the Python - launch latency BETWEEN consecutive kernels in run_fn (~100 µs on - Hopper/Blackwell). Graph capture and PDL hide that latency; eager - mode honestly reports it. For per-kernel timing in eager mode, look - at per_kernel.start_us/end_us in --json-detailed output rather than - the span percentiles. Useful when graph capture is undesirable. - - --no-cupti (with or without --cuda-graph) - No in-bench timing — just runs the kernels for an external profiler - (nsys / ncu) to time. In-process CUPTI conflicts with nsys's own - subscriber, so disable ours when wrapping in nsys. Bench output - reports zeros for median/p95/p99; trust the external trace. - -JSON output schema (--json-output PATH) -======================================= - -Designed to be parsed by collect.py / report.py without touching sqlite or -NVTX traces. Future agents: prefer reading this JSON over re-running nsys. - - { - "metadata": {timestamp, cmd, tp_size, warmup, iters, variant, cupti}, - "results": { - "": {median, p95, p99, n, iters_us, [n_writes_per_iter], [per_kernel]} - } - } - -Key format mirrors collect.py's kernel_data.json convention: - incremental/{batch}/{mtp}/{sd}/k{prev_k}/{sweep_parts}/tp{tp} - triton/{batch}/{mtp}/{sd}/tp{tp} - flashinfer/{batch}/{mtp}/{sd}/tp{tp} - - - is normalized: bf16 / fp16 / fp32 / int8 / int16 / fp8. - - is e.g. "M16_W1_S3_SR0_RECT0_WC1" — flags concatenated by - underscore in canonical (M, W, S, pW, pS, H, R, CT, SR, RECT, WC) order. - - All numeric values in microseconds (us). - -Per-record fields: - - median, p95, p99: span statistics (us). Span = max(kernel_end_ns) - - min(kernel_start_ns) across the iter's kernels — same convention as - nsys-derived collect.py used to use. - - n: number of timed iters that contributed. - - iters_us: list of length n, raw per-iter spans. - - n_writes_per_iter: for mix rows, list of length n with the number of - write-path slots in each timed iteration. - - per_kernel: {: {start_us: [...], end_us: [...]}} where - timestamps are RELATIVE to that iter's first kernel start, in us. Lets - you see PDL overlap directly without an external profiler. Only with - --json-detailed. - -Example usage: - # Basic sweep (default = --cupti, just summary stats) - python benchmark_replay_selective_state_update.py \\ - --batch-sizes 1,2,4 --mtp-lengths 1,4,8 --warmup 5 --iters 20 - - # JSON output, summary stats only (compact) - python benchmark_replay_selective_state_update.py \\ - --batch-sizes 16 --mtp-lengths 6 --json-output /tmp/out.json - - # JSON output, full per-iter / per-kernel data (for PDL analysis etc.) - python benchmark_replay_selective_state_update.py \\ - --batch-sizes 16 --mtp-lengths 6 \\ - --json-output /tmp/out.json --json-detailed - - # nsys capture (--no-cupti so our subscriber doesn't conflict) - nsys profile --capture-range=cudaProfilerApi \\ - python benchmark_replay_selective_state_update.py --profile --no-cupti - - # ncu capture (--no-cupti --no-cuda-graph: each kernel replayable solo) - ncu --target-processes all \\ - python benchmark_replay_selective_state_update.py --profile \\ - --no-cupti --no-cuda-graph \\ - --batch-sizes 1 --mtp-lengths 4 --warmup 5 --iters 5 -""" - -import argparse -import atexit -import ctypes -import importlib -import itertools -import json -import multiprocessing as mp -import os -import queue -import statistics -import sys -import threading -import time -from datetime import datetime -from multiprocessing import shared_memory -from pathlib import Path - -import numpy as np -import torch -from einops import repeat - - -def _import_mamba_kernels_fast(): - """Load kernel modules directly (~40s faster than a full tensorrt_llm init). - Use --full-import as the fallback if module dependencies change. - - Strategy: stub the parent packages (tensorrt_llm, tensorrt_llm._torch, - tensorrt_llm._torch.modules) in sys.modules with __path__ set, but do - NOT execute their __init__.py. Then load the leaf kernel modules. - When a kernel body imports e.g. tensorrt_llm._utils.get_sm_version, - Python's machinery resolves it against our stub's __path__ and loads - only _utils.py — skipping the heavy tensorrt_llm package init. - """ - import types - - repo_root = Path(__file__).resolve().parents[5] - trtllm_dir = repo_root / "tensorrt_llm" - mamba_pkg = "tensorrt_llm._torch.modules.mamba" - mamba_dir = trtllm_dir / "_torch" / "modules" / "mamba" - - def _stub_pkg(fqn: str, pkg_dir: Path): - """Register a stub package in sys.modules without running its - __init__.py. Sets __path__ so Python can resolve submodule imports - against the real directory on disk.""" - if fqn in sys.modules: - return - stub = types.ModuleType(fqn) - stub.__path__ = [str(pkg_dir)] - sys.modules[fqn] = stub - - # Stub the parent chain so `from tensorrt_llm._utils import ...` (and - # similar) work without triggering tensorrt_llm/__init__.py. - _stub_pkg("tensorrt_llm", trtllm_dir) - _stub_pkg("tensorrt_llm._torch", trtllm_dir / "_torch") - _stub_pkg("tensorrt_llm._torch.modules", trtllm_dir / "_torch" / "modules") - - def _load(mod_name: str, file_name: str): - fqn = f"{mamba_pkg}.{mod_name}" if mod_name else mamba_pkg - if fqn in sys.modules: - return sys.modules[fqn] - spec = importlib.util.spec_from_file_location( - fqn, - mamba_dir / file_name, - submodule_search_locations=[str(mamba_dir)] if file_name == "__init__.py" else [], - ) - mod = importlib.util.module_from_spec(spec) - sys.modules[fqn] = mod - spec.loader.exec_module(mod) - return mod - - # 1. Package __init__ (defines PAD_SLOT_ID = -1) - _load("", "__init__.py") - # 2. softplus helper (used by both kernel modules) - _load("softplus", "softplus.py") - # 3. The actual kernels - replay_mod = _load("replay_selective_state_update", "replay_selective_state_update.py") - checkpoint_mod = _load("checkpointing_state_update_refactored", "checkpointing_state_update_refactored.py") - base_mod = _load("selective_state_update", "selective_state_update.py") - conv1d_mod = _load("causal_conv1d_triton", "causal_conv1d_triton.py") - - return ( - replay_mod.replay_selective_state_update, - checkpoint_mod.checkpointing_state_update, - base_mod.selective_state_update, - conv1d_mod.causal_conv1d_update, - ) - - -def _import_mamba_kernels_full(): - """Import via the standard tensorrt_llm package (slow but safe).""" - from tensorrt_llm._torch.modules.mamba.causal_conv1d_triton import causal_conv1d_update - from tensorrt_llm._torch.modules.mamba.checkpointing_state_update_refactored import ( - checkpointing_state_update, - ) - from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( - replay_selective_state_update, - ) - from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update - - return ( - replay_selective_state_update, - checkpointing_state_update, - selective_state_update, - causal_conv1d_update, - ) - - -# Use fast import by default; --full-import parsed later but we need the -# functions at module level. Check sys.argv early. -if "--full-import" in sys.argv: - ( - replay_selective_state_update, - checkpointing_state_update, - selective_state_update, - causal_conv1d_update, - ) = _import_mamba_kernels_full() -else: - try: - ( - replay_selective_state_update, - checkpointing_state_update, - selective_state_update, - causal_conv1d_update, - ) = _import_mamba_kernels_fast() - except Exception as e: # noqa: BLE001 - exit loudly; don't hide a fast-import regression - print( - f"ERROR: fast import failed ({type(e).__name__}: {e})\n" - "Re-run with --full-import for the slow but stable path, " - "then file a bug or fix _import_mamba_kernels_fast.", - file=sys.stderr, - ) - sys.exit(1) - - -_VARIANT_FNS = { - "replay": lambda: replay_selective_state_update, - "checkpointing": lambda: checkpointing_state_update, -} - -# Model config defaults (Nemotron-3-Super-120B full model). -# --tp-size divides nheads and ngroups to get the per-GPU slice. -# TP=1: nheads=128, ngroups=8 -# TP=4: nheads=32, ngroups=2 -# TP=8: nheads=16, ngroups=1 (default) -NHEADS = 128 -HEAD_DIM = 64 -D_STATE = 128 -NGROUPS = 8 -TP_SIZE = 8 # default; overridden by --tp-size - -# L2 flush buffer: ~128 MB — larger than L2 on A100/H100/B200 -_L2_FLUSH_SIZE = 32 * 1024 * 1024 # float32 elements → 128 MB -_l2_flush: torch.Tensor | None = None - - -def _init_l2_flush() -> None: - global _l2_flush - _l2_flush = torch.empty(_L2_FLUSH_SIZE, dtype=torch.float32, device="cuda") - - -def _flush_l2() -> None: - """Evict L2 by writing to a large buffer then synchronising.""" - assert _l2_flush is not None - _l2_flush.fill_(0.0) - torch.cuda.synchronize() - - -def _resolve_prev_ks(args, mtp_len: int) -> list[int]: - """Resolve prev_k values for one mtp_len cell. - - Two input modes (mutually exclusive in spirit; absolute wins if both given): - --prev-tokens-int "0,10,11,16" → use literal integers, clamped to - [0, max_window] (where max_window is the cache T-axis capacity). - --prev-tokens-fracs "0,0.5,1.0" → fractions of mtp_len, clamped to - [0, mtp_len] (current behavior). - - For replay-style checkpointing the cache holds up to max_window old - tokens, so absolute integers are the right knob. Fractions are kept - for back-compat with prior placeholder runs. - """ - upper = getattr(args, "max_window", 0) or mtp_len - if getattr(args, "prev_tokens_int", None): - return sorted(set(max(0, min(upper, int(v))) for v in args.prev_tokens_int)) - return sorted( - set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs) - ) - - -# Tensor construction helpers - -# Module-level cache for tensor buffers shared across cells. Keyed by all -# the "fixed" dimensions (state_dtype, act_dtype, max_window, mtp_len, -# nheads, head_dim, d_state, ngroups). Within a key, the batch dim grows -# in place: if a new cell requests a batch <= cached max_batch, we return -# views (slices) of the existing tensors; if batch > cached max_batch, we -# realloc at the new batch (which becomes the new max). Tensors never shrink. -# -# Rationale: torch.randn/zeros for these tensor shapes at b=512 takes -# ~10-30ms per call. At ~895 cells/min with 5 different batch sizes, -# we were re-allocating every cell. Caching saves the bulk of that per-cell -# overhead, raising GPU util in the timing phase. -# -# Reset state lives in caller (state_work = state0.copy_), so cached state0 -# is purely a reference whose contents stay fixed once allocated. This is -# fine: it's only read by the reset path. -_TENSOR_CACHE: dict = {} - - -def _build_tensors( - batch: int, - mtp_len: int, - state_dtype: torch.dtype, - act_dtype: torch.dtype, - nheads: int, - head_dim: int, - d_state: int, - ngroups: int, - max_window: int | None = None, -): - """ - Build all tensors for one benchmark configuration. - - nheads/ngroups are already TP-split (i.e. full_nheads // tp_size). - - Returns: - state0 : (batch, nheads, head_dim, d_state) – initial SSM state - x, dt, B, C : (batch, mtp_len, ...) – token inputs for both kernels - A, dt_bias, D : SSM parameters (float32, tie_hdim strides) - prev_tokens : (batch,) - out_incr : pre-allocated output for replay kernel (batch, mtp_len, nheads, head_dim) - out_base : pre-allocated output for baseline kernel (batch, mtp_len, nheads, head_dim) - intermediate_states_buffer: for baseline kernel (batch, mtp_len, nheads, head_dim, d_state) - """ - device = "cuda" - - # Cache lookup — grow batch in place if needed; else return views. - cache_key = (state_dtype, act_dtype, max_window, mtp_len, - nheads, head_dim, d_state, ngroups) - cached = _TENSOR_CACHE.get(cache_key) - if cached is not None and cached["max_batch"] >= batch: - # Hit — return slices for current batch. - b = batch - return ( - cached["state0"][:b], - cached["state_scales0"][:b] if cached["state_scales0"] is not None else None, - cached["old_x"][:b], - cached["old_B"][:b], - cached["old_dt"][:b], - cached["old_dA_cumsum"][:b], - cached["cache_buf_idx"][:b], - cached["x"][:b], - cached["dt"][:b], - cached["B"][:b], - cached["C"][:b], - cached["A"], - cached["dt_bias"], - cached["D"], - cached["prev_tokens"][:b], - cached["slot_perm_buf"][:b], - cached["out_incr"][:b], - cached["out_base"][:b], - cached["intermediate_states_buffer"][:b], - cached["xbc_input"][:b], - cached["conv_state"][:b], - cached["conv_weight"], - cached["conv_bias"], - cached["d_inner"], - cached["conv_dim"], - ) - - # Miss or grow. Allocate at new max_batch (existing data, if any, is - # released — caller code re-fills via reset paths anyway). Rebind - # `batch` locally to alloc_batch so the existing allocation code below - # uses the larger size; keep request_batch for the final slice. - request_batch = batch - alloc_batch = batch if cached is None else max(batch, cached["max_batch"]) - batch = alloc_batch - - torch.manual_seed(42) - - # --- SSM parameters (float32, tie_hdim strides) --- - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) # stride(-1)=0, stride(-2)=0 - - dt_bias_base = torch.randn(nheads, device=device, dtype=torch.float32) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) # stride(-1)=0 - - D_base = torch.randn(nheads, device=device, dtype=torch.float32) - D = repeat(D_base, "h -> h p", p=head_dim) - - # --- SSM state --- - # Quantized dtypes need their own initializer (torch.randn doesn't accept - # int) and a parallel fp32 scales tensor (per-(head, dim) channel decode - # scale, broadcast over dstate). Quant state is filled with realistic- - # range values via fp32 → quant; scales are derived consistently so the - # initial state isn't garbage on dequant. - _QUANT_BENCH = { - torch.int8: 127.0, - torch.int16: 32767.0, - torch.float8_e4m3fn: 448.0, - } - if state_dtype in _QUANT_BENCH: - quant_max = _QUANT_BENCH[state_dtype] - state_fp32 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - amax = state_fp32.abs().amax(dim=-1) # (batch, nheads, head_dim) - encode_scale = quant_max / amax.clamp(min=1e-30) - state_scales0 = (1.0 / encode_scale).to(torch.float32) # decode scale - scaled = state_fp32 * encode_scale.unsqueeze(-1) - if state_dtype == torch.float8_e4m3fn: - state0 = scaled.clamp(-quant_max, quant_max).to(state_dtype) - else: - state0 = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) - else: - state0 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=state_dtype - ) - state_scales0 = None - - # --- Cache tensors for replay kernel --- - # max_window is the cache T-axis capacity; defaults to mtp_len (the - # placeholder/degenerate case where every step is a checkpoint step). - # For real replay-style checkpointing, max_window > mtp_len. - cache_T = max_window if max_window is not None else mtp_len - # old_x: single-buffered (cache, max_window, nheads, dim) - old_x = torch.randn(batch, cache_T, nheads, head_dim, device=device, dtype=act_dtype) - # old_B: double-buffered (cache, 2, max_window, ngroups, dstate) - old_B = torch.randn(batch, 2, cache_T, ngroups, d_state, device=device, dtype=act_dtype) - # old_dt: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous - old_dt = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) - # old_dA_cumsum: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous - old_dA_cumsum = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) - # cache_buf_idx: which buffer to read (0 or 1) - cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) - - # --- Token inputs (used by both replay and baseline kernels) --- - x = torch.randn(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - # dt must match D's dtype (fp32) for flashinfer — force it for all paths. - dt_base = torch.randn(batch, mtp_len, nheads, device=device, dtype=torch.float32) - dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) # tie_hdim - B = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) - C = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) - - # prev_tokens placeholder — overwritten per-run - prev_tokens = torch.zeros(batch, device=device, dtype=torch.int32) - # slot_perm placeholder — overwritten per-run by mix pre_iter_fn when - # sort_slots is enabled. Identity by default so cells that don't sort - # (or pure-batch cells) get a meaningful identity perm if the kernel - # ends up reading it (USE_PERM=False makes this path unused). - slot_perm_buf = torch.arange(batch, device=device, dtype=torch.int32) - - out_incr = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - out_base = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - - # intermediate_states_buffer is only consumed by the fp/baseline path; - # for quantized state dtypes we'll skip baselines entirely, so the buffer - # dtype falls back to fp32 to keep selective_state_update happy. - int_buffer_dtype = state_dtype if state_dtype not in _QUANT_BENCH else torch.float32 - intermediate_states_buffer = torch.zeros( - batch, mtp_len, nheads, head_dim, d_state, device=device, dtype=int_buffer_dtype - ) - - # --- Conv1d tensors (for --with-conv1d mode) --- - d_inner = nheads * head_dim - conv_dim = d_inner + 2 * ngroups * d_state - d_conv = 4 # conv kernel width for Nemotron/Mamba2 - - # xbc_input: (batch, conv_dim, mtp_len) — "hot" input from in_proj. - # Match production layout: in_proj output is (batch*mtp_len, conv_dim) - # contiguous, then .view(batch, mtp_len, conv_dim).transpose(1, 2) - # gives strides (mtp_len*conv_dim, 1, conv_dim) — NOT the standard - # (conv_dim*mtp_len, mtp_len, 1) of a freshly allocated 3D tensor. - # Conv1d preserves input strides in its output, so downstream split - # + view inherits the correct layout without needing .contiguous(). - xbc_input_flat = torch.randn(batch * mtp_len, conv_dim, device=device, dtype=act_dtype) - xbc_input = xbc_input_flat.view(batch, mtp_len, conv_dim).transpose(1, 2) - # conv_state: (batch, conv_dim, d_conv) — "cold" cache - conv_state = torch.randn(batch, conv_dim, d_conv, device=device, dtype=act_dtype) - # conv_weight: (conv_dim, d_conv) — parameter - conv_weight = torch.randn(conv_dim, d_conv, device=device, dtype=act_dtype) - # conv_bias: (conv_dim,) — parameter - conv_bias = torch.randn(conv_dim, device=device, dtype=act_dtype) - - # Store full-batch buffers in cache and return slices at request_batch. - _TENSOR_CACHE[cache_key] = { - "max_batch": alloc_batch, - "state0": state0, - "state_scales0": state_scales0, - "old_x": old_x, - "old_B": old_B, - "old_dt": old_dt, - "old_dA_cumsum": old_dA_cumsum, - "cache_buf_idx": cache_buf_idx, - "x": x, - "dt": dt, - "B": B, - "C": C, - "A": A, - "dt_bias": dt_bias, - "D": D, - "prev_tokens": prev_tokens, - "slot_perm_buf": slot_perm_buf, - "out_incr": out_incr, - "out_base": out_base, - "intermediate_states_buffer": intermediate_states_buffer, - "xbc_input": xbc_input, - "conv_state": conv_state, - "conv_weight": conv_weight, - "conv_bias": conv_bias, - "d_inner": d_inner, - "conv_dim": conv_dim, - } - rb = request_batch - return ( - state0[:rb], - state_scales0[:rb] if state_scales0 is not None else None, - old_x[:rb], - old_B[:rb], - old_dt[:rb], - old_dA_cumsum[:rb], - cache_buf_idx[:rb], - x[:rb], - dt[:rb], - B[:rb], - C[:rb], - A, - dt_bias, - D, - prev_tokens[:rb], - slot_perm_buf[:rb], - out_incr[:rb], - out_base[:rb], - intermediate_states_buffer[:rb], - xbc_input[:rb], - conv_state[:rb], - conv_weight, - conv_bias, - d_inner, - conv_dim, - ) - - -# ============================================================================= -# CUPTI in-process kernel timing -# -# Self-contained module-in-a-file. Reads kernel start/end timestamps directly -# from the GPU profiling fabric via CUPTI's Activity API (1 ns -# resolution), avoiding two pitfalls of the cuda-events path: -# -# 1. cudaEvent.elapsed_time() resolution (~0.5 us) is too coarse for the -# short kernels we care about, especially with PDL + cuda graphs at -# small batch — events recorded inside a graph have proven noisy. -# 2. nsys is the only known accurate alternative, but the -# profile-export-sqlite-parse pipeline is heavy and out-of-process. -# -# This is functionally equivalent to wrapping each cell in nsys, except it -# runs in the same benchmark process and sends raw activity buffers to a -# parser process instead of materializing Python objects in the CUPTI callback. -# ============================================================================= - - -# Substring match: kernels run_fn launches that we want to time. Mirrors -# the parser in scripts/.../collect.py so cupti and nsys-based outputs agree. -_CUPTI_KEEP_KERNEL_SUBSTRINGS = ( - "_replay_precompute", - "_checkpointing_precompute", - "_rectangle_precompute", - "_dynamic_precompute", - "_replay_state_update", - "_checkpointing_main", - "_rectangle_main", - "_dynamic_main", - "_persistent_main", - "selective_scan_update", - "selective_state_update", - "causal_conv1d_update", -) - - -def _kernels_per_iter_incremental( - mode: str, - with_conv1d: bool, - *, - persistent_skip_empty: bool = True, -) -> int: - """Expected number of CUPTI-tracked kernels per iter for the incremental - kernel chain, given the dispatch mode and the conv1d flag. - - Used to validate CUPTI record counts (no auto-inference — silent - mis-timing is the failure mode we're guarding against). - - `persistent_skip_empty=True` (today's behavior): the - `mode='persistent_main'` launch helper host-early-outs when its half - is empty (n_writes=0 or n_writes=batch in pure scenarios), so only - one of the two persistent_main_kernel launches actually fires per - iter. With `persistent_skip_empty=False` (future no-eo mode), both - halves always launch and K bumps by 1. - - `persistent_dynamic` always launches 1 main; not affected by the flag. - """ - if mode == "monolithic": - k = 2 # precomp + main - elif mode == "dynamic": - k = 2 # dynamic_precomp + dynamic_main - elif mode == "maindl": - k = 3 # 1 dynamic_precomp + 2 mains (write + nowrite) - elif mode in ("doublelaunch", "dlgrouped"): - k = 4 # 2 precomp + 2 main - elif mode == "dl_write_only": - k = 2 # 1 precomp + 1 main (write only) - elif mode == "persistent_dynamic": - k = 2 # 1 dynamic_precomp + 1 persistent_main - elif mode == "persistent_main": - k = 2 if persistent_skip_empty else 3 # see docstring - else: - raise ValueError(f"_kernels_per_iter_incremental: unknown mode {mode!r}") - if with_conv1d: - k += 1 - return k - - -def _kernels_per_iter_baseline(with_conv1d: bool) -> int: - """Expected kernels per iter for triton / flashinfer baselines. - - Both baselines run a single state-update kernel; `--with-conv1d` - prepends one conv1d kernel. - """ - return 2 if with_conv1d else 1 - - -_LIBCUPTI_CANDIDATES = ( - os.environ.get("CUPTI_LIBRARY_PATH"), - "/usr/local/lib/python3.12/dist-packages/nvidia/cu13/lib/libcupti.so.13", - "libcupti.so.13", - "libcupti.so", -) -_CUPTI_SUCCESS = 0 -_CUPTI_ERROR_MAX_LIMIT_REACHED = 12 -_CUPTI_ERROR_INVALID_KIND = 21 -_CUPTI_ACTIVITY_KIND_KERNEL = 3 -_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL = 10 -_CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER = 5 -_CUPTI_HOST_BUFFER_BYTES = 1024 * 1024 -_CUPTI_HOST_BUFFER_COUNT = 16 - -# Multiprocessing start method for compile-warmup + CUPTI parser children. -# Set in __main__ from --mp-start-method. "spawn" (default) is robust; each -# child re-imports torch/triton/etc (~15s). "forkserver" preloads once and -# forks cheaply (~1s/child) — see __main__ block for the preload setup. -_MP_START_METHOD = "spawn" -_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE = 1 -_DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX = 4 - - -def _load_libcupti() -> ctypes.CDLL: - errors = [] - for candidate in _LIBCUPTI_CANDIDATES: - if not candidate: - continue - try: - return ctypes.CDLL(candidate) - except OSError as exc: - errors.append(f"{candidate}: {exc}") - raise ImportError("Unable to load libcupti: " + "; ".join(errors)) - - -class _CuptiActivityKernel11Prefix(ctypes.Structure): - _pack_ = 1 - _fields_ = [ - ("kind", ctypes.c_int), - ("cache_config", ctypes.c_uint8), - ("shared_memory_config", ctypes.c_uint8), - ("registers_per_thread", ctypes.c_uint16), - ("partitioned_global_cache_requested", ctypes.c_int), - ("partitioned_global_cache_executed", ctypes.c_int), - ("start", ctypes.c_uint64), - ("end", ctypes.c_uint64), - ("completed", ctypes.c_uint64), - ("device_id", ctypes.c_uint32), - ("context_id", ctypes.c_uint32), - ("stream_id", ctypes.c_uint32), - ("grid_x", ctypes.c_int32), - ("grid_y", ctypes.c_int32), - ("grid_z", ctypes.c_int32), - ("block_x", ctypes.c_int32), - ("block_y", ctypes.c_int32), - ("block_z", ctypes.c_int32), - ("static_shared_memory", ctypes.c_int32), - ("dynamic_shared_memory", ctypes.c_int32), - ("local_memory_per_thread", ctypes.c_uint32), - ("local_memory_total", ctypes.c_uint32), - ("correlation_id", ctypes.c_uint32), - ("grid_id", ctypes.c_int64), - ("name", ctypes.c_void_p), - ("reserved0", ctypes.c_void_p), - ("queued", ctypes.c_uint64), - ("submitted", ctypes.c_uint64), - ("launch_type", ctypes.c_uint8), - ("is_shared_memory_carveout_requested", ctypes.c_uint8), - ("shared_memory_carveout_requested", ctypes.c_uint8), - ("padding", ctypes.c_uint8), - ("shared_memory_executed", ctypes.c_uint32), - ("graph_node_id", ctypes.c_uint64), - ] - - -def _configure_cupti_get_next_record(libcupti) -> None: - libcupti.cuptiActivityGetNextRecord.argtypes = [ - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.POINTER(ctypes.c_void_p), - ] - libcupti.cuptiActivityGetNextRecord.restype = ctypes.c_int - - -def _parse_cupti_buffer_ptr(libcupti, buffer_ptr: int, valid_size: int, *, include_names: bool): - records = [] - zero_ts_count = 0 - zero_ts_names: dict[str, int] = {} - record_ptr = ctypes.c_void_p(None) - while True: - result = libcupti.cuptiActivityGetNextRecord( - ctypes.c_void_p(buffer_ptr), - valid_size, - ctypes.byref(record_ptr), - ) - if result == _CUPTI_SUCCESS: - kind = ctypes.cast(record_ptr, ctypes.POINTER(ctypes.c_int)).contents.value - if kind not in (_CUPTI_ACTIVITY_KIND_KERNEL, _CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL): - continue - kernel = ctypes.cast(record_ptr, ctypes.POINTER(_CuptiActivityKernel11Prefix)).contents - name = None - if include_names: - if kernel.name: - name = ctypes.string_at(kernel.name).decode("utf-8", errors="replace") - else: - name = "?" - if kernel.start == 0 or kernel.end == 0: - zero_ts_count += 1 - if name is not None: - zero_ts_names[name] = zero_ts_names.get(name, 0) + 1 - continue - if include_names: - records.append(( - name, - int(kernel.start), - int(kernel.end), - int(kernel.correlation_id), - 0, - int(kernel.graph_node_id), - int(kernel.stream_id), - )) - else: - records.append(( - int(kernel.start), - int(kernel.end), - int(kernel.correlation_id), - int(kernel.graph_node_id), - int(kernel.stream_id), - )) - elif result == _CUPTI_ERROR_MAX_LIMIT_REACHED: - break - elif result == _CUPTI_ERROR_INVALID_KIND: - break - else: - raise RuntimeError(f"cuptiActivityGetNextRecord failed with CUptiResult={result}") - return records, zero_ts_count, zero_ts_names - - -def _apply_cupti_filter_plan(numeric_records, filter_plan): - if not filter_plan: - return [ - (None, start, end, corr, 0, graph_node_id, stream_id) - for start, end, corr, graph_node_id, stream_id in sorted(numeric_records) - ] - - filtered = [] - replay_idx = 0 - record_idx = 0 - for start, end, corr, graph_node_id, stream_id in sorted(numeric_records): - if replay_idx >= len(filter_plan): - break - records_per_replay, ordinal_names = filter_plan[replay_idx] - if record_idx < len(ordinal_names): - name = ordinal_names[record_idx] - if name is not None: - filtered.append((name, start, end, corr, 0, graph_node_id, stream_id)) - record_idx += 1 - if record_idx >= records_per_replay: - replay_idx += 1 - record_idx = 0 - return filtered - - -def _cupti_parser_worker(input_queue, output_queue, ready_event) -> None: - libcupti = _load_libcupti() - _configure_cupti_get_next_record(libcupti) - shared_blocks: dict[str, shared_memory.SharedMemory] = {} - records_by_generation: dict[int, list[tuple[int, int, int, int, int]]] = {} - zero_ts_by_generation: dict[int, int] = {} - ready_event.set() - while True: - item = input_queue.get() - if item is None: - break - kind = item[0] - if kind == "buffer": - _, generation, buffer_id, name, valid_size = item - shm = shared_blocks.get(name) - if shm is None: - shm = shared_memory.SharedMemory(name=name) - shared_blocks[name] = shm - shared_char = ctypes.c_char.from_buffer(shm.buf) - try: - parser_ptr = ctypes.addressof(shared_char) - records, zero_ts_count, _ = _parse_cupti_buffer_ptr( - libcupti, - parser_ptr, - valid_size, - include_names=False, - ) - records_by_generation.setdefault(generation, []).extend(records) - zero_ts_by_generation[generation] = zero_ts_by_generation.get(generation, 0) + zero_ts_count - ctypes.memset(parser_ptr, 0, len(shm.buf)) - except Exception as exc: # pragma: no cover - diagnostic worker path - output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) - finally: - del shared_char - output_queue.put({"kind": "buffer_done", "generation": generation, "buffer_id": buffer_id}) - elif kind == "finish": - if len(item) == 4: - _, generation, filter_plan, stats_request = item - else: - _, generation, filter_plan = item - stats_request = None - try: - raw_records = records_by_generation.pop(generation, []) - zero_ts_count = zero_ts_by_generation.pop(generation, 0) - filtered_records = _apply_cupti_filter_plan(raw_records, filter_plan) - stats = None - parser_stats_ms = 0.0 - stats_ready = stats_request is not None - if stats_request is not None: - stats_start_s = time.perf_counter() - stats = _stats_from_cupti_records( - filtered_records, - int(stats_request["warmup"]), - int(stats_request["iters"]), - str(stats_request["tag"]), - int(stats_request["expected_K"]), - zero_ts_count=zero_ts_count, - zero_ts_names={}, - include_details=bool(stats_request.get("include_details", True)), - ) - parser_stats_ms = 1000.0 * (time.perf_counter() - stats_start_s) - filtered_records = [] - output_queue.put({ - "kind": "finish_done", - "generation": generation, - "records": filtered_records, - "zero_ts_count": zero_ts_count, - "zero_ts_names": {}, - "raw_record_count": len(raw_records), - "stats": stats, - "stats_ready": stats_ready, - "parser_stats_ms": parser_stats_ms, - }) - except Exception as exc: # pragma: no cover - diagnostic worker path - output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) - else: - output_queue.put({"kind": "error", "generation": -1, "error": f"unknown parser message {kind!r}"}) - for shm in shared_blocks.values(): - shm.close() - - -class CuptiKernelTimer: - """Raw CUPTI Activity timer with out-of-process parsing for timed runs. - - CUPTI's callback gives us raw activity buffers. The callback only hands - shared-memory buffer metadata to a parser process, so the main process - avoids the cupti-python per-record object creation cost during the timed - path. A single local calibration replay may parse names in-process to - build an ordinal filter plan for a just-captured CUDA graph. - """ - - _instance = None - _import_error = None - - _request_callback_type = ctypes.CFUNCTYPE( - None, - ctypes.POINTER(ctypes.c_void_p), - ctypes.POINTER(ctypes.c_size_t), - ctypes.POINTER(ctypes.c_size_t), - ) - _complete_callback_type = ctypes.CFUNCTYPE( - None, - ctypes.c_void_p, - ctypes.c_uint32, - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.c_size_t, - ) - - @classmethod - def get(cls) -> "CuptiKernelTimer": - if cls._instance is not None: - return cls._instance - if cls._import_error is not None: - raise cls._import_error - try: - cls._instance = cls() - return cls._instance - except ImportError as exc: # pragma: no cover - env-dependent - cls._import_error = exc - raise - - def __init__(self) -> None: - self._libcupti = _load_libcupti() - self._configure_functions() - self._lock = threading.Lock() - self._shared_buffers: dict[int, shared_memory.SharedMemory] = {} - self._buffer_id_by_ptr: dict[int, int] = {} - self._free_buffer_ids: list[int] = [] - self._local_completed: list[tuple[int, int]] = [] - self._mode = "drop" - self._generation = 0 - self._finish_results: dict[int, dict] = {} - self._parser_errors: list[str] = [] - self._filter_plan = () - self._last_start_timing: dict[str, float] = {} - self._last_stop_timing: dict[str, float] = {} - self._current_flush_period_ms = 0 - self._mp_ctx = mp.get_context(_MP_START_METHOD) - # Retry parser-process spawn: concurrent bench instances on the same - # node race on POSIX named semaphores in /dev/shm — child can die in - # pickle.load with FileNotFoundError in SemLock._rebuild before - # signalling ready_event. Detect early-dead child via is_alive() so - # we don't waste the full timeout, and retry up to 3x with jitter. - last_err = None - for _spawn_attempt in range(3): - self._parse_input_queue = self._mp_ctx.Queue() - self._parse_output_queue = self._mp_ctx.Queue() - ready_event = self._mp_ctx.Event() - self._parse_process = self._mp_ctx.Process( - target=_cupti_parser_worker, - args=(self._parse_input_queue, self._parse_output_queue, ready_event), - ) - self._parse_process.start() - deadline = time.time() + 30.0 - spawn_ok = False - while time.time() < deadline: - if ready_event.wait(timeout=0.5): - spawn_ok = True - break - if not self._parse_process.is_alive(): - break - if spawn_ok: - last_err = None - break - last_err = (f"attempt {_spawn_attempt + 1}: " - f"alive={self._parse_process.is_alive()}, " - f"exitcode={self._parse_process.exitcode}") - try: - if self._parse_process.is_alive(): - self._parse_process.terminate() - self._parse_process.join(timeout=2.0) - except Exception: - pass - time.sleep(0.5 + 0.5 * _spawn_attempt) - if last_err is not None: - raise RuntimeError( - f"CUPTI parser process did not initialize after 3 attempts: {last_err}" - ) - - self._set_zeroed_host_buffer_attr() - for _ in range(_CUPTI_HOST_BUFFER_COUNT): - self._free_buffer_ids.append(self._allocate_shared_buffer()) - - self._request_callback = self._request_callback_type(self._request_buffer) - self._complete_callback = self._complete_callback_type(self._complete_buffer) - self._check(self._libcupti.cuptiActivityRegisterCallbacks( - self._request_callback, - self._complete_callback, - )) - self._check(self._libcupti.cuptiActivityEnable(_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL)) - atexit.register(self.close) - - def _configure_functions(self) -> None: - self._libcupti.cuptiActivityRegisterCallbacks.argtypes = [ - self._request_callback_type, - self._complete_callback_type, - ] - self._libcupti.cuptiActivityRegisterCallbacks.restype = ctypes.c_int - self._libcupti.cuptiActivityEnable.argtypes = [ctypes.c_int] - self._libcupti.cuptiActivityEnable.restype = ctypes.c_int - self._libcupti.cuptiActivityFlushAll.argtypes = [ctypes.c_uint32] - self._libcupti.cuptiActivityFlushAll.restype = ctypes.c_int - self._libcupti.cuptiActivityFlushPeriod.argtypes = [ctypes.c_uint32] - self._libcupti.cuptiActivityFlushPeriod.restype = ctypes.c_int - self._libcupti.cuptiActivitySetAttribute.argtypes = [ - ctypes.c_int, - ctypes.POINTER(ctypes.c_size_t), - ctypes.c_void_p, - ] - self._libcupti.cuptiActivitySetAttribute.restype = ctypes.c_int - _configure_cupti_get_next_record(self._libcupti) - - def _set_zeroed_host_buffer_attr(self) -> None: - value_obj = ctypes.c_uint8(1) - size_obj = ctypes.c_size_t(ctypes.sizeof(value_obj)) - result = self._libcupti.cuptiActivitySetAttribute( - _CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER, - ctypes.byref(size_obj), - ctypes.byref(value_obj), - ) - if result != _CUPTI_SUCCESS: - print( - "[WARN] CUPTI zeroed host-buffer attribute failed; " - f"continuing with default CUPTI buffer handling (CUptiResult={result}).", - file=sys.stderr, - ) - - def _check(self, result: int) -> None: - if result != _CUPTI_SUCCESS: - raise RuntimeError(f"CUPTI call failed with CUptiResult={result}") - - def _allocate_shared_buffer(self) -> int: - buffer_id = len(self._shared_buffers) - shm = shared_memory.SharedMemory(create=True, size=_CUPTI_HOST_BUFFER_BYTES) - shared_char = ctypes.c_char.from_buffer(shm.buf) - try: - ptr = ctypes.addressof(shared_char) - finally: - del shared_char - if ptr % 8 != 0: - shm.close() - shm.unlink() - raise RuntimeError("CUPTI shared-memory activity buffer was not 8-byte aligned") - self._shared_buffers[buffer_id] = shm - self._buffer_id_by_ptr[ptr] = buffer_id - return buffer_id - - def _buffer_ptr(self, buffer_id: int) -> int: - shm = self._shared_buffers[buffer_id] - shared_char = ctypes.c_char.from_buffer(shm.buf) - try: - return ctypes.addressof(shared_char) - finally: - del shared_char - - def _request_buffer(self, buffer, size, max_num_records) -> None: - with self._lock: - if self._free_buffer_ids: - buffer_id = self._free_buffer_ids.pop() - else: - buffer_id = self._allocate_shared_buffer() - ptr = self._buffer_ptr(buffer_id) - buffer[0] = ptr - size[0] = _CUPTI_HOST_BUFFER_BYTES - max_num_records[0] = 0 - - def _complete_buffer(self, context, stream_id, buffer, size, valid_size) -> None: - del context, stream_id, size - buffer_ptr = int(buffer) - valid_size_int = int(valid_size) - with self._lock: - mode = self._mode - generation = self._generation - buffer_id = self._buffer_id_by_ptr[buffer_ptr] - if valid_size_int == 0 or mode == "drop": - self._free_buffer_ids.append(buffer_id) - return - if mode == "local": - self._local_completed.append((buffer_id, valid_size_int)) - return - shm = self._shared_buffers[buffer_id] - self._parse_input_queue.put(("buffer", generation, buffer_id, shm.name, valid_size_int)) - - def _handle_parser_result(self, result: dict) -> None: - kind = result.get("kind") - if kind == "buffer_done": - with self._lock: - self._free_buffer_ids.append(int(result["buffer_id"])) - elif kind == "finish_done": - self._finish_results[int(result["generation"])] = result - elif kind == "error": - self._parser_errors.append(str(result.get("error"))) - - def _drain_parser_results(self) -> None: - while True: - try: - result = self._parse_output_queue.get_nowait() - except queue.Empty: - break - self._handle_parser_result(result) - - def is_generation_ready(self, generation: int) -> bool: - self._drain_parser_results() - return generation in self._finish_results or bool(self._parser_errors) - - def _flush(self, flag: int) -> None: - self._check(self._libcupti.cuptiActivityFlushAll(flag)) - - def _set_flush_period_ms(self, period_ms: int) -> None: - if period_ms == self._current_flush_period_ms: - return - self._check(self._libcupti.cuptiActivityFlushPeriod(period_ms)) - self._current_flush_period_ms = period_ms - - def _begin( - self, - mode: str, - filter_plan=(), - flush_period_ms: int = 0, - collect_timing: bool = False, - ) -> int: - start_timing: dict[str, float] = {} - with self._lock: - self._mode = "drop" - phase_start_s = time.perf_counter() if collect_timing else 0.0 - self._flush(1) - if collect_timing: - start_timing["forced_flush_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) - phase_start_s = time.perf_counter() if collect_timing else 0.0 - self._drain_parser_results() - if collect_timing: - start_timing["drain_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) - with self._lock: - self._generation += 1 - generation = self._generation - self._mode = mode - self._local_completed = [] - self._filter_plan = filter_plan - if flush_period_ms > 0: - phase_start_s = time.perf_counter() if collect_timing else 0.0 - self._set_flush_period_ms(flush_period_ms) - if collect_timing: - start_timing["period_enable_ms"] = 1000.0 * ( - time.perf_counter() - phase_start_s - ) - self._last_start_timing = start_timing - return generation - - def capture_names(self, replay_fn) -> tuple[list[tuple], int, dict]: - """Run a small calibration replay and parse kernel names locally.""" - self._begin("local") - replay_fn() - torch.cuda.synchronize() - self._flush(0) - records: list[tuple] = [] - zero_ts_count = 0 - zero_ts_names: dict[str, int] = {} - with self._lock: - completed = list(self._local_completed) - self._local_completed = [] - self._mode = "drop" - for buffer_id, valid_size in completed: - ptr = self._buffer_ptr(buffer_id) - recs, zeros, zero_names = _parse_cupti_buffer_ptr( - self._libcupti, - ptr, - valid_size, - include_names=True, - ) - records.extend(recs) - zero_ts_count += zeros - for name, count in zero_names.items(): - zero_ts_names[name] = zero_ts_names.get(name, 0) + count - ctypes.memset(ptr, 0, _CUPTI_HOST_BUFFER_BYTES) - with self._lock: - self._free_buffer_ids.append(buffer_id) - records.sort(key=lambda r: r[1]) - return records, zero_ts_count, zero_ts_names - - def start( - self, - filter_plan=(), - flush_period_ms: int = 0, - collect_timing: bool = False, - ) -> None: - self._begin("parser", filter_plan, flush_period_ms, collect_timing) - - def stop_async( - self, - collect_timing: bool = False, - stats_request: dict | None = None, - ) -> tuple[int, dict[str, float]]: - stop_timing: dict[str, float] = {} - generation = self._generation - phase_start_s = time.perf_counter() if collect_timing else 0.0 - self._set_flush_period_ms(0) - if collect_timing: - stop_timing["period_disable_ms"] = 1000.0 * ( - time.perf_counter() - phase_start_s - ) - phase_start_s = time.perf_counter() if collect_timing else 0.0 - self._flush(0) - if collect_timing: - stop_timing["flush_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) - with self._lock: - self._mode = "drop" - filter_plan = self._filter_plan - self._parse_input_queue.put(("finish", generation, filter_plan, stats_request)) - self._last_stop_timing = stop_timing - return generation, stop_timing - - def wait_for_generation_result( - self, - generation: int, - stop_timing: dict[str, float] | None = None, - collect_timing: bool = False, - ) -> dict: - if stop_timing is None: - stop_timing = {} - phase_start_s = time.perf_counter() if collect_timing else 0.0 - deadline = time.perf_counter() + 10.0 - while time.perf_counter() < deadline: - result = self._finish_results.pop(generation, None) - if result is not None: - if collect_timing: - stop_timing["parser_wait_ms"] = 1000.0 * ( - time.perf_counter() - phase_start_s - ) - stop_timing["total_ms"] = ( - stop_timing.get("period_disable_ms", 0.0) - + stop_timing.get("flush_ms", 0.0) - + stop_timing["parser_wait_ms"] - ) - self._last_stop_timing = stop_timing - return result - timeout_s = max(0.0, min(0.01, deadline - time.perf_counter())) - try: - parser_result = self._parse_output_queue.get(timeout=timeout_s) - except queue.Empty: - continue - self._handle_parser_result(parser_result) - if self._parser_errors: - raise RuntimeError("CUPTI parser process failed: " + "; ".join(self._parser_errors)) - raise TimeoutError("Timed out waiting for CUPTI parser process") - - def wait_for_generation( - self, - generation: int, - stop_timing: dict[str, float] | None = None, - collect_timing: bool = False, - ) -> tuple[list[tuple], int, dict, int]: - result = self.wait_for_generation_result(generation, stop_timing, collect_timing) - return ( - list(result["records"]), - int(result["zero_ts_count"]), - dict(result["zero_ts_names"]), - int(result["raw_record_count"]), - ) - - def stop(self, collect_timing: bool = False) -> tuple[list[tuple], int, dict, int]: - generation, stop_timing = self.stop_async(collect_timing) - return self.wait_for_generation(generation, stop_timing, collect_timing) - - def last_start_timing(self) -> dict[str, float]: - return dict(self._last_start_timing) - - def last_stop_timing(self) -> dict[str, float]: - return dict(self._last_stop_timing) - - def close(self) -> None: - parse_process = getattr(self, "_parse_process", None) - if parse_process is not None and parse_process.is_alive(): - self._parse_input_queue.put(None) - parse_process.join(timeout=5.0) - if parse_process.is_alive(): - parse_process.terminate() - parse_process.join(timeout=1.0) - for shm in getattr(self, "_shared_buffers", {}).values(): - try: - shm.close() - shm.unlink() - except FileNotFoundError: - pass - - -# ============================================================================= -# Timing helpers -# ============================================================================= - - -def _stats_from_spans(spans_us: list[float]) -> dict: - """Compute median / p95 / p99 / n from a per-iter span list.""" - s = sorted(spans_us) - return { - "median": statistics.median(s), - "p95": s[int(0.95 * len(s))], - "p99": s[int(0.99 * len(s))], - "n": len(s), - } - - -def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, - zero_ts_count: int = 0, - zero_ts_names: dict | None = None, - include_details: bool = True): - """Bin a flat CUPTI kernel record stream into per-iter spans + per-kernel - relative timestamps. Used by both graph and eager CUPTI paths. - - `records` are tuples (name, start_ns, end_ns, ...) — see CuptiKernelTimer. - `expected_K` is the kernels-per-iter count the caller declares; we - validate the CUPTI total matches `expected_K * (warmup + iters)` exactly. - On mismatch we dump per-name record counts so missing or extra kernels - are obvious (most common cause: a new dispatch mode whose kernels lack - a matching entry in `_CUPTI_KEEP_KERNEL_SUBSTRINGS`, silently filtering - them out). - """ - records = [ - r for r in records - if r[0] is not None and any(s in r[0] for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS) - ] - records.sort(key=lambda r: r[1]) # by start_ns - - total = len(records) - expected_iters = warmup + iters - expected_total = expected_K * expected_iters - if total != expected_total: - from collections import Counter - name_counts = dict(Counter(r[0] for r in records)) - # Non-fatal: skip this cell instead of killing the whole sweep. - # Mismatch may be a CUPTI dropped-records issue (rare configs), - # not necessarily a K-table bug. Log so the user can investigate - # the specific cell post-hoc; return None so the caller can skip - # writing a JSON row. - zero_msg = "" - if zero_ts_count: - zero_msg = ( - f" + {zero_ts_count} records with start/end=0 " - f"(dropped by callback, breakdown {zero_ts_names}). " - f"Total observed kernel records (timed + zero-ts) = " - f"{total + zero_ts_count} / {expected_total}." - ) - print( - f"[WARN] CUPTI capture mismatch for {tag!r}: expected " - f"{expected_K} kernels/iter × {expected_iters} iters " - f"(warmup+iters) = {expected_total} records, got {total}. " - f"Kernel record counts: {name_counts}.{zero_msg} SKIPPING cell.", - file=sys.stderr, - flush=True, - ) - # Per-record dump: (name, start_ns_rel, end_ns_rel, corr_id, graph_id, stream_id). - # Times relative to first record so absolute ns isn't drowning output. - # Limit dump to first 30 records to avoid flooding logs at high K. - if records: - t0_ns = records[0][1] - for i, r in enumerate(records[:30]): - # r = (name, start_ns, end_ns, corr_id, graph_id, graph_node_id, stream_id) - rel_start = (r[1] - t0_ns) / 1000.0 # us - rel_end = (r[2] - t0_ns) / 1000.0 - print( - f" rec[{i:3d}] name={r[0]!r} start={rel_start:.2f}us " - f"end={rel_end:.2f}us corr={r[3]} graph={r[4]} stream={r[6]}", - file=sys.stderr, - flush=True, - ) - if len(records) > 30: - print(f" ... ({len(records) - 30} more records elided)", - file=sys.stderr, flush=True) - return None - K = expected_K - timed = records[warmup * K:] - - spans_us: list[float] = [] - per_kernel: dict[str, dict[str, list[float]]] = {} - for i in range(iters): - chunk = timed[i * K:(i + 1) * K] - iter_start_ns = min(r[1] for r in chunk) - iter_end_ns = max(r[2] for r in chunk) - spans_us.append((iter_end_ns - iter_start_ns) / 1000.0) - if include_details: - for r in chunk: - name = r[0] - slot = per_kernel.setdefault(name, {"start_us": [], "end_us": []}) - slot["start_us"].append((r[1] - iter_start_ns) / 1000.0) - slot["end_us"].append((r[2] - iter_start_ns) / 1000.0) - - out = _stats_from_spans(spans_us) - out["iters_us"] = spans_us - if include_details: - out["per_kernel"] = per_kernel - return out - - -_PRE_GRAPH_WARMUP_ITERS = 1 -_CUPTI_FILTER_PLAN_CACHE: dict[tuple, tuple[int, tuple[str | None, ...]]] = {} - - -class _HostTiming: - def __init__(self, enabled: bool) -> None: - self.enabled = enabled - self.values: dict[str, float | int | bool] = {} - self._total_start_s = time.perf_counter() if enabled else 0.0 - self._phase_start_s = 0.0 - - def start(self) -> None: - if self.enabled: - self._phase_start_s = time.perf_counter() - - def stop(self, key: str) -> None: - if self.enabled: - self.values[key] = 1000.0 * (time.perf_counter() - self._phase_start_s) - - def add(self, key: str, value: float | int | bool) -> None: - if self.enabled: - self.values[key] = value - - def stop_total(self) -> None: - if self.enabled: - self.values["total_ms"] = 1000.0 * (time.perf_counter() - self._total_start_s) - - def attach(self, stats: dict | None) -> None: - if self.enabled and stats is not None: - stats["host_timing"] = self.values - - -class _PendingCuptiStats: - - def __init__( - self, - timer: CuptiKernelTimer, - generation: int, - stop_timing: dict[str, float], - host_timing: _HostTiming, - *, - warmup: int, - iters: int, - tag: str, - expected_K: int, - expected_raw_record_count: int, - ) -> None: - self._timer = timer - self._generation = generation - self._stop_timing = stop_timing - self._host_timing = host_timing - self._warmup = warmup - self._iters = iters - self._tag = tag - self._expected_K = expected_K - self._expected_raw_record_count = expected_raw_record_count - - def is_ready(self) -> bool: - return self._timer.is_generation_ready(self._generation) - - def resolve(self) -> dict | None: - result = self._timer.wait_for_generation_result( - self._generation, - self._stop_timing, - collect_timing=self._host_timing.enabled, - ) - for key, value in self._timer.last_stop_timing().items(): - self._host_timing.add(f"cupti_stop_{key}", value) - raw_record_count = int(result["raw_record_count"]) - if raw_record_count != self._expected_raw_record_count: - print( - f"[WARN] CUPTI raw-record mismatch for {self._tag!r}: expected " - f"{self._expected_raw_record_count}, got {raw_record_count}. SKIPPING cell.", - file=sys.stderr, - ) - return None - - if result.get("stats_ready"): - stats = result.get("stats") - self._host_timing.add("stats_ms", 0.0) - self._host_timing.add("parser_stats_ms", float(result.get("parser_stats_ms", 0.0))) - else: - self._host_timing.start() - stats = _stats_from_cupti_records( - list(result["records"]), - self._warmup, - self._iters, - self._tag, - self._expected_K, - zero_ts_count=int(result["zero_ts_count"]), - zero_ts_names=dict(result["zero_ts_names"]), - ) - self._host_timing.stop("stats_ms") - self._host_timing.attach(stats) - return stats - - -def _target_name_or_none(name: str | None) -> str | None: - if name is None: - return None - if any(s in name for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS): - return name - return None - - -def _capture_group_graph( - args, - run_fn, - reset_fn, - group_iters: int, - graph_pre_iter_fn=None, -) -> torch.cuda.CUDAGraph: - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - for j in range(group_iters): - if graph_pre_iter_fn is not None: - graph_pre_iter_fn(j) - reset_fn() - if args.l2_flush: - _l2_flush.fill_(0.0) - run_fn() - return graph - - -def _graph_group_iters(args, total_iters: int, pre_iter_fn, pre_iter_group_factory) -> int: - """Pick the graph-group size unconditionally; the caller is expected to - round total_iters up to a multiple of this so all iters fit in clean - replays. Sample arrays are pre-padded at allocation (see _sample_pnat - call site) so the per-replay window can index past the user-requested - iter count by up to group_iters-1 extra samples. - """ - if pre_iter_fn is not None and pre_iter_group_factory is None: - # Per-iter callback without a group-factory: can't batch. - return 1 - requested = getattr(args, "cuda_graph_group_iters", None) - if requested is None: - return ( - _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX - if pre_iter_group_factory is not None - else _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE - ) - return max(1, int(requested)) - - -def _get_cupti_filter_plan(timer: CuptiKernelTimer, graph, cache_key: tuple | None, - group_iters: int) -> tuple[int, tuple[str | None, ...]]: - full_cache_key = None if cache_key is None else (cache_key, group_iters) - if full_cache_key is not None: - cached = _CUPTI_FILTER_PLAN_CACHE.get(full_cache_key) - if cached is not None: - return cached - - records, zero_ts_count, zero_ts_names = timer.capture_names(graph.replay) - if zero_ts_count: - print( - f"[WARN] CUPTI calibration saw {zero_ts_count} zero-timestamp records " - f"(breakdown {zero_ts_names}); continuing with nonzero records.", - file=sys.stderr, - ) - ordinal_names = tuple(_target_name_or_none(r[0]) for r in records) - target_count = sum(name is not None for name in ordinal_names) - if target_count == 0: - raise RuntimeError("CUPTI calibration did not find any target kernel records") - plan = (len(records), ordinal_names) - if full_cache_key is not None: - _CUPTI_FILTER_PLAN_CACHE[full_cache_key] = plan - return plan - - -def _time_kernel_cuda_graph( - args, - run_fn, - reset_fn, - tag: str, - *, - expected_K: int, - pre_iter_fn=None, - pre_iter_group_factory=None, - iters_override: int | None = None, - cupti_plan_key: tuple | None = None, -) -> dict: - """CUDA-graph CUPTI timer (graph-per-iter design). - - Captures one CUDA graph holding a small group of logical iterations - (per-iter setup + reset + l2_flush + run_fn) and replays it enough - times to cover `warmup + iters`. - - Why graph-per-iter (vs the older "one giant graph holding all iters" - design): instantiating a CUDA graph is expensive — proportional to - graph size — so a single small graph instantiated once is much - cheaper than one big graph instantiated for each cell of a sweep. - Replays are cheap regardless. - - Mix cells use a per-replay device window: an outside-graph copy loads - the next group of PNAT/n_writes samples, then graph-captured per-iter - copies update kernel inputs before each reset + L2 flush + run. - - Pre-graph eager warmup: forces PyTorch's caching allocator - + Triton's autotune cache to settle before capture so the graph - doesn't bake in init-only allocations. - - ``iters_override`` (if not None) overrides ``args.iters`` for this - call. Used to give mix scenarios a higher iter count than pure - (more iters = more independent mix draws averaged in). - """ - host_timing = _HostTiming(bool(getattr(args, "host_timing", False))) - timer = CuptiKernelTimer.get() - warmup = args.warmup - iters = iters_override if iters_override is not None else args.iters - - # Pre-graph eager warmup: full per-iter chain once. This settles - # Triton/PyTorch setup and wrapper-side intermediate allocations; - # skipping it risks lazy work leaking into graph capture. - warmup_iters = _PRE_GRAPH_WARMUP_ITERS - host_timing.add("pre_graph_warmup_iters", warmup_iters) - host_timing.start() - for _ in range(warmup_iters): - reset_fn() - if pre_iter_fn is not None: - pre_iter_fn(0) - run_fn() - if warmup_iters > 0: - torch.cuda.synchronize() - host_timing.stop("pre_graph_warmup_ms") - - total_iters = warmup + iters - group_iters = _graph_group_iters(args, total_iters, pre_iter_fn, pre_iter_group_factory) - # Args are rounded at argparse-time so warmup+iters/mix_iters are already - # multiples of the relevant group_iters. Assert here to catch any caller - # bypassing argparse. - assert total_iters % group_iters == 0, ( - f"total_iters={total_iters} not a multiple of group_iters={group_iters}; " - f"args.warmup/iters/mix_iters should be rounded post-argparse." - ) - pre_replay_fn = None - graph_pre_iter_fn = None - if pre_iter_group_factory is not None and group_iters > 1: - pre_replay_fn, graph_pre_iter_fn = pre_iter_group_factory(group_iters) - - # Reset just before capture so warmup state changes don't bleed in. - host_timing.start() - reset_fn() - torch.cuda.synchronize() - host_timing.stop("pre_capture_reset_ms") - - # Capture a small group of identical logical iterations. Mix/pre_iter - # cells can group when they provide a graph-side pre-iter updater backed - # by a per-replay device window. - host_timing.start() - g = _capture_group_graph(args, run_fn, reset_fn, group_iters, graph_pre_iter_fn) - host_timing.stop("graph_capture_ms") - - if pre_replay_fn is not None: - host_timing.start() - pre_replay_fn(0) - torch.cuda.synchronize() - host_timing.stop("graph_preload_ms") - - plan_cache_key = None if cupti_plan_key is None else (cupti_plan_key, group_iters) - host_timing.add("cupti_plan_cached", ( - plan_cache_key is not None and plan_cache_key in _CUPTI_FILTER_PLAN_CACHE - )) - host_timing.start() - records_per_replay, ordinal_names = _get_cupti_filter_plan( - timer, - g, - cupti_plan_key, - group_iters, - ) - host_timing.stop("cupti_plan_ms") - target_count = sum(name is not None for name in ordinal_names) - expected_targets_per_replay = expected_K * group_iters - if target_count != expected_targets_per_replay: - print( - f"[WARN] CUPTI calibration mismatch for {tag!r}: expected " - f"{expected_targets_per_replay} target records in a {group_iters}-iter graph replay, " - f"got {target_count} target records out of {records_per_replay} total records.", - file=sys.stderr, - ) - - # Time: replay the grouped graph enough times to cover warmup+iters. - # Mix cells preload one device window per replay on the same stream. - # CUPTI records every kernel launch; _stats_from_cupti_records - # validates against expected_K and slices warmup off the front. - graph_replays = total_iters // group_iters - filter_plan = ((records_per_replay, ordinal_names),) * graph_replays - cupti_flush_period_ms = max(0, int(getattr(args, "cupti_flush_period_ms", 0))) - host_timing.start() - timer.start( - filter_plan, - flush_period_ms=cupti_flush_period_ms, - collect_timing=host_timing.enabled, - ) - host_timing.stop("cupti_start_ms") - for key, value in timer.last_start_timing().items(): - host_timing.add(f"cupti_start_{key}", value) - torch.cuda.nvtx.range_push(tag) - host_timing.start() - for i in range(graph_replays): - if pre_replay_fn is not None: - pre_replay_fn(i) - elif pre_iter_fn is not None: - pre_iter_fn(i) - g.replay() - host_timing.stop("graph_enqueue_ms") - host_timing.start() - torch.cuda.synchronize() - host_timing.stop("graph_sync_ms") - torch.cuda.nvtx.range_pop() - expected_raw_record_count = records_per_replay * graph_replays - host_timing.start() - if int(getattr(args, "cupti_defer_depth", 1)) > 1: - generation, stop_timing = timer.stop_async( - collect_timing=host_timing.enabled, - stats_request={ - "warmup": warmup, - "iters": iters, - "tag": tag, - "expected_K": expected_K, - "include_details": bool(getattr(args, "json_detailed", False)), - }, - ) - host_timing.stop("cupti_stop_ms") - for key, value in timer.last_stop_timing().items(): - host_timing.add(f"cupti_stop_{key}", value) - host_timing.stop_total() - host_timing.add("graph_group_iters", group_iters) - host_timing.add("graph_replays", graph_replays) - host_timing.add("cupti_records_per_replay", records_per_replay) - host_timing.add("cupti_target_records_per_replay", target_count) - host_timing.add("cupti_raw_records_expected", expected_raw_record_count) - host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) - return _PendingCuptiStats( - timer, - generation, - stop_timing, - host_timing, - warmup=warmup, - iters=iters, - tag=tag, - expected_K=expected_K, - expected_raw_record_count=expected_raw_record_count, - ) - - records, zero_ts_count, zero_ts_names, raw_record_count = timer.stop( - collect_timing=host_timing.enabled, - ) - host_timing.stop("cupti_stop_ms") - for key, value in timer.last_stop_timing().items(): - host_timing.add(f"cupti_stop_{key}", value) - if raw_record_count != expected_raw_record_count: - print( - f"[WARN] CUPTI raw-record mismatch for {tag!r}: expected " - f"{records_per_replay} total records/replay × {graph_replays} replays " - f"= {expected_raw_record_count}, got {raw_record_count}. SKIPPING cell.", - file=sys.stderr, - ) - return None - - host_timing.start() - stats = _stats_from_cupti_records( - records, - warmup, - iters, - tag, - expected_K, - zero_ts_count=zero_ts_count, - zero_ts_names=zero_ts_names, - include_details=bool(getattr(args, "json_detailed", False)), - ) - host_timing.stop("stats_ms") - host_timing.stop_total() - host_timing.add("graph_group_iters", group_iters) - host_timing.add("graph_replays", graph_replays) - host_timing.add("cupti_records_per_replay", records_per_replay) - host_timing.add("cupti_target_records_per_replay", target_count) - host_timing.add("cupti_raw_records", raw_record_count) - host_timing.add("cupti_raw_records_expected", expected_raw_record_count) - host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) - host_timing.attach(stats) - return stats - - -def _time_kernel_eager( - args, - run_fn, - reset_fn, - tag: str, - *, - expected_K: int, - pre_iter_fn=None, - iters_override: int | None = None, - cupti_plan_key: tuple | None = None, -) -> dict: - """Non-graph CUPTI timer (for ncu wrapping, debugging, etc.). - - Each iter runs serially with sync between, but kernel start/end still - come from CUPTI — same accuracy as the graph path, just slower per-iter - (extra Python + sync overhead). - """ - host_timing = _HostTiming(bool(getattr(args, "host_timing", False))) - timer = CuptiKernelTimer.get() - warmup = args.warmup - iters = iters_override if iters_override is not None else args.iters - - del cupti_plan_key - - def _run_eager_loop(): - torch.cuda.nvtx.range_push(tag) - # Unified warmup+iters loop; CUPTI filters by warmup count internally. - for i in range(warmup + iters): - reset_fn() - if args.l2_flush: - _flush_l2() # includes synchronize - if pre_iter_fn is not None: - pre_iter_fn(i) - run_fn() - torch.cuda.synchronize() - torch.cuda.nvtx.range_pop() - - host_timing.start() - records, zero_ts_count, zero_ts_names = timer.capture_names(_run_eager_loop) - host_timing.stop("timed_loop_and_cupti_parse_ms") - - host_timing.start() - stats = _stats_from_cupti_records( - records, - warmup, - iters, - tag, - expected_K, - zero_ts_count=zero_ts_count, - zero_ts_names=zero_ts_names, - include_details=bool(getattr(args, "json_detailed", False)), - ) - host_timing.stop("stats_ms") - host_timing.stop_total() - host_timing.attach(stats) - return stats - - -def _run_kernel_untimed(args, run_fn, reset_fn, tag: str) -> dict: - """No in-bench timing: just run the kernels for an external profiler - (nsys / ncu) to time externally. Returns a stats dict full of zeros so - downstream code (table, JSON) doesn't break. - - Note: pre_iter_fn / iters_override aren't plumbed here yet — mix-mode - benchmarking relies on CUPTI. Add when a use-case lands. - """ - warmup = args.warmup - iters = args.iters - - if args.cuda_graph: - # Eager warmup before capture (Triton autotune) - reset_fn(); run_fn(); torch.cuda.synchronize() - reset_fn(); torch.cuda.synchronize() - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - for _ in range(warmup + iters): - reset_fn() - if args.l2_flush: - _l2_flush.fill_(0.0) - run_fn() - torch.cuda.synchronize() - torch.cuda.nvtx.range_push(tag) - g.replay() - torch.cuda.synchronize() - torch.cuda.nvtx.range_pop() - else: - torch.cuda.nvtx.range_push(tag) - for _ in range(warmup + iters): - reset_fn() - if args.l2_flush: - _flush_l2() - run_fn() - torch.cuda.synchronize() - torch.cuda.nvtx.range_pop() - - spans_us = [0.0] * iters - out = _stats_from_spans(spans_us) - out["iters_us"] = spans_us - out["per_kernel"] = {} - return out - - -def _time_kernel( - args, run_fn, reset_fn, tag: str, - *, - expected_K: int, - pre_iter_fn=None, - pre_iter_group_factory=None, - iters_override: int | None = None, - cupti_plan_key: tuple | None = None, -) -> dict: - """Dispatch to graph-CUPTI / eager-CUPTI / no-timer path. - - --cupti: in-process CUPTI Activity API timing (default). Use --no-cupti - when running under nsys (in-process CUPTI conflicts with nsys's own - subscriber); the bench then runs the kernels for nsys to time externally. - - `expected_K` is the kernels-per-iter count the caller declares - (computed via _kernels_per_iter_*). CUPTI paths validate against it - explicitly; the no-timer fallback ignores it (no records to validate). - """ - if not getattr(args, "cupti", True): - if pre_iter_fn is not None: - raise RuntimeError( - "_time_kernel: pre_iter_fn requires CUPTI (mix-mode); " - "got --no-cupti. Re-run with CUPTI on or plumb pre_iter_fn " - "through _run_kernel_untimed." - ) - return _run_kernel_untimed(args, run_fn, reset_fn, tag) - if args.cuda_graph: - return _time_kernel_cuda_graph( - args, run_fn, reset_fn, tag, - expected_K=expected_K, - pre_iter_fn=pre_iter_fn, - pre_iter_group_factory=pre_iter_group_factory, - iters_override=iters_override, - cupti_plan_key=cupti_plan_key, - ) - return _time_kernel_eager( - args, run_fn, reset_fn, tag, - expected_K=expected_K, - pre_iter_fn=pre_iter_fn, - iters_override=iters_override, - cupti_plan_key=cupti_plan_key, - ) - - -# Per-config benchmark (consolidated baseline + replay) - - -def _warm_one_config(args, cfg, baseline_fn) -> None: - """Module-level worker for the compile-warmup process pool. - - Module-level so ProcessPoolExecutor can pickle it (nested functions - aren't picklable). Each worker process holds its own GIL → no - serialization between concurrent compiles. - - ``cfg`` is a tuple of (outer_cfg, inner_overrides_or_list): - * outer_cfg = (batch, mtp_len, prev_ks, state_dtype, act_dtype, - sr_mode, rect, write_ckpt, mode, - sort_slots, reverse_nowrite, hardcode_sort) - * inner_overrides_or_list = dict of args attribute name -> value-string, - OR a list of such dicts. In the list form (CPS-grouped task) the - worker compiles each entry sequentially within the same process so - Triton's in-process kernel cache catches value-spec hits across - related entries (e.g. CPS={1,2} and {4,8} each form a `div_by_16` - spec bucket; the second compile in a bucket short-circuits). - - ``baseline_fn`` is optional — when ``None``, only the checkpointing - kernel is warmed (the baseline-selection kernel can be warmed once in - the parent if needed). This lets us avoid pickling C-extension - function references across processes. - """ - outer_cfg, inner_overrides_or_list = cfg - overrides_list = (inner_overrides_or_list - if isinstance(inner_overrides_or_list, list) - else [inner_overrides_or_list]) - (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, - rect, write_ckpt, mode, sort_slots, reverse_nowrite, hardcode_sort) = outer_cfg - import argparse as _ap - for inner_overrides in overrides_list: - # Fresh clone per entry: prevents knob-value leakage between - # consecutive cells in a CPS-grouped task (entries may set - # different non-CPS knobs in degenerate edge cases). - args_copy = _ap.Namespace(**vars(args)) - for k, v in inner_overrides.items(): - setattr(args_copy, k, v) - _bench_config( - args_copy, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, - sr_mode=sr_mode, rectangle_for_nowrite=rect, - write_checkpoint=write_ckpt, mode=mode, - sort_slots=sort_slots, reverse_nowrite=reverse_nowrite, - hardcode_sort=hardcode_sort, - warmup_only=True, - ) - - -def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, - baseline_fn, max_workers: int) -> None: - _cw_t0 = time.perf_counter() - def _cw(label: str) -> None: - dt = time.perf_counter() - _cw_t0 - print(f"[compile-warmup] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) - _cw("entered _compile_warmup_phase") - """Parallel compile-warmup using a ProcessPoolExecutor with `spawn` - start method. - - Each worker process holds its own GIL and its own CUDA context, so - Triton compiles (Python AST/codegen + LLVM/ptxas) run truly in - parallel. Previous ThreadPoolExecutor design hit GIL contention - in the Python codegen phase, capping throughput at ~1-2 cores even - with 28 threads (observed: 4 R threads vs 28 in pool). - - Compiled binaries land in Triton's on-disk cache (TRITON_CACHE_DIR - or default ~/.triton/cache). Workers share the cache via filesystem - — first to write any given (kernel_source × constexpr_set) hash - wins; concurrent writes to the SAME hash are wasteful but not - corrupting. - - spawn start method avoids inheriting parent CUDA state (which is - unsafe after fork on Linux with active CUDA contexts). Per-worker - import + CUDA init costs ~10s, amortized over each worker's many - compiles. baseline_fn is intentionally NOT passed to workers to - avoid pickling complications; the parent compiles the baseline - kernel itself before launching the pool when applicable. - """ - from concurrent.futures import ProcessPoolExecutor - import multiprocessing - - sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) - - rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) - write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) - modes_list = getattr(args, "modes_list", ["monolithic"]) - sort_list = getattr(args, "sort_slots_list", [False]) - rev_list = getattr(args, "reverse_nowrite_list", [False]) - hsort_list = getattr(args, "hardcode_sort_list", [False]) - - # Compile-warmup task enumeration: outer × inner cartesian. - # CRITICAL: only enumerate axes that change the kernel's COMPILE signature. - # Drop runtime axes (batch, prev_k) that produce identical kernel hashes — - # otherwise we'd pay ~50-100ms of bench setup per redundant cache-hit task. - # - # Batches collapsed to first only: batch is a runtime int passed to the - # kernel, not a constexpr; all batches share the same compiled kernel. - # prev_k is already a list passed into _bench_config (not enumerated here). - configs = [] - _compile_batches = batch_sizes[:1] # collapse runtime axis - for batch in _compile_batches: - for mtp_len in mtp_lengths: - prev_ks = _resolve_prev_ks(args, mtp_len) - for state_dtype in state_dtypes: - for act_dtype in act_dtypes: - for sr_mode in sr_modes_list: - for mode in modes_list: - effective_write_modes = ( - write_modes_list if mode == "monolithic" else [True] - ) - for write_ckpt in effective_write_modes: - if mode == "monolithic": - effective_rect_list = ( - [False] if write_ckpt else rect_list - ) - else: - effective_rect_list = rect_list - for rect in effective_rect_list: - # Note: "persistent_main" is included in - # the dl-family for sort/hsort sweep - # eligibility — it consumes the same - # slot_perm and benefits from the same - # write-first clustering. It additionally - # requires _n_writes (count of write - # slots) which the bench computes from - # the pure-scenario PNAT (mix scenarios - # not yet supported for persistent_main). - is_dl_family = mode in ( - "doublelaunch", "dlgrouped", "maindl", - "dl_write_only", "persistent_main", - "persistent_dynamic", - ) - # Match the timed-run skip: sort=1 only - # makes sense when there's a mix scenario. - can_sort = is_dl_family and (args.mix_csv is not None) - effective_sort_list = ( - sort_list if can_sort else [False] - ) - effective_hsort_list = ( - hsort_list if can_sort else [False] - ) - for sort_slots in effective_sort_list: - effective_rev_list = ( - rev_list if sort_slots else [False] - ) - for reverse_nowrite in effective_rev_list: - for hardcode_sort in effective_hsort_list: - if sort_slots and hardcode_sort: - continue - configs.append(( - batch, mtp_len, prev_ks, - state_dtype, act_dtype, - sr_mode, rect, write_ckpt, mode, - sort_slots, reverse_nowrite, - hardcode_sort, - )) - - # Enumerate inner-knob signatures. Two paths: - # (1) --cell-list mode (preferred when set): pull exactly the cells - # that will be timed from args._cell_list_set. No synthetic - # cartesian — we only pre-compile what will run. - # (2) Sweep-args mode: cartesian over split-aware axes that read - # BOTH unsplit (args.X) and per-half (args.X_write/_nowrite) - # knob settings. Older code read only args.X and silently - # enumerated 1 inner combo when callers set only the per-half - # versions (all cell-list usage, plus any --block-size-m-write/ - # _nowrite CLI invocation), causing massive in-process JIT - # compile tax for persistent_main especially. - # - # In BOTH paths we GROUP tasks by non-CPS signature so each worker - # process compiles all CPS values for its group sequentially. - # NUM_PERSISTENT = CPS * num_sms is a runtime int but Triton auto- - # specializes on `div_by_16`, partitioning {CPS=1,2} (132,264) from - # {CPS=4,8} (528,1056) into two distinct compiled variants. By - # keeping all CPS variants for one (M,W,S,LS,TMA,...) signature in - # the same worker, the second compile in each spec bucket hits the - # in-process Triton cache (no disk-cache round trip). - - def _ps(val): - if val is None or (isinstance(val, str) and not val): - return [None] - if isinstance(val, str): - return [v.strip() for v in val.split(",") if v.strip()] - return [val] - - def _split_or_pair(shared_attr, w_attr, nw_attr): - """Return list of (write_val, nowrite_val) strings. - - Reads shared (args.X), write-side (args.X_write), and nowrite-side - (args.X_nowrite) values. If both per-half attrs are None, emits - tied pairs (v,v) over the shared values. If either per-half is - set, cartesian-iterates per-half values, falling back to shared - for whichever side is None. - """ - w = _ps(getattr(args, w_attr, None)) - nw = _ps(getattr(args, nw_attr, None)) - s = _ps(getattr(args, shared_attr, None)) - if w == [None] and nw == [None]: - return [(v, v) for v in s] - if w == [None]: - w = s - if nw == [None]: - nw = s - return [(a, b) for a in w for b in nw] - - _cw(f"built {len(configs)} outer configs") - cell_set = getattr(args, "_cell_list_set", set()) - cell_keys = getattr(args, "_cell_list_keys", ()) - # CPS keys are runtime ints (kernel value-specializes on `div_by_16`); - # cells differing only on CPS values can SHARE a worker so the second - # CPS value in a div_by_16 bucket hits the in-process Triton cache. - _cps_keys = ("cta_per_sm_write", "cta_per_sm_nowrite", "cta_per_sm") - - if cell_set: - # ============== CELL-LIST PATH ============== - # Build tasks DIRECTLY from cells. Each cell carries its OWN - # outer-axis values (RECT, MODE, SR, WC, SORT, REVN, HSORT) so we - # pair each cell with its specific outer config — NOT the union- - # cartesian of all cells' outer values. Previously the OUTER × - # CELL cartesian doubled task count when a cell-list spanned both - # RECT=0 and RECT=1 (or any other outer-axis split); half the - # tasks then failed the cell-list filter inside the worker and - # wasted dispatch overhead. This path is O(|unique cell groups|). - from collections import defaultdict as _dd - cell_groups: dict = _dd(list) - for tup in cell_set: - d = dict(zip(cell_keys, tup)) - cell_outer = ( - "SR" if d.get("SR", 0) else "RN", # sr_mode - bool(d.get("RECT", 0)), # rect - bool(d.get("WC", 1)), # write_ckpt - d.get("MODE", "monolithic"), # mode - bool(d.get("SORT", 0)), # sort_slots - bool(d.get("REVN", 0)), # reverse_nowrite - bool(d.get("HSORT", 0)), # hardcode_sort - ) - inner = {} - for k, v in d.items(): - if k in _CELL_LIST_KEY_TO_ARG: - inner[_CELL_LIST_KEY_TO_ARG[k]] = str(v) - non_cps_sig = tuple(sorted((k, v) for k, v in inner.items() if k not in _cps_keys)) - cell_groups[(cell_outer, non_cps_sig)].append(inner) - - # CLI-runtime axes (batch/mtp/dtype) are NOT in cell-list — they - # come from CLI args and cartesian here (typically just 1 combo). - cli_outers = [] - for _b in _compile_batches: - for _m in mtp_lengths: - _pk = _resolve_prev_ks(args, _m) - for _sd in state_dtypes: - for _ad in act_dtypes: - cli_outers.append((_b, _m, _pk, _sd, _ad)) - - tasks = [] - for cli_outer in cli_outers: - for (cell_outer, _sig), inner_list in cell_groups.items(): - outer_cfg = (*cli_outer, *cell_outer) - tasks.append((outer_cfg, inner_list)) - n_groups = len(cell_groups) - n_total_cells = sum(len(g) for g in cell_groups.values()) - n_outer_used = len(cli_outers) - else: - # ============== SWEEP-ARGS PATH ============== - # Build inner_dicts via cartesian over knob axes, then cross with - # the `configs` outer cartesian. Existing behavior. - m_pairs = _split_or_pair("block_size_m", "block_size_m_write", "block_size_m_nowrite") - w_pairs = _split_or_pair("num_warps", "num_warps_write", "num_warps_nowrite") - ns_pairs = _split_or_pair("num_stages", "num_stages_write", "num_stages_nowrite") - cps_pairs = _split_or_pair("cta_per_sm", "cta_per_sm_write", "cta_per_sm_nowrite") - ls_pairs = _split_or_pair("num_loop_stages", "num_loop_stages_write", "num_loop_stages_nowrite") - pw_vals = _ps(args.precompute_num_warps) - ps_vals = _ps(args.precompute_num_stages) - h_vals = _ps(args.heads_per_block) - mr_vals = _ps(args.maxnreg) - ct_vals = _ps(args.num_ctas) - fl_vals = _ps(args.flatten) - wsp_vals = _ps(args.warp_specialize) - trl_vals = _ps(args.use_tma_rect_load) - twl_vals = _ps(args.use_tma_replay_write_load) - tnl_vals = _ps(args.use_tma_replay_nowrite_load) - tws_vals = _ps(args.use_tma_replay_write_store) - import itertools as _it - inner_dicts = [] - for ((mw, mnw), (ww, wnw), (sw, snw), (cw, cnw), (lw, lnw), - pw, ps_, h, mr, ct, fl, wsp, - trl, twl, tnl, tws) in _it.product( - m_pairs, w_pairs, ns_pairs, cps_pairs, ls_pairs, - pw_vals, ps_vals, h_vals, mr_vals, ct_vals, - fl_vals, wsp_vals, - trl_vals, twl_vals, tnl_vals, tws_vals): - d = {} - for k, v in ( - ("block_size_m_write", mw), - ("block_size_m_nowrite", mnw), - ("num_warps_write", ww), - ("num_warps_nowrite", wnw), - ("num_stages_write", sw), - ("num_stages_nowrite", snw), - ("cta_per_sm_write", cw), - ("cta_per_sm_nowrite", cnw), - ("num_loop_stages_write", lw), - ("num_loop_stages_nowrite", lnw), - ("precompute_num_warps", pw), - ("precompute_num_stages", ps_), - ("heads_per_block", h), - ("maxnreg", mr), - ("num_ctas", ct), - ("flatten", fl), - ("warp_specialize", wsp), - ("use_tma_rect_load", trl), - ("use_tma_replay_write_load", twl), - ("use_tma_replay_nowrite_load", tnl), - ("use_tma_replay_write_store", tws), - ): - if v is not None: - d[k] = str(v) - inner_dicts.append(d) - - groups: dict = {} - for d in inner_dicts: - sig = tuple(sorted((k, v) for k, v in d.items() if k not in _cps_keys)) - groups.setdefault(sig, []).append(d) - tasks = [] - for outer in configs: - for sig, group in groups.items(): - tasks.append((outer, group)) - n_groups = len(groups) - n_total_cells = sum(len(g) for g in groups.values()) - n_outer_used = len(configs) - - # Shuffle ACROSS tasks (preserve within-group CPS sequence for in-process - # cache adjacency — within-group order is intentional, not shuffled). - import random as _r - _r.shuffle(tasks) - - _cw(f"built {len(tasks)} tasks covering {n_total_cells} cells in {n_groups} groups") - print(f"[compile-warmup] {len(tasks)} compile tasks " - f"({n_outer_used} outer × {n_groups} cell-groups " - f"covering {n_total_cells} cells, CPS-grouped" - + (", per-cell outer" if cell_set else "") - + f") across {max_workers} processes (ProcessPoolExecutor, {_MP_START_METHOD} start)") - t0 = time.perf_counter() - - ctx = multiprocessing.get_context(_MP_START_METHOD) - errors = [] - _cw("about to create ProcessPoolExecutor") - with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as ex: - _cw("ProcessPoolExecutor created, about to submit tasks") - # baseline_fn=None: workers compile only the checkpointing kernel. - # Baseline kernels (if any) get compiled lazily in the parent during - # the timing phase — usually just one extra compile, negligible. - futures = { - ex.submit(_warm_one_config, args, task, None): task - for task in tasks - } - _cw(f"submitted {len(futures)} tasks, waiting for results") - _n_done = 0 - for fut in futures: - try: - fut.result() - except Exception as e: - errors.append((futures[fut], e)) - _n_done += 1 - # Progress beacons at 10/25/50/75/100% to gauge effective parallelism. - if _n_done in (max(1, len(futures)//10), - max(1, len(futures)//4), - max(1, len(futures)//2), - max(1, (3*len(futures))//4), - len(futures)): - _cw(f"{_n_done}/{len(futures)} tasks complete") - - if errors: - for cfg, e in errors: - print(f"[compile-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", - file=sys.stderr) - raise errors[0][1] - - print(f"[compile-warmup] done in {time.perf_counter() - t0:.1f}s") - - -def _bench_config( - args, - batch: int, - mtp_len: int, - prev_ks: list[int], - state_dtype: torch.dtype, - act_dtype: torch.dtype, - baseline_fn, - sr_mode: str = "RN", - rectangle_for_nowrite: bool = False, - write_checkpoint: bool = True, - mode: str = "monolithic", - mix_samples_cpu=None, - mix_label: str = "", - sort_slots: bool = False, - reverse_nowrite: bool = False, - perm_samples_cpu=None, - hardcode_sort: bool = False, - mix_samples_sorted_cpu=None, - warmup_only: bool = False, -) -> None: - """ - Benchmark one (batch, mtp_len, dtype) configuration. - - Runs the baseline kernel (if baseline_fn is not None) followed by the - replay kernel for each prev_k value. Tensors are built once and - shared across all runs in this config. - - When ``warmup_only`` is True, calls each kernel exactly once instead of - timing it. Used by the parallel-warmup phase to populate Triton's - persistent compile cache across all configs concurrently. No timing - output is produced. - """ - state_dtype_name = str(state_dtype).split(".")[-1] - act_dtype_name = str(act_dtype).split(".")[-1] - - ( - state0, - state_scales0, - old_x0, - old_B0, - old_dt0, - old_dA_cumsum0, - cache_buf_idx0, - x, - dt, - B, - C, - A, - dt_bias, - D, - prev_tokens, - slot_perm_buf, - out_incr, - out_base, - intermediate_states_buffer, - xbc_input0, - conv_state0, - conv_weight, - conv_bias, - d_inner, - conv_dim, - ) = _build_tensors( - batch, - mtp_len, - state_dtype, - act_dtype, - args.tp_nheads, - args.head_dim, - args.d_state, - args.tp_ngroups, - max_window=getattr(args, "max_window", None) or None, - ) - - nheads = args.tp_nheads - ngroups = args.tp_ngroups - head_dim = args.head_dim - d_state = args.d_state - with_conv1d = getattr(args, "with_conv1d", False) - use_philox = (sr_mode == "SR") - variant_fn = _VARIANT_FNS[args.variant]() - - # SR rounding: allow fp16 and the quantized dtypes (int8/int16/fp8). - # bf16/fp32 SR is not supported (no PTX path for bf16; fp32 doesn't need - # rounding). When sweeping --sr-modes RN,SR over a mixed dtype set, - # silently skip the SR cell for unsupported dtypes — the RN cell still - # prints, and other dtypes still get their SR row. - rand_seed = None - _SR_SUPPORTED = ( - torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn, - ) - if use_philox: - if state_dtype not in _SR_SUPPORTED: - return - rand_seed = torch.randint(0, 2**62, (1,), device="cuda", dtype=torch.int64) - - is_quantized = state_dtype in (torch.int8, torch.int16, torch.float8_e4m3fn) - - state_work = state0.clone() - state_scales_work = state_scales0.clone() if state_scales0 is not None else None - old_x_work = old_x0.clone() - old_B_work = old_B0.clone() - old_dt_work = old_dt0.clone() - old_dA_cumsum_work = old_dA_cumsum0.clone() - cache_buf_idx_work = cache_buf_idx0.clone() - xbc_input_work = xbc_input0.clone() - conv_state_work = conv_state0.clone() - - def _reset(): - state_work.copy_(state0) - if state_scales_work is not None: - state_scales_work.copy_(state_scales0) - old_x_work.copy_(old_x0) - old_B_work.copy_(old_B0) - old_dt_work.copy_(old_dt0) - old_dA_cumsum_work.copy_(old_dA_cumsum0) - cache_buf_idx_work.copy_(cache_buf_idx0) - if with_conv1d: - conv_state_work.copy_(conv_state0) - - def _reset_conv1d_realistic(): - """Realistic reset: cold cache, L2 flush, then hot in_proj output.""" - # 1. Reset cold state (cache tensors, SSM state) - state_work.copy_(state0) - if state_scales_work is not None: - state_scales_work.copy_(state_scales0) - old_x_work.copy_(old_x0) - old_B_work.copy_(old_B0) - old_dt_work.copy_(old_dt0) - old_dA_cumsum_work.copy_(old_dA_cumsum0) - cache_buf_idx_work.copy_(cache_buf_idx0) - conv_state_work.copy_(conv_state0) - # 2. L2 flush (evicts cold state from cache) - if _l2_flush is not None: - _l2_flush.fill_(0.0) - # 3. Write hot tensors (simulates in_proj output landing in L2) - xbc_input_work.copy_(xbc_input0) - - # Silently skip the baseline row for any (baseline, state_dtype, SR) - # combo it can't run. Better than erroring on a partial sweep — our - # kernel rows still print. Compatibility: - # * Quantized states (int8 / int16 / fp8): no baseline supports them. - # * Triton baseline (selective_state_update): no rand_seed kwarg. - # * flashinfer baseline: rand_seed only on fp16 state. - def _baseline_supports() -> bool: - if baseline_fn is None: - return False - if is_quantized: - return False - if use_philox: - if args.baseline == "triton": - return False - if args.baseline == "flashinfer" and state_dtype != torch.float16: - return False - return True - - if baseline_fn is not None and not _baseline_supports(): - if not warmup_only: - sr_tag = " + SR" if use_philox else "" - print( - f"# Skipping {args.baseline} baseline for " - f"state_dtype={state_dtype_name}{sr_tag} (unsupported)." - ) - baseline_fn = None - - show_kernel_col = baseline_fn is not None - - def _conv1d_split(xbc_in, conv_st, launch_dependent_kernels=False): - """Run conv1d update and split output into (x, B, C) views. - - The input tensor's strides are preserved through conv1d and the - transpose+view chain. With the production-matching layout - (contiguous (batch*T, conv_dim) viewed as (batch, conv_dim, T)), - the output after transpose+view has stride(-1)==1 and - stride(1)==dim, satisfying both our kernel and flashinfer. - """ - xbc_result = causal_conv1d_update( - xbc_in, - conv_st, - conv_weight, - conv_bias, - activation="silu", - launch_dependent_kernels=launch_dependent_kernels, - ) - xbc_flat = xbc_result.transpose(1, 2).view(batch * mtp_len, conv_dim) - x_flat, B_flat, C_flat = torch.split( - xbc_flat, [d_inner, ngroups * d_state, ngroups * d_state], dim=-1 - ) - x_conv = x_flat.view(batch, mtp_len, nheads, head_dim) - B_conv = B_flat.view(batch, mtp_len, ngroups, d_state) - C_conv = C_flat.view(batch, mtp_len, ngroups, d_state) - return x_conv, B_conv, C_conv - - # --- Baseline --- - if baseline_fn is not None: - tag = f"base_b{batch}_mtp{mtp_len}_s{state_dtype_name}_a{act_dtype_name}" - - philox_kwargs = {} - if rand_seed is not None and args.baseline == "flashinfer": - philox_kwargs = {"rand_seed": rand_seed, "philox_rounds": args.philox_rounds} - - if with_conv1d: - - def _run_baseline(): - x_conv, B_conv, C_conv = _conv1d_split(xbc_input_work, conv_state_work) - baseline_fn( - state_work, - x=x_conv, - dt=dt, - A=A, - B=B_conv, - C=C_conv, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - out=out_base, - disable_state_update=True, - intermediate_states_buffer=intermediate_states_buffer, - cache_steps=mtp_len, - **philox_kwargs, - ) - else: - - def _run_baseline(): - baseline_fn( - state_work, - x=x, - dt=dt, - A=A, - B=B, - C=C, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - out=out_base, - disable_state_update=True, - intermediate_states_buffer=intermediate_states_buffer, - cache_steps=mtp_len, - **philox_kwargs, - ) - - reset_fn = _reset_conv1d_realistic if with_conv1d else _reset - if warmup_only: - reset_fn() - _run_baseline() - torch.cuda.synchronize() - else: - stats = _time_kernel( - args, _run_baseline, reset_fn, tag, - expected_K=_kernels_per_iter_baseline(with_conv1d), - cupti_plan_key=( - "baseline", - args.baseline, - batch, - mtp_len, - state_dtype_name, - act_dtype_name, - with_conv1d, - bool(args.l2_flush), - bool(args.external_pdl), - bool(use_philox), - _kernels_per_iter_baseline(with_conv1d), - ), - ) - - _submit_result_job( - args, - stats, - show_kernel_col=show_kernel_col, - kernel_name=args.baseline, - batch=batch, - mtp_len=mtp_len, - prev_k="N/A", - state_dtype_name=state_dtype_name, - act_dtype_name=act_dtype_name, - skipped_tag=tag, - ) - - # --- Sweep parameter parsing (invariant across prev_k) --- - def _parse_sweep(val): - if val is None: - return [None] - return [int(v) for v in val.split(",")] - - block_size_m_values = _parse_sweep(args.block_size_m) - num_warps_values = _parse_sweep(args.num_warps) - num_stages_values = _parse_sweep(args.num_stages) - precompute_num_warps_values = _parse_sweep(args.precompute_num_warps) - precompute_num_stages_values = _parse_sweep(args.precompute_num_stages) - heads_per_block_values = _parse_sweep(args.heads_per_block) - maxnreg_values = _parse_sweep(args.maxnreg) - num_ctas_values = _parse_sweep(args.num_ctas) - # Persistent-only sweep dims; ignored when the cell's mode != persistent_main. - cta_per_sm_values = _parse_sweep(args.cta_per_sm) - num_loop_stages_values = _parse_sweep(args.num_loop_stages) - flatten_values = _parse_sweep(args.flatten) - warp_specialize_values = _parse_sweep(args.warp_specialize) - # Per-main split-knob sweeps. Default = same as the shared sweep (so each - # combo is tied). When set independently, the inner loop sweeps the - # cross-product (write × nowrite); --skip-diagonal drops the tied subset. - def _split_or_share(split_csv, shared_values): - return _parse_sweep(split_csv) if split_csv else shared_values - block_size_m_write_values = _split_or_share(args.block_size_m_write, block_size_m_values) - block_size_m_nowrite_values = _split_or_share(args.block_size_m_nowrite, block_size_m_values) - num_warps_write_values = _split_or_share(args.num_warps_write, num_warps_values) - num_warps_nowrite_values = _split_or_share(args.num_warps_nowrite, num_warps_values) - num_stages_write_values = _split_or_share(args.num_stages_write, num_stages_values) - num_stages_nowrite_values = _split_or_share(args.num_stages_nowrite, num_stages_values) - cta_per_sm_write_values = _split_or_share(args.cta_per_sm_write, cta_per_sm_values) - cta_per_sm_nowrite_values = _split_or_share(args.cta_per_sm_nowrite, cta_per_sm_values) - num_loop_stages_write_values = _split_or_share(args.num_loop_stages_write, num_loop_stages_values) - num_loop_stages_nowrite_values = _split_or_share(args.num_loop_stages_nowrite, num_loop_stages_values) - # Whether any *_write / *_nowrite knob was independently set — used by - # --skip-diagonal to know if the cross-product is non-trivial. Without - # any split, the per-main values == shared values and skip-diagonal is - # a no-op (which is correct). - _any_split = any(getattr(args, name) for name in ( - "block_size_m_write", "block_size_m_nowrite", - "num_warps_write", "num_warps_nowrite", - "num_stages_write", "num_stages_nowrite", - "cta_per_sm_write", "cta_per_sm_nowrite", - "num_loop_stages_write", "num_loop_stages_nowrite", - )) - # TMA toggles — independent 0/1 sweep per path. The skip-dupe at the - # top of the inner loop body collapses cells where a flag's path is - # unreachable, so e.g. monolithic + WC=True only runs the value=0 - # cells for nowrite-load and rect-load. - use_tma_rect_load_values = _parse_sweep(args.use_tma_rect_load) - use_tma_replay_write_load_values = _parse_sweep(args.use_tma_replay_write_load) - use_tma_replay_nowrite_load_values = _parse_sweep(args.use_tma_replay_nowrite_load) - use_tma_replay_write_store_values = _parse_sweep(args.use_tma_replay_write_store) - - # --- Replay kernel --- - # Cache T-axis capacity (for prev_k validity check on the nowrite path). - max_window = getattr(args, "max_window", 0) or mtp_len - - # Build the list of scenarios to time. A scenario is one cell in the - # output: pure-mode scenarios fill prev_tokens with one constant before - # the timing loop; mix-mode scenarios feed a pre-baked per-iter samples - # tensor, with the per-iter copy captured inside the CUDA graph. Pure - # and mix can coexist in one call so a single nsys trace covers both. - scenarios = [] - if not (getattr(args, "mix_only", False) and mix_samples_cpu is not None): - for prev_k in prev_ks: - # On the nowrite path, new tokens append at [prev_k, prev_k+T) of - # the active buffer, so prev_k+T must fit within max_window. - # mode != monolithic dispatches per-slot from PNAT, so any - # prev_k <= max_window is valid for those modes. - if mode == "monolithic" and not write_checkpoint and prev_k + mtp_len > max_window: - continue - scenarios.append({ - "label": f"k{prev_k}", - "print_label": prev_k, - "fill": prev_k, - "pre_iter": None, - "iters": None, # use args.iters - }) - # Mix scenario: skip on monolithic (mono on mixed PNAT corrupts the - # wrong-mode slots). Persistent_main + mix is now supported: bench - # pre-bakes both a per-iter PNAT samples tensor and a per-iter - # n_writes samples tensor; grouped graph capture copies window rows - # into kernel-input tensors (PNAT and n_writes_dev) before each - # in-graph L2 flush, so the timed kernels read PNAT cold. - if mix_samples_cpu is not None and mode != "monolithic": - device = state_work.device - # Hardcode-sort: per-iter prev_tokens are CPU-sorted write-first. - # Kernel runs USE_PERM=False but the EO gate sees clustered modes. - # Output is scrambled (we don't permute x/B/C/dt to match) but - # timing is meaningful — isolates clustering benefit from the - # per-program perm-load overhead in --sort-slots. - src = mix_samples_sorted_cpu if (hardcode_sort and mix_samples_sorted_cpu is not None) else mix_samples_cpu - samples_gpu = torch.from_numpy(src).to(device=device, dtype=torch.int32) - - # For persistent_main + mix: pre-compute the per-iter n_writes - # (count of slots needing the write path = PNAT+T > max_window) - # and the (1,) scratch the kernel reads from. Both halves of - # persistent_main always launch in mix scenarios (host can't - # cheaply read n_writes per iter without a sync), so the kernel's - # slot-range derivation must be correct from device n_writes. - # persistent_dynamic doesn't need n_writes (the kernel ignores - # n_writes_dev when IS_DYNAMIC=True via Triton DCE), but we still - # allocate a sentinel scratch so the wrapper API is uniform. - n_writes_samples_gpu = None - n_writes_dev_mix = None - # n_writes per iter = number of slots that overflow the window. - # Computed for ALL mix scenarios so the JSON output (--json-detailed) - # can pair each iter's span_us with its mix composition for downstream - # analysis (group iters by # writes → per-bucket median → analytic - # expectation under the steady-state PNAT distribution). - n_writes_per_iter_all = ((src + mtp_len) > max_window).sum(axis=1).astype(np.int32) - if mode in ("persistent_main", "persistent_dynamic"): - n_writes_samples_gpu = torch.from_numpy(n_writes_per_iter_all).to( - device=device, dtype=torch.int32 - ) - n_writes_dev_mix = torch.zeros(1, dtype=torch.int32, device=device) - - # Build _mix_pre_iter — the closure that runs OUTSIDE the captured - # graph between replays. Updates: prev_tokens (always), - # slot_perm_buf (when sort_slots), n_writes_dev_mix (persistent). - perm_samples_gpu = None - if sort_slots and perm_samples_cpu is not None: - perm_samples_gpu = torch.from_numpy(perm_samples_cpu).to( - device=device, dtype=torch.int32 - ) - if n_writes_samples_gpu is not None: - def _mix_pre_iter(i, _s=samples_gpu, _ps=perm_samples_gpu, - _ns=n_writes_samples_gpu, _pt=prev_tokens, - _pm=slot_perm_buf, _nw=n_writes_dev_mix): - _pt.copy_(_s[i]) - _pm.copy_(_ps[i]) - _nw.copy_(_ns[i:i+1]) - else: - def _mix_pre_iter(i, _s=samples_gpu, _ps=perm_samples_gpu, - _pt=prev_tokens, _pm=slot_perm_buf): - _pt.copy_(_s[i]) - _pm.copy_(_ps[i]) - else: - if n_writes_samples_gpu is not None: - def _mix_pre_iter(i, _s=samples_gpu, _ns=n_writes_samples_gpu, - _pt=prev_tokens, _nw=n_writes_dev_mix): - _pt.copy_(_s[i]) - _nw.copy_(_ns[i:i+1]) - else: - def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): - _pt.copy_(_s[i]) - - def _mix_pre_iter_group_factory( - group_iters, - _s=samples_gpu, - _ps=perm_samples_gpu, - _ns=n_writes_samples_gpu, - _pt=prev_tokens, - _pm=slot_perm_buf, - _nw=n_writes_dev_mix, - ): - sample_window = torch.empty( - (group_iters, _s.shape[1]), device=_s.device, dtype=_s.dtype, - ) - perm_window = ( - torch.empty((group_iters, _ps.shape[1]), device=_ps.device, dtype=_ps.dtype) - if _ps is not None else None - ) - nw_window = ( - torch.empty((group_iters,), device=_ns.device, dtype=_ns.dtype) - if _ns is not None else None - ) - - def _pre_replay(replay_idx): - start = replay_idx * group_iters - end = start + group_iters - sample_window.copy_(_s[start:end]) - if perm_window is not None: - perm_window.copy_(_ps[start:end]) - if nw_window is not None: - nw_window.copy_(_ns[start:end]) - - def _graph_pre_iter(j): - _pt.copy_(sample_window[j]) - if perm_window is not None: - _pm.copy_(perm_window[j]) - if nw_window is not None: - _nw.copy_(nw_window[j:j + 1]) - - return _pre_replay, _graph_pre_iter - - # Mix iters override: if --mix-iters set, use it; else use args.iters. - mix_iters = getattr(args, "mix_iters", None) - scenarios.append({ - "label": f"mix{mix_label}", - "print_label": "mix", - "fill": None, - "pre_iter": _mix_pre_iter, - "pre_iter_group_factory": _mix_pre_iter_group_factory, - "iters": mix_iters, # None => use args.iters - # Pass through to _run_incr so the wrapper receives _n_writes_dev - # (mix scenarios) instead of _n_writes (pure scenarios). - "n_writes_dev": n_writes_dev_mix, - # Full per-iter n_writes array (size = warmup + iters). Used by - # the JSON-detailed output to pair each iter's span with its - # mix composition for post-hoc bucketing analysis. - "n_writes_per_iter": n_writes_per_iter_all, - }) - - # Pure scenarios don't pre-allocate n_writes_dev; mix scenarios do. - # Default empty-halves skip: True for pure (host knows n_writes, - # production-equivalent host-skip), False for mix (host can't read - # device n_writes per iter without sync, must always launch both). - for scn in scenarios: - scenario_n_writes_dev = scn.get("n_writes_dev") # None for pure - scenario_skip_empty = scenario_n_writes_dev is None - if scn["fill"] is not None: - prev_tokens.fill_(scn["fill"]) - prev_k_for_print = scn["print_label"] - scenario_pre_iter = scn["pre_iter"] - scenario_pre_iter_group_factory = scn.get("pre_iter_group_factory") - scenario_iters = scn.get("iters") # None => use args.iters - tag = f"incr_b{batch}_mtp{mtp_len}_{scn['label']}_s{state_dtype_name}_a{act_dtype_name}" - - # Iteration over per-cell knob combos. - # When NO per-main split is requested (_any_split=False), each row in - # the cross-product gives the same value to both write_main and - # nowrite_main (current behavior — backward-compat). When ANY split - # IS requested, we iterate the write and nowrite axes independently - # (cross-product blowup is the user's responsibility — they typically - # pair this with --skip-diagonal to drop the tied subset). - if _any_split: - _iter_axes = ( - block_size_m_write_values, block_size_m_nowrite_values, - num_warps_write_values, num_warps_nowrite_values, - num_stages_write_values, num_stages_nowrite_values, - precompute_num_warps_values, - precompute_num_stages_values, - heads_per_block_values, - maxnreg_values, num_ctas_values, - cta_per_sm_write_values, cta_per_sm_nowrite_values, - num_loop_stages_write_values, num_loop_stages_nowrite_values, - flatten_values, warp_specialize_values, - use_tma_rect_load_values, - use_tma_replay_write_load_values, - use_tma_replay_nowrite_load_values, - use_tma_replay_write_store_values, - ) - else: - # Tied: one value per shared knob. Wrap in single-element list for - # uniform iteration; the body sets w/nw both to the shared value. - _iter_axes = ( - block_size_m_values, [None], - num_warps_values, [None], - num_stages_values, [None], - precompute_num_warps_values, - precompute_num_stages_values, - heads_per_block_values, - maxnreg_values, num_ctas_values, - cta_per_sm_values, [None], - num_loop_stages_values, [None], - flatten_values, warp_specialize_values, - use_tma_rect_load_values, - use_tma_replay_write_load_values, - use_tma_replay_nowrite_load_values, - use_tma_replay_write_store_values, - ) - # Iteration source: when --cell-list is active AND this is the main - # timing path (not a compile-warmup worker), iterate the cell set - # DIRECTLY (one yield per cell). The earlier design iterated the - # full inner cartesian and filtered each iteration via membership in - # args._cell_list_set — that's O(cartesian) which blows up to - # billions of iterations when the cell-list spans wide split-knob - # values (CPS, LS, M, W, S each contributing a Wx*Wnw factor on top - # of TMA flags), producing 50+ min of CPU spin per bench call before - # any actual timing. Direct iteration is O(|cell_list|). - # - # IMPORTANT exception for workers (warmup_only=True): _warm_one_config - # clamps args.*_write/_nowrite via inner_overrides to single values, - # making the cartesian 1×1×...×1 = 1 iter, which is exactly the one - # cell that worker was given. If we used cell-list-direct iteration - # here, every worker would iterate ALL 2884 cells instead of just - # its assigned one — turning compile-warmup into 28-way duplication. - # (Observed: 256 tasks in 233s under that bug vs ~18s correct.) - if getattr(args, "_cell_list_set", None) and not warmup_only: - def _gen_from_cell_list(): - keys = args._cell_list_keys - for tup in args._cell_list_set: - d = dict(zip(keys, tup)) - yield ( - d.get("Mw"), d.get("Mnw"), - d.get("Ww"), d.get("Wnw"), - d.get("Sw"), d.get("Snw"), - d.get("pW"), d.get("pS"), - d.get("H"), - d.get("R"), d.get("CT"), - d.get("CPSw"), d.get("CPSnw"), - d.get("LSw"), d.get("LSnw"), - d.get("FL"), d.get("WS"), - d.get("TMARL"), d.get("TMAWL"), - d.get("TMANL"), d.get("TMAWS"), - ) - _iter_source = _gen_from_cell_list() - else: - _iter_source = itertools.product(*_iter_axes) - - for ( - block_size_m_w, - block_size_m_nw, - num_warps_w, - num_warps_nw, - num_stages_w, - num_stages_nw, - precompute_num_warps, - precompute_num_stages, - heads_per_block, - maxnreg, - num_ctas, - cta_per_sm_w, - cta_per_sm_nw, - num_loop_stages_w, - num_loop_stages_nw, - flatten, - warp_specialize, - use_tma_rect_load, - use_tma_replay_write_load, - use_tma_replay_nowrite_load, - use_tma_replay_write_store, - ) in _iter_source: - # When tied, _nw values were placeholder None; fill from _w (the - # shared value). When split, _w and _nw came from independent lists. - if not _any_split: - block_size_m_nw = block_size_m_w - num_warps_nw = num_warps_w - num_stages_nw = num_stages_w - cta_per_sm_nw = cta_per_sm_w - num_loop_stages_nw = num_loop_stages_w - # Skip-diagonal: when split is on, drop the tied subset (same as a - # prior shared-knob sweep would cover). - if _any_split and args.skip_diagonal and ( - block_size_m_w == block_size_m_nw and - num_warps_w == num_warps_nw and - num_stages_w == num_stages_nw and - cta_per_sm_w == cta_per_sm_nw and - num_loop_stages_w == num_loop_stages_nw - ): - continue - # Backward-compat aliases used by the existing body below. When - # tied, these are simply the shared value. When split, the - # _write copy is used for sweep_tag and grouping (a stable choice - # so the tag is unique per (write, nowrite) combo). - block_size_m = block_size_m_w - num_warps = num_warps_w - num_stages = num_stages_w - cta_per_sm = cta_per_sm_w - num_loop_stages = num_loop_stages_w - # Skip-dupe for TMA flag sweeps: a flag whose code path isn't - # reachable in this cell produces identical timing for value=0 - # and value=1. We canonicalize by skipping value=1 cells when - # the flag's path is unreachable. Path reachability rules: - # * write path (replay write-load + write-store): mono+WC=True, - # OR any non-monolithic mode. - # * rect path (rect-load): rectangle_for_nowrite=True AND a - # nowrite path exists in this mode (mono+WC=False, OR any - # non-monolithic mode). - # * replay-nowrite path (nowrite-load): nowrite path exists - # AND rect isn't taking it: mono+WC=False+rect=False, OR - # any non-monolithic mode with rect=False. - _is_mono = (mode == "monolithic") - _write_path = (_is_mono and write_checkpoint) or (not _is_mono) - _rect_path = rectangle_for_nowrite and ( - (not _is_mono) or (_is_mono and not write_checkpoint) - ) - _replay_nowrite_path = ( - (_is_mono and not write_checkpoint and not rectangle_for_nowrite) - or ((not _is_mono) and not rectangle_for_nowrite) - ) - def _set(v): # flag set to a non-zero sweep value - return v is not None and v != 0 - if (_set(use_tma_rect_load) and not _rect_path - or _set(use_tma_replay_write_load) and not _write_path - or _set(use_tma_replay_nowrite_load) and not _replay_nowrite_path - or _set(use_tma_replay_write_store) and not _write_path): - continue - - # Pre-allocate n_writes_dev tensor OUTSIDE the captured graph for - # persistent modes in pure scenarios. Mix scenarios already have - # `scenario_n_writes_dev` pre-allocated. The wrapper's fallback - # `torch.tensor([...], device=...)` allocation would invalidate - # the CUDA-graph capture stream — must allocate here, before the - # `_run_incr` lambda (which is what gets captured) is defined. - # For persistent_dynamic the kernel ignores the value (IS_DYNAMIC - # DCE's the load); we still need a valid pointer. For - # persistent_main pure, the value is constant per cell so we set - # it once here. - _n_writes_dev_pure: torch.Tensor | None = None - _host_n_writes_pure: int | None = None - if mode in ("persistent_main", "persistent_dynamic") and scenario_n_writes_dev is None: - _n_writes_dev_pure = torch.zeros(1, dtype=torch.int32, device=state_work.device) - if mode == "persistent_main": - scn_fill = scn["fill"] - is_write_scenario_local = (scn_fill + mtp_len) > max_window - _host_n_writes_pure = batch if is_write_scenario_local else 0 - _n_writes_dev_pure.fill_(_host_n_writes_pure) - - def _run_incr( - block_size_m=block_size_m, - num_warps=num_warps, - num_stages=num_stages, - precompute_num_warps=precompute_num_warps, - precompute_num_stages=precompute_num_stages, - heads_per_block=heads_per_block, - maxnreg=maxnreg, - num_ctas=num_ctas, - cta_per_sm=cta_per_sm, - num_loop_stages=num_loop_stages, - flatten=flatten, - warp_specialize=warp_specialize, - use_tma_rect_load=use_tma_rect_load, - use_tma_replay_write_load=use_tma_replay_write_load, - use_tma_replay_nowrite_load=use_tma_replay_nowrite_load, - use_tma_replay_write_store=use_tma_replay_write_store, - ): - if with_conv1d: - x_call, B_call, C_call = _conv1d_split( - xbc_input_work, conv_state_work, launch_dependent_kernels=args.external_pdl - ) - extra_kwargs = {"launch_with_pdl": args.external_pdl} - else: - x_call, B_call, C_call = x, B, C - extra_kwargs = {} - # write_checkpoint is only meaningful for the checkpointing - # variant; replay variant ignores the kwarg. state_scales - # is also checkpointing-only (replay kernel doesn't quantize). - if args.variant == "checkpointing": - extra_kwargs["write_checkpoint"] = write_checkpoint - extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite - extra_kwargs["mode"] = mode - if sort_slots: - extra_kwargs["slot_perm"] = slot_perm_buf - # reverse_nowrite is meaningful in two ways: - # - with slot_perm: walk the perm tail-first - # - without slot_perm (hardcode-sort): walk pid_b - # itself tail-first via the REVERSE_PERM constexpr - if sort_slots or (hardcode_sort and reverse_nowrite): - extra_kwargs["reverse_nowrite"] = reverse_nowrite - if state_scales_work is not None: - extra_kwargs["state_scales"] = state_scales_work - if use_tma_rect_load: # 1 → True, 0/None → False - extra_kwargs["_use_tma_rect_load"] = True - if use_tma_replay_write_load: - extra_kwargs["_use_tma_replay_write_load"] = True - if use_tma_replay_nowrite_load: - extra_kwargs["_use_tma_replay_nowrite_load"] = True - if use_tma_replay_write_store: - extra_kwargs["_use_tma_replay_write_store"] = True - # persistent_main needs n_writes (count of write-mode - # slots in the pre-sorted batch) as a host-side int. - # Pure scenarios: every slot has the same PNAT, so - # n_writes is either 0 (all nowrite) or batch (all - # write) depending on whether PNAT+T overflows the - # window. Mix scenarios are skipped earlier. - if mode in ("persistent_main", "persistent_dynamic"): - # Per-cell sweep values for persistent-only knobs. - # Apply to both persistent variants. _parse_sweep - # returns [None] when the user didn't pass the flag, - # in which case we leave the wrapper's defaults. - if cta_per_sm is not None: - extra_kwargs["_cta_per_sm"] = cta_per_sm - if num_loop_stages is not None: - extra_kwargs["_num_loop_stages"] = num_loop_stages - if flatten is not None: - extra_kwargs["_flatten"] = bool(flatten) - if warp_specialize is not None: - extra_kwargs["_warp_specialize"] = bool(warp_specialize) - if mode in ("persistent_main", "persistent_dynamic"): - # persistent_main + mix REQUIRES sort: the kernel - # partitions slots [0, n_writes) = write half, - # [n_writes, batch) = nowrite half. This only holds - # if PNAT is monotone (writes first), which sort - # provides via either: - # sort_slots=1 → USE_PERM reads slot_perm to remap - # hardcode_sort=1 → PNAT itself is CPU-pre-sorted - # persistent_dynamic doesn't need sort (per-slot - # runtime dispatch); persistent_main pure scenarios - # are trivially sorted (homogeneous PNAT). - if (mode == "persistent_main" - and scenario_n_writes_dev is not None - and not (sort_slots or hardcode_sort)): - raise AssertionError( - "persistent_main + mix requires sort_slots=1 " - "or hardcode_sort=1 — kernel partitions slots " - "by index, which is only valid when PNAT is " - "monotone (writes first). Without sort, the " - "partition silently mismatches actual slot " - "modes. Re-run with --sort-slots 1 or " - "--hardcode-sort 1." - ) - # n_writes plumbing: pure scenarios pass an int - # (host knows the value, can host-skip empty halves); - # mix scenarios pass a (1,) device tensor updated - # per iter by the benchmark pre-iter path. - # _persistent_skip_empty_halves=False on mix so both - # halves always launch (kernel uses device n_writes - # to derive its slot range). - if scenario_n_writes_dev is not None: - # Mix path: caller-allocated tensor, updated per - # iter by scenario_pre_iter outside capture. - extra_kwargs["_n_writes_dev"] = scenario_n_writes_dev - extra_kwargs["_persistent_skip_empty_halves"] = False - elif mode == "persistent_main": - # Pure: caller pre-allocated `_n_writes_dev_pure` - # outside this lambda (so the alloc doesn't land - # inside the captured graph). Pass both the - # tensor and the host int so the wrapper can use - # host-skip when `_persistent_skip_empty_halves`. - extra_kwargs["_n_writes"] = _host_n_writes_pure - extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure - extra_kwargs["_persistent_skip_empty_halves"] = scenario_skip_empty - elif mode == "persistent_dynamic": - # persistent_dynamic pure: kernel ignores n_writes - # via IS_DYNAMIC DCE, but the wrapper needs a - # valid (1,) tensor pointer. Pass the pre-allocated - # zero tensor to avoid any in-capture alloc. - extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure - variant_fn( - state_work, - old_x_work, - old_B_work, - old_dt_work, - old_dA_cumsum_work, - cache_buf_idx_work, - prev_tokens, - x=x_call, - dt=dt, - A=A, - B=B_call, - C=C_call, - out=out_incr, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=None, - rand_seed=rand_seed, - philox_rounds=args.philox_rounds, - use_internal_pdl=args.internal_pdl, - _block_size_m=block_size_m, - _num_warps=num_warps, - _num_stages=num_stages, - _precompute_num_warps=precompute_num_warps, - _precompute_num_stages=precompute_num_stages, - _heads_per_block=heads_per_block, - _maxnreg=maxnreg, - _num_ctas=num_ctas, - # Per-main overrides (None = tied to shared above; explicit - # only when the inner loop is iterating split axes). - _block_size_m_write=block_size_m_w if _any_split else None, - _block_size_m_nowrite=block_size_m_nw if _any_split else None, - _num_warps_write=num_warps_w if _any_split else None, - _num_warps_nowrite=num_warps_nw if _any_split else None, - _num_stages_write=num_stages_w if _any_split else None, - _num_stages_nowrite=num_stages_nw if _any_split else None, - _cta_per_sm_write=cta_per_sm_w if _any_split else None, - _cta_per_sm_nowrite=cta_per_sm_nw if _any_split else None, - _num_loop_stages_write=num_loop_stages_w if _any_split else None, - _num_loop_stages_nowrite=num_loop_stages_nw if _any_split else None, - **extra_kwargs, - ) - - parts = [] - # When tied (not _any_split), emit the shared single-value tag - # (M=8 etc). When split, emit explicit Mw / Mnw tags so cells - # with the same shared value but different per-main values get - # unique JSON keys. - def _emit_split(name_w, name_nw, val_w, val_nw): - if val_w is None and val_nw is None: - return - if not _any_split or val_w == val_nw: - parts.append(f"{name_w[:-1]}={val_w}") # strip the 'w' suffix - else: - parts.append(f"{name_w}={val_w}") - parts.append(f"{name_nw}={val_nw}") - _emit_split("Mw", "Mnw", block_size_m_w, block_size_m_nw) - _emit_split("Ww", "Wnw", num_warps_w, num_warps_nw) - _emit_split("Sw", "Snw", num_stages_w, num_stages_nw) - if precompute_num_warps is not None: - parts.append(f"pW={precompute_num_warps}") - if precompute_num_stages is not None: - parts.append(f"pS={precompute_num_stages}") - if heads_per_block is not None: - parts.append(f"H={heads_per_block}") - if maxnreg is not None: - parts.append(f"R={maxnreg}") - if num_ctas is not None: - parts.append(f"CT={num_ctas}") - # Persistent-only knobs (only meaningful when MODE=persistent_main; - # printed unconditionally so output rows are uniformly comparable - # across modes when the user passed these sweeps). - _emit_split("CPSw", "CPSnw", cta_per_sm_w, cta_per_sm_nw) - _emit_split("LSw", "LSnw", num_loop_stages_w, num_loop_stages_nw) - if flatten is not None: - parts.append(f"FL={flatten}") - if warp_specialize is not None: - parts.append(f"WS={warp_specialize}") - # TMA sweep tags. Four wrapper-level flags map to three - # kernel-level constexprs (rect-load and replay-nowrite-load - # share `USE_TMA_LOAD_NOWRITE`, picked by the wrapper based on - # RECTANGLE). TMARL specifically gates the rectangle path's - # state load; TMANL specifically gates the replay-style - # nowrite path's state load. Distinct because their measured - # perf profiles differ (see CHECKPOINTING_DESIGN.md item #17: - # rect TMA is "not a win" while replay-nowrite TMA is the - # biggest measured win at int8 b>=64). - if use_tma_rect_load is not None: - parts.append(f"TMARL={use_tma_rect_load}") # rect path load - if use_tma_replay_write_load is not None: - parts.append(f"TMAWL={use_tma_replay_write_load}") # replay-write load - if use_tma_replay_nowrite_load is not None: - parts.append(f"TMANL={use_tma_replay_nowrite_load}") # replay-NOWRITE load (NOT rect) - if use_tma_replay_write_store is not None: - parts.append(f"TMAWS={use_tma_replay_write_store}") # replay-write store - parts.append(f"SR={1 if use_philox else 0}") - parts.append(f"RECT={1 if rectangle_for_nowrite else 0}") - parts.append(f"WC={1 if write_checkpoint else 0}") - parts.append(f"MODE={mode}") - parts.append(f"SORT={1 if sort_slots else 0}") - parts.append(f"REVN={1 if reverse_nowrite else 0}") - parts.append(f"HSORT={1 if hardcode_sort else 0}") - sweep_suffix = (" " + ",".join(parts)) if parts else "" - sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") - - reset_fn = _reset_conv1d_realistic if with_conv1d else _reset - # --cell-list filter: only time cells whose (canonical-knob-values) - # tuple is in the loaded set. Robust to bench gaining new knobs - # (old cell-list files keep working: any keys they don't list - # become wildcards that retain CLI defaults). - if args._cell_list_keys: - _tup = _current_cell_tuple(args, locals()) - if _tup is None or _tup not in args._cell_list_set: - continue - # Resume from JSONL: skip cells already recorded. Built the same - # way _print_row builds JSON keys; must stay in sync. - done_keys = getattr(args, "_done_keys", None) - if done_keys: - # One key per scenario (k=6, k=11, mix) — skip the whole cell - # only if ALL of its scenarios are already done. We don't - # know which scenarios will be emitted here without - # re-evaluating the inner scenario loop; conservatively skip - # only when the prev_k_for_print's specific key is done. - _resume_key = _build_json_key( - args.variant, batch, mtp_len, prev_k_for_print, - state_dtype_name, sweep_suffix, args.tp_size, - ) - if _resume_key in done_keys: - continue - if warmup_only: - reset_fn() - if scenario_pre_iter is not None: - scenario_pre_iter(0) - _run_incr() - torch.cuda.synchronize() - else: - # Inline retry: CUPTI sometimes loses records under PDL + - # high cell count; retrying the SAME cell often catches it - # because the failure is transient at the kernel-launch level. - # Per --cupti-retry budget. On final failure, append tag to - # the skipped list for an external rerun in a fresh process. - defer_results = ( - args.cuda_graph - and getattr(args, "cupti", True) - and int(getattr(args, "cupti_defer_depth", 1)) > 1 - ) - retry_budget = 0 if defer_results else max(0, getattr(args, "cupti_retry", 1)) - stats = None - expected_K = _kernels_per_iter_incremental( - mode, with_conv1d=with_conv1d, - persistent_skip_empty=scenario_skip_empty, - ) - plan_key = ( - "incremental", - args.variant, - mode, - batch, - mtp_len, - state_dtype_name, - act_dtype_name, - with_conv1d, - bool(args.l2_flush), - bool(args.external_pdl), - bool(args.internal_pdl), - bool(use_philox), - bool(rectangle_for_nowrite), - bool(write_checkpoint), - bool(sort_slots), - bool(reverse_nowrite), - bool(hardcode_sort), - scenario_pre_iter is not None, - expected_K, - ) - for attempt in range(retry_budget + 1): - stats = _time_kernel( - args, _run_incr, reset_fn, sweep_tag, - expected_K=expected_K, - pre_iter_fn=scenario_pre_iter, - pre_iter_group_factory=scenario_pre_iter_group_factory, - iters_override=scenario_iters, - cupti_plan_key=plan_key, - ) - if stats is not None: - break - if attempt < retry_budget: - print( - f"[retry] CUPTI mismatch on {sweep_tag!r}; " - f"retrying ({attempt + 1}/{retry_budget})", - file=sys.stderr, - flush=True, - ) - if stats is None: - args._skipped_cells.append(sweep_tag) - - # Attach n_writes_per_iter when it is needed for scoring. - # For pure scenarios, n_writes is constant: 0 (nowrite) or - # batch (write), determined by scn["fill"] + mtp_len > max_window. - # For mix, scn carries the precomputed per-iter array. - per_iter_nw = None - if stats is not None and ( - getattr(args, "json_detailed", False) or scn["fill"] is None - ): - eff_iters = scenario_iters if scenario_iters is not None else args.iters - if scn["fill"] is not None: - if getattr(args, "json_detailed", False): - # Pure scenario: constant n_writes for every iter. - is_write = (scn["fill"] + mtp_len > max_window) - per_iter_nw = [batch if is_write else 0] * eff_iters - else: - # Mix scenario: slice off warmup, keep timed iters. - nw_full = scn.get("n_writes_per_iter") - if nw_full is not None: - per_iter_nw = nw_full[args.warmup:args.warmup + eff_iters].tolist() - else: - per_iter_nw = None - - if stats is not None: - _submit_result_job( - args, - stats, - show_kernel_col=show_kernel_col, - kernel_name=args.variant, - batch=batch, - mtp_len=mtp_len, - prev_k=prev_k_for_print, - state_dtype_name=state_dtype_name, - act_dtype_name=act_dtype_name, - sweep_suffix=sweep_suffix, - per_iter_nw=per_iter_nw, - skipped_tag=sweep_tag, - ) - - -# Map full torch dtype name → short tag used in JSON keys (matches collect.py). -_DTYPE_SHORT = { - "float32": "fp32", "bfloat16": "bf16", "float16": "fp16", - "int8": "int8", "int16": "int16", "float8_e4m3fn": "fp8", -} - - -def _build_json_key( - kernel_name, batch, mtp_len, prev_k, state_dtype_name, sweep_suffix, tp_size -): - """Build a key matching collect.py's kernel_data.json convention: - - incremental/{batch}/{mtp}/{sd}/k{k}/{sweep_parts}/tp{tp} - triton/{batch}/{mtp}/{sd}/tp{tp} - flashinfer/{batch}/{mtp}/{sd}/tp{tp} - - `kernel_name` is what _print_row receives: variant name for the timed - kernel (replay/checkpointing) or baseline name for the baseline row. - Variant rows collapse to "incremental" — the variant choice is captured - by the sweep flags collect.py would otherwise apply via --variant. - """ - if kernel_name in ("replay", "checkpointing"): - kind = "incremental" - else: - kind = kernel_name # "triton" / "flashinfer" - - sd = _DTYPE_SHORT.get(state_dtype_name, state_dtype_name) - parts = [kind, str(batch), str(mtp_len), sd] - if prev_k != "N/A": - parts.append(f"k{prev_k}") - if sweep_suffix: - # sweep_suffix format: " M=4,W=1,S=1,SR=0,RECT=0,WC=1" - # collect.py format: "M4_W1_S1_SR0_RECT0_WC0" - # Strip leading/trailing whitespace, drop '=', commas → underscores. - parts.append( - sweep_suffix.strip().replace("=", "").replace(",", "_") - ) - parts.append(f"tp{tp_size}") - return "/".join(parts) - - -def _print_row( - show_kernel_col, - kernel_name, - batch, - mtp_len, - prev_k, - state_dtype_name, - act_dtype_name, - stats, - sweep_suffix="", - tp_size=None, - json_detailed=False, - jsonl_path=None, - jsonl_host=None, - jsonl_gpu=None, -): - """Print one summary row and append the result to the JSONL sidecar. - - `stats` is a dict from _time_kernel: {median, p95, p99, n, iters_us, - [n_writes_per_iter], [per_kernel]}. The summary table only shows the - headline percentiles. JSONL captures the compact per-iter spans + - n_writes_per_iter by default; with json_detailed=True it also captures - per-kernel data. - - When `jsonl_path` is provided, appends one JSON line per row to the - JSONL sidecar (crash-safe incremental persistence; lets a killed sweep - resume from the last completed cell on rerun, even across hosts). Open - per-write because `args` is pickled to ProcessPoolExecutor workers and - file handles aren't picklable. JSONL is the canonical artifact — the - bench no longer writes a final `.json` summary; use `jsonl_to_json.py` - if a one-shot `.json` snapshot is needed. - """ - kernel_col = f"{kernel_name:>11} | " if show_kernel_col else "" - print( - f"| {kernel_col}{batch:>5} | {mtp_len:>7} | {str(prev_k):>6} | " - f"{state_dtype_name:>11} | {act_dtype_name:>9} | " - f"{stats['median']:>9.2f} | {stats['p95']:>7.2f} | {stats['p99']:>7.2f} |" - f"{sweep_suffix}" - ) - if jsonl_path is not None: - key = _build_json_key( - kernel_name, batch, mtp_len, prev_k, state_dtype_name, - sweep_suffix, tp_size, - ) - if json_detailed: - row_stats = stats - else: - row_stats = { - k: stats[k] - for k in ("median", "p95", "p99", "n", "iters_us", "n_writes_per_iter") - if k in stats - } - if "host_timing" in stats: - row_stats["host_timing"] = stats["host_timing"] - # Append to JSONL sidecar if a path is set (incremental persistence). - # Open per-write because args is pickled to ProcessPoolExecutor - # workers, and file handles aren't picklable. A clean SIGTERM or - # Python exception will leave the file consistent up to the last - # newline; catastrophic kills can leave a partial last line, which - # the resume reader tolerates via json.JSONDecodeError pass. - if jsonl_path is not None: - # Wall-clock timestamp (float seconds since UNIX epoch) at write - # time. Lets post-hoc analysis diff consecutive rows to derive - # per-cell wall budget and identify startup-bound vs steady-state - # segments (cells/sec, downtime between bench invocations) without - # needing to instrument the bench's outer loops separately. - import time as _time - rec = {"key": key, "stats": row_stats, "t": _time.time()} - if jsonl_host is not None: - rec["host"] = jsonl_host - if jsonl_gpu is not None: - rec["gpu"] = jsonl_gpu - with open(jsonl_path, "a") as f: - f.write(json.dumps(rec) + "\n") - - -def _finish_result_job(args, job: dict) -> None: - result = job["result"] - if isinstance(result, _PendingCuptiStats): - stats = result.resolve() - else: - stats = result - - if stats is None: - skipped_tag = job.get("skipped_tag") - if skipped_tag is not None: - args._skipped_cells.append(skipped_tag) - return - - per_iter_nw = job.get("per_iter_nw") - if per_iter_nw is not None: - stats["n_writes_per_iter"] = per_iter_nw - - _print_row( - job["show_kernel_col"], - job["kernel_name"], - job["batch"], - job["mtp_len"], - job["prev_k"], - job["state_dtype_name"], - job["act_dtype_name"], - stats, - job.get("sweep_suffix", ""), - tp_size=args.tp_size, - json_detailed=getattr(args, "json_detailed", False), - jsonl_path=getattr(args, "_jsonl_path", None), - jsonl_host=getattr(args, "_jsonl_host", None), - jsonl_gpu=getattr(args, "_jsonl_gpu", None), - ) - - -def _drain_pending_results(args, *, force: bool = False) -> None: - pending_results = getattr(args, "_pending_results", None) - if not pending_results: - return - - max_pending = max(1, int(getattr(args, "cupti_defer_depth", 1))) - while pending_results: - first_result = pending_results[0]["result"] - should_block = force or len(pending_results) >= max_pending - if ( - not should_block - and isinstance(first_result, _PendingCuptiStats) - and not first_result.is_ready() - ): - break - job = pending_results.pop(0) - _finish_result_job(args, job) - - -def _submit_result_job( - args, - result, - *, - show_kernel_col, - kernel_name, - batch, - mtp_len, - prev_k, - state_dtype_name, - act_dtype_name, - sweep_suffix="", - per_iter_nw=None, - skipped_tag=None, -) -> None: - job = { - "result": result, - "show_kernel_col": show_kernel_col, - "kernel_name": kernel_name, - "batch": batch, - "mtp_len": mtp_len, - "prev_k": prev_k, - "state_dtype_name": state_dtype_name, - "act_dtype_name": act_dtype_name, - "sweep_suffix": sweep_suffix, - "per_iter_nw": per_iter_nw, - "skipped_tag": skipped_tag, - } - if isinstance(result, _PendingCuptiStats): - args._pending_results.append(job) - _drain_pending_results(args) - else: - _finish_result_job(args, job) - - -# Cell-list mode — canonical knob-key mapping to argparse args + local -# loop variable. See _load_cell_list_into_args / inner-loop filter. -# -# Each entry: cell-key → (args attribute name, comma-separated string flag) -# For split (write/nowrite) knobs, we use Xw / Xnw keys. Tied forms (M, W, -# S, CPS, LS) accepted on load and expanded to their w/nw variants. -_CELL_LIST_KEY_TO_ARG = { - "Mw": "block_size_m_write", - "Mnw": "block_size_m_nowrite", - "Ww": "num_warps_write", - "Wnw": "num_warps_nowrite", - "Sw": "num_stages_write", - "Snw": "num_stages_nowrite", - "CPSw": "cta_per_sm_write", - "CPSnw": "cta_per_sm_nowrite", - "LSw": "num_loop_stages_write", - "LSnw": "num_loop_stages_nowrite", - "pW": "precompute_num_warps", - "pS": "precompute_num_stages", - "H": "heads_per_block", - "R": "maxnreg", - "CT": "num_ctas", - "FL": "flatten", - "WS": "warp_specialize", - "TMARL": "use_tma_rect_load", - "TMAWL": "use_tma_replay_write_load", - "TMANL": "use_tma_replay_nowrite_load", - "TMAWS": "use_tma_replay_write_store", - "RECT": "rectangle_for_nowrite", - "WC": "write_modes", - "SORT": "sort_slots", - "REVN": "reverse_nowrite", - "HSORT": "hardcode_sort", - # MODE and SR get special handling (string values): - # MODE → args.modes (single mode name) - # SR → args.sr_modes ("RN" if 0, "SR" if 1) -} - -# Split-knob tied form: "M" expands to both "Mw" and "Mnw". -_CELL_LIST_TIED_EXPANSIONS = { - "M": ("Mw", "Mnw"), - "W": ("Ww", "Wnw"), - "S": ("Sw", "Snw"), - "CPS": ("CPSw", "CPSnw"), - "LS": ("LSw", "LSnw"), -} - - -def _normalize_cell(cell: dict) -> dict: - """Expand tied-form keys (M, W, S, CPS, LS) to their w/nw variants. - Returns a new dict with only canonical split-or-plain keys. - """ - out = dict(cell) - for tied, (w_key, nw_key) in _CELL_LIST_TIED_EXPANSIONS.items(): - if tied in out: - v = out.pop(tied) - out.setdefault(w_key, v) - out.setdefault(nw_key, v) - return out - - -def _load_cell_list_into_args(args) -> None: - """Read --cell-list JSON, normalize, override args.* knob ranges, and - populate args._cell_list_keys + args._cell_list_set for the inner-loop - filter. Errors out if cells aren't uniform (different key sets). - """ - with open(args.cell_list) as f: - raw = json.load(f) - if not isinstance(raw, list): - sys.exit(f"--cell-list: expected JSON list, got {type(raw).__name__}") - cells = [_normalize_cell(c) for c in raw] - if not cells: - print("[cell-list] empty list — nothing to time", file=sys.stderr) - return - # All cells must share the same key set (uniform schema) - keys0 = frozenset(cells[0].keys()) - for i, c in enumerate(cells[1:], start=1): - if frozenset(c.keys()) != keys0: - sys.exit( - f"--cell-list: cells must have uniform key sets; cell[0] " - f"has {sorted(keys0)} but cell[{i}] has {sorted(c.keys())}" - ) - - # Auto-cover: collect per-knob value set across all cells - cover: dict = {} - for c in cells: - for k, v in c.items(): - cover.setdefault(k, set()).add(v) - # Apply overrides - for key, vals in cover.items(): - if key in _CELL_LIST_KEY_TO_ARG: - arg_name = _CELL_LIST_KEY_TO_ARG[key] - vals_str = ",".join(str(v) for v in sorted(vals)) - setattr(args, arg_name, vals_str) - elif key == "MODE": - args.modes = ",".join(sorted({str(v) for v in vals})) - elif key == "SR": - args.sr_modes = ",".join(sorted({"SR" if v else "RN" for v in vals})) - else: - print(f"[cell-list] WARNING: unknown key {key!r} in cells; " - f"will not override any args.* attribute (the value will " - f"still be matched in the filter if a matching local var " - f"is in scope)", file=sys.stderr) - - # Canonical key order (sorted) for tuple matching in the inner loop - args._cell_list_keys = tuple(sorted(keys0)) - args._cell_list_set = { - tuple(c[k] for k in args._cell_list_keys) for c in cells - } - print(f"[cell-list] loaded {len(cells)} cells with keys " - f"{list(args._cell_list_keys)}; overrode args.* to auto-cover", - file=sys.stderr) - - -# Maps cell-list key → name of the local variable in _bench_config's inner -# loop. Used to extract the "current cell" tuple for the filter check. -# Keep in sync with the loop-variable names; the filter is lenient about -# missing names (it picks them up from the inner scope at runtime). -_CELL_LIST_KEY_TO_LOCAL = { - "Mw": "block_size_m_w", - "Mnw": "block_size_m_nw", - "Ww": "num_warps_w", - "Wnw": "num_warps_nw", - "Sw": "num_stages_w", - "Snw": "num_stages_nw", - "CPSw": "cta_per_sm_w", - "CPSnw": "cta_per_sm_nw", - "LSw": "num_loop_stages_w", - "LSnw": "num_loop_stages_nw", - "pW": "precompute_num_warps", - "pS": "precompute_num_stages", - "H": "heads_per_block", - "R": "maxnreg", - "CT": "num_ctas", - "FL": "flatten", - "WS": "warp_specialize", - "TMARL": "use_tma_rect_load", - "TMAWL": "use_tma_replay_write_load", - "TMANL": "use_tma_replay_nowrite_load", - "TMAWS": "use_tma_replay_write_store", - "RECT": "rectangle_for_nowrite", - "WC": "write_checkpoint", - "MODE": "mode", - "SORT": "sort_slots", - "REVN": "reverse_nowrite", - "HSORT": "hardcode_sort", - "SR": "use_philox", -} - - -def _current_cell_tuple(args, locals_dict: dict) -> tuple | None: - """Build the (key1=val1, key2=val2, ...) tuple for the current inner-loop - iteration, matching args._cell_list_keys' order. Used by the inner-loop - filter to check membership in args._cell_list_set. Returns None if any - expected local is missing (the bench evolved a knob name — caller skips). - """ - if not args._cell_list_keys: - return None - vals = [] - for k in args._cell_list_keys: - local_name = _CELL_LIST_KEY_TO_LOCAL.get(k, k) - if local_name not in locals_dict: - return None - v = locals_dict[local_name] - # Coerce bools to ints to match cell-list JSON (1/0) - if isinstance(v, bool): - v = int(v) - vals.append(v) - return tuple(vals) - - -# Main benchmark loop - - -def _run_benchmark(args) -> None: - # Phase-timing markers — emit timestamped checkpoints so a captured-stdout - # run can later attribute wall time to setup vs compile-warmup vs prewarm - # vs timing. Single-line format makes log-grepping trivial. - _phase_t0 = time.perf_counter() - def _phase(label: str) -> None: - dt = time.perf_counter() - _phase_t0 - print(f"[phase] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) - _phase("enter _run_benchmark") - - # Pending-results FIFO for srxl's deferred CUPTI parsing pipeline. Each - # entry holds a _PendingCuptiStats handle; _drain_pending_results pulls - # ready entries and routes them to _print_row (which appends to JSONL). - args._pending_results = [] - - # JSONL incremental sidecar. Path = `.jsonl`. Each completed - # cell appends one line `{"key": , "stats": {...}, "host": }` - # to this file as it finishes timing. On startup we read this sidecar (if - # present) and populate _done_keys so a killed bench can resume without - # redoing already-timed cells. Crash-safe by construction: append-only - # writes survive SIGTERM/SIGKILL/reboot mid-sweep. - # - # Resume is host-blind: _done_keys includes records from any host, so a - # bench restarted on a different node fills in the missing cells without - # redoing cells already covered elsewhere. Cross-host *timings* aren't - # directly comparable, but each JSONL record carries its `host` stamp so - # the analyzer can group/compare per host. This bench no longer writes a - # final `.json` summary — the JSONL is the canonical artifact; use the - # `jsonl_to_json.py` helper if a one-shot `.json` snapshot is needed. - # - # Note: we store only paths/strings on `args` because args is pickled to - # ProcessPoolExecutor workers during compile-warmup, and file handles - # (TextIOWrapper) aren't picklable. _print_row open-appends per cell. - args._jsonl_path = None - args._done_keys: set[str] = set() - args._jsonl_host = None # hostname stamp for the current run - args._jsonl_gpu = None # GPU device id stamp (current process visibility) - if getattr(args, "json_output", None): - import socket - args._jsonl_host = socket.gethostname() - # Capture GPU id once at startup. Used by the oracle-cache layer in - # search_driver to attribute timings to a specific (host, gpu) pair - # for cross-process pruning. os.environ['CUDA_VISIBLE_DEVICES'] - # is the right source pre-torch-init (it's what the harness sets); - # post-init we could use torch.cuda.current_device() but we keep it - # to env to avoid forcing a CUDA init at this point in startup. - args._jsonl_gpu = os.environ.get("CUDA_VISIBLE_DEVICES", "") - args._jsonl_path = args.json_output + ".jsonl" - # Read existing JSONL if present: load every record's key into the - # skip set regardless of host (gap-fill on a new node). - if os.path.exists(args._jsonl_path): - n_loaded = 0 - host_counts: dict[str, int] = {} - with open(args._jsonl_path) as f: - for line in f: - line = line.strip() - if not line: - continue - try: - rec = json.loads(line) - except json.JSONDecodeError: - # Tolerate partial last line from a crash mid-write. - continue - k = rec.get("key") - if k is None: - continue - args._done_keys.add(k) - n_loaded += 1 - rec_host = rec.get("host") - if rec_host: - host_counts[rec_host] = host_counts.get(rec_host, 0) + 1 - if n_loaded: - host_summary = ", ".join( - f"{h}={n}" for h, n in sorted(host_counts.items()) - ) if host_counts else "(no host stamps)" - print( - f"[resume] {args._jsonl_path}: loaded {n_loaded} prior " - f"cell results across hosts [{host_summary}]; sweep will " - f"skip them. New cells stamp host={args._jsonl_host}.", - file=sys.stderr, - ) - - # Sidecar metadata: cmd, host, tp_size, variant, cupti, etc. Written - # once at startup; helps later analysis identify how this JSONL was - # produced even though there's no top-level .json wrapper anymore. - meta_path = args.json_output + ".meta.json" - meta_payload = { - "timestamp": datetime.now().isoformat(), - "host": args._jsonl_host, - "cmd": " ".join(sys.argv), - "tp_size": getattr(args, "tp_size", None), - "warmup": getattr(args, "warmup", None), - "iters": getattr(args, "iters", None), - "variant": getattr(args, "variant", None), - "cupti": getattr(args, "cupti", False), - } - # Append to a list so successive runs (gap-fill, retry) keep history. - existing_meta = [] - if os.path.exists(meta_path): - try: - with open(meta_path) as f: - existing_meta = json.load(f) - if not isinstance(existing_meta, list): - existing_meta = [existing_meta] - except (OSError, json.JSONDecodeError): - existing_meta = [] - existing_meta.append(meta_payload) - # Bench is sometimes invoked with --json-output pointing into a dir - # the caller hasn't created (subprocess driver, search loop, etc.). - # Ensure the dir exists before writing the meta sidecar OR the JSONL. - os.makedirs(os.path.dirname(os.path.abspath(meta_path)), exist_ok=True) - tmp = meta_path + ".tmp" - with open(tmp, "w") as f: - json.dump(existing_meta, f, indent=2) - os.replace(tmp, meta_path) - - # Skipped cells accumulator — populated by _bench_config when CUPTI capture - # mismatch causes a cell to be skipped. Written to args.skipped_output - # (or derived from json_output) at end of run. - args._skipped_cells = [] - - # Cell-list filter (replaces the old --retry-cells tag-string filter). - # When set, the sweep iterates ONLY the cells described in the list. - # - # Each entry in the JSON file is a dict of canonical knob keys → values, - # using the same names that appear in the sweep_tag (Mw/Mnw, Ww/Wnw, - # Sw/Snw, pW, pS, H, R, CT, CPSw/CPSnw, LSw/LSnw, FL, WS, TMARL, - # TMAWL, TMANL, TMAWS, SR, RECT, WC, MODE, SORT, REVN, HSORT). Each - # cell may also use the tied forms M / W / S / CPS / LS (single value - # applied to both write and nowrite halves). - # - # On load we: - # - Override the bench's CLI knob args (`args.block_size_m_write`, - # etc.) with the union of values present across all cells per knob, - # so the cartesian iteration auto-covers the list. - # - Build `args._cell_list_keys` (the canonical key order used by - # every cell — must be uniform across the list) and - # `args._cell_list_set` (frozen tuples for O(1) membership check - # inside the inner loop). - # - # In the inner loop, we build the current iteration's tuple and skip - # cells not in the set. Dict-matching is robust to bench gaining new - # knobs (old cell-list files keep working — newly-added knobs simply - # aren't matched on, so they retain CLI defaults). - # Cell-list state may already have been populated by main() (so that - # the args.*_list derivations downstream see the override). Default to - # empty if not. - _phase(f"done loading _done_keys ({len(args._done_keys)} entries)") - - if not hasattr(args, "_cell_list_keys"): - args._cell_list_keys: tuple = () - args._cell_list_set: set = set() - if getattr(args, "cell_list", None): - _load_cell_list_into_args(args) - _phase(f"done loading cell-list ({len(args._cell_list_set)} cells)") - - assert args.nheads % args.tp_size == 0, ( - f"nheads ({args.nheads}) must be divisible by tp_size ({args.tp_size})" - ) - assert args.ngroups % args.tp_size == 0, ( - f"ngroups ({args.ngroups}) must be divisible by tp_size ({args.tp_size})" - ) - args.tp_nheads = args.nheads // args.tp_size - args.tp_ngroups = args.ngroups // args.tp_size - - batch_sizes = [int(x) for x in args.batch_sizes.split(",")] - mtp_lengths = [int(x) for x in args.mtp_lengths.split(",")] - - dtype_map = { - "bf16": torch.bfloat16, - "fp32": torch.float32, - "fp16": torch.float16, - "int8": torch.int8, - "int16": torch.int16, - "fp8": torch.float8_e4m3fn, - } - state_dtypes = [dtype_map[s] for s in args.state_dtypes.split(",")] - act_dtypes = [dtype_map[s] for s in args.act_dtypes.split(",")] - - # Resolve baseline function - if args.baseline == "flashinfer": - from flashinfer.mamba import selective_state_update as baseline_fn - elif args.baseline == "triton": - baseline_fn = selective_state_update - else: - baseline_fn = None - - # --with-conv1d uses its own realistic L2 flush (cold cache flush then - # hot in_proj write). Override the generic l2_flush to avoid double-flushing. - if args.with_conv1d: - args.l2_flush = False - _init_l2_flush() # still needed for the realistic reset's flush step - elif args.l2_flush: - _init_l2_flush() - - _phase("about to enter compile-warmup") - if args.compile_threads > 0: - _compile_warmup_phase( - args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, - baseline_fn, max_workers=args.compile_threads, - ) - _phase("returned from compile-warmup") - - # Pre-warm the per-(state_dtype, act_dtype, mtp_len, ...) tensor cache at - # the largest requested batch size. Without this, the timing loop would - # progressively grow the cache as it encounters larger batches (e.g., - # iterate 1 -> 16 -> 64 -> 128 -> 512 = 5 separate growth allocations, - # each freeing the previous buffers). Pre-warming at max-batch up front - # makes every subsequent timing cell a view-slice (zero alloc cost). - _max_batch = max(batch_sizes) - for state_dtype in state_dtypes: - for act_dtype in act_dtypes: - for mtp_len in mtp_lengths: - _build_tensors( - _max_batch, mtp_len, state_dtype, act_dtype, - args.tp_nheads, args.head_dim, args.d_state, args.tp_ngroups, - max_window=getattr(args, "max_window", None) or None, - ) - _phase("done tensor prewarm — entering timing") - - if args.profile: - torch.cuda.cudart().cudaProfilerStart() - - # Print header - if baseline_fn is not None: - print( - f"| {'kernel':>11} | {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " - f"{'state_dtype':>11} | {'act_dtype':>9} | " - f"{'median_us':>9} | {'p95_us':>7} | {'p99_us':>7} |" - ) - print( - f"|{'-' * 13}|{'-' * 7}|{'-' * 9}|{'-' * 8}|" - f"{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" - ) - else: - print( - f"| {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " - f"{'state_dtype':>11} | {'act_dtype':>9} | " - f"{'median_us':>9} | {'p95_us':>7} | {'p99_us':>7} |" - ) - print( - f"|{'-' * 7}|{'-' * 9}|{'-' * 8}|{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" - ) - - sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) - rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) - write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) - modes_list = getattr(args, "modes_list", ["monolithic"]) - sort_list = getattr(args, "sort_slots_list", [False]) - rev_list = getattr(args, "reverse_nowrite_list", [False]) - hsort_list = getattr(args, "hardcode_sort_list", [False]) - - # Pre-load AL distribution for mix mode (if --mix-csv set). - mix_al = None - mix_label = "" - if args.mix_csv is not None: - from pathlib import Path as _Path - from checkpoint_mix_sim import load_al_distribution as _load_al - mix_label = _Path(args.mix_csv).stem - # T (= mtp_len) varies per cell; load once with the LARGEST mtp so - # we have enough columns; the loader normalizes the dist anyway. - mix_al = _load_al(_Path(args.mix_csv), T=max(mtp_lengths), column=args.mix_csv_column) - - for batch in batch_sizes: - for mtp_len in mtp_lengths: - # Resolve prev_k fractions → clamped integers in [0, mtp_len] - prev_ks = _resolve_prev_ks(args, mtp_len) - - # Pre-generate mix samples once per (batch, mtp_len) cell so all - # tuning configs see the same per-iter prev_tokens vectors — - # tuning differences become signal, mix-noise is shared. - # Size the sample buffer for the LARGER of args.iters and - # args.mix_iters since mix scenarios use mix_iters. - mix_samples_cpu = None - perm_samples_cpu = None # per-iter slot perm sorted write-first - mix_samples_sorted_cpu = None # per-iter prev_tokens, write-first - if mix_al is not None: - from checkpoint_mix_sim import sample_steady_state_pnat as _sample_pnat - _max_window = getattr(args, "max_window", 0) or mtp_len - _max_iters = max(args.iters, getattr(args, "mix_iters", None) or args.iters) - mix_samples_cpu = _sample_pnat( - mix_al, T=mtp_len, window=_max_window, batch=batch, - K=args.warmup + _max_iters, seed=args.mix_seed, - ) - if any(sort_list) or any(hsort_list): - # write-first stable argsort: kind='stable' preserves - # original-slot order within each mode group. - write_mask = ( - mix_samples_cpu + mtp_len > _max_window - ).astype(np.int8) # 1 = write, 0 = nowrite - perm_idx = np.argsort( - -write_mask, kind="stable", axis=-1 - ).astype(np.int32) - if any(sort_list): - perm_samples_cpu = perm_idx - if any(hsort_list): - # Apply the perm to the prev_tokens samples themselves. - # Result row i = mix_samples_cpu[i] reordered such - # that write-mode entries come first. - mix_samples_sorted_cpu = np.take_along_axis( - mix_samples_cpu, perm_idx, axis=-1 - ).astype(mix_samples_cpu.dtype) - - for state_dtype in state_dtypes: - for act_dtype in act_dtypes: - for sr_mode in sr_modes_list: - for mode in modes_list: - # Non-monolithic modes ignore write_checkpoint - # (per-slot from PNAT) — collapse the sweep so we - # don't duplicate identical cells. - effective_write_modes = ( - write_modes_list if mode == "monolithic" else [True] - ) - for write_ckpt in effective_write_modes: - # Rectangle is meaningful for: nowrite cells in - # monolithic; always for dynamic / doublelaunch - # (constexpr knob). - if mode == "monolithic": - effective_rect_list = ( - [False] if write_ckpt else rect_list - ) - else: - effective_rect_list = rect_list - for rect in effective_rect_list: - # Sort/reverse only meaningful for the - # dl-family early-out kernels AND only - # against the mix scenario (the actual - # sort experiment). Pure k= scenarios - # under sort=1 would just run a - # USE_PERM=True kernel against an - # identity perm — same data point as - # sort=0 + extra compile. Skip sort=1 - # when no mix is configured; mono / - # dynamic also skip sort=1; reverse=1 - # with sort=0 is a no-op (skip). - # Note: "persistent_main" is included in - # the dl-family for sort/hsort sweep - # eligibility — it consumes the same - # slot_perm and benefits from the same - # write-first clustering. It additionally - # requires _n_writes (count of write - # slots) which the bench computes from - # the pure-scenario PNAT (mix scenarios - # not yet supported for persistent_main). - is_dl_family = mode in ( - "doublelaunch", "dlgrouped", "maindl", - "dl_write_only", "persistent_main", - "persistent_dynamic", - ) - can_sort = ( - is_dl_family and mix_samples_cpu is not None - ) - effective_sort_list = ( - sort_list if can_sort else [False] - ) - effective_hsort_list = ( - hsort_list if can_sort else [False] - ) - for sort_slots in effective_sort_list: - for hardcode_sort in effective_hsort_list: - # sort_slots and hardcode_sort - # are alternative experiments - # for the same idea — skip the - # combined cell to avoid double - # interpretation. - if sort_slots and hardcode_sort: - continue - # rev=1 is meaningful with EITHER - # sort_slots=1 (perm-based) or - # hardcode_sort=1 (raw pid_b - # subtraction in unsorted-perm - # path). rev=1 with both 0 is - # a no-op. - effective_rev_list = ( - rev_list if (sort_slots or hardcode_sort) else [False] - ) - for reverse_nowrite in effective_rev_list: - _bench_config( - args, batch, mtp_len, - prev_ks, state_dtype, - act_dtype, baseline_fn, - sr_mode=sr_mode, - rectangle_for_nowrite=rect, - write_checkpoint=write_ckpt, - mode=mode, - mix_samples_cpu=mix_samples_cpu, - mix_label=mix_label, - sort_slots=sort_slots, - reverse_nowrite=reverse_nowrite, - perm_samples_cpu=perm_samples_cpu, - hardcode_sort=hardcode_sort, - mix_samples_sorted_cpu=mix_samples_sorted_cpu, - ) - - _drain_pending_results(args, force=True) - - if args.profile: - torch.cuda.cudart().cudaProfilerStop() - - # JSONL is the canonical artifact (written incrementally per cell with - # host stamps). No clean-exit `.json` write — use `jsonl_to_json.py` to - # materialize a snapshot when an analyzer wants one. - if args.json_output and args._jsonl_path is not None: - print(f"\nJSONL results: {args._jsonl_path} " - f"(meta sidecar: {args.json_output}.meta.json)") - - # Write the skipped-cells sidecar. Caller can convert this list to a - # --cell-list JSON (one dict per skipped cell) to drive a retry pass in - # a fresh process. - skipped_path = getattr(args, "skipped_output", None) - if skipped_path is None and args.json_output: - # Derive default: foo.json -> foo.skipped.json - skipped_path = args.json_output.rsplit(".", 1)[0] + ".skipped.json" - if skipped_path is not None and args._skipped_cells: - payload = { - "metadata": { - "timestamp": datetime.now().isoformat(), - "cmd": " ".join(sys.argv), - "skipped_count": len(args._skipped_cells), - }, - "skipped": args._skipped_cells, - } - tmp = skipped_path + ".tmp" - with open(tmp, "w") as f: - json.dump(payload, f, indent=2) - os.replace(tmp, skipped_path) - print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " - f"tags written to: {skipped_path}", file=sys.stderr) - elif args._skipped_cells: - # No output path but there are skipped cells — emit a stderr summary. - print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " - f"first 5: {args._skipped_cells[:5]}", file=sys.stderr) - - -# CLI - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Benchmark replay_selective_state_update Triton kernel", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - parser.add_argument( - "--nheads", - type=int, - default=NHEADS, - help="Full-model nheads (divided by --tp-size for per-GPU slice)", - ) - parser.add_argument( - "--ngroups", - type=int, - default=NGROUPS, - help="Full-model ngroups (divided by --tp-size for per-GPU slice)", - ) - parser.add_argument( - "--head-dim", type=int, default=HEAD_DIM, help="Head dimension (not TP-split)" - ) - parser.add_argument( - "--d-state", type=int, default=D_STATE, help="SSM state dimension (not TP-split)" - ) - parser.add_argument( - "--tp-size", - type=int, - default=TP_SIZE, - help="Tensor parallel size; divides nheads and ngroups", - ) - parser.add_argument( - "--batch-sizes", default="1,2,4,8", help="Comma-separated decode batch sizes" - ) - parser.add_argument( - "--mtp-lengths", - default="1,2,4,8", - help="Comma-separated per-request sequence lengths (num_draft_tokens + 1 target)", - ) - parser.add_argument( - "--state-dtypes", - default="fp32", - help="Comma-separated state dtypes: fp16,bf16,fp32,int8,int16,fp8. " - "Quantized dtypes (int8/int16/fp8) require the checkpointing variant " - "and skip baselines (selective_state_update doesn't accept them).", - ) - parser.add_argument( - "--act-dtypes", - default="bf16", - help="Comma-separated activation dtypes for x/B/C/dt: fp32,bf16", - ) - parser.add_argument("--warmup", type=int, default=4, - help="Number of warmup iterations. Default aligns with " - "the graph group-iters (default 4 for mix scenarios) so " - "warmup + iters / mix-iters lands on a clean multiple " - "without per-args rounding overhead. Earlier default of " - "20 was overkill for steady-state warming.") - parser.add_argument("--iters", type=int, default=100, help="Number of timed iterations") - parser.add_argument( - "--compile-threads", - type=int, - default=64, - help="Number of THREADS used in the compile-warmup phase (one call " - "per (batch, mtp_len, prev_k, dtype, sweep) cell, parallelized over " - "N threads). Triton compile releases the GIL, so threads compile " - "in parallel and populate the persistent cache for free hits during " - "the sequential timed phase. 0 disables the phase. Default 64.", - ) - parser.add_argument( - "--mp-start-method", - choices=("spawn", "forkserver"), - default="spawn", - help="multiprocessing start method for compile-warmup workers AND " - "the CUPTI parser child process. 'spawn' (default) is robust but " - "each child re-imports the bench module (~15s torch+triton import " - "cost). 'forkserver' starts a server once, preloads the bench " - "module ONCE, then forks children cheaply (~1s each). When 4 " - "benches run concurrently with --compile-threads 26 each, spawn " - "still incurs 4*26=104 imports per round; forkserver cuts this to " - "4 (one per server).", - ) - parser.add_argument( - "--profile", - action="store_true", - help="Wrap timed region in cudaProfilerStart/Stop (for ncu --target-processes all)", - ) - parser.add_argument( - "--l2-flush", - action=argparse.BooleanOptionalAction, - default=True, - help="L2 eviction between iterations", - ) - parser.add_argument( - "--cuda-graph", - action=argparse.BooleanOptionalAction, - default=True, - help="Capture all warmup + timed iterations in a " - "single CUDA graph with per-iteration events " - "inside the graph, eliminating all host overhead.", - ) - parser.add_argument( - "--cuda-graph-group-iters", - type=int, - default=None, - help="Capture this many logical benchmark iterations per graph " - "replay when warmup + iters is divisible by this value. Default " - f"auto-selects {_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE} for pure " - f"cells and {_DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX} for mix cells. " - "Mix cells use a per-replay device window so they can group " - "iterations too.", - ) - parser.add_argument( - "--cupti", - action=argparse.BooleanOptionalAction, - default=True, - help="Time kernels via CUPTI Activity API (1 ns from the GPU " - "profiling fabric); per-iter span = max(kernel_end) - " - "min(kernel_start). Default ON. --no-cupti disables in-bench " - "timing entirely (kernels still run, but median/p95/p99 are zero) " - "— use when wrapping the bench in nsys/ncu, where the external " - "profiler provides timings and our CUPTI subscriber would conflict.", - ) - parser.add_argument( - "--cupti-flush-period-ms", - type=int, - default=0, - help="If >0, ask CUPTI to periodically flush activity buffers during " - "the timed CUDA-graph region. This can overlap raw-buffer parsing with " - "long timed cells; 0 leaves flushing explicit at the end of each cell.", - ) - parser.add_argument( - "--cupti-defer-depth", - type=int, - default=4, - help="Maximum number of CUDA-graph CUPTI timing results that may be " - "left for the parser process while the main process starts later cells. " - "1 preserves synchronous per-cell parsing and inline retry behavior.", - ) - parser.add_argument( - "--json-output", - default=None, - help="If set, write per-cell results to this JSON file in the " - "shape consumed by collect.py / report.py. See the 'JSON output " - "schema' section at the top of this file.", - ) - parser.add_argument( - "--json-detailed", - action=argparse.BooleanOptionalAction, - default=False, - help="When --json-output is set, also include per_kernel " - "(per-iter relative start/end timestamps for each kernel). Compact " - "JSON always includes iters_us, and mix rows include " - "n_writes_per_iter. Default off keeps records compact.", - ) - parser.add_argument( - "--host-timing", - action=argparse.BooleanOptionalAction, - default=False, - help="Attach benchmark host-side phase timings to JSON/JSONL results. " - "Useful for diagnosing benchmark overhead, but it adds roughly 1 KB " - "per compact JSONL row and several perf_counter calls per cell.", - ) - parser.add_argument( - "--cupti-retry", - type=int, - default=1, - help="On CUPTI capture mismatch (kernel record count != expected), " - "retry the cell this many times in-process before giving up. CUPTI " - "gets racy after thousands of cells in one process (PDL + small " - "kernels occasionally lose records); a single retry usually catches " - "transient cases. Set 0 to disable and skip on first mismatch.", - ) - parser.add_argument( - "--skipped-output", - default=None, - help="Path to write the list of cells that failed CUPTI capture even " - "after --cupti-retry retries (JSON list of sweep_tag strings). " - "Default: derived from --json-output by replacing .json with " - ".skipped.json.", - ) - parser.add_argument( - "--cell-list", - default=None, - help="Path to a JSON list of cell dicts (one per cell to time). " - "Each dict has canonical knob keys → values: Mw, Mnw, Ww, Wnw, Sw, " - "Snw, pW, pS, H, R, CT, CPSw, CPSnw, LSw, LSnw, FL, WS, TMARL, " - "TMAWL, TMANL, TMAWS, SR, RECT, WC, MODE, SORT, REVN, HSORT (tied " - "forms M / W / S / CPS / LS are also accepted and auto-expanded). " - "When set, bench's CLI knob ranges are auto-overridden to the " - "per-knob union across all cells, and the inner-loop filter skips " - "any iteration whose knob-value tuple isn't in the list. All cells " - "must share the same key set (uniform schema).", - ) - parser.add_argument( - "--prev-tokens-fracs", - default="0,0.5,1.0", - type=lambda s: [float(x) for x in s.split(",")], - help="Fractions of mtp_len to use as prev_num_accepted_tokens " - "for the replay kernel sweep. Values are rounded " - "and clamped to [0, mtp_len].", - ) - parser.add_argument( - "--baseline", - default=None, - nargs="?", - const="triton", - choices=[None, "triton", "flashinfer"], - help="Baseline to benchmark alongside the replay kernel. " - "'triton': native Triton selective_state_update. " - "'flashinfer': flashinfer selective_state_update (same signature). " - "Pass --baseline alone for 'triton'. Default: no baseline.", - ) - parser.add_argument( - "--output", - default=None, - help="Path to save results (file or directory). " - "If a directory, writes benchmark_replay_.txt inside it.", - ) - parser.add_argument( - "--block-size-m", - type=str, - default=None, - help="Override BLOCK_SIZE_M: single value or comma-separated sweep (e.g. '4,8,16,32').", - ) - parser.add_argument( - "--num-warps", - type=str, - default=None, - help="Override num_warps: single value or comma-separated sweep (e.g. '1,2,4').", - ) - parser.add_argument( - "--internal-pdl", - action=argparse.BooleanOptionalAction, - default=True, - help="Internal PDL between precompute and main kernels (default: on).", - ) - parser.add_argument( - "--num-stages", - type=str, - default=None, - help="Override num_stages for the main kernel (comma-separated sweep).", - ) - parser.add_argument( - "--block-size-m-write", type=str, default=None, - help="Sweep BLOCK_SIZE_M for the WRITE main only (overrides --block-size-m " - "for the write half). Tied to --block-size-m if unset.", - ) - parser.add_argument( - "--block-size-m-nowrite", type=str, default=None, - help="Sweep BLOCK_SIZE_M for the NOWRITE main only. Tied to --block-size-m if unset.", - ) - parser.add_argument( - "--num-warps-write", type=str, default=None, - help="Sweep num_warps for the WRITE main only. Tied to --num-warps if unset.", - ) - parser.add_argument( - "--num-warps-nowrite", type=str, default=None, - help="Sweep num_warps for the NOWRITE main only. Tied to --num-warps if unset.", - ) - parser.add_argument( - "--num-stages-write", type=str, default=None, - help="Sweep num_stages for the WRITE main only. Tied to --num-stages if unset.", - ) - parser.add_argument( - "--num-stages-nowrite", type=str, default=None, - help="Sweep num_stages for the NOWRITE main only. Tied to --num-stages if unset.", - ) - parser.add_argument( - "--cta-per-sm-write", type=str, default=None, - help="Sweep cta_per_sm for the WRITE persistent_main only. Tied to --cta-per-sm if unset.", - ) - parser.add_argument( - "--cta-per-sm-nowrite", type=str, default=None, - help="Sweep cta_per_sm for the NOWRITE persistent_main only. Tied to --cta-per-sm if unset.", - ) - parser.add_argument( - "--num-loop-stages-write", type=str, default=None, - help="Sweep num_loop_stages for the WRITE persistent_main only. Tied to --num-loop-stages if unset.", - ) - parser.add_argument( - "--num-loop-stages-nowrite", type=str, default=None, - help="Sweep num_loop_stages for the NOWRITE persistent_main only. Tied to --num-loop-stages if unset.", - ) - parser.add_argument( - "--skip-diagonal", action=argparse.BooleanOptionalAction, default=False, - help="When sweeping any per-main *_write / *_nowrite knobs, skip cells " - "where ALL splittable knobs satisfy write_value == nowrite_value (i.e. " - "the 'diagonal' that's already covered by a prior shared-knob sweep). " - "Useful for incremental sweeps that extend earlier results without redoing " - "the tied-knob cells.", - ) - parser.add_argument( - "--precompute-num-warps", - type=str, - default=None, - help="Override num_warps for precompute kernel (comma-separated sweep).", - ) - parser.add_argument( - "--precompute-num-stages", - type=str, - default=None, - help="Override num_stages for precompute kernel (comma-separated sweep).", - ) - parser.add_argument( - "--max-window", - type=int, - default=16, - help="Cache T-axis capacity (max replay buffer length). Default 16 " - "matches Nemotron-3-Super-120B production. Pass 0 to fall back to " - "mtp_len (degenerate every-step-checkpoint case, mostly unused).", - ) - parser.add_argument( - "--prev-tokens-int", - type=lambda s: [int(x) for x in s.split(",")] if s else None, - default=None, - help="Absolute prev_num_accepted_tokens values to test, comma-separated " - "(e.g. '0,10,11,16'). Clamped to [0, max_window]. When set, " - "overrides --prev-tokens-fracs.", - ) - parser.add_argument( - "--write-checkpoint", - action=argparse.BooleanOptionalAction, - default=True, - help="Whether the checkpointing kernel should write the post-replay " - "state to HBM. True = checkpoint step (default). False = " - "non-checkpoint step (skip state HBM write + Philox). No effect on " - "the replay variant. Ignored if --write-modes is set.", - ) - parser.add_argument( - "--write-modes", - type=str, - default=None, - help="Comma-separated 0/1 values to sweep both write modes in a " - "single nsys process — for apples-to-apples comparison of write " - "vs nowrite (replay) vs nowrite (rectangle) within one timeline. " - "Skips silently for (write=False, prev_k+T>max_window) combos. " - "When set, overrides --write-checkpoint.", - ) - parser.add_argument( - "--with-conv1d", - action="store_true", - help="Include conv1d kernel before replay SSM. " - "Uses realistic L2 flush: cold caches flushed, hot in_proj output " - "kept warm. Measures conv1d → precompute → main span.", - ) - parser.add_argument( - "--external-pdl", - action=argparse.BooleanOptionalAction, - default=True, - help="External PDL: conv1d launches dependents, precompute waits. " - "Only relevant with --with-conv1d. --no-external-pdl disables.", - ) - parser.add_argument( - "--heads-per-block", - type=str, - default=None, - help="Override HEADS_PER_BLOCK for precompute kernel (comma-separated sweep).", - ) - parser.add_argument( - "--maxnreg", - type=str, - default=None, - help="Override maxnreg for the main kernel (comma-separated sweep).", - ) - parser.add_argument( - "--num-ctas", - type=str, - default=None, - help="Override num_ctas for the main kernel (comma-separated sweep).", - ) - parser.add_argument( - "--cta-per-sm", - type=str, - default=None, - help="CTAs per SM in the 1D persistent grid for mode=persistent_main " - "(comma-separated sweep). num_persistent = cta_per_sm × NUM_SMS. " - "Default = 1 (one CTA per SM). Replaces the old --num-persistent. " - "Ignored for non-persistent_main modes.", - ) - parser.add_argument( - "--num-loop-stages", - type=str, - default=None, - help="num_stages on the inner tl.range(...) persistent loop for " - "mode=persistent_main (comma-separated sweep). Default = 2. Note: " - "this is loop-level, NOT the kernel-arg num_stages (which only " - "pipelines dot-feeding loads). Watch Triton issue #8259 — " - "num_stages>1 + flatten=True can corrupt stores in non-dot kernels. " - "Ignored for non-persistent_main modes.", - ) - parser.add_argument( - "--flatten", - type=str, - default=None, - help="`flatten` arg on tl.range(...) for mode=persistent_main " - "(comma-separated 0/1 sweep). Default = 1. Ignored for " - "non-persistent_main modes.", - ) - parser.add_argument( - "--warp-specialize", - type=str, - default=None, - help="`warp_specialize` arg on tl.range(...) for mode=persistent_main " - "(comma-separated 0/1 sweep). Default = 0. Triton 3.6 only " - "supports it on simple matmul loops; our scan loop probably won't " - "pattern-match — exposed as a knob for sweep experiments. Requires " - "num_warps >= 4 if 1. Ignored for non-persistent_main modes.", - ) - parser.add_argument( - "--sr-modes", - type=str, - default="RN", - help="Comma-separated rounding modes to sweep: any combination of " - "{RN, SR}. SR (stochastic rounding) is silently skipped for state " - "dtypes that don't support it (bf16, fp32). Default 'RN' matches " - "legacy --philox-rounding=False behavior.", - ) - parser.add_argument( - "--rectangle-for-nowrite", - type=str, - default="0", - help="Comma-separated 0/1 values: 0 = replay-style nowrite kernel, " - "1 = dedicated rectangle nowrite kernel. Sweep both with '0,1' to " - "compare in one invocation. Silently no-op for write cells (the " - "write path always uses replay-style). Only applies to the " - "checkpointing variant.", - ) - parser.add_argument( - "--use-tma-rect-load", - type=str, - default=None, - help="Comma-separated 0/1 sweep. Use TMA (host-built tensor " - "descriptor) for state load in the rectangle nowrite path. " - "Cells where the rect path isn't reachable (e.g. mode=monolithic " - "+ WC=True) skip the value=1 case as a dupe.", - ) - parser.add_argument( - "--use-tma-replay-write-load", - type=str, - default=None, - help="Comma-separated 0/1 sweep. TMA state LOAD in replay main " - "when WC=True. Independent from nowrite-load and rect TMA — see " - "CHECKPOINTING_DESIGN.md item #17 for measured perf.", - ) - parser.add_argument( - "--use-tma-replay-nowrite-load", - type=str, - default=None, - help="Comma-separated 0/1 sweep. TMA state LOAD in replay main " - "when WC=False. Design doc reports the largest win on this path " - "(int8 b>=64: -8 to -12%%).", - ) - parser.add_argument( - "--use-tma-replay-write-store", - type=str, - default=None, - help="Comma-separated 0/1 sweep. TMA state STORE in replay main " - "(WC=True path only — no-op for WC=False). Independent from all " - "load TMA flags.", - ) - parser.add_argument( - "--modes", - type=str, - default="monolithic", - help="Comma-separated dispatch modes to sweep, any of " - "{monolithic,dynamic,doublelaunch}. monolithic = today's behavior " - "(one kernel pair, write_checkpoint applied to whole batch); " - "dynamic = single kernel pair that dispatches per-slot at runtime " - "based on PNAT (rectangle_for_nowrite picks RECTANGLE constexpr); " - "doublelaunch = two kernel pairs launched in sequence with " - "EARLY_OUT=True, each handling slots whose mode matches it. " - "Only applies to the checkpointing variant; non-monolithic modes " - "ignore --write-modes (per-slot from PNAT).", - ) - parser.add_argument( - "--mix-csv", - type=str, - default=None, - help="Path to AL histogram CSV (cols: AL, count). When set, an " - "additional 'mix' cell is emitted per (batch, mtp, dtype, sr, " - "mode, RECT, M, W, ...) combo where prev_tokens varies per iter, " - "drawn from the steady-state PNAT distribution induced by the " - "AL histogram. Mix cells run only on dynamic and doublelaunch " - "modes (mono on a mixed batch corrupts wrong-mode slots). " - "Each iteration of the captured CUDA graph has a different " - "pre-baked prev_tokens vector; warmup iters use distinct samples " - "from the timed iters so nsys-included warmup leaks don't bias.", - ) - parser.add_argument( - "--mix-csv-column", - type=int, - default=1, - help="Column index (0-based) in the AL histogram CSV for the " - "count/probability column. Default 1 (second column).", - ) - parser.add_argument( - "--mix-seed", - type=int, - default=42, - help="RNG seed for the steady-state PNAT sampler. Same seed " - "across runs => same per-slot samples for reproducible " - "comparisons.", - ) - parser.add_argument( - "--sort-slots", - type=str, - default="0", - help="Comma-separated 0/1. When 1, mix scenarios pre-sort slots " - "write-first (write slots at the head of slot_perm, nowrite at the " - "tail) and the dl-family kernels read pid_b through that perm — " - "clusters early-outs at one end of the grid. Only meaningful for " - "doublelaunch/dlgrouped/maindl with mix scenarios; mono/dynamic " - "and pure-batch cells skip sort=1.", - ) - parser.add_argument( - "--reverse-nowrite", - type=str, - default="0", - help="Comma-separated 0/1. When 1 (and --sort-slots 1), the " - "nowrite-side kernels in dlgrouped/doublelaunch/maindl walk the " - "perm in reverse so both halves of the dl chain front-load real " - "work. reverse=1 with sort=0 is skipped (no perm to reverse).", - ) - parser.add_argument( - "--hardcode-sort", - type=str, - default="0", - help="Comma-separated 0/1. When 1, the per-iter prev_tokens " - "samples are pre-sorted write-first OFFLINE (CPU-side) before " - "the timed region — kernel runs unchanged (USE_PERM=False) but " - "the EO gate sees sorted PNAT so early-outs cluster naturally. " - "Zero per-program load cost vs --sort-slots; output is " - "scrambled (we don't permute x/B/C/dt) but timing is meaningful. " - "Used to isolate whether clustering helps independent of the " - "perm-load overhead in the sort-slots path.", - ) - parser.add_argument( - "--mix-iters", - type=int, - default=None, - help="Iteration count override for mix scenarios (each iter is a " - "different per-slot prev_tokens draw). Default (None) uses " - "--iters. Mix scenarios benefit from more iters since each " - "iter samples a different mix; pure scenarios don't.", - ) - parser.add_argument( - "--mix-only", - action=argparse.BooleanOptionalAction, - default=False, - help="When --mix-csv is set, emit only mix scenarios and skip the " - "pure prev_k sibling scenarios. Default: false.", - ) - parser.add_argument( - "--philox-rounding", - action="store_true", - help="DEPRECATED — equivalent to --sr-modes SR. Retained for " - "backward compatibility; use --sr-modes for new scripts. fp16 SR " - "and fp8 SR require sm_100a (Blackwell B200+).", - ) - parser.add_argument( - "--philox-rounds", - type=int, - default=5, - help="Number of Philox PRNG rounds. Default 5 matches the " - "Nemotron-3-Super-120B production config (mamba_ssm_philox_rounds=5 " - "in examples/configs and tests/integration/perf configs). The " - "wrapper's generic fallback default is 10; callers without explicit " - "config see 10. Only consulted when --philox-rounding is enabled.", - ) - parser.add_argument( - "--variant", - choices=["replay", "checkpointing"], - default="replay", - help="Which kernel to time as the 'replay' row. 'replay' = today's " - "kernel (selective_state_update.py:replay). 'checkpointing' = " - "checkpointing_state_update.py. Both share the same wrapper signature.", - ) - parser.add_argument( - "--full-import", - action="store_true", - help="Use standard tensorrt_llm import path instead of fast direct " - "module loading. Slower (~40s startup) but guaranteed correct " - "if the fast path breaks due to package changes.", - ) - args = parser.parse_args() - if args.mix_only and args.mix_csv is None: - parser.error("--mix-only requires --mix-csv") - - # Round iter counts up so warmup + iters (and warmup + mix_iters) are clean - # multiples of the graph group-iters used downstream. Default mix group is - # 4, default pure group is 2. An explicit --cuda-graph-group-iters can - # request a larger group. We round to the max of the two so all scenarios - # in a single run (pure + mix) share a clean total_iters. The cost is at - # most (group-1) extra iters per scenario — negligible — and the win is - # that graph_group_iters never falls back to 1 (which caused ~5x slowdown - # in observed benchmark walls). - _group_for_rounding = max( - _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX, - _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE, - getattr(args, "cuda_graph_group_iters", None) or 0, - ) - def _round_iters_to_group(name, val): - total = args.warmup + val - if total % _group_for_rounding == 0: - return val - new_total = ((total + _group_for_rounding - 1) // _group_for_rounding) * _group_for_rounding - new_val = new_total - args.warmup - print(f"[bench] rounding --{name} {val} → {new_val} so warmup+{name} " - f"({new_total}) is a multiple of graph group_iters={_group_for_rounding}", - file=sys.stderr) - return new_val - args.iters = _round_iters_to_group("iters", args.iters) - if getattr(args, "mix_iters", None): - args.mix_iters = _round_iters_to_group("mix-iters", args.mix_iters) - - # Cell-list (if any) must be applied BEFORE the post-argparse string→list - # derivations below — those build args.*_list from args.* strings, so a - # cell-list override of e.g. args.modes='maindl' needs to land before - # args.modes_list is computed. The function populates args._cell_list_keys - # and args._cell_list_set, plus overrides args.* knob strings to the - # per-knob union of values across the listed cells. - if getattr(args, "cell_list", None): - _load_cell_list_into_args(args) - - # Backward-compat: --philox-rounding implies --sr-modes SR if --sr-modes - # was left at the default. If both are set explicitly, error. - sr_modes_default = (args.sr_modes == "RN") - if args.philox_rounding: - if not sr_modes_default and args.sr_modes != "SR": - parser.error( - "--philox-rounding (deprecated) is incompatible with explicit " - f"--sr-modes={args.sr_modes!r}. Use --sr-modes SR (or " - "RN,SR) instead and drop --philox-rounding." - ) - args.sr_modes = "SR" - - sr_modes = [m.strip() for m in args.sr_modes.split(",") if m.strip()] - for m in sr_modes: - if m not in ("RN", "SR"): - parser.error(f"--sr-modes value must be RN or SR, got {m!r}") - args.sr_modes_list = sr_modes - - rect_modes = [v.strip() for v in args.rectangle_for_nowrite.split(",") if v.strip()] - rect_list = [] - for v in rect_modes: - if v not in ("0", "1"): - parser.error(f"--rectangle-for-nowrite value must be 0 or 1, got {v!r}") - rect_list.append(v == "1") - args.rectangle_for_nowrite_list = rect_list - - sort_modes = [v.strip() for v in (args.sort_slots or "0").split(",") if v.strip()] - sort_list = [] - for v in sort_modes: - if v not in ("0", "1"): - parser.error(f"--sort-slots value must be 0 or 1, got {v!r}") - sort_list.append(v == "1") - args.sort_slots_list = sort_list - - rev_modes = [v.strip() for v in (args.reverse_nowrite or "0").split(",") if v.strip()] - rev_list = [] - for v in rev_modes: - if v not in ("0", "1"): - parser.error(f"--reverse-nowrite value must be 0 or 1, got {v!r}") - rev_list.append(v == "1") - args.reverse_nowrite_list = rev_list - - hsort_modes = [v.strip() for v in (args.hardcode_sort or "0").split(",") if v.strip()] - hsort_list = [] - for v in hsort_modes: - if v not in ("0", "1"): - parser.error(f"--hardcode-sort value must be 0 or 1, got {v!r}") - hsort_list.append(v == "1") - args.hardcode_sort_list = hsort_list - - if args.write_modes is not None: - wm = [v.strip() for v in args.write_modes.split(",") if v.strip()] - write_list = [] - for v in wm: - if v not in ("0", "1"): - parser.error(f"--write-modes value must be 0 or 1, got {v!r}") - write_list.append(v == "1") - args.write_modes_list = write_list - else: - args.write_modes_list = [args.write_checkpoint] - - modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] - valid_modes = { - "monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl", - "dl_write_only", "persistent_main", "persistent_dynamic", - } - for m in modes_raw: - if m not in valid_modes: - parser.error( - f"--modes value must be one of {sorted(valid_modes)}, got {m!r}" - ) - args.modes_list = modes_raw or ["monolithic"] - return args - - -class _Tee: - """Write to both stdout and a file simultaneously.""" - - def __init__(self, path: str): - parent = os.path.dirname(path) - if parent: - os.makedirs(parent, exist_ok=True) - self._file = open(path, "w") # noqa: SIM115 - self._stdout = sys.stdout - - def write(self, data): - self._stdout.write(data) - self._file.write(data) - - def flush(self): - self._stdout.flush() - self._file.flush() - - def close(self): - self._file.close() - - -if __name__ == "__main__": - _args = _parse_args() - - # Configure multiprocessing start method early — must be before any - # mp.get_context() that uses the chosen method. For forkserver, also - # add this file's dir to sys.path so the forkserver can import this - # module by basename for preload (otherwise it tries to import - # __main__, which is a different beast across processes). - if _args.mp_start_method == "forkserver": - sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - mp.set_start_method("forkserver", force=True) - try: - mp.set_forkserver_preload([ - "benchmark_replay_selective_state_update", - ]) - except Exception as _e: - print(f"[warn] set_forkserver_preload failed: {_e!r}; " - f"forks will still work but pay full import cost", - file=sys.stderr) - _MP_START_METHOD = _args.mp_start_method - - _out_path = None - if _args.output != "-": - _ts = datetime.now().strftime("%Y%m%d_%H%M%S") - _fname = f"benchmark_replay_{_ts}.txt" - if _args.output is None: - _out_path = os.path.expanduser(f"~/nemo_logs/{_fname}") - elif os.path.isdir(_args.output) or _args.output.endswith("/"): - _out_path = os.path.join(_args.output, _fname) - else: - _out_path = _args.output - - if _out_path: - _tee = _Tee(_out_path) - sys.stdout = _tee - print(f"# benchmark_replay_selective_state_update {datetime.now().isoformat()}") - print(f"# cmd: {' '.join(sys.argv)}") - - try: - _run_benchmark(_args) - finally: - if _out_path: - sys.stdout = _tee._stdout - _tee.close() - print(f"\nResults saved to: {_out_path}") diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_refactored.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_refactored.py deleted file mode 100644 index 2e87b9963715..000000000000 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_refactored.py +++ /dev/null @@ -1,2230 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - -import math - -import pytest -import torch -import torch.nn.functional as F -import triton -import triton.language as tl -from einops import repeat - -from tensorrt_llm._torch.modules.mamba.checkpointing_state_update_refactored import ( - _stochastic_round_int8_packed, - _stochastic_round_int16_packed, - checkpointing_state_update, -) -from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update -from tensorrt_llm._utils import get_sm_version - -# Philox stochastic rounding uses PTX cvt.rs.f16x2.f32 which requires sm >= 100. -_skip_pre_sm100 = pytest.mark.skipif( - get_sm_version() < 100, reason="Philox stochastic rounding needs sm >= 100" -) - -# Configs derived from NVIDIA-Nemotron-3-Super-120B-A12B Mamba2 parameters -# (nheads=128, headdim=64, d_state=128, ngroups=8) with TP split applied: -# TP=8: nheads=16, ngroups=1 — primary production config -# TP=4: nheads=32, ngroups=2 — exercises ngroups>1 (grouped B/C path) -_CONFIGS = [ - # (nheads, head_dim, d_state, ngroups) - (16, 64, 128, 1), # TP=8 production config - (32, 64, 128, 2), # TP=4, ngroups>1 (more heads than B/C groups) -] - -# Quantized state dtypes and their representable-magnitude limits (== QUANT_MAX -# in the kernel). fp8_e4m3fn cells require SM 89+ for the fp32↔fp8 cvt PTX -# instructions; SR variants of fp16/fp8 additionally need SM 100+. -_QUANT_MAX_BY_DTYPE = { - torch.int8: 127.0, - torch.int16: 32767.0, - torch.float8_e4m3fn: 448.0, -} - - -def _quantize_state(state_fp32: torch.Tensor, state_dtype: torch.dtype, quant_max: float): - """Quantize fp32 state to (state_quant, decode_scale) using the same - per-(head, dim) channel scheme the kernel does on store. decode_scale = - max_abs_per_channel / quant_max (= 1/encode_scale). - """ - amax = state_fp32.abs().amax(dim=-1) # (cache, nheads, head_dim) - encode_scale = quant_max / amax.clamp(min=1e-30) - decode_scale = 1.0 / encode_scale - scaled = state_fp32 * encode_scale.unsqueeze(-1) - if state_dtype == torch.float8_e4m3fn: - # Native cast does RN at the fp8 grid; explicit round() would destroy - # sub-integer precision (matches the kernel's fp8 RN path). - state_quant = scaled.clamp(-quant_max, quant_max).to(state_dtype) - else: - state_quant = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) - return state_quant, decode_scale - - -def _dequantize_state(state_quant: torch.Tensor, decode_scale: torch.Tensor): - return state_quant.to(torch.float32) * decode_scale.unsqueeze(-1) - - -def _maybe_skip_dtype(state_dtype, use_sr): - """Skip on insufficient SM. fp8 e4m3fn (any) needs SM 89+; fp16/fp8 SR - needs SM 100+; int8/int16 (RN or SR) runs anywhere.""" - if state_dtype == torch.float8_e4m3fn and get_sm_version() < 89: - pytest.skip("fp8_e4m3fn requires SM 89+ (Ada Lovelace / Hopper / Blackwell)") - if use_sr and state_dtype in (torch.float16, torch.float8_e4m3fn) and get_sm_version() < 100: - pytest.skip(f"{state_dtype} stochastic rounding requires SM 100+ (Blackwell B200+)") - - -@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) -@pytest.mark.parametrize( - "state_dtype", - [ - torch.float16, - torch.bfloat16, - torch.float32, - torch.int8, - torch.int16, - torch.float8_e4m3fn, - ], - ids=["fp16", "bf16", "fp32", "int8", "int16", "fp8"], -) -@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) -@pytest.mark.parametrize( - "T", [6, 10, 16, 27, 32, 55], ids=["T6", "T10", "T16", "T27", "T32", "T55"] -) -@pytest.mark.parametrize( - "write_checkpoint,rectangle_for_nowrite", - [ - (True, False), # write path (rectangle_for_nowrite is ignored) - (False, False), # nowrite path via replay-style kernels - (False, True), # nowrite path via dedicated rectangle kernels - ], - ids=["write", "no_write_replay", "no_write_rectangle"], -) -@pytest.mark.parametrize( - "mode", - ["monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl"], - ids=["monolithic", "dynamic", "doublelaunch", "dlgrouped", "maindl"], -) -def test_checkpointing_state_update( - nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T, - write_checkpoint, rectangle_for_nowrite, mode, -): - """ - Verify that: - checkpointing_state_update(state0, old_caches, k, new_x, ...) - produces the same output as: - selective_state_update(state_after_k_old_tokens, new_x, ...) - and writes state_after_k_old_tokens back to the state tensor. - - Quantized state dtypes (int8/int16/fp8) follow the same flow with - a per-(head, dim) channel decode-scale tensor; comparison is done - via dequant(state, scales) against the fp32 reference. - """ - _maybe_skip_dtype(state_dtype, use_sr=False) - - quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) - is_quantized = quant_max > 0.0 - - batch = 2 - device = "cuda" - dtype = torch.bfloat16 # input activations are bf16 - assert nheads % ngroups == 0 - - # Cache T-axis size (max_window). Use the kernel's BLOCK_SIZE_T as the - # ceiling — this is what the wrapper allows and enables PNAT-aware writes - # at [PNAT, PNAT+T) for no-checkpoint mode. For T=6 that's 16 (production - # max_window); for larger T it scales with np2(T). - max_window = max(triton.next_power_of_2(T), 16) - - if paged_cache: - cache_size = 4 - state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) - else: - cache_size = batch - state_batch_indices = None - - torch.manual_seed(42) - - # A: (nheads, head_dim, d_state) with stride(-2)=0, stride(-1)=0 [tie_hdim] - A_base = -torch.rand(nheads, device=device) - 0.5 # float32, negative - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - - # dt_bias: (nheads, head_dim) with stride(-1)=0 [tie_hdim] - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - - # D: (nheads, head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - # Initial SSM state (cache_size slots). Quantized dtypes need a separate - # init: derive scales from a fp32 source so the quantized state isn't - # garbage on dequant. ref_input_state is what the fp32 reference run - # sees — for non-quant it's state0 (cast to fp32 inside reference); for - # quant it's the lossy dequant of state0 (matches what the kernel sees - # internally on load). - if is_quantized: - state0_fp32 = torch.randn( - cache_size, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - state0, state0_scales = _quantize_state(state0_fp32, state_dtype, quant_max) - ref_input_state = _dequantize_state(state0, state0_scales) - else: - state0 = torch.randn( - cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype - ) - state0_scales = None - ref_input_state = state0.float() - - # Old inputs: up to `max_window` tokens per batch request, so the test - # loop can probe PNAT > T-1 (which the prior T-token setup couldn't - # reach). step1_T = max_window covers the full PNAT range we sweep. - step1_T = max_window - x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) - dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) # stride(-1)=0 - B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - - # Capture intermediate SSM states using selective_state_update across - # all step1_T positions — gives us reference states for k ∈ [0, step1_T]. - states_buffer_f32 = torch.zeros( - cache_size, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - cache_idx_for_capture = ( - state_batch_indices - if paged_cache - else torch.arange(batch, device=device, dtype=torch.int32) - ) - out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_input_state.clone(), - x1, - dt1, - A, - B1, - C1, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=cache_idx_for_capture, - intermediate_states_buffer=states_buffer_f32, - cache_steps=step1_T, - out=out1, - disable_state_update=True, - ) - - # Build cache tensors for the replay kernel. - # old_x: (cache, max_window, nheads, dim) bf16 — single-buffered - # old_B: (cache, 2, max_window, ngroups, dstate) bf16 — double-buffered - # old_dt: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous - # old_dA_cumsum: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous - # cache_buf_idx: random 0s and 1s to verify indexing correctness - old_x = torch.zeros(cache_size, max_window, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) - cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) - - # Fill each slot's active buffer (= cache_buf_idx) with step 1's data at - # positions [0:step1_T) = [0:max_window). Whole buffer covered so PNAT - # values up to max_window are exercised. Inactive buffer has random - # garbage to catch indexing bugs. - slots = state_batch_indices if paged_cache else slice(None) - old_x[slots, :step1_T] = x1 - - # Compute processed dt and dA_cumsum for step 1 - dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) - dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) - - # Write to each slot's active buffer based on its cache_buf_idx - slot_indices = state_batch_indices.tolist() if paged_cache else list(range(cache_size)) - for i, slot in enumerate(slot_indices): - buf = cache_buf_idx[slot].item() - batch_idx = i # maps slot back to the batch index - old_B[slot, buf, :step1_T] = B1[batch_idx] - old_dt[slot, buf, :, :step1_T] = dt1[batch_idx].T # (step1_T, nheads) → (nheads, step1_T) - old_dA_cumsum[slot, buf, :, :step1_T] = dA_cumsum1[batch_idx].T - - # Main loop: test each k (number of old tokens replayed). - # write_checkpoint=False (nowrite): k ∈ [0, max_window-T] — new tokens - # append at [k, k+T) of the active buffer; need k+T ≤ max_window. - # write_checkpoint=True (write): k ∈ [max_window-T+1, max_window] — - # new tokens land in the staging buffer at [0, T); k > max_window-T - # captures the overflow case that triggers a checkpoint in production. - # Combined sweep covers the full k ∈ [0, max_window] with the - # appropriate boundary handling per mode. - if write_checkpoint: - k_lo = max(0, max_window - T + 1) - k_hi = max_window + 1 # exclusive - else: - k_lo = 0 - k_hi = max_window - T + 1 # exclusive - for k in range(k_lo, k_hi): - torch.manual_seed(k + 100) - - x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) - B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - # Reference (fp32, starting from the same lossy-or-not state the - # kernel sees). - ref_state_f32 = ref_input_state.clone() - if k > 0: - ref_state_f32[slots] = states_buffer_f32[slots, k - 1] - - ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state_f32, - x2, - dt2, - A, - B2, - C2, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=(state_batch_indices if paged_cache else None), - out=ref_out, - ) - - # Replay kernel — clone caches into mutable working copies that we - # can inspect AFTER the call to verify cache postconditions. - test_state = state0.clone() - test_scales = state0_scales.clone() if is_quantized else None - prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) - test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - old_x_w = old_x.clone() - old_B_w = old_B.clone() - old_dt_w = old_dt.clone() - old_dA_cumsum_w = old_dA_cumsum.clone() - # cache_buf_idx stays at its random values — each slot reads from its own buffer - - checkpointing_state_update( - test_state, - old_x_w, - old_B_w, - old_dt_w, - old_dA_cumsum_w, - cache_buf_idx.clone(), - prev_tokens, - x=x2, - dt=dt2, - A=A, - B=B2, - C=C2, - out=test_out, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=state_batch_indices, - state_scales=test_scales, - write_checkpoint=write_checkpoint, - rectangle_for_nowrite=rectangle_for_nowrite, - mode=mode, - ) - - # Tolerance rationale: the replay kernel uses bf16 tl.dot for four - # matmuls (dB_scaled @ old_x, C @ state, CB_scaled @ x, and C @ B in - # precompute). The reference (selective_state_update) and flashinfer - # baseline use fp32 element-wise MACs. The bf16 input casts lose the - # dt_bias/A-derived bits that the baselines keep — per-element rounding, - # not accumulating. Prefill (ssd_chunk_scan) does identical bf16 tl.dot - # casts, so we match prefill precision exactly. Empirical: max ~1.0 at - # T<=16, ~2.0 at T=32-55; mean ~0.014; <0.02% of elements exceed 0.5. - # State dtype (fp16/bf16/fp32) doesn't shift the error — bf16 dot - # inputs dominate, not state storage. - # - # Quantized states add a per-element state quant error eps that - # propagates through C @ state in the output dot. With dstate=128 - # and C ~ N(0,1), the output channel std from this noise is roughly - # eps * sqrt(128/3) ≈ 6.5 * eps. Stack with the bf16 baseline: - # out_atol = bf16_atol + 6.5 * eps_max - # where eps_max is the worst-case per-element error at the - # post-replay state magnitude (T=55 → amax ≈ 23). - # - # Per-element error (eps_max for T=55): - # int8 (uniform grid): amax/(2*127) ≈ 0.091 - # int16 (uniform grid): amax/(2*32767) ≈ 3.5e-4 - # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 (worst-case - # cell at top of channel; smaller for - # smaller-magnitude elements) - out_atol = ( - {torch.int8: 1.6, torch.int16: 1.05, torch.float8_e4m3fn: 4.0}[state_dtype] - if is_quantized else 1.0 - ) - out_rtol = ( - {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] - if is_quantized else 2e-2 - ) - out_diff = (test_out.float() - ref_out.float()).abs() - out_max = out_diff.max().item() - out_mean = out_diff.mean().item() - try: - torch.testing.assert_close( - test_out, ref_out, rtol=out_rtol, atol=out_atol, - msg=f"Output mismatch at k={k}", - ) - except AssertionError: - print( - f"k={k} out: max={out_max:.4f} mean={out_mean:.4f} " - f"nan={torch.isnan(test_out).any().item()} " - f"inf={torch.isinf(test_out).any().item()}" - ) - raise - - # State expectation depends on write_checkpoint: - # True → kernel writes the post-replay state; expect the - # selective_state_update reference's state at step k-1. - # False → kernel skips the HBM store; state must be UNCHANGED - # from the input (state0; for quant, scales also unchanged). - if is_quantized: - if write_checkpoint: - # Compare via dequant against the fp32 reference state. - expected_fp32 = ( - ref_input_state[slots] if k == 0 else states_buffer_f32[slots, k - 1] - ) - actual_fp32 = _dequantize_state(test_state[slots], test_scales[slots]) - # State diff = bf16_replay_error + quant_error (per element). - # The bf16 component is the SAME error source the non-quant - # test absorbs in its atol=1.0 baseline (replay's tl.dot is - # bf16-input fp32-accum; per-element error ~ 2^-7 * amax, - # empirically ≤ ~0.2 at T=55 amax≈23). Quant adds: - # int8: amax/(2*127) ≈ 0.091 worst-case - # int16: amax/(2*32767) ≈ 3.5e-4 (negligible vs bf16) - # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 worst-case - # Atol = bf16_baseline (1.0) + quant_eps_max. - state_atol = { - torch.int8: 1.1, torch.int16: 1.0, torch.float8_e4m3fn: 2.5, - }[state_dtype] - state_rtol = { - torch.int8: 5e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 1e-1, - }[state_dtype] - try: - torch.testing.assert_close( - actual_fp32, expected_fp32, - rtol=state_rtol, atol=state_atol, - msg=f"State mismatch at k={k} dtype={state_dtype}", - ) - except AssertionError: - diff = (actual_fp32 - expected_fp32).abs() - print( - f"k={k} state(dequant): max={diff.max().item():.4f} " - f"mean={diff.mean().item():.4f}" - ) - raise - # Scales sanity (fp32, finite, positive). - assert test_scales.dtype == torch.float32 - assert torch.isfinite(test_scales[slots]).all(), ( - f"state_scales has non-finite values at k={k}" - ) - assert (test_scales[slots] > 0).all(), ( - f"state_scales has non-positive values at k={k}" - ) - else: - # No write: raw quant state and scales unchanged. Use - # torch.equal for byte-level equality (dtype-agnostic; works - # for int8 / int16 / fp8 alike). - assert torch.equal(test_state[slots], state0[slots]), ( - f"Quant state changed at k={k} write_checkpoint=False" - ) - assert torch.equal(test_scales[slots], state0_scales[slots]), ( - f"State scales changed at k={k} write_checkpoint=False" - ) - else: - if write_checkpoint: - expected_state = ( - state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) - ) - else: - expected_state = state0[slots] - state_diff = (test_state[slots].float() - expected_state.float()).abs() - state_max = state_diff.max().item() - state_mean = state_diff.mean().item() - try: - torch.testing.assert_close( - test_state[slots], - expected_state, - rtol=2e-2, - atol=1.0 if write_checkpoint else 0.0, - msg=f"State mismatch at k={k} (write_checkpoint={write_checkpoint})", - ) - except AssertionError: - print( - f"k={k} state: max={state_max:.4f} mean={state_mean:.4f} " - f"nan={torch.isnan(test_state).any().item()} " - f"inf={torch.isinf(test_state).any().item()}" - ) - raise - - # --- Cache postconditions --- - # Compute step 2's processed values (what the kernel should have - # stored at [write_offset : write_offset+T) of write_buf): - # write_buf = (1 - active_buf) if write_checkpoint else active_buf - # write_offset = 0 if write_checkpoint else k - # Untouched cache regions must equal their pre-call snapshots - # (old_x / old_B / old_dt / old_dA_cumsum captured before the call). - dt2_proc = F.softplus(dt2_base.float() + dt_bias_base.float()[None, None, :]) # (B,T,H) - dA_cumsum2 = torch.cumsum(A_base.float()[None, None, :] * dt2_proc, dim=1) - write_offset = 0 if write_checkpoint else k - - for batch_idx, slot in enumerate(slot_indices): - active = cache_buf_idx[slot].item() - wb = (1 - active) if write_checkpoint else active - - # --- old_x (single-buffered): write at [write_offset : +T) of slot --- - written_x = old_x_w[slot, write_offset : write_offset + T] - torch.testing.assert_close( - written_x, x2[batch_idx], rtol=0, atol=0, - msg=f"old_x written region wrong at k={k} write={write_checkpoint}", - ) - # Untouched ranges of old_x[slot] - if write_offset > 0: - torch.testing.assert_close( - old_x_w[slot, :write_offset], old_x[slot, :write_offset], - rtol=0, atol=0, - msg=f"old_x [0:{write_offset}) modified at k={k} write={write_checkpoint}", - ) - if write_offset + T < max_window: - torch.testing.assert_close( - old_x_w[slot, write_offset + T:], old_x[slot, write_offset + T:], - rtol=0, atol=0, - msg=f"old_x [{write_offset+T}:) modified at k={k} write={write_checkpoint}", - ) - - # --- old_B (double-buffered): write at write_buf, [write_offset:+T) --- - torch.testing.assert_close( - old_B_w[slot, wb, write_offset : write_offset + T], - B2[batch_idx], rtol=0, atol=0, - msg=f"old_B written region wrong at k={k} write={write_checkpoint}", - ) - # Other-buffer (= 1-wb) untouched - torch.testing.assert_close( - old_B_w[slot, 1 - wb], old_B[slot, 1 - wb], - rtol=0, atol=0, - msg=f"old_B inactive buffer modified at k={k} write={write_checkpoint}", - ) - - # --- old_dt (double-buffered, fp32, layout (heads, T)): --- - torch.testing.assert_close( - old_dt_w[slot, wb, :, write_offset : write_offset + T], - dt2_proc[batch_idx].T, - rtol=1e-4, atol=1e-4, - msg=f"old_dt written region wrong at k={k} write={write_checkpoint}", - ) - torch.testing.assert_close( - old_dt_w[slot, 1 - wb], old_dt[slot, 1 - wb], - rtol=0, atol=0, - msg=f"old_dt inactive buffer modified at k={k} write={write_checkpoint}", - ) - - # --- old_dA_cumsum (double-buffered, fp32, layout (heads, T)): --- - torch.testing.assert_close( - old_dA_cumsum_w[slot, wb, :, write_offset : write_offset + T], - dA_cumsum2[batch_idx].T, - rtol=1e-4, atol=1e-4, - msg=f"old_dA_cumsum written region wrong at k={k} write={write_checkpoint}", - ) - torch.testing.assert_close( - old_dA_cumsum_w[slot, 1 - wb], old_dA_cumsum[slot, 1 - wb], - rtol=0, atol=0, - msg=f"old_dA_cumsum inactive buf modified at k={k} write={write_checkpoint}", - ) - - -@pytest.mark.parametrize( - "mode,rectangle_for_nowrite", - [ - ("dynamic", False), - ("dynamic", True), - ("doublelaunch", False), - ("doublelaunch", True), - ("dlgrouped", False), - ("dlgrouped", True), - ("maindl", False), - ("maindl", True), - ], - ids=[ - "dynamic_replay", - "dynamic_rectangle", - "doublelaunch_replay", - "doublelaunch_rectangle", - "dlgrouped_replay", - "dlgrouped_rectangle", - "maindl_replay", - "maindl_rectangle", - ], -) -def test_checkpointing_state_update_mixed_mode(mode, rectangle_for_nowrite): - """ - Mixed-mode dispatch: a batch where some slots have PNAT triggering - write and others triggering nowrite, exercising the per-slot dispatch - of mode={dynamic, doublelaunch}. - - Setup: 4 slots, max_window=16, T=6. - pnat_per_slot = [3, 10, 12, 16] - slots 0, 1: nowrite (pnat + T <= max_window) - slots 2, 3: write (pnat + T > max_window) - - Reference: per-slot post-replay state computed from the captured - state evolution, then selective_state_update for the new step. - Output and post-replay state are verified per-slot. - """ - nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 - T = 6 - max_window = 16 - batch = 4 - cache_size = batch - device = "cuda" - dtype = torch.bfloat16 - state_dtype = torch.bfloat16 - - # PNAT mix: write threshold is pnat + T > max_window → pnat >= 11. - pnat_per_slot = torch.tensor([3, 10, 12, 16], device=device, dtype=torch.int32) - # Per-slot dispatch destinations under each mode (for the postcondition checks). - pnat_means_write = (pnat_per_slot + T > max_window).tolist() # [F, F, T, T] - - torch.manual_seed(42) - - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - state0 = torch.randn( - cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype - ) - ref_input_state = state0.float() - - # Old inputs spanning the full window (so any pnat 0..max_window is exercised). - step1_T = max_window - x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) - dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) - B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - - # Capture per-step intermediate SSM states across the window. - states_buffer_f32 = torch.zeros( - cache_size, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) - out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_input_state.clone(), - x1, dt1_input, A, B1, C1, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=cache_idx_for_capture, - intermediate_states_buffer=states_buffer_f32, - cache_steps=step1_T, - out=out1, - disable_state_update=True, - ) - - # Build the cache tensors with old data on each slot's active buffer. - old_x = torch.zeros(cache_size, max_window, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn( - cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 - ) - old_dA_cumsum = torch.randn( - cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 - ) - cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) - - old_x[:, :step1_T] = x1 - dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) - dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) - for i in range(cache_size): - buf = cache_buf_idx[i].item() - old_B[i, buf, :step1_T] = B1[i] - old_dt[i, buf, :, :step1_T] = dt1_processed[i].T - old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T - - # New-step inputs. - torch.manual_seed(123) - x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) - B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - # Reference: per-slot post-replay state, then selective_state_update. - ref_state_f32 = ref_input_state.clone() - for i in range(cache_size): - k_i = pnat_per_slot[i].item() - if k_i > 0: - ref_state_f32[i] = states_buffer_f32[i, k_i - 1] - ref_state_after_replay = ref_state_f32.clone() - - ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state_f32, - x2, dt2, A, B2, C2, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, - out=ref_out, - ) - - # Kernel call. - test_state = state0.clone() - test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - old_x_w = old_x.clone() - old_B_w = old_B.clone() - old_dt_w = old_dt.clone() - old_dA_cumsum_w = old_dA_cumsum.clone() - cache_buf_idx_w = cache_buf_idx.clone() - - checkpointing_state_update( - test_state, - old_x_w, old_B_w, old_dt_w, old_dA_cumsum_w, - cache_buf_idx_w, - pnat_per_slot, - x=x2, dt=dt2, A=A, B=B2, C=C2, - out=test_out, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, - rectangle_for_nowrite=rectangle_for_nowrite, - mode=mode, - # write_checkpoint is ignored in dynamic / doublelaunch. - ) - - # Output: every slot must match its reference (bf16 atol consistent - # with existing tests). - torch.testing.assert_close( - test_out.float(), ref_out.float(), - atol=1.0, rtol=0.05, - msg=f"Output mismatch (mode={mode}, rect={rectangle_for_nowrite})", - ) - - # State postconditions per slot: - # write slots (pnat + T > max_window): state in HBM is the post-replay - # fp32 reference (cast back to state_dtype with bf16 atol). - # nowrite slots: HBM state is unchanged (still state0 bitwise). - for i in range(cache_size): - if pnat_means_write[i]: - torch.testing.assert_close( - test_state[i].float(), - ref_state_after_replay[i].float(), - atol=1.0, rtol=0.05, - msg=f"Write slot {i}: state mismatch (mode={mode})", - ) - else: - torch.testing.assert_close( - test_state[i], state0[i], - rtol=0, atol=0, - msg=f"Nowrite slot {i}: state HBM was modified (mode={mode})", - ) - - -@pytest.mark.parametrize( - "mode,rectangle_for_nowrite,reverse_nowrite", - [ - ("doublelaunch", False, False), - ("doublelaunch", False, True), - ("doublelaunch", True, False), - ("doublelaunch", True, True), - ("dlgrouped", False, False), - ("dlgrouped", False, True), - ("dlgrouped", True, False), - ("dlgrouped", True, True), - ("maindl", False, False), - ("maindl", False, True), - ("maindl", True, False), - ("maindl", True, True), - ], - ids=lambda v: str(v), -) -def test_checkpointing_state_update_sorted_dispatch(mode, rectangle_for_nowrite, reverse_nowrite): - """ - Sort-driven dispatch: caller pre-sorts slots write-first via slot_perm. - Verifies the perm-aware kernels remap pid_b correctly so each slot's - work lands at the right grid program — i.e. slot S's output and HBM - state still match the reference under any permutation. - - Setup mirrors test_checkpointing_state_update_mixed_mode but with - slot_perm = [2, 3, 0, 1] (write slots 2, 3 first; nowrite 0, 1 after). - reverse_nowrite=True walks the perm tail-first on the nowrite-side, - so e.g. rectangle main reads perm[B-1-pid_grid] = [1, 0, 3, 2]. - """ - nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 - T = 6 - max_window = 16 - batch = 4 - device = "cuda" - dtype = torch.bfloat16 - - pnat_per_slot = torch.tensor([3, 10, 12, 16], device=device, dtype=torch.int32) - pnat_means_write = (pnat_per_slot + T > max_window).tolist() # [F, F, T, T] - # Write-first perm: indices 2, 3 (write) then 0, 1 (nowrite). - slot_perm = torch.tensor([2, 3, 0, 1], device=device, dtype=torch.int32) - - torch.manual_seed(42) - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=dtype) - ref_input_state = state0.float() - - step1_T = max_window - x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) - dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) - B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - - states_buffer_f32 = torch.zeros( - batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) - out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_input_state.clone(), - x1, dt1_input, A, B1, C1, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=cache_idx_for_capture, - intermediate_states_buffer=states_buffer_f32, - cache_steps=step1_T, - out=out1, - disable_state_update=True, - ) - - old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) - cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) - - old_x[:, :step1_T] = x1 - dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) - dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) - for i in range(batch): - buf = cache_buf_idx[i].item() - old_B[i, buf, :step1_T] = B1[i] - old_dt[i, buf, :, :step1_T] = dt1_processed[i].T - old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T - - torch.manual_seed(123) - x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) - B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - ref_state_f32 = ref_input_state.clone() - for i in range(batch): - k_i = pnat_per_slot[i].item() - if k_i > 0: - ref_state_f32[i] = states_buffer_f32[i, k_i - 1] - ref_state_after_replay = ref_state_f32.clone() - - ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state_f32, x2, dt2, A, B2, C2, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, out=ref_out, - ) - - test_state = state0.clone() - test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - checkpointing_state_update( - test_state, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), - cache_buf_idx.clone(), - pnat_per_slot, - x=x2, dt=dt2, A=A, B=B2, C=C2, - out=test_out, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, - rectangle_for_nowrite=rectangle_for_nowrite, - mode=mode, - slot_perm=slot_perm, - reverse_nowrite=reverse_nowrite, - ) - - torch.testing.assert_close( - test_out.float(), ref_out.float(), - atol=1.0, rtol=0.05, - msg=f"Output mismatch (mode={mode}, rect={rectangle_for_nowrite}, rev={reverse_nowrite})", - ) - - for i in range(batch): - if pnat_means_write[i]: - torch.testing.assert_close( - test_state[i].float(), ref_state_after_replay[i].float(), - atol=1.0, rtol=0.05, - msg=f"Write slot {i}: state mismatch (mode={mode}, rev={reverse_nowrite})", - ) - else: - torch.testing.assert_close( - test_state[i], state0[i], rtol=0, atol=0, - msg=f"Nowrite slot {i}: state HBM modified (mode={mode}, rev={reverse_nowrite})", - ) - - -@pytest.mark.parametrize( - "scenario,pnat_per_slot_list,n_writes_expected,slot_perm_list", - [ - # All-write: every slot has PNAT triggering write - # (PNAT + T > max_window). No permutation needed. - ("all_write", [12, 13, 14, 15], 4, [0, 1, 2, 3]), - # All-nowrite: every slot fits in the window. n_writes = 0. - ("all_nowrite", [3, 4, 5, 6], 0, [0, 1, 2, 3]), - # Mixed (write-first sorted via slot_perm): physical slots 2, 3 - # are writes; physical slots 0, 1 are nowrites. slot_perm - # remaps grid pid_b 0..3 to physical slots 2, 3, 0, 1 — so the - # first n_writes=2 grid programs hit write slots and the rest - # hit nowrite slots. - ("mixed_sorted", [3, 10, 12, 16], 2, [2, 3, 0, 1]), - ], - ids=["all_write", "all_nowrite", "mixed_sorted"], -) -def test_checkpointing_state_update_persistent_main( - scenario, pnat_per_slot_list, n_writes_expected, slot_perm_list, -): - """ - Persistent-CTA main kernel: 1D-grid kernel that loops over - (slot, M-tile, head) work units via tl.range. Caller pre-sorts - slots write-first and passes _n_writes (count of write slots) so - the kernel can split the persistent loop into write and nowrite - halves with the right WRITE_CHECKPOINT constexpr each time. - - Setup mirrors test_checkpointing_state_update_sorted_dispatch - (same fixed seeds, same input shapes) so the reference state - evolution is identical and we can compare per-slot output and - HBM-state postconditions to the same reference. - - Cases: - - all_write (n_writes=B): every slot exercises the - WRITE_CHECKPOINT=True branch of the persistent loop. - - all_nowrite (n_writes=0): every slot exercises the - WRITE_CHECKPOINT=False branch. Verifies the kernel handles - the "write half is empty" launch (n_slots=0 → early return). - - mixed_sorted: slots [2, 3] are writes, slots [0, 1] are - nowrites. slot_perm = [2, 3, 0, 1]. Persistent kernel - should call its impl with pid_b ∈ {2, 3} for the write half - and pid_b ∈ {0, 1} for the nowrite half, even though the - grid pid_b_grid is 0..n_slots-1 in each. - """ - nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 - T = 6 - max_window = 16 - batch = 4 - device = "cuda" - dtype = torch.bfloat16 - - pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) - pnat_means_write = (pnat_per_slot + T > max_window).tolist() - slot_perm = torch.tensor(slot_perm_list, device=device, dtype=torch.int32) - # Sanity: caller-supplied n_writes must match the actual count of - # write slots in the post-perm order. - write_count = sum(pnat_means_write) - assert write_count == n_writes_expected, ( - f"test setup error: expected {n_writes_expected} writes, got {write_count}" - ) - - torch.manual_seed(42) - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - state0 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=dtype - ) - ref_input_state = state0.float() - - step1_T = max_window - x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) - dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) - B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - - states_buffer_f32 = torch.zeros( - batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) - out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_input_state.clone(), - x1, dt1_input, A, B1, C1, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=cache_idx_for_capture, - intermediate_states_buffer=states_buffer_f32, - cache_steps=step1_T, - out=out1, - disable_state_update=True, - ) - - old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn( - batch, 2, nheads, max_window, device=device, dtype=torch.float32 - ) - cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) - - old_x[:, :step1_T] = x1 - dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) - dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) - for i in range(batch): - buf = cache_buf_idx[i].item() - old_B[i, buf, :step1_T] = B1[i] - old_dt[i, buf, :, :step1_T] = dt1_processed[i].T - old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T - - torch.manual_seed(123) - x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) - B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - ref_state_f32 = ref_input_state.clone() - for i in range(batch): - k_i = pnat_per_slot[i].item() - if k_i > 0: - ref_state_f32[i] = states_buffer_f32[i, k_i - 1] - ref_state_after_replay = ref_state_f32.clone() - - ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state_f32, x2, dt2, A, B2, C2, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, out=ref_out, - ) - - test_state = state0.clone() - test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - checkpointing_state_update( - test_state, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), - cache_buf_idx.clone(), - pnat_per_slot, - x=x2, dt=dt2, A=A, B=B2, C=C2, - out=test_out, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, - mode="persistent_main", - slot_perm=slot_perm, - _n_writes=n_writes_expected, - ) - - torch.testing.assert_close( - test_out.float(), ref_out.float(), - atol=1.0, rtol=0.05, - msg=f"Output mismatch (scenario={scenario})", - ) - - for i in range(batch): - if pnat_means_write[i]: - torch.testing.assert_close( - test_state[i].float(), ref_state_after_replay[i].float(), - atol=1.0, rtol=0.05, - msg=f"Write slot {i}: state mismatch (scenario={scenario})", - ) - else: - torch.testing.assert_close( - test_state[i], state0[i], rtol=0, atol=0, - msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", - ) - - -@pytest.mark.parametrize( - "scenario,pnat_per_slot_list", - [ - # All-write: every slot has PNAT triggering write (PNAT + T > max_window). - ("all_write", [12, 13, 14, 15]), - # All-nowrite: every slot fits in the window. - ("all_nowrite", [3, 4, 5, 6]), - # Mixed: some slots write, some nowrite. No pre-sort needed; the - # dynamic kernel dispatches per-slot at runtime via PNAT load. - ("mixed_unsorted", [3, 12, 10, 15]), - ], - ids=["all_write", "all_nowrite", "mixed_unsorted"], -) -def test_checkpointing_state_update_persistent_dynamic( - scenario, pnat_per_slot_list, -): - """ - Persistent-dynamic kernel: 1D persistent-CTA grid covering the full - batch, with runtime per-slot WRITE_CHECKPOINT branch derived from - each slot's PNAT. Single launch, no half-split, no n_writes needed, - no slot_perm needed (handles unsorted batches natively). - - Same setup as test_checkpointing_state_update_persistent_main; we - verify all three scenarios — including a mixed-unsorted batch the - persistent_main kernel can't handle without pre-sorting. - """ - nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 - T = 6 - max_window = 16 - batch = 4 - device = "cuda" - dtype = torch.bfloat16 - - pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) - pnat_means_write = (pnat_per_slot + T > max_window).tolist() - - torch.manual_seed(42) - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - state0 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=dtype - ) - ref_input_state = state0.float() - - step1_T = max_window - x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) - dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) - B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - - states_buffer_f32 = torch.zeros( - batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) - out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_input_state.clone(), - x1, dt1_input, A, B1, C1, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=cache_idx_for_capture, - intermediate_states_buffer=states_buffer_f32, - cache_steps=step1_T, - out=out1, - disable_state_update=True, - ) - - old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn( - batch, 2, nheads, max_window, device=device, dtype=torch.float32 - ) - cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) - - old_x[:, :step1_T] = x1 - dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) - dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) - for i in range(batch): - buf = cache_buf_idx[i].item() - old_B[i, buf, :step1_T] = B1[i] - old_dt[i, buf, :, :step1_T] = dt1_processed[i].T - old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T - - torch.manual_seed(123) - x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) - B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - ref_state_f32 = ref_input_state.clone() - for i in range(batch): - k_i = pnat_per_slot[i].item() - if k_i > 0: - ref_state_f32[i] = states_buffer_f32[i, k_i - 1] - ref_state_after_replay = ref_state_f32.clone() - - ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state_f32, x2, dt2, A, B2, C2, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, out=ref_out, - ) - - test_state = state0.clone() - test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - checkpointing_state_update( - test_state, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), - cache_buf_idx.clone(), - pnat_per_slot, - x=x2, dt=dt2, A=A, B=B2, C=C2, - out=test_out, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, - mode="persistent_dynamic", - ) - - torch.testing.assert_close( - test_out.float(), ref_out.float(), - atol=1.0, rtol=0.05, - msg=f"Output mismatch (scenario={scenario})", - ) - - for i in range(batch): - if pnat_means_write[i]: - torch.testing.assert_close( - test_state[i].float(), ref_state_after_replay[i].float(), - atol=1.0, rtol=0.05, - msg=f"Write slot {i}: state mismatch (scenario={scenario})", - ) - else: - torch.testing.assert_close( - test_state[i], state0[i], rtol=0, atol=0, - msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", - ) - - -@pytest.mark.parametrize( - "scenario,pnat_per_slot_list,n_writes_expected,slot_perm_list", - [ - # Same input shape as the mixed_sorted case in - # test_checkpointing_state_update_persistent_main, but exercises - # the device-tensor n_writes plumbing + skip_empty_halves=False - # path that the bench's mix-mode benchmarking depends on. - ("mixed_sorted", [3, 10, 12, 16], 2, [2, 3, 0, 1]), - # Boundary: n_writes=batch — when skip_empty_halves=False, the - # nowrite half launches with an empty slot range and the kernel - # must do nothing useful (tl.range covers 0 iterations). - ("all_write_noskip", [12, 13, 14, 15], 4, [0, 1, 2, 3]), - # Boundary: n_writes=0 — write half launches with empty range. - ("all_nowrite_noskip", [3, 4, 5, 6], 0, [0, 1, 2, 3]), - ], - ids=["mixed_sorted", "all_write_noskip", "all_nowrite_noskip"], -) -def test_checkpointing_state_update_persistent_main_device_n_writes( - scenario, pnat_per_slot_list, n_writes_expected, slot_perm_list, -): - """ - Persistent_main with the device-tensor n_writes plumbing. - - The bench's mix-mode benchmarking captures a single CUDA graph and - replays it many times with varying per-iter n_writes. To do that the - wrapper takes `_n_writes_dev` (a (1,) int32 device tensor) instead of - `_n_writes` (host int), and the kernel reads `n_writes` from device - memory at entry — same captured pointer across replays, value can - change between replays via an outside-graph copy. Also exercises - `_persistent_skip_empty_halves=False`: both halves of persistent_main - always launch, even when one half has no work (mix scenarios can't - cheaply know n_writes host-side per iter to skip). - - Verifies: - 1. Kernel reads device n_writes correctly (output matches reference). - 2. Empty-half launches don't corrupt state (skip_empty_halves=False - + n_writes=0 / =batch boundary cases). - 3. slot_perm + USE_PERM still works through the device-n_writes path. - """ - nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 - T = 6 - max_window = 16 - batch = 4 - device = "cuda" - dtype = torch.bfloat16 - - pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) - pnat_means_write = (pnat_per_slot + T > max_window).tolist() - slot_perm = torch.tensor(slot_perm_list, device=device, dtype=torch.int32) - write_count = sum(pnat_means_write) - assert write_count == n_writes_expected, ( - f"test setup error: expected {n_writes_expected} writes, got {write_count}" - ) - - # Device-tensor n_writes (the new path). Caller mutates between iters - # in mix-mode benchmarking; we only run one iter here so a single fill. - n_writes_dev = torch.tensor([n_writes_expected], device=device, dtype=torch.int32) - - torch.manual_seed(42) - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - state0 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=dtype - ) - ref_input_state = state0.float() - - step1_T = max_window - x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) - dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) - B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - - states_buffer_f32 = torch.zeros( - batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) - out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_input_state.clone(), - x1, dt1_input, A, B1, C1, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=cache_idx_for_capture, - intermediate_states_buffer=states_buffer_f32, - cache_steps=step1_T, - out=out1, - disable_state_update=True, - ) - - old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn( - batch, 2, nheads, max_window, device=device, dtype=torch.float32 - ) - cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) - - old_x[:, :step1_T] = x1 - dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) - dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) - for i in range(batch): - buf = cache_buf_idx[i].item() - old_B[i, buf, :step1_T] = B1[i] - old_dt[i, buf, :, :step1_T] = dt1_processed[i].T - old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T - - torch.manual_seed(123) - x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) - B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - ref_state_f32 = ref_input_state.clone() - for i in range(batch): - k_i = pnat_per_slot[i].item() - if k_i > 0: - ref_state_f32[i] = states_buffer_f32[i, k_i - 1] - ref_state_after_replay = ref_state_f32.clone() - - ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state_f32, x2, dt2, A, B2, C2, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, out=ref_out, - ) - - test_state = state0.clone() - test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - checkpointing_state_update( - test_state, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), - cache_buf_idx.clone(), - pnat_per_slot, - x=x2, dt=dt2, A=A, B=B2, C=C2, - out=test_out, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, - mode="persistent_main", - slot_perm=slot_perm, - # NEW PATHS: - _n_writes_dev=n_writes_dev, # device tensor (not host int) - _persistent_skip_empty_halves=False, # both halves always launch - ) - - torch.testing.assert_close( - test_out.float(), ref_out.float(), - atol=1.0, rtol=0.05, - msg=f"Output mismatch (scenario={scenario})", - ) - - for i in range(batch): - if pnat_means_write[i]: - torch.testing.assert_close( - test_state[i].float(), ref_state_after_replay[i].float(), - atol=1.0, rtol=0.05, - msg=f"Write slot {i}: state mismatch (scenario={scenario})", - ) - else: - torch.testing.assert_close( - test_state[i], state0[i], rtol=0, atol=0, - msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", - ) - - -@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) -@pytest.mark.parametrize( - "state_dtype", - [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], - ids=["fp16", "int8", "int16", "fp8"], -) -@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) -@pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) -def test_checkpointing_state_update_philox(state_dtype, nheads, head_dim, d_state, ngroups, paged_cache, T): - """ - Verify that Philox stochastic rounding produces correct results across - all SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). - - Runs our kernel twice with identical inputs — once without rand_seed - (deterministic RN), once with rand_seed (Philox SR) — and confirms: - - Outputs are within bf16-dot tolerance (state perturbation ≤ 1 ULP). - - State dtype is preserved. - - State difference is bounded by ~1 ULP of the chosen grid. - """ - _maybe_skip_dtype(state_dtype, use_sr=True) - - quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) - is_quantized = quant_max > 0.0 - - batch = 2 - device = "cuda" - dtype = torch.bfloat16 - assert nheads % ngroups == 0 - - if paged_cache: - cache_size = 4 - state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) - else: - cache_size = batch - state_batch_indices = None - - torch.manual_seed(42) - - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - if is_quantized: - state0_fp32 = torch.randn( - cache_size, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - state0, state0_scales = _quantize_state(state0_fp32, state_dtype, quant_max) - else: - state0 = torch.randn( - cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype - ) - state0_scales = None - - # Cache tensors - old_x = torch.randn(cache_size, T, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) - - # New token inputs - x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) - B = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - prev_tokens = torch.full((cache_size,), T // 2, device=device, dtype=torch.int32) - - common_kwargs = dict( - x=x, - dt=dt, - A=A, - B=B, - C=C, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=state_batch_indices, - ) - - # --- Run without rounding (deterministic RN store) --- - state_no_round = state0.clone() - scales_no_round = state0_scales.clone() if is_quantized else None - out_no_round = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - checkpointing_state_update( - state_no_round, - old_x.clone(), - old_B.clone(), - old_dt.clone(), - old_dA_cumsum.clone(), - cache_buf_idx.clone(), - prev_tokens, - out=out_no_round, - state_scales=scales_no_round, - **common_kwargs, - ) - - # --- Run with Philox rounding --- - rand_seed = torch.tensor([12345], device=device, dtype=torch.int64) - state_rounded = state0.clone() - scales_rounded = state0_scales.clone() if is_quantized else None - out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - checkpointing_state_update( - state_rounded, - old_x.clone(), - old_B.clone(), - old_dt.clone(), - old_dA_cumsum.clone(), - cache_buf_idx.clone(), - prev_tokens, - out=out_rounded, - rand_seed=rand_seed, - philox_rounds=10, - state_scales=scales_rounded, - **common_kwargs, - ) - - # Outputs should be nearly identical — rounding only perturbs the - # post-replay state by ±1 ULP before the output phase reads it. - # Out_atol = bf16_baseline + 6.5 * per_elem_ULP_after_dequant: - # non-quant fp16: fp16 ULP at typical magnitude is tiny → 1.0 - # int8: amax/127 ≈ 23/127 → 6.5*0.18 ≈ 1.2 + bf16_baseline - # int16: amax/32767 ≈ 7e-4 → ~bf16_baseline only - # fp8: amax/14 ≈ 23/14 → 6.5*1.6 ≈ 10.7 + bf16_baseline - out_atol = ( - {torch.int8: 1.5, torch.int16: 1.0, torch.float8_e4m3fn: 6.0}[state_dtype] - if is_quantized else 1.0 - ) - out_rtol = ( - {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] - if is_quantized else 2e-2 - ) - torch.testing.assert_close( - out_rounded, out_no_round, rtol=out_rtol, atol=out_atol, - msg=f"Output diverged with Philox rounding ({state_dtype})", - ) - - # State dtype preserved. - assert state_rounded.dtype == state_dtype - - # State diff between RN and SR is bounded by 1 quant cell per element. - # Per-channel decode_scale varies by 10x+ across channels (amax depends - # on randn extremes), so a single flat atol can't bound it accurately — - # use per-channel ULP-aware comparison. - slots = state_batch_indices if paged_cache else slice(None) - if is_quantized: - rounded_fp32 = _dequantize_state(state_rounded[slots], scales_rounded[slots]) - no_round_fp32 = _dequantize_state(state_no_round[slots], scales_no_round[slots]) - diff = (rounded_fp32 - no_round_fp32).abs() - # Per-element bound = max(decode_scale_no_round, decode_scale_rounded). - # decode_scale is shape (cache, nheads, dim); broadcast over dstate. - scale_bound = torch.maximum( - scales_no_round[slots], scales_rounded[slots] - ).unsqueeze(-1) - # int8 / int16: 1 cell after dequant = decode_scale exactly. - # fp8_e4m3: variable grid; the largest cell within a channel scaled - # to fit ±448 is at the channel's max-magnitude element, where the - # cell is ~32x larger than the average. Bound = decode_scale * 32. - # Apply a 1.5x slack pad for floating-point compare quirks at the - # exact-cell boundary. - cell_pad = ( - 32.0 if state_dtype == torch.float8_e4m3fn else 1.0 - ) - bound = scale_bound * (cell_pad * 1.5) - if not (diff <= bound).all(): - offenders = (diff > bound).sum().item() - n_total = diff.numel() - pytest.fail( - f"State RN-SR diff exceeds 1 cell per element for " - f"{offenders}/{n_total} elements ({state_dtype}). " - f"max_diff={diff.max().item():.4g}, " - f"max_bound={bound.max().item():.4g}." - ) - else: - # fp16 ULP depends on magnitude — rtol absorbs that. - torch.testing.assert_close( - state_rounded[slots], - state_no_round[slots], - rtol=2e-3, - atol=0.2, - msg=f"State diverged with Philox rounding ({state_dtype})", - ) - - -@pytest.mark.parametrize( - "state_dtype", - [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], - ids=["fp16", "int8", "int16", "fp8"], -) -def test_philox_rounding_unbiased(state_dtype): - """ - Verify that Philox stochastic rounding is unbiased across all - SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). - - Captures the true fp32 post-replay state by running with fp32 storage, - then runs the kernel with the target dtype + Philox SR. Compares the - SR rounding residual against the deterministic-RN residual: SR should - have mean residual closer to zero than RN, since RN has a systematic - round-to-nearest-even bias and SR is unbiased by construction. - - Uses a large batch (16) for ~2M state elements — plenty of statistics. - """ - _maybe_skip_dtype(state_dtype, use_sr=True) - - quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) - is_quantized = quant_max > 0.0 - - nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 - batch, T = 16, 6 - device = "cuda" - dtype = torch.bfloat16 - - torch.manual_seed(42) - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - # fp32 reference state — replay produces values that don't fit cleanly - # in the target dtype's grid, exposing the rounding bias. - state0_fp32 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - - old_x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(batch, 2, T, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) - cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) - - x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt_val = repeat(dt_base, "b t h -> b t h p", p=head_dim) - B = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - prev_tokens = torch.full((batch,), T, device=device, dtype=torch.int32) - - common_kwargs = dict( - x=x, dt=dt_val, A=A, B=B, C=C, D=D, dt_bias=dt_bias, dt_softplus=True, - ) - - # 1. fp32 state — captures true post-replay fp32 state. - state_fp32 = state0_fp32.clone() - out_fp32 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - checkpointing_state_update( - state_fp32, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), - cache_buf_idx.clone(), prev_tokens, out=out_fp32, **common_kwargs, - ) - - # 2. Target dtype + Philox SR. For quant we also need scales (derived - # from the same per-channel amax used by the kernel on store). - rand_seed = torch.tensor([99999], device=device, dtype=torch.int64) - if is_quantized: - state_rounded, scales_rounded = _quantize_state(state0_fp32, state_dtype, quant_max) - else: - state_rounded = state0_fp32.to(state_dtype) - scales_rounded = None - out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - checkpointing_state_update( - state_rounded, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), - cache_buf_idx.clone(), prev_tokens, out=out_rounded, - rand_seed=rand_seed, philox_rounds=10, - state_scales=scales_rounded, - **common_kwargs, - ) - - # Compute residuals. For non-quant: stochastic_residual = SR(fp32) - - # fp32, deterministic_residual = RN(fp32) - fp32. For quant: dequant - # both, comparing in fp32. - if is_quantized: - fp32_vals = state_fp32.flatten() - stochastic_residual = ( - _dequantize_state(state_rounded, scales_rounded).flatten() - fp32_vals - ) - # Deterministic reference: do the same per-channel quant on the - # captured fp32 state, then dequant. This is what the kernel would - # have produced with rand_seed=None. - det_quant, det_scales = _quantize_state(state_fp32, state_dtype, quant_max) - deterministic_residual = ( - _dequantize_state(det_quant, det_scales).flatten() - fp32_vals - ) - else: - fp32_vals = state_fp32.flatten() - stochastic_residual = state_rounded.float().flatten() - fp32_vals - deterministic_residual = fp32_vals.to(state_dtype).float() - fp32_vals - - # Only consider elements where rounding matters (non-zero residual possible). - nonzero_mask = deterministic_residual.abs() > 0 - num_nonzero = nonzero_mask.sum().item() - assert num_nonzero > 1000, f"Too few roundable elements: {num_nonzero}" - - stochastic_mean = stochastic_residual[nonzero_mask].mean().item() - stochastic_std = stochastic_residual[nonzero_mask].std().item() - deterministic_mean = deterministic_residual[nonzero_mask].mean().item() - - # SE-based bias check. An unbiased estimator's sample mean has standard - # error SE = std / sqrt(n). We require |sr_mean| < K*SE (K=4 ≈ ~3.2e-5 - # one-sided false-positive rate). This auto-calibrates per dtype: - # * int16: residual std ~1e-4 → SE ~9e-8 (very tight bound) - # * int8: residual std ~3e-2 → SE ~2e-5 - # * fp8: residual std ~1e-1 → SE ~9e-5 (loosest, magnitude-driven) - # The previous fixed-1e-5 threshold was below SE for int8/fp8 and would - # always fail by chance. Note the |sr|<|det| fallback was also dropped: - # on Gaussian (symmetric) inputs RN's bias is ~0 by symmetry, so SR vs RN - # is just two unbiased estimators racing — unreliable as a unbias test. - se_sr = stochastic_std / (num_nonzero ** 0.5) - K = 4 - assert abs(stochastic_mean) < K * se_sr, ( - f"SR mean exceeds {K}*SE (likely biased) ({state_dtype}): " - f"stochastic_mean={stochastic_mean:.3e}, " - f"SE={se_sr:.3e} (K*SE={K * se_sr:.3e}), " - f"deterministic_mean={deterministic_mean:.3e} (for reference), " - f"n_elements={num_nonzero}" - ) - - -# HEADS_PER_BLOCK > 1 test. The default heuristic only picks HPB > 1 at large -# total_heads (>= 256-512), which the main test with batch=2 never reaches. -# This test overrides _heads_per_block to exercise the two-loop structure in -# the precompute kernel (store-then-reload of per-head dt/dA_cumsum). -# Configs: (nheads=16, ngroups=1) and (nheads=32, ngroups=2) both have -# heads_per_group=16. The heuristic caps HPB at min(2|4, hpg), so HPB=2, 4. -@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) -@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) -@pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) -def test_checkpointing_heads_per_block( - nheads, - head_dim, - d_state, - ngroups, - state_dtype, - T, - heads_per_block, -): - # PDL flags use wrapper defaults; trimming the parametrize keeps this - # suite fast. Coverage of {launch_with_pdl, use_internal_pdl} variations - # lives in the dedicated correctness tests above (test_checkpointing_state_update). - batch = 8 - """ - Verify checkpointing_state_update produces correct results when - _heads_per_block > 1, exercising the precompute kernel's two-loop - structure (store per-head dt/dA_cumsum in loop 1, reload in loop 2). - """ - device = "cuda" - dtype = torch.bfloat16 - - if nheads % heads_per_block != 0: - pytest.skip(f"nheads ({nheads}) not divisible by heads_per_block ({heads_per_block})") - if heads_per_block > nheads // ngroups: - pytest.skip( - f"heads_per_block ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})" - ) - - torch.manual_seed(42) - - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - cache_size = batch - state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) - - x1 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) - B1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - states_buffer_f32 = torch.zeros( - cache_size, T, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) - out1 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - state0.clone(), - x1, - dt1, - A, - B1, - C1, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=cache_idx_for_capture, - intermediate_states_buffer=states_buffer_f32, - cache_steps=T, - out=out1, - disable_state_update=True, - ) - - old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) - - old_x[:] = x1 - dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) - dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) - - for slot in range(cache_size): - buf = cache_buf_idx[slot].item() - old_B[slot, buf] = B1[slot] - old_dt[slot, buf] = dt1[slot].T - old_dA_cumsum[slot, buf] = dA_cumsum1[slot].T - - k = T - torch.manual_seed(123) - - x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) - B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - ref_state_f32 = state0.float().clone() - ref_state_f32[:] = states_buffer_f32[:, k - 1] - ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state_f32, - x2, - dt2, - A, - B2, - C2, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=None, - out=ref_out, - ) - - test_state = state0.clone() - prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) - test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - - checkpointing_state_update( - test_state, - old_x.clone(), - old_B.clone(), - old_dt.clone(), - old_dA_cumsum.clone(), - cache_buf_idx.clone(), - prev_tokens, - x=x2, - dt=dt2, - A=A, - B=B2, - C=C2, - out=test_out, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=None, - _heads_per_block=heads_per_block, - ) - - torch.testing.assert_close( - test_out, - ref_out, - rtol=2e-2, - atol=1.0, - msg=f"Output mismatch with HPB={heads_per_block}, T={T}, " - f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", - ) - - expected_state = states_buffer_f32[:, k - 1].to(state_dtype) - torch.testing.assert_close( - test_state, - expected_state, - rtol=2e-2, - atol=1.0, - msg=f"State mismatch with HPB={heads_per_block}, T={T}, " - f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", - ) - - -# HPB > 1 multi-step test. Production chains decode steps; bugs in -# buffer ordering or stale cache values accumulate across steps and can -# be invisible in a single-step test. -@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) -@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("T", [6, 16], ids=["T6", "T16"]) -@pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) -@pytest.mark.parametrize("paged_cache", [False, True], ids=["contig", "paged"]) -def test_checkpointing_heads_per_block_multistep( - nheads, head_dim, d_state, ngroups, state_dtype, T, heads_per_block, paged_cache -): - """ - Chain N decode steps with HPB > 1 and verify each step's output matches - a fresh reference. A bug that mixes up WRITE/READ buffers, writes wrong - data to cache, or races in the two-loop structure would accumulate - across steps. - """ - batch = 2 - device = "cuda" - dtype = torch.bfloat16 - n_steps = 8 - - if nheads % heads_per_block != 0: - pytest.skip(f"nheads ({nheads}) not divisible by HPB ({heads_per_block})") - if heads_per_block > nheads // ngroups: - pytest.skip(f"HPB ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})") - - torch.manual_seed(42) - - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - if paged_cache: - cache_size = 4 - state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) - slots = state_batch_indices - else: - cache_size = batch - state_batch_indices = None - slots = slice(None) - - all_x = [] - all_dt = [] - all_B = [] - all_C = [] - for step in range(n_steps): - torch.manual_seed(1000 + step) - all_x.append(torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype)) - dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - all_dt.append(repeat(dt_base, "b t h -> b t h p", p=head_dim)) - all_B.append(torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype)) - all_C.append(torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype)) - - torch.manual_seed(999) - state_init = torch.randn( - cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype - ) - - ref_state = state_init.float().clone() - ref_outs = [] - ref_slots = ( - state_batch_indices - if paged_cache - else torch.arange(batch, device=device, dtype=torch.int32) - ) - for step in range(n_steps): - out_step = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state, - all_x[step], - all_dt[step], - A, - all_B[step], - all_C[step], - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=ref_slots, - out=out_step, - ) - ref_outs.append(out_step) - - test_state = state_init.clone() - old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.zeros(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - old_dA_cumsum = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) - - for step in range(n_steps): - k = T if step > 0 else 0 - prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) - test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - - checkpointing_state_update( - test_state, - old_x, - old_B, - old_dt, - old_dA_cumsum, - cache_buf_idx, - prev_tokens, - x=all_x[step], - dt=all_dt[step], - A=A, - B=all_B[step], - C=all_C[step], - out=test_out, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=state_batch_indices, - _heads_per_block=heads_per_block, - ) - - if paged_cache: - cache_buf_idx[slots] = 1 - cache_buf_idx[slots] - else: - cache_buf_idx[:] = 1 - cache_buf_idx - - torch.testing.assert_close( - test_out, - ref_outs[step], - rtol=2e-2, - atol=2.0, - msg=f"Output mismatch at step {step} with HPB={heads_per_block}, " - f"T={T}, nheads={nheads}, ngroups={ngroups}, " - f"state_dtype={state_dtype}, paged_cache={paged_cache}", - ) - - -# ----- SR grid-bracket tests (fp8 and fp16) ----- -# -# Verify that each PTX SR output lands on the destination dtype's grid as -# a bracket neighbour of the fp32 input. Catches byte-order traps in the -# inline-asm source-register specifier: -# * fp8: cvt.rs.satfinite.e4m3x4.f32 with pack=4, asm "{$4,$3,$2,$1}" -# * fp16: cvt.rs.f16x2.f32 with pack=2, asm "$0, $2, $1, $3" -# The unbiased test (test_philox_rounding_unbiased) wouldn't catch a -# shuffle: outputs that are still on-grid but swapped within a pack still -# average correctly. Only the per-element bracket check exposes it. -# -# Both kernels are inline copies of the production helpers — kept here so -# the test exercises the exact PTX form independent of wrapper changes. - - -@triton.jit -def _packed_int8_sr_kernel(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): - offs = tl.arange(0, BLOCK) - x = tl.load(x_ptr + offs) - rand = tl.load(rand_ptr + (offs // 4)) - y = _stochastic_round_int8_packed(x, rand, offs) - tl.store(out_ptr + offs, y.to(tl.int8)) - - -@triton.jit -def _packed_int16_sr_kernel(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): - offs = tl.arange(0, BLOCK) - x = tl.load(x_ptr + offs) - rand = tl.load(rand_ptr + (offs // 2)) - y = _stochastic_round_int16_packed(x, rand, offs) - tl.store(out_ptr + offs, y.to(tl.int16)) - - -def _bitrev_int(x: int, bits: int) -> int: - out = 0 - for _ in range(bits): - out = (out << 1) | (x & 1) - x >>= 1 - return out - - -def _rand_words(rand: torch.Tensor) -> list[int]: - return [int(v) & 0xFFFFFFFF for v in rand.cpu().tolist()] - - -def test_packed_int_sr_matches_reference(): - device = "cuda" - n = 1024 - offs = torch.arange(n, device=device, dtype=torch.float32) - x = ((offs % 37) - 18.0) + (((offs * 13.0) % 97.0) + 0.3) / 128.0 - - torch.manual_seed(42) - rand_i8 = torch.randint(-(2**31), 2**31, (n // 4,), device=device, dtype=torch.int32) - out_i8 = torch.empty(n, device=device, dtype=torch.int8) - _packed_int8_sr_kernel[(1,)](x, rand_i8, out_i8, BLOCK=n) - - x_cpu = x.cpu().tolist() - rand_i8_words = _rand_words(rand_i8) - ref_i8 = [] - for i, value in enumerate(x_cpu): - word = rand_i8_words[i // 4] - low = word & 0x0000FFFF - high = (word >> 16) & 0x0000FFFF - pos = i & 3 - if pos == 0: - rand16 = low - elif pos == 1: - rand16 = _bitrev_int(low, 16) - elif pos == 2: - rand16 = high - else: - rand16 = _bitrev_int(high, 16) - ref_i8.append(math.floor(value + rand16 / float(1 << 16))) - - torch.testing.assert_close( - out_i8.cpu().to(torch.int16), - torch.tensor(ref_i8, dtype=torch.int16), - rtol=0, - atol=0, - ) - - rand_i16 = torch.randint(-(2**31), 2**31, (n // 2,), device=device, dtype=torch.int32) - out_i16 = torch.empty(n, device=device, dtype=torch.int16) - _packed_int16_sr_kernel[(1,)](x, rand_i16, out_i16, BLOCK=n) - - rand_i16_words = _rand_words(rand_i16) - ref_i16 = [] - for i, value in enumerate(x_cpu): - word = rand_i16_words[i // 2] - rand_bits = word if (i & 1) == 0 else _bitrev_int(word, 32) - rand24 = rand_bits & 0x00FFFFFF - ref_i16.append(math.floor(value + rand24 / float(1 << 24))) - - torch.testing.assert_close( - out_i16.cpu(), - torch.tensor(ref_i16, dtype=torch.int16), - rtol=0, - atol=0, - ) - - -@triton.jit -def _bracket_kernel_fp8(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): - offs = tl.arange(0, BLOCK) - x = tl.load(x_ptr + offs) - rand = tl.load(rand_ptr + offs) - y = tl.inline_asm_elementwise( - asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", - constraints="=r,r,r,r,r,r,r,r,r", - args=(x, rand), - dtype=tl.float8e4nv, - is_pure=True, - pack=4, - ) - tl.store(out_ptr + offs, y) - - -@triton.jit -def _bracket_kernel_fp16(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): - offs = tl.arange(0, BLOCK) - x = tl.load(x_ptr + offs) - rand = tl.load(rand_ptr + offs) - y = tl.inline_asm_elementwise( - asm="""{ - cvt.rs.f16x2.f32 $0, $2, $1, $3; - }""", - constraints=("=r,r,r,r,r"), - args=(x, rand), - dtype=tl.float16, - is_pure=True, - pack=2, - ) - tl.store(out_ptr + offs, y) - - -_BRACKET_KERNEL = { - torch.float8_e4m3fn: _bracket_kernel_fp8, - torch.float16: _bracket_kernel_fp16, -} - - -def _build_finite_grid(dtype: torch.dtype, device: str) -> torch.Tensor: - """Reinterpret all bit patterns of ``dtype`` as floats; return sorted - unique finite values (drops ±inf, NaNs).""" - if dtype == torch.float8_e4m3fn: - ints = torch.arange(256, dtype=torch.uint8, device=device) - full = ints.view(torch.float8_e4m3fn).to(torch.float32) - elif dtype == torch.float16: - # int16 view of all 65536 patterns (covers fp16 normals + subnormals - # + ±inf + NaN; we filter to finite below). - ints = torch.arange(65536, dtype=torch.int32, device=device).to(torch.int16) - full = ints.view(torch.float16).to(torch.float32) - else: - raise ValueError(f"Unsupported bracket-test dtype: {dtype}") - return full[torch.isfinite(full)].sort()[0].unique() - - -def _build_bracket_inputs(dtype: torch.dtype, n: int, device: str) -> torch.Tensor: - """Test inputs spanning the dtype's grid range. Includes on-grid points - so we exercise the no-rounding case; for fp8 also includes overflow to - test saturation (PTX `cvt.rs.satfinite.e4m3x4.f32` clamps in-op). - - fp16 inputs are kept inside the finite range — `cvt.rs.f16x2.f32` does - NOT have a `satfinite` modifier and produces ±inf for OOR inputs (not - a saturate-to-±max). The kernel only ever sees in-range fp32 state in - practice (state_amax is always ≪ fp16_max), so the test mirrors that. - """ - grid = _build_finite_grid(dtype, device) - g_min, g_max = grid[0].item(), grid[-1].item() - x = torch.empty(n, device=device, dtype=torch.float32) - if dtype == torch.float8_e4m3fn: - # 1.5x range exercises saturation; satfinite handles it in-op. - x.uniform_(g_min * 1.5, g_max * 1.5) - else: # fp16: four magnitude bands, all within finite range. - x[: n // 4].uniform_(-1.0, 1.0) - x[n // 4 : n // 2].uniform_(-100, 100) - x[n // 2 : 3 * n // 4].uniform_(-1000, 1000) - x[3 * n // 4 :].uniform_(g_min * 0.99, g_max * 0.99) - return x, grid - - -@_skip_pre_sm100 -@pytest.mark.parametrize( - "state_dtype", - [torch.float8_e4m3fn, torch.float16], - ids=["fp8", "fp16"], -) -def test_sr_grid_bracket(state_dtype): - """Verify SR PTX outputs each lie on the destination grid as a bracket - neighbour of the fp32 input.""" - device = "cuda" - n = 1024 # multiple of both pack=4 (fp8) and pack=2 (fp16) - - torch.manual_seed(42) - x, grid_finite = _build_bracket_inputs(state_dtype, n, device) - g_min, g_max = grid_finite[0].item(), grid_finite[-1].item() - - # Bracket [lo, hi] in the destination grid for each input. For - # out-of-range inputs the bracket is the saturating endpoint pair. - x_clamped = x.clamp(g_min, g_max) - idx = torch.searchsorted(grid_finite, x_clamped, right=False).clamp( - min=1, max=len(grid_finite) - 1 - ) - lo = grid_finite[idx - 1] - hi = grid_finite[idx] - # For x exactly on grid, idx points at it; lo = grid[i-1], hi = x — the - # bracket allows out==hi (=x) which is what RN-on-grid produces. - - kernel = _BRACKET_KERNEL[state_dtype] - - for seed in range(4): - torch.manual_seed(seed) - # int32 for raw random bits — PTX takes the bit pattern, sign - # interpretation doesn't matter. - rand = torch.randint(-(2**31), 2**31, (n,), device=device, dtype=torch.int32) - out = torch.empty(n, device=device, dtype=state_dtype) - kernel[(1,)](x, rand, out, BLOCK=n) - out_fp32 = out.to(torch.float32) - - on_grid = (out_fp32 == lo) | (out_fp32 == hi) - if not on_grid.all(): - offenders = ~on_grid - n_off = offenders.sum().item() - sample = ( - x[offenders][:5].tolist(), - lo[offenders][:5].tolist(), - hi[offenders][:5].tolist(), - out_fp32[offenders][:5].tolist(), - ) - pytest.fail( - f"{state_dtype} SR output not on grid bracket for {n_off}/{n} " - f"elements (seed={seed}). x={sample[0]} lo={sample[1]} " - f"hi={sample[2]} out={sample[3]}. Likely the PTX byte-order " - "bug (cvt.rs source-register order)." - ) - From 6a8d3c4e2369d4188047741eedf4d19fd0f92971 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Sun, 17 May 2026 23:21:37 -0700 Subject: [PATCH 53/89] =?UTF-8?q?mamba=5Fcheckpointing:=20=5Fslim=20copies?= =?UTF-8?q?=20(kernel/test/bench)=20=E2=80=94=20persistent=5Fdynamic=20and?= =?UTF-8?q?=20persistent=5Fmain=20only,=20dropping=20monolithic/dynamic/ma?= =?UTF-8?q?indl/doublelaunch/dlgrouped/dl=5Fwrite=5Fonly=20modes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update_slim.py | 2993 ++++++++++ ...mark_replay_selective_state_update_slim.py | 4847 +++++++++++++++++ .../test_checkpointing_state_update_slim.py | 1915 +++++++ 3 files changed, 9755 insertions(+) create mode 100644 tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py create mode 100644 tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py create mode 100644 tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_slim.py diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py new file mode 100644 index 000000000000..73047aec3240 --- /dev/null +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py @@ -0,0 +1,2993 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +# +# Adapted from: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +# SPDX-FileCopyrightText: Copyright contributors to the sglang project +# +# Copyright (c) 2024, Tri Dao, Albert Gu. +# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID +from tensorrt_llm._utils import get_sm_version + +from .softplus import softplus + + +# Lazy global allocator for Triton TMA tensor descriptors. Required by any +# host- or device-built tensor_descriptor; without it Triton raises at first +# launch. See TMA backlog item #17 / scratch experiment notes. +_TMA_ALLOCATOR_SET = False + + +def _ensure_tma_allocator() -> None: + global _TMA_ALLOCATOR_SET + if _TMA_ALLOCATOR_SET: + return + + def _alloc_fn(size, alignment, stream): + # Triton expects an int8 buffer of `size` bytes; alignment is enforced + # by the allocator returning a buffer satisfying it (PyTorch's + # cudaMalloc-backed tensors are 256B-aligned, so we're fine). + return torch.empty(size, device="cuda", dtype=torch.int8) + + triton.set_allocator(_alloc_fn) + _TMA_ALLOCATOR_SET = True + + +@triton.jit +def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 pair → fp16x2 using Philox random bits. + + Uses PTX cvt.rs.f16x2.f32 which rounds each fp32 value to fp16 using + the random bits to break ties, avoiding systematic rounding bias that + accumulates over many decode steps with fp16 state. + + Adapted from flashinfer (Apache-2.0, vLLM/mamba lineage). + """ + return tl.inline_asm_elementwise( + asm="""{ + cvt.rs.f16x2.f32 $0, $2, $1, $3; + }""", + constraints=("=r,r,r,r,r"), + args=(x, rand), + dtype=tl.float16, + is_pure=True, + pack=2, + ) + + +@triton.jit +def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 quad → fp8 e4m3 using Philox random bits. + + Uses PTX cvt.rs.satfinite.e4m3x4.f32 which combines stochastic rounding + and saturating cast in a single op (output is final fp8, no separate + clamp needed). The reversed source-register order {$4,$3,$2,$1} is + load-bearing — PTX packs leftmost source into the high byte but Triton's + pack=4 is little-endian, so the natural {$1,$2,$3,$4} order would + silently shuffle every group of 4 contiguous outputs. + + Requires SM_100a+ (Blackwell B200). Caller must gate at the wrapper + level — this kernel does not check. + + Adapted from vLLM PR #40012 (Apache-2.0). + """ + return tl.inline_asm_elementwise( + asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", + constraints="=r,r,r,r,r,r,r,r,r", + args=(x, rand), + dtype=tl.float8e4nv, + is_pure=True, + pack=4, + ) + + +@triton.jit +def _bitrev32(x: tl.tensor) -> tl.tensor: + return tl.inline_asm_elementwise( + asm="brev.b32 $0, $1;", + constraints="=r,r", + args=(x,), + dtype=tl.uint32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _stochastic_round_int8_packed( + x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor +) -> tl.tensor: + """Stochastic rounding for int8 using one random uint32 per 4 values.""" + low = rand & 0x0000FFFF + high = (rand >> 16) & 0x0000FFFF + low_rev = _bitrev32(low) >> 16 + high_rev = _bitrev32(high) >> 16 + rand_pos = offs_n & 3 + rand16 = tl.where( + rand_pos == 0, + low, + tl.where(rand_pos == 1, low_rev, tl.where(rand_pos == 2, high, high_rev)), + ) + rand01 = rand16.to(tl.float32) * (1.0 / float(1 << 16)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +@triton.jit +def _stochastic_round_int16_packed( + x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor +) -> tl.tensor: + """Stochastic rounding for int16 using one random uint32 per 2 values.""" + rand_bits = tl.where((offs_n & 1) == 0, rand, _bitrev32(rand)) + rand01 = (rand_bits & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +# Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, +# old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. +# Grid: (batch, nheads // HEADS_PER_BLOCK). + + +@triton.jit() +def _replay_precompute_impl( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers (both buffers reachable via stride_*_dbuf). This + # kernel writes to either the active (= cache_buf_idx) or inactive + # (= 1 - cache_buf_idx) buffer depending on WRITE_CHECKPOINT — see + # comment block at top of kernel body. + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + # Double-buffer index (per cache slot) — selects this step's "active" + # buffer (= where the historical inputs for this step live). + cache_buf_idx_ptr, + # Per-request accepted-tokens count (already-cached old tokens at + # [0, PNAT) of the active buffer; new tokens this step go after them + # on no-checkpoint steps). + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides: (cache, 2, T, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Checkpointing flag — selects target buffer + offset for new-token + # cache writes. See "Cache write semantics" block below. + # Runtime (not constexpr): the only WRITE_CHECKPOINT-dependent code in + # this body is the write_buf/write_offset selection, which is plain + # arithmetic — no constexpr-shaped tile or whole-block gate. Letting + # it be runtime lets the dynamic dispatch kernel call us once with the + # per-slot needs_write flag instead of inlining two specializations. + write_checkpoint, +): + pid_b = tl.program_id(axis=0) + pid_hg = tl.program_id(axis=1) # head-group index + first_head = pid_hg * HEADS_PER_BLOCK + + # Resolve cache index for writes + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + + # --- Cache write semantics --- + # cache_buf_idx names this step's "active" buffer — the one with the + # historical inputs at [0, PNAT). The other buffer is "staging". + # + # Where do we write new tokens this step? + # WRITE_CHECKPOINT=False (no overflow): append to ACTIVE buffer at + # offset [PNAT : PNAT+T). Caller does NOT flip cache_buf_idx + # afterward; PNAT_next = PNAT + accepted. [0, PNAT) preserved. + # WRITE_CHECKPOINT=True (would overflow): write to STAGING buffer at + # [0, T). Caller flips cache_buf_idx afterward; next step's + # active = the one we just wrote. PNAT_next = accepted. Old + # data in the previous active buffer is folded into state via + # the replay update and discarded. This matches today's replay + # kernel behavior exactly. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if write_checkpoint: + write_buf = 1 - buf_active + write_offset = 0 + else: + write_buf = buf_active + write_offset = prev_num_accepted_tokens + + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + t_mask = offs_t < T + n_mask = offs_n < dstate + + # Causal mask is shared across all heads (depends only on offs_t) + causal_mask = offs_t[:, None] >= offs_t[None, :] + valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] + + # --- Vectorized pre-wait phase across HEADS_PER_BLOCK heads --- + # Compute dt, dA_cumsum, decay_vec as (H, T) tiles. Pre-compute + # scale_combo = decay_matrix * dt[:, None, :] as an (H, T, T) tile that + # stays in registers across gdc_wait — eliminates the post-wait reload + # of dt + dA_cumsum and the per-head loop. + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) + + # Load dt (H, T) + dt_addrs = ( + dt_ptr + pid_b * stride_dt_batch + + heads_block[:, None] * stride_dt_head + + offs_t[None, :] * stride_dt_T + ) + dt = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias[:, None] + if DT_SOFTPLUS: + dt = softplus(dt) + + A = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) # (H,) + dA_cumsum = tl.cumsum(A[:, None] * dt, axis=1) # (H, T) + decay_vec = tl.exp(dA_cumsum) # (H, T) + + # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of write_buf. + old_dt_addrs = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block[:, None] * stride_old_dt_head + + (write_offset + offs_t)[None, :] * stride_old_dt_T + ) + tl.store(old_dt_addrs, dt, mask=t_mask[None, :]) + + old_dA_cumsum_addrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block[:, None] * stride_old_dA_cumsum_head + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T + ) + tl.store(old_dA_cumsum_addrs, dA_cumsum, mask=t_mask[None, :]) + + # decay_vec scratch — always at offs_t. + decay_vec_addrs = ( + decay_vec_ptr + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) + tl.store(decay_vec_addrs, decay_vec, mask=t_mask[None, :]) + + # scale_combo (H, T, T) = exp(dA_cumsum[h, t1] - dA_cumsum[h, t2]) * dt[h, t2] + # Stays live across gdc_wait — used post-wait to compute CB_scaled. + decay_matrix = tl.exp(dA_cumsum[:, :, None] - dA_cumsum[:, None, :]) # (H, T, T) + scale_combo = decay_matrix * dt[:, None, :] # (H, T, T) + + # --- Wait for upstream kernel (external PDL) before loading B and C --- + # All dt processing above is independent of conv1d outputs. + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- + group_idx = first_head // nheads_ngroups_ratio + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_all = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_all = tl.load( + B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + # Compute raw CB once — shared across all heads in this block + raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) + + # Store B to cache at [write_offset : write_offset+T) of write_buf. + if first_head % nheads_ngroups_ratio == 0: + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + write_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_all, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # --- Vectorized post-wait phase: scale_combo (H, T, T) is still live in + # registers from pre-wait; multiply by raw_CB (T, T), apply causal mask, + # store as one (H, T, T) tile. --- + CB_scaled_block = tl.where( + valid_mask[None, :, :], + raw_CB[None, :, :] * scale_combo, + 0.0, + ) # (H, T, T) + cb_scaled_addrs = ( + cb_scaled_ptr + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_t[None, None, :] * stride_cb_j + ) # (H, T, T) + cb_store_mask = ( + (offs_t[None, :, None] < BLOCK_SIZE_T) + & (offs_t[None, None, :] < BLOCK_SIZE_T) + ) + tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) + + +# Replay-style precompute kernel. Thin wrapper around _replay_precompute_impl +# that carries the @triton.heuristics for constexpr derivation; called from +# the Python wrapper on the replay-style path (write or replay-nowrite). +@triton.jit() +def _rectangle_precompute_impl( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle + decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) + # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite + # path: read from buf_active at [0, PNAT), write new tokens at + # [PNAT, PNAT+T) of buf_active (same buffer). + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # rectangle K-axis bound + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides (rectangle: (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides: (cache, 2, T_max, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T_max) + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T_max) + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, +): + pid_b = tl.program_id(axis=0) + pid_hg = tl.program_id(axis=1) + first_head = pid_hg * HEADS_PER_BLOCK + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_buf = buf_active + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout: old at [0, PNAT) (mask is_old_k); + # new at [MAX-T, MAX) at compile-time shift K_NEW_SHIFT = MAX - T. + # PNAT + T <= MAX is guaranteed on the nowrite path → no overlap. + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis (output rows) + offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis (rectangle input cols) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + t_mask = offs_t < T + n_mask = offs_n < dstate + + # K-axis masks (approach C: runtime PNAT-offset instead of K_NEW_SHIFT) + # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Loop 1: per-head dt processing. dt → dt_processed → dA_cumsum → + # decay_vec_new (= exp(cumAdt_new)). Stored to write_buf for next step. + # decay_vec_full (= total_decay * decay_vec_new) is finalized in loop 2 + # once total_decay is loaded; loop 1 stores raw decay_vec_new to scratch. + for h_local in range(HEADS_PER_BLOCK): + head_idx = first_head + h_local + + dt_base = dt_ptr + pid_b * stride_dt_batch + head_idx * stride_dt_head + dt = tl.load(dt_base + offs_t * stride_dt_T, mask=t_mask, other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + head_idx * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias + if DT_SOFTPLUS: + dt = softplus(dt) + + A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) + dA_cumsum = tl.cumsum(A * dt, axis=0) + + # Store dt and dA_cumsum to write_buf at [write_offset, write_offset+T) + # for next step's replay/rectangle use. + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + tl.store( + old_dt_base + (write_offset + offs_t) * stride_old_dt_T, + dt, + mask=t_mask, + ) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + tl.store( + old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, + dA_cumsum, + mask=t_mask, + ) + + # ---- Hoisted: cache-only loads independent of conv1d ---- + # old_B (group-level, BLOCK_K × BLOCK_DSTATE = ~8KB tile) and the + # decay_vec_full per-head pre-compute (which writes to DRAM and doesn't + # need cross-gdc_wait variables) are issued BEFORE gdc_wait so their + # HBM latency overlaps with conv1d. Per-head factor_dt/exp_diff stay + # below gdc_wait — they need cross-iteration spans, which Triton can't + # express without a DRAM round-trip; the per-head LOADS in the post- + # wait loop are small and cheap, so leave them. + group_idx = first_head // nheads_ngroups_ratio + + # Group-level: old B from active buffer at [0, PNAT) of the K-axis. + old_B_read_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + buf_active * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_load = tl.load( + old_B_read_base + + safe_old_k[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + mask=is_old_k[:, None] & n_mask[None, :], + other=0.0, + ) + + # Pre-wait: vectorized across HEADS_PER_BLOCK heads. Compute decay_vec_full + # (H, T) and combo = factor_dt * exp_diff (H, T, K). Store decay_vec_full; + # combo_block stays in registers across gdc_wait — used directly post-wait + # to compute rect_CB_scaled without a global memory roundtrip. + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) + + # Per-head bases (H,) — broadcast with offs_k or offs_t for 2D loads. + old_dt_read_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_active * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_read_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + ) + old_dt_write_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_write_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + ) + + # (H, K) loads at [0, PNAT) — old data from previous step. + hk_mask = is_old_k[None, :] # (1, K) + old_dt_all = tl.load( + old_dt_read_h[:, None] + safe_old_k[None, :] * stride_old_dt_T, + mask=hk_mask, other=0.0, + ).to(tl.float32) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, + mask=hk_mask, other=0.0, + ).to(tl.float32) + # (H,) scalar-per-head: total_dA_cumsum at prev_k_idx. + total_dA_cumsum = tl.load( + old_dA_cumsum_read_h + prev_k_idx * stride_old_dA_cumsum_T + ).to(tl.float32) + # (H, T) loads at [PNAT, PNAT+T) — this step's dA_cumsum_new from loop 1. + ht_mask = t_mask[None, :] # (1, T) + dA_cumsum_new = tl.load( + old_dA_cumsum_write_h[:, None] + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T, + mask=ht_mask, other=0.0, + ).to(tl.float32) + # (H, K) loads at K_NEW_SHIFT-shifted positions for new tokens. + hkn_mask = is_new_k[None, :] + dt_at_kn = tl.load( + old_dt_write_h[:, None] + + (write_offset + safe_k_new)[None, :] * stride_old_dt_T, + mask=hkn_mask, other=0.0, + ).to(tl.float32) + dA_cumsum_at_kn = tl.load( + old_dA_cumsum_write_h[:, None] + + (write_offset + safe_k_new)[None, :] * stride_old_dA_cumsum_T, + mask=hkn_mask, other=0.0, + ).to(tl.float32) + + # decay_vec_full = total_decay * exp(cumAdt_new). (H, T). + total_decay = tl.where( + prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0 + ) # (H,) + decay_vec_full_block = total_decay[:, None] * tl.exp(dA_cumsum_new) # (H, T) + decay_vec_addrs = ( + decay_vec_ptr + + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) # (H, T) + tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) + + # combo_block = factor_dt * exp_diff — (H, T, K). Stays in registers + # across gdc_wait. + factor_dt = tl.where(is_old_k[None, :], old_dt_all, dt_at_kn) # (H, K) + s_k = tl.where( + is_old_k[None, :], + total_dA_cumsum[:, None] - old_dA_cumsum_all, + -dA_cumsum_at_kn, + ) # (H, K) + # exp_diff (H, T, K) = exp(s_k (H, 1, K) + dA_cumsum_new (H, T, 1)). + exp_diff = tl.exp(s_k[:, None, :] + dA_cumsum_new[:, :, None]) + combo_block = factor_dt[:, None, :] * exp_diff # (H, T, K) + + # ---- gdc_wait: from here on we depend on conv1d's outputs ---- + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # Conv1d outputs: B and C + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + B_new_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_all = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_new_orig = tl.load( + B_new_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_new_shifted = tl.load( + B_new_base + safe_k_new[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=is_new_k[:, None] & n_mask[None, :], + other=0.0, + ) + # Disjoint masks: old at [0, PNAT), new at [K_NEW_SHIFT, K_NEW_SHIFT+T). + B_combined = old_B_load + B_new_shifted + raw_rect_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_combined).to(tl.bfloat16)) + + # Append new B to cache at [PNAT, PNAT+T) of write_buf (once per group). + if first_head % nheads_ngroups_ratio == 0: + old_B_write_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + write_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_write_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_new_orig, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # Causal mask (BLOCK_SIZE_T × BLOCK_SIZE_K, shared across heads). + # Approach C: new tokens at runtime [PNAT, PNAT+T) instead of K_NEW_SHIFT. + t_idx_2d = offs_t[:, None] + k_idx_2d = offs_k[None, :] + is_old_k_2d = k_idx_2d < prev_num_accepted_tokens + k_new_idx_2d = k_idx_2d - prev_num_accepted_tokens + is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) + causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] + + # Post-wait vectorized: combo_block (H, T, K) is still live in registers. + # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store as + # one (H, T, K) tile. + rect_CB_scaled_block = tl.where( + causal_combined[None, :, :], + raw_rect_CB[None, :, :] * combo_block, + 0.0, + ) # (H, T, K) + cb_scaled_addrs = ( + cb_scaled_ptr + + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_k[None, None, :] * stride_cb_j + ) # (H, T, K) + cb_store_mask_3d = ( + (offs_t[None, :, None] < BLOCK_SIZE_T) + & (offs_k[None, None, :] < BLOCK_SIZE_K) + ) # (1, T, K) → broadcasts to (H, T, K) + tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) + + +# Rectangle precompute kernel. Thin wrapper around _rectangle_precompute_impl +# that carries the @triton.heuristics for constexpr derivation; called from +# the Python wrapper on the rectangle nowrite path. +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _dynamic_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. + RECTANGLE: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + pid_b = tl.program_id(axis=0) + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH + # write_checkpoint is now runtime in replay precompute, so a single + # call site handles both write and nowrite for the replay branch. + # Take rectangle only when RECTANGLE is True AND this slot doesn't + # need write; everything else funnels into replay. + if needs_write_runtime or not RECTANGLE: + _replay_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + needs_write_runtime, + ) + else: + _rectangle_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + ) + + +# Main kernel: tl.dot replay + precomputed CB output. +# Grid: (cdiv(dim, M), batch, nheads). + + +@triton.jit() +def _persistent_main_impl( + # Per-work-unit indices (computed by the persistent wrapper). + # `pid_b` is the post-perm slot index (caller has already applied any + # slot permutation and slot_offset). + pid_m, + pid_b, + pid_h, + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or + # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor + # USE_TMA_STORE is enabled (kernel ignores it via constexpr). + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + # IS_DYNAMIC: kept in the signature for caller-side bookkeeping (the + # outer _persistent_main_kernel still inspects it to decide the slot- + # IS_DYNAMIC: when True (persistent_dynamic), is_write is per-slot from + # PNAT. When False (persistent_main), is_write is constexpr from + # WRITE_CHECKPOINT. See also WC_IS_CONSTEXPR below. + IS_DYNAMIC: tl.constexpr, + # WC_IS_CONSTEXPR: when True, force is_write = WRITE_CHECKPOINT (constexpr) + # regardless of IS_DYNAMIC. Callers in RECT=1 use this in the is_w=True + # arm of _persistent_main_kernel (we know all slots that reach this call + # need is_write=True because is_w was the PNAT-derived runtime check, and + # this arm only fires when is_w is True). Passing WRITE_CHECKPOINT=True + # as a literal at the call site + WC_IS_CONSTEXPR=True here lets the inner + # body DCE the nowrite path under IS_DYNAMIC=True too — same codegen + # quality as persistent_main mode (-3.7% measured at b=1024 dyn-shape). + # When False (RECT=0 callers, where both write and nowrite slots are + # dispatched to ONE call), use the original runtime is_write under + # IS_DYNAMIC=True; avoids the binary-doubling regression that two + # specialized calls would cause. + WC_IS_CONSTEXPR: tl.constexpr = False, + # TMA flags — picked inside body based on is_write. When is_write is + # constexpr (either IS_DYNAMIC=False or WC_IS_CONSTEXPR=True), the + # use_tma_load = USE_TMA_LOAD_WRITE if is_write else USE_TMA_LOAD_NOWRITE + # ternary constexpr-folds and only one TMA load form survives. + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # IS_DYNAMIC: kernel-mode label, used by the OUTER _persistent_main_kernel + # to decide slot-range derivation and outer is_w dispatch strategy + # (constexpr WC for persistent_main; runtime is_w split -> 2 specialized + # impl calls for persistent_dynamic). Inside this impl, IS_DYNAMIC is + # NOT consulted at runtime -- WRITE_CHECKPOINT is the only constexpr that + # gates the write/nowrite codegen, in BOTH modes. + + # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized + # state dtype (int8 / int16 / float8e4nv) and only those. + tl.static_assert( + (QUANT_MAX > 0.0) + == ( + (state_ptr.dtype.element_ty == tl.int8) + or (state_ptr.dtype.element_ty == tl.int16) + or (state_ptr.dtype.element_ty == tl.float8e4nv) + ), + "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", + ) + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + # Resolve is_write: see WC_IS_CONSTEXPR / IS_DYNAMIC docs in the param + # list above. Three cases: + # - WC_IS_CONSTEXPR=True (RECT=1 is_w=True arm callers): use WC + # constexpr. Caller knows the slot needs write; inner DCEs nowrite + # paths. Avoids the binary-doubling overhead that calling the impl + # twice would cause, while still constexpr-DCEing the nowrite half. + # - IS_DYNAMIC=True (RECT=0 caller, persistent_dynamic): runtime + # branch on PNAT. Both write and nowrite codegen live in one body + # (no bloat) — same as the pre-refactor behavior. + # - IS_DYNAMIC=False (persistent_main): WC constexpr from caller. + if WC_IS_CONSTEXPR: + is_write: tl.constexpr = WRITE_CHECKPOINT + elif IS_DYNAMIC: + is_write = (prev_num_accepted_tokens + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_write = WRITE_CHECKPOINT + if is_write: + write_buf = 1 - active_buf # noqa: F841 + write_offset = 0 + else: + write_buf = active_buf # noqa: F841 + write_offset = prev_num_accepted_tokens + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # Load state. state_tma_descriptor is a host-built tensor_descriptor + # over a flat (cache*nheads*dim, dstate) view of state when any TMA + # path is enabled; raw `state_ptr` is the underlying tensor and is + # always passed. state_ptrs / state_ptr_raw are the raw-pointer view + # used for !TMA load and store paths. offs_y is the flat row index + # for TMA load/store; computed unconditionally (cheap int math; DCE'd + # when no TMA path is reachable). + state_mask = m_mask[:, None] & n_mask[None, :] + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + # Load state. Branch on is_write (constexpr = WRITE_CHECKPOINT in BOTH + # modes after the outer-dispatch refactor), then constexpr-pick TMA-vs- + # tl.load per side. Outer `if` DCE's, only the matching side's + # constexpr-gated load survives -- same compile-time picking for both + # persistent_main and persistent_dynamic (the latter dispatches at the + # outer kernel level so each impl instance sees a constexpr WC). + if is_write: + if USE_TMA_LOAD_WRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + else: + if USE_TMA_LOAD_NOWRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, + other=1.0, + ).to(tl.float32) + state = state * decode_scale[:, None] + + # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) + group_idx = pid_h // nheads_ngroups_ratio + + old_window_mask = offs_window < prev_num_accepted_tokens + + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + active_buf * stride_old_dt_dbuf + + pid_h * stride_old_dt_head + ) + old_dt_all = tl.load( + old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 + ).to(tl.float32) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + active_buf * stride_old_dA_cumsum_dbuf + + pid_h * stride_old_dA_cumsum_head + ) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, + mask=old_window_mask, other=0.0, + ).to(tl.float32) + + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( + tl.float32 + ) + + coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all + + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + old_x_all = tl.load( + old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + mask=old_window_mask[:, None] & m_mask[None, :], + other=0.0, + ) + + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + active_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_all = tl.load( + old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + mask=old_window_mask[:, None] & n_mask[None, :], + other=0.0, + ).to(tl.float32) + + dB_scaled = coeff[:, None] * old_B_all + + total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) + state *= total_decay + + state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) + + if is_write: + if USE_RS_ROUNDING: + # Generate random tensor for stochastic rounding. The amount of + # randomness needed depends on the SR codegen path: + # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) + # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) + # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs + # int16 SR (24b + bitrev32): 1 b32 per 2 outputs + # The PTX cvt.rs.* instructions consume a single 32-bit random + # and split the bits internally for 2 or 4 conversions. Generate + # only what's actually consumed and broadcast to fill the unused + # slots — saves Philox rounds proportionally. + if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: + RAND_DIVISOR: tl.constexpr = 4 # fp8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: + RAND_DIVISOR: tl.constexpr = 4 # int8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: + RAND_DIVISOR: tl.constexpr = 2 # int16 SR + elif QUANT_MAX == 0.0: + RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) + else: + RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized + + rand_seed = tl.load(rand_seed_ptr) + base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head + # Number of unique randoms per row = dstate / RAND_DIVISOR. + # randint4x emits 4 randoms per offset, so use that / 4 offsets. + offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) + rand_offsets_q = ( + base_rand + + offs_m[:, None] * stride_state_dim + + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) + ) # (M, dstate / (4*RAND_DIVISOR)) + if PHILOX_ROUNDS > 0: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + else: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) + r01 = tl.join(r0, r1) + r23 = tl.join(r2, r3) + r0123 = tl.join(r01, r23) + rand_compact = tl.reshape( + r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) + ) + # Broadcast each unique rand to RAND_DIVISOR adjacent positions. + # Pack-group (pack=2 fp16 / pack=4 fp8) consumes adjacent positions; + # the unique rand lands at the asm's read slot; duplicates feed + # the dead slots. Triton's broadcast_to is stride-0 in IR. + if RAND_DIVISOR > 1: + rand_3d = rand_compact[:, :, None] + rand_3d = tl.broadcast_to( + rand_3d, + (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), + ) + rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + else: + rand = rand_compact + + if QUANT_MAX > 0.0: + amax = tl.max(tl.abs(state), axis=1) + encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) + decode_scale = 1.0 / encode_scale + state_scales_ptrs = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + + offs_m * stride_state_scales_dim + ) + tl.store(state_scales_ptrs, decode_scale, mask=m_mask) + state_q = state * encode_scale[:, None] + if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): + _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_fp8sr) + else: + tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) + else: + if USE_RS_ROUNDING: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized SR fall-through expects int8 or int16; " + "fp8 SR is handled by the prior branch.", + ) + if state_ptrs.dtype.element_ty == tl.int8: + state_q = _stochastic_round_int8_packed( + state_q, rand, offs_n[None, :] + ) + else: + state_q = _stochastic_round_int16_packed( + state_q, rand, offs_n[None, :] + ) + elif state_ptrs.dtype.element_ty != tl.float8e4nv: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized RN with explicit round() expects int8 or int16.", + ) + state_q = tl.extra.cuda.libdevice.round(state_q) + state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) + _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_cast) + else: + tl.store(state_ptrs, _state_q_cast, mask=state_mask) + elif USE_RS_ROUNDING: + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "Non-quantized SR only supports fp16 state.", + ) + _state_sr = _stochastic_round_fp16x2(state, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_sr) + else: + tl.store(state_ptrs, _state_sr, mask=state_mask) + else: + _state_cast = state.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_cast) + else: + tl.store(state_ptrs, _state_cast, mask=state_mask) + + # Phase 2: Output using precomputed CB_scaled and decay_vec + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + x_all = tl.load( + x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_base + + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_all, + mask=t_mask[:, None] & m_mask[None, :], + ) + x_all = x_all.to(tl.float32) + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + tl.float32 + ) + + init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) + out_all = init_out + cb_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# `_persistent_rectangle_impl`: rectangle nowrite path for the persistent +# kernel. Body is a copy of `_rectangle_main_impl` with `pid_m`/`pid_b`/`pid_h` +# lifted to args (same pattern as `_persistent_main_impl` vs `_replay_main_impl`). +# Called only for nowrite slots when the kernel runs with RECTANGLE=True. +# Dropped from the rect impl: LAUNCH_DEPENDENT_KERNELS / USE_PERM / REVERSE_PERM +# (kernel-level, signalled once at top); slot_perm_ptr (kernel resolves perm). +@triton.jit() +def _persistent_rectangle_impl( + # Per-work-unit indices (computed by the persistent wrapper). + pid_m, + pid_b, + pid_h, + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as + # replay path). Used when USE_TMA_LOAD; ignored otherwise. + state_tma_descriptor, + state_scales_ptr, # only consulted when QUANT_MAX > 0 + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides (rectangle (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + QUANT_MAX: tl.constexpr, + USE_TMA_LOAD: tl.constexpr = False, +): + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout (matches precompute). + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_k = tl.arange(0, BLOCK_SIZE_K) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # K-axis masks (approach C: PNAT-runtime offset, matches precompute). + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Load state. Quant scale hoist: defer `* decode_scale` post-matmul. + if USE_TMA_LOAD: + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state = state_tma_descriptor.load([offs_y, 0]) + else: + state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, other=1.0, + ).to(tl.float32) + else: + state = state.to(tl.float32) + + # Group / pointer offset setup + group_idx = pid_h // nheads_ngroups_ratio + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + # Hoist: old_x doesn't depend on conv1d/precompute; load before gdc_wait. + old_x_load = tl.load( + old_x_base + + safe_old_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + mask=is_old_k[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + x_K = tl.load( + x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=is_new_k[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_base + + offs_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_K, + mask=is_new_k[:, None] & m_mask[None, :], + ) + + x_K_f32 = x_K.to(tl.float32) + x_combined = old_x_load + x_K_f32 + + if HAS_D or HAS_Z: + sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) + x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) + else: + x_all = x_K_f32 # placeholder; unused + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec_full = tl.load( + decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 + ).to(tl.float32) + + state_out = ( + tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) + * decay_vec_full[:, None] + ) + if QUANT_MAX > 0.0: + state_out = state_out * decode_scale[None, :] + + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) + + out_all = state_out + token_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# Persistent main kernel: 1D grid, persistent CTA loop. +# Heuristics mirror those of `_checkpointing_main_kernel`. +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"NUM_PID_M_BLOCKS": lambda args: triton.cdiv(args["dim"], args["BLOCK_SIZE_M"])} +) +@triton.jit() +def _persistent_main_kernel( + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view. + # Shared across BOTH the replay path (consumed by _persistent_main_impl + # when USE_TMA_LOAD_*/STORE) AND the rectangle path (consumed by + # _persistent_rectangle_impl when USE_TMA_LOAD) — same descriptor, same + # block_shape, just gated by separate constexprs per impl. Wrapper sets + # this to a TensorDescriptor when ANY of the three TMA flags is on, else + # to `state_ptr` (raw); each impl ignores it via its own constexpr when + # not consuming it. + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + rand_seed_ptr, + pad_slot_id, + # Persistent-loop work-distribution scalars. Caller pre-sorts the batch + # write-first; the kernel uses (n_writes, batch_total, WRITE_CHECKPOINT) + # to derive its own slot range. Write half processes [0, n_writes), + # nowrite half processes [n_writes, batch_total). + # + # n_writes_ptr is a device pointer to a (1,) int32 tensor. Reading + # from device memory (rather than taking a Python int kernel arg) is + # required so mix-mode benchmarking can vary n_writes per iter inside + # a captured CUDA graph — the source tensor's contents change, the + # pointer doesn't. Cost: one int load per kernel launch (~negligible). + # When IS_DYNAMIC=True the value is unused (Triton DCEs the load). + n_writes_ptr, # int32 *: device-side count of write-mode slots + batch_total, # int32: total slot count + nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + USE_PERM: tl.constexpr, + # NUM_PERSISTENT: runtime int (not constexpr). Used ONLY as the loop + # stride in `tl.range(pid, total_work, NUM_PERSISTENT, ...)`. Making it + # runtime collapses the cta_per_sm tuning dim from the kernel's compile + # signature: 8 CPS values used to mean 8x recompiles; now they share one + # compiled kernel. Work decomposition (pid_m, pid_b_local, pid_h) does + # NOT depend on NUM_PERSISTENT — it uses constexpr NUM_PID_M_BLOCKS and + # runtime n_slots_local — so loop unrolling and flatten=/num_stages=/ + # warp_specialize= optimizations on `tl.range` operate independently of + # the stride value. + NUM_PERSISTENT, + NUM_LOOP_STAGES: tl.constexpr, + NUM_PID_M_BLOCKS: tl.constexpr, + FLATTEN: tl.constexpr, + WARP_SPECIALIZE: tl.constexpr, + IS_DYNAMIC: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr = 16, # rectangle K-axis (heuristic-derived) + RECTANGLE: tl.constexpr = False, # when True, dispatch nowrite slots to _persistent_rectangle_impl + # 3 TMA toggles per the 3 live paths per-compilation: + # USE_TMA_LOAD_WRITE — replay-style state load when is_write + # USE_TMA_LOAD_NOWRITE — nowrite-path state load (rect when RECTANGLE, + # else replay-nowrite) + # USE_TMA_STORE — replay-style state store (only fires on write + # path; no-op when not is_write) + # Wrapper picks USE_TMA_LOAD_NOWRITE = _use_tma_rect_load (if rectangle) + # or _use_tma_replay_nowrite_load (if not). + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # PDL signal: fire once at kernel entry (not per work unit). + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + # Load runtime n_writes from device memory. Read once at kernel entry; + # used only by the !IS_DYNAMIC slot-range derivation below. Triton + # DCEs the load when IS_DYNAMIC=True (n_writes is dead there). + n_writes = tl.load(n_writes_ptr) + + # Derive this kernel's slot range. Two modes: + # IS_DYNAMIC=False (persistent_main): caller pre-sorts and splits halves; + # slot range is [0, n_writes) when WRITE_CHECKPOINT else [n_writes, batch_total) + # IS_DYNAMIC=True (persistent_dynamic): single launch covers full batch; + # each work-item dispatches via runtime PNAT check inside the impl. + if IS_DYNAMIC: + slot_lo = 0 + slot_hi = batch_total + else: + if WRITE_CHECKPOINT: + slot_lo = 0 + slot_hi = n_writes + else: + slot_lo = n_writes + slot_hi = batch_total + n_slots_local = slot_hi - slot_lo + + pid = tl.program_id(axis=0) + total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads + + # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) + # with pid_m varying fastest (M-tile cache locality on state load), then + # slot, then head — mirrors the existing 3D grid's axis ordering + # (axis=0 fastest = pid_m). + for tile_id in tl.range( + pid, total_work, NUM_PERSISTENT, + flatten=FLATTEN, num_stages=NUM_LOOP_STAGES, warp_specialize=WARP_SPECIALIZE, + ): + pid_m = tile_id % NUM_PID_M_BLOCKS + pid_b_local = (tile_id // NUM_PID_M_BLOCKS) % n_slots_local + pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) + # Translate local slot index → global slot index. When USE_PERM is + # set, the caller-provided slot_perm gives the original slot index + # for the post-sort position. + pid_b_grid = pid_b_local + slot_lo + if USE_PERM: + pid_b = tl.load(slot_perm_ptr + pid_b_grid) + else: + pid_b = pid_b_grid + + # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle + # impl. Pre-load cache_batch_idx + PNAT here only for the RECTANGLE + # path's branch decision. Both impls re-load and handle pad_slot_id + # internally (Triton's L1 cache makes the duplicate loads ~free). + if RECTANGLE: + if HAS_CACHE_BATCH_INDICES: + cbi_pre = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + is_pad = cbi_pre == pad_slot_id + else: + cbi_pre = pid_b.to(tl.int64) + is_pad = False + if not is_pad: + pnat_pre = tl.load(prev_num_accepted_tokens_ptr + cbi_pre) + if IS_DYNAMIC: + is_w = (pnat_pre + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_w = WRITE_CHECKPOINT + if is_w: + # Pass WRITE_CHECKPOINT=True constexpr to specialize this + # impl call for the write path. Under IS_DYNAMIC=True, the + # kernel-level WRITE_CHECKPOINT is False (launcher default), + # but the OUTER is_w branch we are inside narrows the + # runtime path to writes-only, so we override to True here + # so the impl's constexpr-gated `if is_write:` blocks DCE + # to the write-only codegen. Under IS_DYNAMIC=False + # (persistent_main), the kernel-level WRITE_CHECKPOINT is + # itself True for this half (write half launches with + # WC=True), and the outer is_w = WRITE_CHECKPOINT = True + # constexpr-folds; passing literal True here is consistent + # and constexpr-equivalent. + _persistent_main_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, + True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) + True, # WC_IS_CONSTEXPR — force inner to use WC constexpr + # 3 TMA flags: write-load fires here (we're in the + # is_write branch), nowrite-load is dead (no slot + # reaches it), store fires (write path). + USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + ) + else: + # Rectangle nowrite: pass state_ptr (raw, always) + + # state_tma_descriptor (the single unified descriptor — + # same memory replay paths use). Rect impl gates use + # of the descriptor via its USE_TMA_LOAD constexpr. + _persistent_rectangle_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, + LAUNCH_WITH_PDL, QUANT_MAX, + USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle + ) + # else: pad slot — skip both impls (both would early-return anyway) + else: + # No rectangle path — single _persistent_main_impl call covers + # both write and nowrite slots via WC constexpr (non-dynamic) or + # runtime is_write (IS_DYNAMIC=True). Pass all 3 TMA flags; + # impl picks USE_TMA_LOAD_WRITE vs USE_TMA_LOAD_NOWRITE based on + # its computed is_write — constexpr-folds when is_write is + # constexpr (non-dyn), runtime branch when IS_DYNAMIC=True. + # (Reverted from outer two-call dispatch: that doubled the + # compiled body size under IS_DYNAMIC=True and regressed RECT=0 + # perf by ~+24%.) + _persistent_main_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, + WRITE_CHECKPOINT, IS_DYNAMIC, + False, # WC_IS_CONSTEXPR=False — RECT=0 has both write/nowrite slots in one call + USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + ) + + +# ============================================================================ +# Python wrapper +# ============================================================================ + + +_QUANT_MAX_BY_DTYPE = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, +} + + +def checkpointing_state_update( + state: torch.Tensor, + old_x: torch.Tensor, + old_B: torch.Tensor, + old_dt: torch.Tensor, + old_dA_cumsum: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + out: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + pad_slot_id: int = PAD_SLOT_ID, + rand_seed: torch.Tensor | None = None, + philox_rounds: int = 10, + state_scales: torch.Tensor | None = None, + launch_with_pdl=False, + use_internal_pdl=True, + write_checkpoint: bool = True, + rectangle_for_nowrite: bool = False, + mode: str = "persistent_dynamic", + # Slot permutation: int32 (batch,) tensor mapping grid program_id -> + # original slot index. When provided, the persistent kernels read pid_b + # through this perm so callers can pre-sort slots write-first (required + # for `persistent_main`, optional perf hint for `persistent_dynamic`). + # None => identity perm. + slot_perm: torch.Tensor | None = None, + _block_size_m: int | None = None, + _num_warps: int | None = None, + _num_stages: int | None = None, + _precompute_num_warps: int | None = None, + _precompute_num_stages: int | None = None, + _heads_per_block: int | None = None, + _maxnreg: int | None = None, + _num_ctas: int | None = None, + # Per-main knobs (override shared values for one half of the dl-family / + # persistent_main launches). Default None = tied to the shared value + # (backward compat). The two main kernels (write vs nowrite) have + # different per-slot work — write does a state shift + store, nowrite + # just appends — so the optimum (M, W, S, H) can differ. Precompute + # knobs are intentionally NOT split: shared precompute wins (cheaper + # launch, hotter precompute outputs in L2). Persistent CPS / LS knobs + # are also split per-main since the two persistent_main launches have + # different grid sizes. + _block_size_m_write: int | None = None, + _block_size_m_nowrite: int | None = None, + _num_warps_write: int | None = None, + _num_warps_nowrite: int | None = None, + _num_stages_write: int | None = None, + _num_stages_nowrite: int | None = None, + # Note: heads_per_block / precompute_num_warps are NOT split — they only + # affect the precompute kernel, which is shared across write/nowrite. + # TMA state-tensor toggles — 4 independent paths (see CHECKPOINTING_DESIGN.md + # item #17 for measured perf profiles). Each is False=raw load/store, True= + # use a host-built TMA tensor_descriptor for that path. + _use_tma_rect_load: bool = False, # rect kernel's state load (nowrite-only) + _use_tma_replay_write_load: bool = False, # replay-style state load when WC=True + _use_tma_replay_write_store: bool = False, # replay-style state store when WC=True + _use_tma_replay_nowrite_load: bool = False, # replay-style state load when WC=False + # Persistent-mode bench kwargs (only consulted when mode == "persistent_main"): + # _n_writes : int — count of write-mode slots in the (pre-sorted) batch. + # Required when mode == "persistent_main"; the persistent kernel uses + # it as a runtime int32 to compute total_work for write/nowrite halves. + # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally + # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. Default = 1. + # _num_loop_stages : int — `num_stages` arg on the inner `tl.range(...)` + # persistent loop. Note: this is loop-level, NOT the kernel-arg + # `num_stages` (which only pipelines dot-feeding loads). Default 2. + # _flatten : bool — `flatten` arg on `tl.range(...)`. Default True + # (the canonical Triton 3.6 persistent idiom). + # _warp_specialize : bool — `warp_specialize` arg on `tl.range(...)`. + # Default False. Triton 3.6 only supports it on simple matmul loops; + # our scan loop probably won't pattern-match — but exposed as a knob + # for sweep experiments. Requires num_warps >= 4 if True. + _n_writes: int | None = None, + # Optional pre-allocated (1,) int32 device tensor for the persistent + # kernel's n_writes input. Bench passes this in mix scenarios so the + # captured CUDA graph can read varying n_writes per iter without + # re-capture. When None and `_n_writes` is provided, we allocate a + # scratch tensor and fill from `_n_writes` (pure scenarios). + _n_writes_dev: torch.Tensor | None = None, + # When True, persistent_main host-skips empty-half launches (n_writes=0 + # or =batch in pure scenarios). Default True preserves today's behavior. + # Set False to always launch both halves — used by mix scenarios (where + # host can't cheaply read n_writes per iter) and for fair K-consistent + # comparisons. + _persistent_skip_empty_halves: bool = True, + _cta_per_sm: int | None = None, + _num_loop_stages: int | None = None, + _flatten: bool | None = None, + _warp_specialize: bool | None = None, + # Per-main persistent-specific knobs. Same rationale as the BLOCK_SIZE_M + # split above: the two persistent_main launches (write half vs nowrite + # half) have different grid sizes and per-work-item costs, so they may + # want different cta_per_sm / num_loop_stages. + _cta_per_sm_write: int | None = None, + _cta_per_sm_nowrite: int | None = None, + _num_loop_stages_write: int | None = None, + _num_loop_stages_nowrite: int | None = None, +): + """ + Replay SSM state update with precomputed CB and tl.dot fast-forward. + + Two-kernel architecture: + 1. Precompute kernel: computes CB_scaled and decay_vec from B, C, dt, A. + Writes processed dt/dA_cumsum/B to double-buffered cache for next step. + 2. Main kernel: replays old tokens via tl.dot fast-forward on cached data, + then computes output using precomputed CB_scaled and new x/C inputs. + + PDL (Programmatic Dependent Launch) chain: + conv1d → (external PDL) → precompute → (internal PDL) → main + External PDL: precompute starts while conv1d is running; gdc_wait() + in precompute blocks until conv1d completes before loading B/C. + Internal PDL: main starts while precompute is running; main's replay + phase uses only cached data from the previous step. gdc_wait() in + main blocks until precompute completes before loading conv1d outputs + (x, C) and precompute outputs (CB_scaled, decay_vec). + + Uses double-buffered cache tensors. cache_buf_idx[slot] indicates which + buffer (0 or 1) to READ from for replay. The WRITE buffer is 1 - read. + Caller must flip cache_buf_idx[slot] after each call. + + Arguments: + state: (cache, nheads, dim, dstate) in-place. After the call, contains + the state after replaying prev_num_accepted_tokens old tokens. + old_x: (cache, T, nheads, dim) bf16 — old x cache (single-buffered). + old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. + old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. + old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. + cache_buf_idx: (cache,) int32 — which buffer to read (0 or 1). + prev_num_accepted_tokens: (cache,) int32. + x: (batch, T, nheads, dim) new token inputs. + dt: (batch, T, nheads, dim) with stride(-1)==0 (tie_hdim). + A: (nheads, dim, dstate) with stride(-1)==0, stride(-2)==0 (tie_hdim). + B: (batch, T, ngroups, dstate). + C: (batch, T, ngroups, dstate). + out: (batch, T, nheads, dim) preallocated output. + D: (nheads, dim) optional feed-through parameter. + z: (batch, T, nheads, dim) optional silu gate. + dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). + state_batch_indices: (batch,) optional cache slot mapping. + rand_seed: optional single-element int64 CUDA tensor for Philox PRNG seed. + When provided, state is stochastically rounded on store. Supported + for state.dtype in (fp16, int8, int16, fp8_e4m3fn); other dtypes + silently use deterministic rounding. fp16+SR and fp8+SR both + require sm_100a (Blackwell B200+) — wrapper asserts this loudly. + philox_rounds: number of Philox PRNG rounds (default 10). + state_scales: required when state.dtype in (int8, int16, fp8_e4m3fn). + Shape (cache_size, nheads, dim), fp32. Per-(head, dim) channel + decode scale (= 1 / encode_scale). The kernel writes scales on + checkpoint steps and reads them on load (broadcast over dstate). + Ignored for non-quantized state dtypes. + launch_with_pdl: enable external PDL (conv1d → precompute chain). + Defaults False; caller opts in when the upstream chain is PDL-safe. + Ignored on hardware that doesn't support PDL (sm < 90). + use_internal_pdl: enable internal PDL (precompute → main overlap). + Defaults True; override for testing only. + Ignored on hardware that doesn't support PDL (sm < 90). + + _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, + _precompute_num_warps, _precompute_num_stages, _heads_per_block, + _maxnreg, _num_ctas) are benchmark-only overrides; production callers + should leave them None to use the heuristic-tuned defaults. + """ + # PDL needs sm >= 90. + if get_sm_version() < 90: + launch_with_pdl = False + use_internal_pdl = False + + # Mode selection: + # mode="persistent_dynamic" (default): single persistent-CTA kernel + # covering the full batch. Each work-item dispatches via runtime + # PNAT check (is_write = (pnat + T) > MAX). No write/nowrite split. + # slot_perm is honored but optional. write_checkpoint is ignored + # (per-slot from PNAT). + # mode="persistent_main": persistent-CTA kernel with two launches + # (write half + nowrite half). Caller must pre-sort slots + # write-first via slot_perm and pass _n_writes / _n_writes_dev so + # the kernel can split the persistent loop into the two halves + # with the right WRITE_CHECKPOINT constexpr each time. RECTANGLE + # constexpr (= rectangle_for_nowrite) picks rect vs replay for the + # nowrite half. write_checkpoint is ignored (per-slot from PNAT). + assert mode in ("persistent_dynamic", "persistent_main"), ( + f"unknown mode {mode!r}; expected 'persistent_dynamic' or 'persistent_main'" + ) + + # --- Hardware support gates --- + # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX + # instructions (Ada Lovelace introduced them; Hopper/Blackwell carry them). + if state.dtype == torch.float8_e4m3fn: + assert get_sm_version() >= 89, ( + "fp8_e4m3fn state requires SM 89+ (Ada Lovelace / Hopper / Blackwell) " + f"for fp32↔fp8 cvt PTX instructions; current SM is {get_sm_version()}." + ) + + # PTX cvt.rs.* (stochastic rounding) family lands on Blackwell only. + # Wrapper fails loud; framework decides fall-back (e.g. drop SR, use RN). + # int8 / int16 SR uses pure-Triton libdevice.floor + uniform noise — no + # PTX SR instruction needed, runs anywhere. + if rand_seed is not None: + if state.dtype == torch.float16: + assert get_sm_version() >= 100, ( + "fp16 stochastic rounding (PTX cvt.rs.f16x2.f32) requires " + f"sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + ) + elif state.dtype == torch.float8_e4m3fn: + assert get_sm_version() >= 100, ( + "fp8 stochastic rounding (PTX cvt.rs.satfinite.e4m3x4.f32) " + f"requires sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + ) + + # --- Unsqueeze inputs to canonical shapes --- + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if x.dim() == 3: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if dt.dim() == 3: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if B.dim() == 3: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if C.dim() == 3: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None: + if z.dim() == 2: + z = z.unsqueeze(1) + if z.dim() == 3: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + if out.dim() == 2: + out = out.unsqueeze(1) + if out.dim() == 3: + out = out.unsqueeze(1) + + cache_size, nheads, dim, dstate = state.shape + batch, T, _, _ = x.shape + ngroups = B.shape[2] + assert nheads % ngroups == 0 + + # --- Quantization plumbing --- + # QUANT_MAX > 0 ⇔ state is int8 / int16 / fp8_e4m3fn. Kernel-entry + # static_assert on the Triton side mirrors this invariant. + quant_max = _QUANT_MAX_BY_DTYPE.get(state.dtype, 0.0) + is_quantized = quant_max > 0.0 + if is_quantized: + assert state_scales is not None, ( + f"state.dtype={state.dtype} requires state_scales tensor " + "(shape (cache_size, nheads, dim), fp32)." + ) + assert state_scales.shape == (cache_size, nheads, dim), ( + f"state_scales shape mismatch: expected {(cache_size, nheads, dim)}, " + f"got {state_scales.shape}." + ) + assert state_scales.dtype == torch.float32, ( + f"state_scales must be fp32, got {state_scales.dtype}." + ) + assert state_scales.device == state.device + + # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the + # placeholder degenerate case max_window = T (every step is a checkpoint + # step). For real replay-style checkpointing, max_window > T and + # `prev_num_accepted_tokens` can be 0..max_window. + max_window = old_x.shape[1] + assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" + # Replay-style code path uses BLOCK_SIZE_T = max(np2(T), 16) for the + # combined T-axis (T_new tile size) and reuses it for window loads. Until + # the heuristic is generalized to track max_window separately, require + # max_window to fit within that tile. + block_size_t = max(triton.next_power_of_2(T), 16) + assert max_window <= block_size_t, ( + f"max_window={max_window} exceeds BLOCK_SIZE_T={block_size_t} " + f"derived from T={T}; extend the heuristic to include max_window." + ) + + assert x.shape == (batch, T, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + assert B.shape == (batch, T, ngroups, dstate) + assert C.shape == B.shape + assert old_x.shape == (cache_size, max_window, nheads, dim) + assert old_B.shape == (cache_size, 2, max_window, ngroups, dstate) + assert old_dt.shape == (cache_size, 2, nheads, max_window) + assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_window) + assert cache_buf_idx.shape == (cache_size,) + assert prev_num_accepted_tokens.shape == (cache_size,) + + tie_hdim = ( + A.stride(-1) == 0 + and A.stride(-2) == 0 + and dt.stride(-1) == 0 + and (dt_bias is None or dt_bias.stride(-1) == 0) + ) + assert tie_hdim + + device = x.device + BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) + # Rectangle K-axis bound = window (max_window). Computed unconditionally + # so the launch sites can refer to it; only used on the rectangle path. + BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) + + # Allocate precomputed intermediates (per-call, not cached). Always + # allocate (T, K) — the largest layout that any path uses. Replay-style + # paths only touch the first T columns; rectangle/dynamic use the full K. + # The few extra unused columns per row are negligible (~6KB per layer at + # production sizes) and let the dispatch helpers share one buffer. + cb_scaled = torch.empty( + batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 + ) + decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) + + z_strides = ( + (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) if z is not None else (0, 0, 0, 0) + ) + + # Kernel tuning: BLOCK_SIZE_M, num_warps, HEADS_PER_BLOCK, precompute_num_warps. + # Dtype-aware heuristic from B200 sweeps (batch 1-512, T=6/32, TP=8, conv1d + + # chained PDL). Keyed on total_heads, BLOCK_SIZE_T, and state dtype; 16-bit + # states prefer different tiles from fp32 due to lower bandwidth. Philox + # gets its own branch — stochastic rounding shifts compute toward CUDA cores, + # so small-batch configs want more warps to hide the extra work. + total_heads = batch * nheads + heads_per_group = nheads // ngroups + state_is_16bit = state.dtype in (torch.float16, torch.bfloat16) + use_philox = rand_seed is not None + if BLOCK_SIZE_T <= 16: + if use_philox and state_is_16bit: + # Philox: more warps at small batch to hide CUDA core work. + # At large batch, converges to non-Philox fp16 config. + if total_heads <= 16: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 4, 4, 4, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + elif state_is_16bit: + if total_heads <= 16: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 4, 1 + elif total_heads <= 64: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 2, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 1, + min(2, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + else: # fp32 state (no Philox — fp32 doesn't need stochastic rounding) + if total_heads <= 32: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 4, 1 + elif total_heads <= 64: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + elif total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 2, 2, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 1, 2, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 2, + min(2, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + else: # T > 16 + if state_is_16bit: + if total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 16, + 1, + 4, + min(2, heads_per_group), + ) + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 1, + min(4, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 4, + min(2, heads_per_group), + ) + else: # fp32 state + if total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 2, + 4, + min(2, heads_per_group), + ) + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 2, + min(4, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 4, + min(2, heads_per_group), + ) + if _block_size_m is not None: + BLOCK_SIZE_M = _block_size_m + if _num_warps is not None: + num_warps = _num_warps + if _heads_per_block is not None: + heads_per_block = _heads_per_block + if _precompute_num_warps is not None: + precompute_num_warps = _precompute_num_warps + + # Per-main knob resolution: each _*_{write,nowrite} arg, if not None, + # overrides the corresponding shared value for ONE main launch only. + # Default (None) = tied to shared value (current behavior). + BLOCK_SIZE_M_WRITE = _block_size_m_write if _block_size_m_write is not None else BLOCK_SIZE_M + BLOCK_SIZE_M_NOWRITE = _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M + NUM_WARPS_WRITE = _num_warps_write if _num_warps_write is not None else num_warps + NUM_WARPS_NOWRITE = _num_warps_nowrite if _num_warps_nowrite is not None else num_warps + NUM_STAGES_WRITE = _num_stages_write if _num_stages_write is not None else _num_stages + NUM_STAGES_NOWRITE = _num_stages_nowrite if _num_stages_nowrite is not None else _num_stages + # Persistent-only per-main: + CTA_PER_SM_WRITE = _cta_per_sm_write if _cta_per_sm_write is not None else _cta_per_sm + CTA_PER_SM_NOWRITE = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _cta_per_sm + NUM_LOOP_STAGES_WRITE = _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages + NUM_LOOP_STAGES_NOWRITE = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages + + HAS_CACHE_BATCH_INDICES = state_batch_indices is not None + + assert nheads % heads_per_block == 0, ( + f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" + ) + assert heads_per_block <= heads_per_group, ( + f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" + ) + + # state_scales pointer + strides: real tensor when quantized, otherwise + # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). + if is_quantized: + state_scales_arg = state_scales + state_scales_strides = ( + state_scales.stride(0), + state_scales.stride(1), + state_scales.stride(2), + ) + else: + state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 + state_scales_strides = (0, 0, 0) + + # Per-path TMA descriptors for state — write-side and nowrite-side. Each + # kernel launch consumes the descriptor whose block_shape[0] matches its + # BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need + # distinct descriptors; otherwise the descriptor's block_shape[0] would + # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot / arithmetic + # on the loaded tile fails shape inference at compile time + # ("Cannot make_shape_compatible: incompatible dimensions"). When Mw == + # Mnw (tied, the common case) the two descriptors are the same object. + # Same memory (state's flat 2D view, shape (cache*nheads*dim, dstate)) + # and same dstate block_shape — only block_shape[0] differs. + # When no TMA flag is on, both variables hold the raw `state` tensor as a + # dummy; kernels never reference it because their constexprs are all + # False (Triton DCEs the dead branches). + # `triton.set_allocator()` must run before any descriptor-using launch. + if (_use_tma_rect_load or _use_tma_replay_write_load + or _use_tma_replay_write_store or _use_tma_replay_nowrite_load): + from triton.tools.tensor_descriptor import TensorDescriptor + _ensure_tma_allocator() + assert state.is_contiguous(), "TMA state requires contiguous state" + assert state.stride(-1) == 1, "TMA state requires inner stride 1" + _state_flat = state.view(-1, state.shape[-1]) + _dstate_pow2 = triton.next_power_of_2(dstate) + state_tma_descriptor_write = TensorDescriptor.from_tensor( + _state_flat, block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], + ) + if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: + state_tma_descriptor_nowrite = state_tma_descriptor_write + else: + state_tma_descriptor_nowrite = TensorDescriptor.from_tensor( + _state_flat, block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], + ) + else: + state_tma_descriptor_write = state # dummy; all consuming constexprs False + state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False + + # Slot permutation — pointer + USE_PERM gate. When the caller provides + # a perm tensor the dl-family launches read pid_b through it; otherwise + # we pass any valid pointer (state_batch_indices) and USE_PERM=False so + # the kernel falls back to pid_grid. + if slot_perm is not None: + assert slot_perm.dtype in (torch.int32, torch.int64), ( + f"slot_perm must be int32/int64, got {slot_perm.dtype}" + ) + assert slot_perm.numel() >= batch, ( + f"slot_perm has {slot_perm.numel()} entries; need >= batch ({batch})" + ) + slot_perm_arg = slot_perm + use_perm = True + else: + # Any valid ptr — gated by USE_PERM=False at compile time. + slot_perm_arg = state_batch_indices if state_batch_indices is not None else state + use_perm = False + + precomp_grid = (batch, nheads // heads_per_block) + d_strides = (D.stride(0), D.stride(1)) if D is not None else (0, 0) + + # ---- Launch helpers (close over locals) ------------------------------- + # Each helper is a thin closure that calls one Triton kernel with the + # full positional + kwarg argument list. Mode-dependent constexprs + # (write_checkpoint, early_out, rectangle) are passed in. + + def launch_dynamic_precompute(rectangle: bool): + _dynamic_precompute_kernel[precomp_grid]( + dt, dt_bias, A, B, C, + cb_scaled, decay_vec, + old_B, old_dt, old_dA_cumsum, + cache_buf_idx, prev_num_accepted_tokens, + state_batch_indices, pad_slot_id, + T, max_window, dstate, nheads // ngroups, + dt.stride(0), dt.stride(1), dt.stride(2), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), B.stride(1), B.stride(2), B.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + dt_softplus, + HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, + LAUNCH_WITH_PDL=launch_with_pdl, + LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, + HEADS_PER_BLOCK=heads_per_block, + RECTANGLE=rectangle, + num_warps=precompute_num_warps, + **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), + launch_pdl=launch_with_pdl, + ) + + # ---- launch_persistent_main ------------------------------------------ + # Persistent-CTA main kernel. Single launch covers `n_slots` slots + # starting at `slot_offset`. Caller invokes twice: once for the write + # half (slot_offset=0, n_slots=n_writes, write_checkpoint=True) and + # once for the nowrite half (slot_offset=n_writes, + # n_slots=batch-n_writes, write_checkpoint=False). Hard-sort + # contract: caller has pre-sorted slots so [0, n_writes) are writes + # and [n_writes, batch) are nowrites. + + # Resolve persistent-mode bench knobs. Defaults: cta_per_sm = 1 + # (one CTA per SM, matches upstream `_p_matmul_ogs.py`); num_loop_stages + # = 2 (matches in-tree `swiglu` precedent for non-dot persistent loops); + # flatten = True (canonical Triton 3.6 idiom); warp_specialize = False. + _num_sms = torch.cuda.get_device_properties(device).multi_processor_count + cta_per_sm_arg = _cta_per_sm if _cta_per_sm else 1 + num_persistent_arg = cta_per_sm_arg * _num_sms + num_loop_stages_arg = _num_loop_stages if _num_loop_stages else 2 + flatten_arg = True if _flatten is None else bool(_flatten) + warp_specialize_arg = False if _warp_specialize is None else bool(_warp_specialize) + # Per-launch work-item count. At small batch, total_work may be < the + # full persistent grid; capping `grid` at `min(NUM_PERSISTENT, total_work)` + # avoids launching empty CTAs that pay setup cost for no work. Correctness: + # the kernel's `tl.range(pid, total_work, NUM_PERSISTENT)` ensures each + # tile_id is covered exactly once across all live pids in [0, grid) when + # grid <= NUM_PERSISTENT (each CTA does 1 tile; loop step >= total_work + # exits immediately) AND when grid == NUM_PERSISTENT (each CTA loops over + # multiple tiles). NUM_PERSISTENT is now a runtime int (see kernel def + # docstring at _persistent_main_kernel) so changing cta_per_sm does NOT + # trigger a new Triton compile — same kernel binary, different loop step. + # (Named UPPERCASE for historical Triton-style consistency only; not + # constexpr.) + _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + + def launch_persistent_main(write_checkpoint: bool, + n_writes_dev: torch.Tensor, + *, + host_n_writes: int | None = None, + skip_empty_halves: bool = True, + launch_dependent_kernels: bool = False, + rectangle: bool = False): + # `n_writes_dev` is a (1,) int32 device tensor; the kernel reads + # the count from device memory. `host_n_writes` is the same value + # known host-side (when available — pure scenarios) and lets us + # skip the launch entirely if its half is empty. In mix scenarios + # the host doesn't know n_writes per iter without a sync, so + # `host_n_writes is None` and `skip_empty_halves` is forced False + # — both halves always launch and the kernel processes whatever + # range device-n_writes implies. + if skip_empty_halves and host_n_writes is not None: + n_slots_for_kernel = host_n_writes if write_checkpoint else (batch - host_n_writes) + if n_slots_for_kernel <= 0: + return + # Per-main knob selection. The two persistent_main launches (write + # half vs nowrite half) get independent BLOCK_SIZE_M / num_warps / + # num_stages / cta_per_sm / num_loop_stages. See the per-main args + # block in the wrapper signature. + _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE + _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE + _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE + _cps = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE + _cps = _cps if _cps else 1 + _nls = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE + _nls = _nls if _nls else 2 + _num_persistent = _cps * _num_sms + _num_pid_m_local = (dim + _bsm - 1) // _bsm + # Grid sizing: cap at min(full persistent grid, actual total_work). + # `n_slots` for this launch is `host_n_writes` (write half) / `batch - + # host_n_writes` (nowrite half) when host knows it (pure); else upper + # bound `batch` for mix scenarios where host can't read n_writes_dev + # without a sync. Upper-bound is fine — the kernel's runtime check + # only iterates actual work; the only cost of overcounting is a few + # extra CTAs. + if host_n_writes is not None: + _n_slots_for_launch = host_n_writes if write_checkpoint else (batch - host_n_writes) + else: + _n_slots_for_launch = batch + _total_work_launch = max(1, _n_slots_for_launch * _num_pid_m_local * nheads) + grid = (min(_num_persistent, _total_work_launch),) + # Per-path TMA descriptor — block_shape[0] must match _bsm. + _desc = (state_tma_descriptor_write if write_checkpoint + else state_tma_descriptor_nowrite) + _persistent_main_kernel[grid]( + state, _desc, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + n_writes_dev, batch, nheads, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + _bsm, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=write_checkpoint, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + NUM_PERSISTENT=_num_persistent, + NUM_LOOP_STAGES=_nls, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=False, + RECTANGLE=rectangle, + # 3 TMA flags. IS_DYNAMIC=False: WC fixed per launch; impl + # constexpr-folds the LOAD pick. When WC=True (write half), + # NOWRITE_LOAD is dummy False; when WC=False, WRITE_LOAD/STORE + # dummy False. NOWRITE_LOAD picks rect-load (RECTANGLE) or + # replay-nowrite-load. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), + USE_TMA_LOAD_NOWRITE=bool( + (_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load) + and not write_checkpoint + ), + USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), + num_warps=_nw, + **({"num_stages": _ns} if _ns else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, + launch_dependent_kernels: bool = False, + rectangle: bool = False): + # Single-launch persistent kernel covering the whole batch with + # runtime per-slot WRITE_CHECKPOINT branch. No half-split, no + # n_writes needed (the kernel ignores n_writes_dev when + # IS_DYNAMIC=True; Triton DCEs the load). is_write is computed + # at runtime per work-item from the loaded PNAT. + # We still pass `n_writes_dev` (the same tensor the persistent_main + # path uses) so the kernel signature is uniform; the value is + # immaterial. + # Grid sizing: cap at total_work (= batch * num_pid_m * nheads) for + # the dynamic case (full-batch coverage); see launch_persistent_main + # comment for correctness rationale. + _total_work_launch = max(1, batch * _num_pid_m * nheads) + grid = (min(num_persistent_arg, _total_work_launch),) + # Persistent-dynamic kernel uses a single BLOCK_SIZE_M (same as the + # wrapper's BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE tied convention), so + # the write-side descriptor matches. Both write and nowrite slots + # in this kernel share that BSM. + _persistent_main_kernel[grid]( + state, state_tma_descriptor_write, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + n_writes_dev, batch, nheads, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=False, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + # persistent_dynamic forces USE_PERM=False regardless of caller- + # provided slot_perm — our pd tuning runs all happened with + # SORT=0 (no slot_perm passed), so honoring slot_perm here would + # silently shift pd to an untimed code path. Revisit if/when + # we benchmark pd with slot_perm. + USE_PERM=False, + NUM_PERSISTENT=num_persistent_arg, + NUM_LOOP_STAGES=num_loop_stages_arg, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=True, + RECTANGLE=rectangle, + # 3 TMA flags. IS_DYNAMIC=True: is_write is runtime per slot; + # impl's load TMA picks per-slot (constexpr ternary becomes a + # runtime branch — both load forms emitted, ~negligible cost). + # NOWRITE_LOAD picks rect-load when RECTANGLE, else + # replay-nowrite-load. STORE only fires on runtime is_write. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load), + USE_TMA_LOAD_NOWRITE=bool( + _use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load + ), + USE_TMA_STORE=bool(_use_tma_replay_write_store), + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + # ---- Mode dispatch ---------------------------------------------------- + with torch.cuda.device(device.index): + if mode == "persistent_dynamic": + # Single-launch persistent kernel covering the full batch. + # Each work-item dispatches via runtime PNAT check (is_write = + # (pnat + T) > MAX). No n_writes/half-split — kernel ignores + # n_writes_dev when IS_DYNAMIC=True (Triton DCEs the load). + # We still need a valid pointer to satisfy the kernel arg + # signature; allocate or reuse `_n_writes_dev`. + n_writes_dev_local = ( + _n_writes_dev if _n_writes_dev is not None + else torch.zeros(1, dtype=torch.int32, device=device) + ) + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_persistent_dynamic_main( + n_writes_dev_local, + launch_dependent_kernels=False, + rectangle=rectangle_for_nowrite, + ) + elif mode == "persistent_main": + # Persistent-CTA main kernel. One shared dynamic_precompute + # (per-slot dispatch at runtime via PNAT) feeds two + # persistent_main launches (write half + nowrite half). + # + # Hard-sort contract: caller has pre-sorted slots host-side so + # PNAT is monotone (writes first). Pass the perm via + # slot_perm + USE_PERM. + # + # n_writes is read by the kernel from a (1,) int32 device + # tensor. The caller can provide: + # * _n_writes_dev only (mix): a pre-filled (1,) int32 tensor + # it updates per iter via pre_iter_fn outside the captured + # graph. host can't cheaply read it without a sync, so both + # halves always launch. + # * _n_writes only (non-graph callers, e.g. unit tests): host + # int. We allocate the scratch tensor on the fly. CANNOT + # be used inside CUDA-graph capture — alloc inside capture + # invalidates the stream. + # * Both (pure under graph capture): caller pre-allocates the + # tensor outside capture and tells us the host value too. + # We skip the internal allocation and apply host-skip when + # _persistent_skip_empty_halves=True. This is the + # production-equivalent path the bench's pure cells take. + if _n_writes_dev is not None: + n_writes_dev_local = _n_writes_dev # no allocation + if _n_writes is not None: + # Caller provided both: pure scenario with pre-allocated + # tensor. Use host_n_writes for the skip-empty fast path. + assert 0 <= _n_writes <= batch, ( + f"_n_writes={_n_writes} must be in [0, batch={batch}]" + ) + host_n_writes_local = _n_writes + skip_empty_local = _persistent_skip_empty_halves + else: + # Mix: host doesn't know n_writes without a sync. + host_n_writes_local = None + skip_empty_local = False + else: + # No pre-allocated tensor. Fall back to on-the-fly alloc + # from _n_writes (host int). NOT graph-capture-safe. + assert _n_writes is not None, ( + "mode='persistent_main' requires either _n_writes " + "(host int, non-graph callers) or _n_writes_dev (device " + "tensor, recommended for graph-capture callers)." + ) + assert 0 <= _n_writes <= batch, ( + f"_n_writes={_n_writes} must be in [0, batch={batch}]" + ) + n_writes_dev_local = torch.tensor( + [_n_writes], dtype=torch.int32, device=device, + ) + host_n_writes_local = _n_writes + skip_empty_local = _persistent_skip_empty_halves + # rectangle_for_nowrite=True: precompute populates cb_scaled + # for the rect path; nowrite half uses the rectangle impl; + # write half always replay-style (rect doesn't apply). + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_persistent_main( + write_checkpoint=True, + n_writes_dev=n_writes_dev_local, + host_n_writes=host_n_writes_local, + skip_empty_halves=skip_empty_local, + launch_dependent_kernels=True, + rectangle=False, # write always replay-style + ) + launch_persistent_main( + write_checkpoint=False, + n_writes_dev=n_writes_dev_local, + host_n_writes=host_n_writes_local, + skip_empty_halves=skip_empty_local, + launch_dependent_kernels=False, + rectangle=rectangle_for_nowrite, + ) + else: + raise ValueError( + f"mode={mode!r} is not supported. Supported modes: " + f"'persistent_dynamic', 'persistent_main'." + ) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py new file mode 100644 index 000000000000..793cbb22c95d --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py @@ -0,0 +1,4847 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Standalone benchmark for replay_selective_state_update (Triton kernel). + +Suitable for nsight-compute (ncu) and nsight-systems (nsys) capture. + +Fixed model config: NVIDIA-Nemotron-3-Super-120B-A12B at TP=8 + nheads=16, head_dim=64, d_state=128, ngroups=1 + +mtp_len is the per-request sequence length processed by replay: in MTP it +equals num_draft_tokens + 1 target token, so --mtp-lengths 6 models 5 drafts ++ 1 target. + +Baseline kernel (--baseline [triton|flashinfer]): + Calls selective_state_update with T=mtp_len tokens and disable_state_update=True, + matching the MTP scoring pass in mamba2_mixer.py exactly. + +Timing methodology +================== + +All in-bench timing comes from CUPTI's Activity API (1 ns kernel +timestamps from the GPU profiling fabric). cudaEvent.elapsed_time() was +removed — its ~0.5 us resolution overshoots CUPTI by ~50% on short kernels +in graphs, and we have no other use for it here. See the CUPTI block +lower in this file for the timer source. + +Three modes: + + --cupti --cuda-graph (default) + Capture a small CUDA graph for the cell, replay it for warmup + timed + iterations, and read kernel start/end from CUPTI. Raw CUPTI buffers are + parsed out-of-process on the timed path, with a cached ordinal plan used + to keep only the kernels we care about. + + --cupti --no-cuda-graph + Eager loop with CUPTI. Per-kernel timestamps are still accurate, but + the per-iter SPAN (max(end) - min(start)) now includes the Python + launch latency BETWEEN consecutive kernels in run_fn (~100 µs on + Hopper/Blackwell). Graph capture and PDL hide that latency; eager + mode honestly reports it. For per-kernel timing in eager mode, look + at per_kernel.start_us/end_us in --json-detailed output rather than + the span percentiles. Useful when graph capture is undesirable. + + --no-cupti (with or without --cuda-graph) + No in-bench timing — just runs the kernels for an external profiler + (nsys / ncu) to time. In-process CUPTI conflicts with nsys's own + subscriber, so disable ours when wrapping in nsys. Bench output + reports zeros for median/p95/p99; trust the external trace. + +JSON output schema (--json-output PATH) +======================================= + +Designed to be parsed by collect.py / report.py without touching sqlite or +NVTX traces. Future agents: prefer reading this JSON over re-running nsys. + + { + "metadata": {timestamp, cmd, tp_size, warmup, iters, variant, cupti}, + "results": { + "": {median, p95, p99, n, iters_us, [n_writes_per_iter], [per_kernel]} + } + } + +Key format mirrors collect.py's kernel_data.json convention: + incremental/{batch}/{mtp}/{sd}/k{prev_k}/{sweep_parts}/tp{tp} + triton/{batch}/{mtp}/{sd}/tp{tp} + flashinfer/{batch}/{mtp}/{sd}/tp{tp} + + - is normalized: bf16 / fp16 / fp32 / int8 / int16 / fp8. + - is e.g. "M16_W1_S3_SR0_RECT0_WC1" — flags concatenated by + underscore in canonical (M, W, S, pW, pS, H, R, CT, SR, RECT, WC) order. + - All numeric values in microseconds (us). + +Per-record fields: + - median, p95, p99: span statistics (us). Span = max(kernel_end_ns) - + min(kernel_start_ns) across the iter's kernels — same convention as + nsys-derived collect.py used to use. + - n: number of timed iters that contributed. + - iters_us: list of length n, raw per-iter spans. + - n_writes_per_iter: for mix rows, list of length n with the number of + write-path slots in each timed iteration. + - per_kernel: {: {start_us: [...], end_us: [...]}} where + timestamps are RELATIVE to that iter's first kernel start, in us. Lets + you see PDL overlap directly without an external profiler. Only with + --json-detailed. + +Example usage: + # Basic sweep (default = --cupti, just summary stats) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 1,2,4 --mtp-lengths 1,4,8 --warmup 5 --iters 20 + + # JSON output, summary stats only (compact) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 16 --mtp-lengths 6 --json-output /tmp/out.json + + # JSON output, full per-iter / per-kernel data (for PDL analysis etc.) + python benchmark_replay_selective_state_update.py \\ + --batch-sizes 16 --mtp-lengths 6 \\ + --json-output /tmp/out.json --json-detailed + + # nsys capture (--no-cupti so our subscriber doesn't conflict) + nsys profile --capture-range=cudaProfilerApi \\ + python benchmark_replay_selective_state_update.py --profile --no-cupti + + # ncu capture (--no-cupti --no-cuda-graph: each kernel replayable solo) + ncu --target-processes all \\ + python benchmark_replay_selective_state_update.py --profile \\ + --no-cupti --no-cuda-graph \\ + --batch-sizes 1 --mtp-lengths 4 --warmup 5 --iters 5 +""" + +import argparse +import atexit +import ctypes +import importlib +import itertools +import json +import multiprocessing as mp +import os +import queue +import statistics +import sys +import threading +import time +from datetime import datetime +from multiprocessing import shared_memory +from pathlib import Path + +import numpy as np +import torch +from einops import repeat + + +def _import_mamba_kernels_fast(): + """Load kernel modules directly (~40s faster than a full tensorrt_llm init). + Use --full-import as the fallback if module dependencies change. + + Strategy: stub the parent packages (tensorrt_llm, tensorrt_llm._torch, + tensorrt_llm._torch.modules) in sys.modules with __path__ set, but do + NOT execute their __init__.py. Then load the leaf kernel modules. + When a kernel body imports e.g. tensorrt_llm._utils.get_sm_version, + Python's machinery resolves it against our stub's __path__ and loads + only _utils.py — skipping the heavy tensorrt_llm package init. + """ + import types + + repo_root = Path(__file__).resolve().parents[5] + trtllm_dir = repo_root / "tensorrt_llm" + mamba_pkg = "tensorrt_llm._torch.modules.mamba" + mamba_dir = trtllm_dir / "_torch" / "modules" / "mamba" + + def _stub_pkg(fqn: str, pkg_dir: Path): + """Register a stub package in sys.modules without running its + __init__.py. Sets __path__ so Python can resolve submodule imports + against the real directory on disk.""" + if fqn in sys.modules: + return + stub = types.ModuleType(fqn) + stub.__path__ = [str(pkg_dir)] + sys.modules[fqn] = stub + + # Stub the parent chain so `from tensorrt_llm._utils import ...` (and + # similar) work without triggering tensorrt_llm/__init__.py. + _stub_pkg("tensorrt_llm", trtllm_dir) + _stub_pkg("tensorrt_llm._torch", trtllm_dir / "_torch") + _stub_pkg("tensorrt_llm._torch.modules", trtllm_dir / "_torch" / "modules") + + def _load(mod_name: str, file_name: str): + fqn = f"{mamba_pkg}.{mod_name}" if mod_name else mamba_pkg + if fqn in sys.modules: + return sys.modules[fqn] + spec = importlib.util.spec_from_file_location( + fqn, + mamba_dir / file_name, + submodule_search_locations=[str(mamba_dir)] if file_name == "__init__.py" else [], + ) + mod = importlib.util.module_from_spec(spec) + sys.modules[fqn] = mod + spec.loader.exec_module(mod) + return mod + + # 1. Package __init__ (defines PAD_SLOT_ID = -1) + _load("", "__init__.py") + # 2. softplus helper (used by both kernel modules) + _load("softplus", "softplus.py") + # 3. The actual kernels + replay_mod = _load("replay_selective_state_update", "replay_selective_state_update.py") + checkpoint_mod = _load("checkpointing_state_update", "checkpointing_state_update.py") + base_mod = _load("selective_state_update", "selective_state_update.py") + conv1d_mod = _load("causal_conv1d_triton", "causal_conv1d_triton.py") + + return ( + replay_mod.replay_selective_state_update, + checkpoint_mod.checkpointing_state_update, + base_mod.selective_state_update, + conv1d_mod.causal_conv1d_update, + ) + + +def _import_mamba_kernels_full(): + """Import via the standard tensorrt_llm package (slow but safe).""" + from tensorrt_llm._torch.modules.mamba.causal_conv1d_triton import causal_conv1d_update + from tensorrt_llm._torch.modules.mamba.checkpointing_state_update_slim import ( + checkpointing_state_update, + ) + from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( + replay_selective_state_update, + ) + from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update + + return ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) + + +# Use fast import by default; --full-import parsed later but we need the +# functions at module level. Check sys.argv early. +if "--full-import" in sys.argv: + ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) = _import_mamba_kernels_full() +else: + try: + ( + replay_selective_state_update, + checkpointing_state_update, + selective_state_update, + causal_conv1d_update, + ) = _import_mamba_kernels_fast() + except Exception as e: # noqa: BLE001 - exit loudly; don't hide a fast-import regression + print( + f"ERROR: fast import failed ({type(e).__name__}: {e})\n" + "Re-run with --full-import for the slow but stable path, " + "then file a bug or fix _import_mamba_kernels_fast.", + file=sys.stderr, + ) + sys.exit(1) + + +_VARIANT_FNS = { + "replay": lambda: replay_selective_state_update, + "checkpointing": lambda: checkpointing_state_update, +} + +# Model config defaults (Nemotron-3-Super-120B full model). +# --tp-size divides nheads and ngroups to get the per-GPU slice. +# TP=1: nheads=128, ngroups=8 +# TP=4: nheads=32, ngroups=2 +# TP=8: nheads=16, ngroups=1 (default) +NHEADS = 128 +HEAD_DIM = 64 +D_STATE = 128 +NGROUPS = 8 +TP_SIZE = 8 # default; overridden by --tp-size + +# L2 flush buffer: ~128 MB — larger than L2 on A100/H100/B200 +_L2_FLUSH_SIZE = 32 * 1024 * 1024 # float32 elements → 128 MB +_l2_flush: torch.Tensor | None = None + + +def _init_l2_flush() -> None: + global _l2_flush + _l2_flush = torch.empty(_L2_FLUSH_SIZE, dtype=torch.float32, device="cuda") + + +def _flush_l2() -> None: + """Evict L2 by writing to a large buffer then synchronising.""" + assert _l2_flush is not None + _l2_flush.fill_(0.0) + torch.cuda.synchronize() + + +def _resolve_prev_ks(args, mtp_len: int) -> list[int]: + """Resolve prev_k values for one mtp_len cell. + + Two input modes (mutually exclusive in spirit; absolute wins if both given): + --prev-tokens-int "0,10,11,16" → use literal integers, clamped to + [0, max_window] (where max_window is the cache T-axis capacity). + --prev-tokens-fracs "0,0.5,1.0" → fractions of mtp_len, clamped to + [0, mtp_len] (current behavior). + + For replay-style checkpointing the cache holds up to max_window old + tokens, so absolute integers are the right knob. Fractions are kept + for back-compat with prior placeholder runs. + """ + upper = getattr(args, "max_window", 0) or mtp_len + if getattr(args, "prev_tokens_int", None): + return sorted(set(max(0, min(upper, int(v))) for v in args.prev_tokens_int)) + return sorted( + set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs) + ) + + +# Tensor construction helpers + +# Module-level cache for tensor buffers shared across cells. Keyed by all +# the "fixed" dimensions (state_dtype, act_dtype, max_window, mtp_len, +# nheads, head_dim, d_state, ngroups). Within a key, the batch dim grows +# in place: if a new cell requests a batch <= cached max_batch, we return +# views (slices) of the existing tensors; if batch > cached max_batch, we +# realloc at the new batch (which becomes the new max). Tensors never shrink. +# +# Rationale: torch.randn/zeros for these tensor shapes at b=512 takes +# ~10-30ms per call. At ~895 cells/min with 5 different batch sizes, +# we were re-allocating every cell. Caching saves the bulk of that per-cell +# overhead, raising GPU util in the timing phase. +# +# Reset state lives in caller (state_work = state0.copy_), so cached state0 +# is purely a reference whose contents stay fixed once allocated. This is +# fine: it's only read by the reset path. +_TENSOR_CACHE: dict = {} + + +def _build_tensors( + batch: int, + mtp_len: int, + state_dtype: torch.dtype, + act_dtype: torch.dtype, + nheads: int, + head_dim: int, + d_state: int, + ngroups: int, + max_window: int | None = None, +): + """ + Build all tensors for one benchmark configuration. + + nheads/ngroups are already TP-split (i.e. full_nheads // tp_size). + + Returns: + state0 : (batch, nheads, head_dim, d_state) – initial SSM state + x, dt, B, C : (batch, mtp_len, ...) – token inputs for both kernels + A, dt_bias, D : SSM parameters (float32, tie_hdim strides) + prev_tokens : (batch,) + out_incr : pre-allocated output for replay kernel (batch, mtp_len, nheads, head_dim) + out_base : pre-allocated output for baseline kernel (batch, mtp_len, nheads, head_dim) + intermediate_states_buffer: for baseline kernel (batch, mtp_len, nheads, head_dim, d_state) + """ + device = "cuda" + + # Cache lookup — grow batch in place if needed; else return views. + cache_key = (state_dtype, act_dtype, max_window, mtp_len, + nheads, head_dim, d_state, ngroups) + cached = _TENSOR_CACHE.get(cache_key) + if cached is not None and cached["max_batch"] >= batch: + # Hit — return slices for current batch. + b = batch + return ( + cached["state0"][:b], + cached["state_scales0"][:b] if cached["state_scales0"] is not None else None, + cached["old_x"][:b], + cached["old_B"][:b], + cached["old_dt"][:b], + cached["old_dA_cumsum"][:b], + cached["cache_buf_idx"][:b], + cached["x"][:b], + cached["dt"][:b], + cached["B"][:b], + cached["C"][:b], + cached["A"], + cached["dt_bias"], + cached["D"], + cached["prev_tokens"][:b], + cached["slot_perm_buf"][:b], + cached["out_incr"][:b], + cached["out_base"][:b], + cached["intermediate_states_buffer"][:b], + cached["xbc_input"][:b], + cached["conv_state"][:b], + cached["conv_weight"], + cached["conv_bias"], + cached["d_inner"], + cached["conv_dim"], + ) + + # Miss or grow. Allocate at new max_batch (existing data, if any, is + # released — caller code re-fills via reset paths anyway). Rebind + # `batch` locally to alloc_batch so the existing allocation code below + # uses the larger size; keep request_batch for the final slice. + request_batch = batch + alloc_batch = batch if cached is None else max(batch, cached["max_batch"]) + batch = alloc_batch + + torch.manual_seed(42) + + # --- SSM parameters (float32, tie_hdim strides) --- + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) # stride(-1)=0, stride(-2)=0 + + dt_bias_base = torch.randn(nheads, device=device, dtype=torch.float32) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) # stride(-1)=0 + + D_base = torch.randn(nheads, device=device, dtype=torch.float32) + D = repeat(D_base, "h -> h p", p=head_dim) + + # --- SSM state --- + # Quantized dtypes need their own initializer (torch.randn doesn't accept + # int) and a parallel fp32 scales tensor (per-(head, dim) channel decode + # scale, broadcast over dstate). Quant state is filled with realistic- + # range values via fp32 → quant; scales are derived consistently so the + # initial state isn't garbage on dequant. + _QUANT_BENCH = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, + } + if state_dtype in _QUANT_BENCH: + quant_max = _QUANT_BENCH[state_dtype] + state_fp32 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + amax = state_fp32.abs().amax(dim=-1) # (batch, nheads, head_dim) + encode_scale = quant_max / amax.clamp(min=1e-30) + state_scales0 = (1.0 / encode_scale).to(torch.float32) # decode scale + scaled = state_fp32 * encode_scale.unsqueeze(-1) + if state_dtype == torch.float8_e4m3fn: + state0 = scaled.clamp(-quant_max, quant_max).to(state_dtype) + else: + state0 = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) + else: + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state_scales0 = None + + # --- Cache tensors for replay kernel --- + # max_window is the cache T-axis capacity; defaults to mtp_len (the + # placeholder/degenerate case where every step is a checkpoint step). + # For real replay-style checkpointing, max_window > mtp_len. + cache_T = max_window if max_window is not None else mtp_len + # old_x: single-buffered (cache, max_window, nheads, dim) + old_x = torch.randn(batch, cache_T, nheads, head_dim, device=device, dtype=act_dtype) + # old_B: double-buffered (cache, 2, max_window, ngroups, dstate) + old_B = torch.randn(batch, 2, cache_T, ngroups, d_state, device=device, dtype=act_dtype) + # old_dt: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous + old_dt = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) + # old_dA_cumsum: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous + old_dA_cumsum = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) + # cache_buf_idx: which buffer to read (0 or 1) + cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) + + # --- Token inputs (used by both replay and baseline kernels) --- + x = torch.randn(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) + # dt must match D's dtype (fp32) for flashinfer — force it for all paths. + dt_base = torch.randn(batch, mtp_len, nheads, device=device, dtype=torch.float32) + dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) # tie_hdim + B = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) + C = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) + + # prev_tokens placeholder — overwritten per-run + prev_tokens = torch.zeros(batch, device=device, dtype=torch.int32) + # slot_perm placeholder — overwritten per-run by mix pre_iter_fn when + # sort_slots is enabled. Identity by default so cells that don't sort + # (or pure-batch cells) get a meaningful identity perm if the kernel + # ends up reading it (USE_PERM=False makes this path unused). + slot_perm_buf = torch.arange(batch, device=device, dtype=torch.int32) + + out_incr = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) + out_base = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) + + # intermediate_states_buffer is only consumed by the fp/baseline path; + # for quantized state dtypes we'll skip baselines entirely, so the buffer + # dtype falls back to fp32 to keep selective_state_update happy. + int_buffer_dtype = state_dtype if state_dtype not in _QUANT_BENCH else torch.float32 + intermediate_states_buffer = torch.zeros( + batch, mtp_len, nheads, head_dim, d_state, device=device, dtype=int_buffer_dtype + ) + + # --- Conv1d tensors (for --with-conv1d mode) --- + d_inner = nheads * head_dim + conv_dim = d_inner + 2 * ngroups * d_state + d_conv = 4 # conv kernel width for Nemotron/Mamba2 + + # xbc_input: (batch, conv_dim, mtp_len) — "hot" input from in_proj. + # Match production layout: in_proj output is (batch*mtp_len, conv_dim) + # contiguous, then .view(batch, mtp_len, conv_dim).transpose(1, 2) + # gives strides (mtp_len*conv_dim, 1, conv_dim) — NOT the standard + # (conv_dim*mtp_len, mtp_len, 1) of a freshly allocated 3D tensor. + # Conv1d preserves input strides in its output, so downstream split + # + view inherits the correct layout without needing .contiguous(). + xbc_input_flat = torch.randn(batch * mtp_len, conv_dim, device=device, dtype=act_dtype) + xbc_input = xbc_input_flat.view(batch, mtp_len, conv_dim).transpose(1, 2) + # conv_state: (batch, conv_dim, d_conv) — "cold" cache + conv_state = torch.randn(batch, conv_dim, d_conv, device=device, dtype=act_dtype) + # conv_weight: (conv_dim, d_conv) — parameter + conv_weight = torch.randn(conv_dim, d_conv, device=device, dtype=act_dtype) + # conv_bias: (conv_dim,) — parameter + conv_bias = torch.randn(conv_dim, device=device, dtype=act_dtype) + + # Store full-batch buffers in cache and return slices at request_batch. + _TENSOR_CACHE[cache_key] = { + "max_batch": alloc_batch, + "state0": state0, + "state_scales0": state_scales0, + "old_x": old_x, + "old_B": old_B, + "old_dt": old_dt, + "old_dA_cumsum": old_dA_cumsum, + "cache_buf_idx": cache_buf_idx, + "x": x, + "dt": dt, + "B": B, + "C": C, + "A": A, + "dt_bias": dt_bias, + "D": D, + "prev_tokens": prev_tokens, + "slot_perm_buf": slot_perm_buf, + "out_incr": out_incr, + "out_base": out_base, + "intermediate_states_buffer": intermediate_states_buffer, + "xbc_input": xbc_input, + "conv_state": conv_state, + "conv_weight": conv_weight, + "conv_bias": conv_bias, + "d_inner": d_inner, + "conv_dim": conv_dim, + } + rb = request_batch + return ( + state0[:rb], + state_scales0[:rb] if state_scales0 is not None else None, + old_x[:rb], + old_B[:rb], + old_dt[:rb], + old_dA_cumsum[:rb], + cache_buf_idx[:rb], + x[:rb], + dt[:rb], + B[:rb], + C[:rb], + A, + dt_bias, + D, + prev_tokens[:rb], + slot_perm_buf[:rb], + out_incr[:rb], + out_base[:rb], + intermediate_states_buffer[:rb], + xbc_input[:rb], + conv_state[:rb], + conv_weight, + conv_bias, + d_inner, + conv_dim, + ) + + +# ============================================================================= +# CUPTI in-process kernel timing +# +# Self-contained module-in-a-file. Reads kernel start/end timestamps directly +# from the GPU profiling fabric via CUPTI's Activity API (1 ns +# resolution), avoiding two pitfalls of the cuda-events path: +# +# 1. cudaEvent.elapsed_time() resolution (~0.5 us) is too coarse for the +# short kernels we care about, especially with PDL + cuda graphs at +# small batch — events recorded inside a graph have proven noisy. +# 2. nsys is the only known accurate alternative, but the +# profile-export-sqlite-parse pipeline is heavy and out-of-process. +# +# This is functionally equivalent to wrapping each cell in nsys, except it +# runs in the same benchmark process and sends raw activity buffers to a +# parser process instead of materializing Python objects in the CUPTI callback. +# ============================================================================= + + +# Substring match: kernels run_fn launches that we want to time. Mirrors +# the parser in scripts/.../collect.py so cupti and nsys-based outputs agree. +_CUPTI_KEEP_KERNEL_SUBSTRINGS = ( + "_replay_precompute", + "_checkpointing_precompute", + "_rectangle_precompute", + "_dynamic_precompute", + "_replay_state_update", + "_checkpointing_main", + "_rectangle_main", + "_dynamic_main", + "_persistent_main", + "selective_scan_update", + "selective_state_update", + "causal_conv1d_update", +) + + +def _kernels_per_iter_incremental( + mode: str, + with_conv1d: bool, + *, + persistent_skip_empty: bool = True, +) -> int: + """Expected number of CUPTI-tracked kernels per iter for the incremental + kernel chain, given the dispatch mode and the conv1d flag. + + Used to validate CUPTI record counts (no auto-inference — silent + mis-timing is the failure mode we're guarding against). + + `persistent_skip_empty=True` (today's behavior): the + `mode='persistent_main'` launch helper host-early-outs when its half + is empty (n_writes=0 or n_writes=batch in pure scenarios), so only + one of the two persistent_main_kernel launches actually fires per + iter. With `persistent_skip_empty=False` (future no-eo mode), both + halves always launch and K bumps by 1. + + `persistent_dynamic` always launches 1 main; not affected by the flag. + """ + if mode == "persistent_dynamic": + k = 2 # 1 dynamic_precomp + 1 persistent_main + elif mode == "persistent_main": + k = 2 if persistent_skip_empty else 3 # see docstring + else: + raise ValueError(f"_kernels_per_iter_incremental: unknown mode {mode!r}") + if with_conv1d: + k += 1 + return k + + +def _kernels_per_iter_baseline(with_conv1d: bool) -> int: + """Expected kernels per iter for triton / flashinfer baselines. + + Both baselines run a single state-update kernel; `--with-conv1d` + prepends one conv1d kernel. + """ + return 2 if with_conv1d else 1 + + +_LIBCUPTI_CANDIDATES = ( + os.environ.get("CUPTI_LIBRARY_PATH"), + "/usr/local/lib/python3.12/dist-packages/nvidia/cu13/lib/libcupti.so.13", + "libcupti.so.13", + "libcupti.so", +) +_CUPTI_SUCCESS = 0 +_CUPTI_ERROR_MAX_LIMIT_REACHED = 12 +_CUPTI_ERROR_INVALID_KIND = 21 +_CUPTI_ACTIVITY_KIND_KERNEL = 3 +_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL = 10 +_CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER = 5 +_CUPTI_HOST_BUFFER_BYTES = 1024 * 1024 +_CUPTI_HOST_BUFFER_COUNT = 16 + +# Multiprocessing start method for compile-warmup + CUPTI parser children. +# Set in __main__ from --mp-start-method. "spawn" (default) is robust; each +# child re-imports torch/triton/etc (~15s). "forkserver" preloads once and +# forks cheaply (~1s/child) — see __main__ block for the preload setup. +_MP_START_METHOD = "spawn" +_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE = 1 +_DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX = 4 + + +def _load_libcupti() -> ctypes.CDLL: + errors = [] + for candidate in _LIBCUPTI_CANDIDATES: + if not candidate: + continue + try: + return ctypes.CDLL(candidate) + except OSError as exc: + errors.append(f"{candidate}: {exc}") + raise ImportError("Unable to load libcupti: " + "; ".join(errors)) + + +class _CuptiActivityKernel11Prefix(ctypes.Structure): + _pack_ = 1 + _fields_ = [ + ("kind", ctypes.c_int), + ("cache_config", ctypes.c_uint8), + ("shared_memory_config", ctypes.c_uint8), + ("registers_per_thread", ctypes.c_uint16), + ("partitioned_global_cache_requested", ctypes.c_int), + ("partitioned_global_cache_executed", ctypes.c_int), + ("start", ctypes.c_uint64), + ("end", ctypes.c_uint64), + ("completed", ctypes.c_uint64), + ("device_id", ctypes.c_uint32), + ("context_id", ctypes.c_uint32), + ("stream_id", ctypes.c_uint32), + ("grid_x", ctypes.c_int32), + ("grid_y", ctypes.c_int32), + ("grid_z", ctypes.c_int32), + ("block_x", ctypes.c_int32), + ("block_y", ctypes.c_int32), + ("block_z", ctypes.c_int32), + ("static_shared_memory", ctypes.c_int32), + ("dynamic_shared_memory", ctypes.c_int32), + ("local_memory_per_thread", ctypes.c_uint32), + ("local_memory_total", ctypes.c_uint32), + ("correlation_id", ctypes.c_uint32), + ("grid_id", ctypes.c_int64), + ("name", ctypes.c_void_p), + ("reserved0", ctypes.c_void_p), + ("queued", ctypes.c_uint64), + ("submitted", ctypes.c_uint64), + ("launch_type", ctypes.c_uint8), + ("is_shared_memory_carveout_requested", ctypes.c_uint8), + ("shared_memory_carveout_requested", ctypes.c_uint8), + ("padding", ctypes.c_uint8), + ("shared_memory_executed", ctypes.c_uint32), + ("graph_node_id", ctypes.c_uint64), + ] + + +def _configure_cupti_get_next_record(libcupti) -> None: + libcupti.cuptiActivityGetNextRecord.argtypes = [ + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_void_p), + ] + libcupti.cuptiActivityGetNextRecord.restype = ctypes.c_int + + +def _parse_cupti_buffer_ptr(libcupti, buffer_ptr: int, valid_size: int, *, include_names: bool): + records = [] + zero_ts_count = 0 + zero_ts_names: dict[str, int] = {} + record_ptr = ctypes.c_void_p(None) + while True: + result = libcupti.cuptiActivityGetNextRecord( + ctypes.c_void_p(buffer_ptr), + valid_size, + ctypes.byref(record_ptr), + ) + if result == _CUPTI_SUCCESS: + kind = ctypes.cast(record_ptr, ctypes.POINTER(ctypes.c_int)).contents.value + if kind not in (_CUPTI_ACTIVITY_KIND_KERNEL, _CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL): + continue + kernel = ctypes.cast(record_ptr, ctypes.POINTER(_CuptiActivityKernel11Prefix)).contents + name = None + if include_names: + if kernel.name: + name = ctypes.string_at(kernel.name).decode("utf-8", errors="replace") + else: + name = "?" + if kernel.start == 0 or kernel.end == 0: + zero_ts_count += 1 + if name is not None: + zero_ts_names[name] = zero_ts_names.get(name, 0) + 1 + continue + if include_names: + records.append(( + name, + int(kernel.start), + int(kernel.end), + int(kernel.correlation_id), + 0, + int(kernel.graph_node_id), + int(kernel.stream_id), + )) + else: + records.append(( + int(kernel.start), + int(kernel.end), + int(kernel.correlation_id), + int(kernel.graph_node_id), + int(kernel.stream_id), + )) + elif result == _CUPTI_ERROR_MAX_LIMIT_REACHED: + break + elif result == _CUPTI_ERROR_INVALID_KIND: + break + else: + raise RuntimeError(f"cuptiActivityGetNextRecord failed with CUptiResult={result}") + return records, zero_ts_count, zero_ts_names + + +def _apply_cupti_filter_plan(numeric_records, filter_plan): + if not filter_plan: + return [ + (None, start, end, corr, 0, graph_node_id, stream_id) + for start, end, corr, graph_node_id, stream_id in sorted(numeric_records) + ] + + filtered = [] + replay_idx = 0 + record_idx = 0 + for start, end, corr, graph_node_id, stream_id in sorted(numeric_records): + if replay_idx >= len(filter_plan): + break + records_per_replay, ordinal_names = filter_plan[replay_idx] + if record_idx < len(ordinal_names): + name = ordinal_names[record_idx] + if name is not None: + filtered.append((name, start, end, corr, 0, graph_node_id, stream_id)) + record_idx += 1 + if record_idx >= records_per_replay: + replay_idx += 1 + record_idx = 0 + return filtered + + +def _cupti_parser_worker(input_queue, output_queue, ready_event) -> None: + libcupti = _load_libcupti() + _configure_cupti_get_next_record(libcupti) + shared_blocks: dict[str, shared_memory.SharedMemory] = {} + records_by_generation: dict[int, list[tuple[int, int, int, int, int]]] = {} + zero_ts_by_generation: dict[int, int] = {} + ready_event.set() + while True: + item = input_queue.get() + if item is None: + break + kind = item[0] + if kind == "buffer": + _, generation, buffer_id, name, valid_size = item + shm = shared_blocks.get(name) + if shm is None: + shm = shared_memory.SharedMemory(name=name) + shared_blocks[name] = shm + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + parser_ptr = ctypes.addressof(shared_char) + records, zero_ts_count, _ = _parse_cupti_buffer_ptr( + libcupti, + parser_ptr, + valid_size, + include_names=False, + ) + records_by_generation.setdefault(generation, []).extend(records) + zero_ts_by_generation[generation] = zero_ts_by_generation.get(generation, 0) + zero_ts_count + ctypes.memset(parser_ptr, 0, len(shm.buf)) + except Exception as exc: # pragma: no cover - diagnostic worker path + output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) + finally: + del shared_char + output_queue.put({"kind": "buffer_done", "generation": generation, "buffer_id": buffer_id}) + elif kind == "finish": + if len(item) == 4: + _, generation, filter_plan, stats_request = item + else: + _, generation, filter_plan = item + stats_request = None + try: + raw_records = records_by_generation.pop(generation, []) + zero_ts_count = zero_ts_by_generation.pop(generation, 0) + filtered_records = _apply_cupti_filter_plan(raw_records, filter_plan) + stats = None + parser_stats_ms = 0.0 + stats_ready = stats_request is not None + if stats_request is not None: + stats_start_s = time.perf_counter() + stats = _stats_from_cupti_records( + filtered_records, + int(stats_request["warmup"]), + int(stats_request["iters"]), + str(stats_request["tag"]), + int(stats_request["expected_K"]), + zero_ts_count=zero_ts_count, + zero_ts_names={}, + include_details=bool(stats_request.get("include_details", True)), + ) + parser_stats_ms = 1000.0 * (time.perf_counter() - stats_start_s) + filtered_records = [] + output_queue.put({ + "kind": "finish_done", + "generation": generation, + "records": filtered_records, + "zero_ts_count": zero_ts_count, + "zero_ts_names": {}, + "raw_record_count": len(raw_records), + "stats": stats, + "stats_ready": stats_ready, + "parser_stats_ms": parser_stats_ms, + }) + except Exception as exc: # pragma: no cover - diagnostic worker path + output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) + else: + output_queue.put({"kind": "error", "generation": -1, "error": f"unknown parser message {kind!r}"}) + for shm in shared_blocks.values(): + shm.close() + + +class CuptiKernelTimer: + """Raw CUPTI Activity timer with out-of-process parsing for timed runs. + + CUPTI's callback gives us raw activity buffers. The callback only hands + shared-memory buffer metadata to a parser process, so the main process + avoids the cupti-python per-record object creation cost during the timed + path. A single local calibration replay may parse names in-process to + build an ordinal filter plan for a just-captured CUDA graph. + """ + + _instance = None + _import_error = None + + _request_callback_type = ctypes.CFUNCTYPE( + None, + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t), + ) + _complete_callback_type = ctypes.CFUNCTYPE( + None, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_size_t, + ) + + @classmethod + def get(cls) -> "CuptiKernelTimer": + if cls._instance is not None: + return cls._instance + if cls._import_error is not None: + raise cls._import_error + try: + cls._instance = cls() + return cls._instance + except ImportError as exc: # pragma: no cover - env-dependent + cls._import_error = exc + raise + + def __init__(self) -> None: + self._libcupti = _load_libcupti() + self._configure_functions() + self._lock = threading.Lock() + self._shared_buffers: dict[int, shared_memory.SharedMemory] = {} + self._buffer_id_by_ptr: dict[int, int] = {} + self._free_buffer_ids: list[int] = [] + self._local_completed: list[tuple[int, int]] = [] + self._mode = "drop" + self._generation = 0 + self._finish_results: dict[int, dict] = {} + self._parser_errors: list[str] = [] + self._filter_plan = () + self._last_start_timing: dict[str, float] = {} + self._last_stop_timing: dict[str, float] = {} + self._current_flush_period_ms = 0 + self._mp_ctx = mp.get_context(_MP_START_METHOD) + # Retry parser-process spawn: concurrent bench instances on the same + # node race on POSIX named semaphores in /dev/shm — child can die in + # pickle.load with FileNotFoundError in SemLock._rebuild before + # signalling ready_event. Detect early-dead child via is_alive() so + # we don't waste the full timeout, and retry up to 3x with jitter. + last_err = None + for _spawn_attempt in range(3): + self._parse_input_queue = self._mp_ctx.Queue() + self._parse_output_queue = self._mp_ctx.Queue() + ready_event = self._mp_ctx.Event() + self._parse_process = self._mp_ctx.Process( + target=_cupti_parser_worker, + args=(self._parse_input_queue, self._parse_output_queue, ready_event), + ) + self._parse_process.start() + deadline = time.time() + 30.0 + spawn_ok = False + while time.time() < deadline: + if ready_event.wait(timeout=0.5): + spawn_ok = True + break + if not self._parse_process.is_alive(): + break + if spawn_ok: + last_err = None + break + last_err = (f"attempt {_spawn_attempt + 1}: " + f"alive={self._parse_process.is_alive()}, " + f"exitcode={self._parse_process.exitcode}") + try: + if self._parse_process.is_alive(): + self._parse_process.terminate() + self._parse_process.join(timeout=2.0) + except Exception: + pass + time.sleep(0.5 + 0.5 * _spawn_attempt) + if last_err is not None: + raise RuntimeError( + f"CUPTI parser process did not initialize after 3 attempts: {last_err}" + ) + + self._set_zeroed_host_buffer_attr() + for _ in range(_CUPTI_HOST_BUFFER_COUNT): + self._free_buffer_ids.append(self._allocate_shared_buffer()) + + self._request_callback = self._request_callback_type(self._request_buffer) + self._complete_callback = self._complete_callback_type(self._complete_buffer) + self._check(self._libcupti.cuptiActivityRegisterCallbacks( + self._request_callback, + self._complete_callback, + )) + self._check(self._libcupti.cuptiActivityEnable(_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL)) + atexit.register(self.close) + + def _configure_functions(self) -> None: + self._libcupti.cuptiActivityRegisterCallbacks.argtypes = [ + self._request_callback_type, + self._complete_callback_type, + ] + self._libcupti.cuptiActivityRegisterCallbacks.restype = ctypes.c_int + self._libcupti.cuptiActivityEnable.argtypes = [ctypes.c_int] + self._libcupti.cuptiActivityEnable.restype = ctypes.c_int + self._libcupti.cuptiActivityFlushAll.argtypes = [ctypes.c_uint32] + self._libcupti.cuptiActivityFlushAll.restype = ctypes.c_int + self._libcupti.cuptiActivityFlushPeriod.argtypes = [ctypes.c_uint32] + self._libcupti.cuptiActivityFlushPeriod.restype = ctypes.c_int + self._libcupti.cuptiActivitySetAttribute.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_void_p, + ] + self._libcupti.cuptiActivitySetAttribute.restype = ctypes.c_int + _configure_cupti_get_next_record(self._libcupti) + + def _set_zeroed_host_buffer_attr(self) -> None: + value_obj = ctypes.c_uint8(1) + size_obj = ctypes.c_size_t(ctypes.sizeof(value_obj)) + result = self._libcupti.cuptiActivitySetAttribute( + _CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER, + ctypes.byref(size_obj), + ctypes.byref(value_obj), + ) + if result != _CUPTI_SUCCESS: + print( + "[WARN] CUPTI zeroed host-buffer attribute failed; " + f"continuing with default CUPTI buffer handling (CUptiResult={result}).", + file=sys.stderr, + ) + + def _check(self, result: int) -> None: + if result != _CUPTI_SUCCESS: + raise RuntimeError(f"CUPTI call failed with CUptiResult={result}") + + def _allocate_shared_buffer(self) -> int: + buffer_id = len(self._shared_buffers) + shm = shared_memory.SharedMemory(create=True, size=_CUPTI_HOST_BUFFER_BYTES) + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + ptr = ctypes.addressof(shared_char) + finally: + del shared_char + if ptr % 8 != 0: + shm.close() + shm.unlink() + raise RuntimeError("CUPTI shared-memory activity buffer was not 8-byte aligned") + self._shared_buffers[buffer_id] = shm + self._buffer_id_by_ptr[ptr] = buffer_id + return buffer_id + + def _buffer_ptr(self, buffer_id: int) -> int: + shm = self._shared_buffers[buffer_id] + shared_char = ctypes.c_char.from_buffer(shm.buf) + try: + return ctypes.addressof(shared_char) + finally: + del shared_char + + def _request_buffer(self, buffer, size, max_num_records) -> None: + with self._lock: + if self._free_buffer_ids: + buffer_id = self._free_buffer_ids.pop() + else: + buffer_id = self._allocate_shared_buffer() + ptr = self._buffer_ptr(buffer_id) + buffer[0] = ptr + size[0] = _CUPTI_HOST_BUFFER_BYTES + max_num_records[0] = 0 + + def _complete_buffer(self, context, stream_id, buffer, size, valid_size) -> None: + del context, stream_id, size + buffer_ptr = int(buffer) + valid_size_int = int(valid_size) + with self._lock: + mode = self._mode + generation = self._generation + buffer_id = self._buffer_id_by_ptr[buffer_ptr] + if valid_size_int == 0 or mode == "drop": + self._free_buffer_ids.append(buffer_id) + return + if mode == "local": + self._local_completed.append((buffer_id, valid_size_int)) + return + shm = self._shared_buffers[buffer_id] + self._parse_input_queue.put(("buffer", generation, buffer_id, shm.name, valid_size_int)) + + def _handle_parser_result(self, result: dict) -> None: + kind = result.get("kind") + if kind == "buffer_done": + with self._lock: + self._free_buffer_ids.append(int(result["buffer_id"])) + elif kind == "finish_done": + self._finish_results[int(result["generation"])] = result + elif kind == "error": + self._parser_errors.append(str(result.get("error"))) + + def _drain_parser_results(self) -> None: + while True: + try: + result = self._parse_output_queue.get_nowait() + except queue.Empty: + break + self._handle_parser_result(result) + + def is_generation_ready(self, generation: int) -> bool: + self._drain_parser_results() + return generation in self._finish_results or bool(self._parser_errors) + + def _flush(self, flag: int) -> None: + self._check(self._libcupti.cuptiActivityFlushAll(flag)) + + def _set_flush_period_ms(self, period_ms: int) -> None: + if period_ms == self._current_flush_period_ms: + return + self._check(self._libcupti.cuptiActivityFlushPeriod(period_ms)) + self._current_flush_period_ms = period_ms + + def _begin( + self, + mode: str, + filter_plan=(), + flush_period_ms: int = 0, + collect_timing: bool = False, + ) -> int: + start_timing: dict[str, float] = {} + with self._lock: + self._mode = "drop" + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._flush(1) + if collect_timing: + start_timing["forced_flush_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._drain_parser_results() + if collect_timing: + start_timing["drain_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) + with self._lock: + self._generation += 1 + generation = self._generation + self._mode = mode + self._local_completed = [] + self._filter_plan = filter_plan + if flush_period_ms > 0: + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._set_flush_period_ms(flush_period_ms) + if collect_timing: + start_timing["period_enable_ms"] = 1000.0 * ( + time.perf_counter() - phase_start_s + ) + self._last_start_timing = start_timing + return generation + + def capture_names(self, replay_fn) -> tuple[list[tuple], int, dict]: + """Run a small calibration replay and parse kernel names locally.""" + self._begin("local") + replay_fn() + torch.cuda.synchronize() + self._flush(0) + records: list[tuple] = [] + zero_ts_count = 0 + zero_ts_names: dict[str, int] = {} + with self._lock: + completed = list(self._local_completed) + self._local_completed = [] + self._mode = "drop" + for buffer_id, valid_size in completed: + ptr = self._buffer_ptr(buffer_id) + recs, zeros, zero_names = _parse_cupti_buffer_ptr( + self._libcupti, + ptr, + valid_size, + include_names=True, + ) + records.extend(recs) + zero_ts_count += zeros + for name, count in zero_names.items(): + zero_ts_names[name] = zero_ts_names.get(name, 0) + count + ctypes.memset(ptr, 0, _CUPTI_HOST_BUFFER_BYTES) + with self._lock: + self._free_buffer_ids.append(buffer_id) + records.sort(key=lambda r: r[1]) + return records, zero_ts_count, zero_ts_names + + def start( + self, + filter_plan=(), + flush_period_ms: int = 0, + collect_timing: bool = False, + ) -> None: + self._begin("parser", filter_plan, flush_period_ms, collect_timing) + + def stop_async( + self, + collect_timing: bool = False, + stats_request: dict | None = None, + ) -> tuple[int, dict[str, float]]: + stop_timing: dict[str, float] = {} + generation = self._generation + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._set_flush_period_ms(0) + if collect_timing: + stop_timing["period_disable_ms"] = 1000.0 * ( + time.perf_counter() - phase_start_s + ) + phase_start_s = time.perf_counter() if collect_timing else 0.0 + self._flush(0) + if collect_timing: + stop_timing["flush_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) + with self._lock: + self._mode = "drop" + filter_plan = self._filter_plan + self._parse_input_queue.put(("finish", generation, filter_plan, stats_request)) + self._last_stop_timing = stop_timing + return generation, stop_timing + + def wait_for_generation_result( + self, + generation: int, + stop_timing: dict[str, float] | None = None, + collect_timing: bool = False, + ) -> dict: + if stop_timing is None: + stop_timing = {} + phase_start_s = time.perf_counter() if collect_timing else 0.0 + deadline = time.perf_counter() + 10.0 + while time.perf_counter() < deadline: + result = self._finish_results.pop(generation, None) + if result is not None: + if collect_timing: + stop_timing["parser_wait_ms"] = 1000.0 * ( + time.perf_counter() - phase_start_s + ) + stop_timing["total_ms"] = ( + stop_timing.get("period_disable_ms", 0.0) + + stop_timing.get("flush_ms", 0.0) + + stop_timing["parser_wait_ms"] + ) + self._last_stop_timing = stop_timing + return result + timeout_s = max(0.0, min(0.01, deadline - time.perf_counter())) + try: + parser_result = self._parse_output_queue.get(timeout=timeout_s) + except queue.Empty: + continue + self._handle_parser_result(parser_result) + if self._parser_errors: + raise RuntimeError("CUPTI parser process failed: " + "; ".join(self._parser_errors)) + raise TimeoutError("Timed out waiting for CUPTI parser process") + + def wait_for_generation( + self, + generation: int, + stop_timing: dict[str, float] | None = None, + collect_timing: bool = False, + ) -> tuple[list[tuple], int, dict, int]: + result = self.wait_for_generation_result(generation, stop_timing, collect_timing) + return ( + list(result["records"]), + int(result["zero_ts_count"]), + dict(result["zero_ts_names"]), + int(result["raw_record_count"]), + ) + + def stop(self, collect_timing: bool = False) -> tuple[list[tuple], int, dict, int]: + generation, stop_timing = self.stop_async(collect_timing) + return self.wait_for_generation(generation, stop_timing, collect_timing) + + def last_start_timing(self) -> dict[str, float]: + return dict(self._last_start_timing) + + def last_stop_timing(self) -> dict[str, float]: + return dict(self._last_stop_timing) + + def close(self) -> None: + parse_process = getattr(self, "_parse_process", None) + if parse_process is not None and parse_process.is_alive(): + self._parse_input_queue.put(None) + parse_process.join(timeout=5.0) + if parse_process.is_alive(): + parse_process.terminate() + parse_process.join(timeout=1.0) + for shm in getattr(self, "_shared_buffers", {}).values(): + try: + shm.close() + shm.unlink() + except FileNotFoundError: + pass + + +# ============================================================================= +# Timing helpers +# ============================================================================= + + +def _stats_from_spans(spans_us: list[float]) -> dict: + """Compute median / p95 / p99 / n from a per-iter span list.""" + s = sorted(spans_us) + return { + "median": statistics.median(s), + "p95": s[int(0.95 * len(s))], + "p99": s[int(0.99 * len(s))], + "n": len(s), + } + + +def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, + zero_ts_count: int = 0, + zero_ts_names: dict | None = None, + include_details: bool = True): + """Bin a flat CUPTI kernel record stream into per-iter spans + per-kernel + relative timestamps. Used by both graph and eager CUPTI paths. + + `records` are tuples (name, start_ns, end_ns, ...) — see CuptiKernelTimer. + `expected_K` is the kernels-per-iter count the caller declares; we + validate the CUPTI total matches `expected_K * (warmup + iters)` exactly. + On mismatch we dump per-name record counts so missing or extra kernels + are obvious (most common cause: a new dispatch mode whose kernels lack + a matching entry in `_CUPTI_KEEP_KERNEL_SUBSTRINGS`, silently filtering + them out). + """ + records = [ + r for r in records + if r[0] is not None and any(s in r[0] for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS) + ] + records.sort(key=lambda r: r[1]) # by start_ns + + total = len(records) + expected_iters = warmup + iters + expected_total = expected_K * expected_iters + if total != expected_total: + from collections import Counter + name_counts = dict(Counter(r[0] for r in records)) + # Non-fatal: skip this cell instead of killing the whole sweep. + # Mismatch may be a CUPTI dropped-records issue (rare configs), + # not necessarily a K-table bug. Log so the user can investigate + # the specific cell post-hoc; return None so the caller can skip + # writing a JSON row. + zero_msg = "" + if zero_ts_count: + zero_msg = ( + f" + {zero_ts_count} records with start/end=0 " + f"(dropped by callback, breakdown {zero_ts_names}). " + f"Total observed kernel records (timed + zero-ts) = " + f"{total + zero_ts_count} / {expected_total}." + ) + print( + f"[WARN] CUPTI capture mismatch for {tag!r}: expected " + f"{expected_K} kernels/iter × {expected_iters} iters " + f"(warmup+iters) = {expected_total} records, got {total}. " + f"Kernel record counts: {name_counts}.{zero_msg} SKIPPING cell.", + file=sys.stderr, + flush=True, + ) + # Per-record dump: (name, start_ns_rel, end_ns_rel, corr_id, graph_id, stream_id). + # Times relative to first record so absolute ns isn't drowning output. + # Limit dump to first 30 records to avoid flooding logs at high K. + if records: + t0_ns = records[0][1] + for i, r in enumerate(records[:30]): + # r = (name, start_ns, end_ns, corr_id, graph_id, graph_node_id, stream_id) + rel_start = (r[1] - t0_ns) / 1000.0 # us + rel_end = (r[2] - t0_ns) / 1000.0 + print( + f" rec[{i:3d}] name={r[0]!r} start={rel_start:.2f}us " + f"end={rel_end:.2f}us corr={r[3]} graph={r[4]} stream={r[6]}", + file=sys.stderr, + flush=True, + ) + if len(records) > 30: + print(f" ... ({len(records) - 30} more records elided)", + file=sys.stderr, flush=True) + return None + K = expected_K + timed = records[warmup * K:] + + spans_us: list[float] = [] + per_kernel: dict[str, dict[str, list[float]]] = {} + for i in range(iters): + chunk = timed[i * K:(i + 1) * K] + iter_start_ns = min(r[1] for r in chunk) + iter_end_ns = max(r[2] for r in chunk) + spans_us.append((iter_end_ns - iter_start_ns) / 1000.0) + if include_details: + for r in chunk: + name = r[0] + slot = per_kernel.setdefault(name, {"start_us": [], "end_us": []}) + slot["start_us"].append((r[1] - iter_start_ns) / 1000.0) + slot["end_us"].append((r[2] - iter_start_ns) / 1000.0) + + out = _stats_from_spans(spans_us) + out["iters_us"] = spans_us + if include_details: + out["per_kernel"] = per_kernel + return out + + +_PRE_GRAPH_WARMUP_ITERS = 1 +_CUPTI_FILTER_PLAN_CACHE: dict[tuple, tuple[int, tuple[str | None, ...]]] = {} + + +class _HostTiming: + def __init__(self, enabled: bool) -> None: + self.enabled = enabled + self.values: dict[str, float | int | bool] = {} + self._total_start_s = time.perf_counter() if enabled else 0.0 + self._phase_start_s = 0.0 + + def start(self) -> None: + if self.enabled: + self._phase_start_s = time.perf_counter() + + def stop(self, key: str) -> None: + if self.enabled: + self.values[key] = 1000.0 * (time.perf_counter() - self._phase_start_s) + + def add(self, key: str, value: float | int | bool) -> None: + if self.enabled: + self.values[key] = value + + def stop_total(self) -> None: + if self.enabled: + self.values["total_ms"] = 1000.0 * (time.perf_counter() - self._total_start_s) + + def attach(self, stats: dict | None) -> None: + if self.enabled and stats is not None: + stats["host_timing"] = self.values + + +class _PendingCuptiStats: + + def __init__( + self, + timer: CuptiKernelTimer, + generation: int, + stop_timing: dict[str, float], + host_timing: _HostTiming, + *, + warmup: int, + iters: int, + tag: str, + expected_K: int, + expected_raw_record_count: int, + ) -> None: + self._timer = timer + self._generation = generation + self._stop_timing = stop_timing + self._host_timing = host_timing + self._warmup = warmup + self._iters = iters + self._tag = tag + self._expected_K = expected_K + self._expected_raw_record_count = expected_raw_record_count + + def is_ready(self) -> bool: + return self._timer.is_generation_ready(self._generation) + + def resolve(self) -> dict | None: + result = self._timer.wait_for_generation_result( + self._generation, + self._stop_timing, + collect_timing=self._host_timing.enabled, + ) + for key, value in self._timer.last_stop_timing().items(): + self._host_timing.add(f"cupti_stop_{key}", value) + raw_record_count = int(result["raw_record_count"]) + if raw_record_count != self._expected_raw_record_count: + print( + f"[WARN] CUPTI raw-record mismatch for {self._tag!r}: expected " + f"{self._expected_raw_record_count}, got {raw_record_count}. SKIPPING cell.", + file=sys.stderr, + ) + return None + + if result.get("stats_ready"): + stats = result.get("stats") + self._host_timing.add("stats_ms", 0.0) + self._host_timing.add("parser_stats_ms", float(result.get("parser_stats_ms", 0.0))) + else: + self._host_timing.start() + stats = _stats_from_cupti_records( + list(result["records"]), + self._warmup, + self._iters, + self._tag, + self._expected_K, + zero_ts_count=int(result["zero_ts_count"]), + zero_ts_names=dict(result["zero_ts_names"]), + ) + self._host_timing.stop("stats_ms") + self._host_timing.attach(stats) + return stats + + +def _target_name_or_none(name: str | None) -> str | None: + if name is None: + return None + if any(s in name for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS): + return name + return None + + +def _capture_group_graph( + args, + run_fn, + reset_fn, + group_iters: int, + graph_pre_iter_fn=None, +) -> torch.cuda.CUDAGraph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for j in range(group_iters): + if graph_pre_iter_fn is not None: + graph_pre_iter_fn(j) + reset_fn() + if args.l2_flush: + _l2_flush.fill_(0.0) + run_fn() + return graph + + +def _graph_group_iters(args, total_iters: int, pre_iter_fn, pre_iter_group_factory) -> int: + """Pick the graph-group size unconditionally; the caller is expected to + round total_iters up to a multiple of this so all iters fit in clean + replays. Sample arrays are pre-padded at allocation (see _sample_pnat + call site) so the per-replay window can index past the user-requested + iter count by up to group_iters-1 extra samples. + """ + if pre_iter_fn is not None and pre_iter_group_factory is None: + # Per-iter callback without a group-factory: can't batch. + return 1 + requested = getattr(args, "cuda_graph_group_iters", None) + if requested is None: + return ( + _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX + if pre_iter_group_factory is not None + else _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE + ) + return max(1, int(requested)) + + +def _get_cupti_filter_plan(timer: CuptiKernelTimer, graph, cache_key: tuple | None, + group_iters: int) -> tuple[int, tuple[str | None, ...]]: + full_cache_key = None if cache_key is None else (cache_key, group_iters) + if full_cache_key is not None: + cached = _CUPTI_FILTER_PLAN_CACHE.get(full_cache_key) + if cached is not None: + return cached + + records, zero_ts_count, zero_ts_names = timer.capture_names(graph.replay) + if zero_ts_count: + print( + f"[WARN] CUPTI calibration saw {zero_ts_count} zero-timestamp records " + f"(breakdown {zero_ts_names}); continuing with nonzero records.", + file=sys.stderr, + ) + ordinal_names = tuple(_target_name_or_none(r[0]) for r in records) + target_count = sum(name is not None for name in ordinal_names) + if target_count == 0: + raise RuntimeError("CUPTI calibration did not find any target kernel records") + plan = (len(records), ordinal_names) + if full_cache_key is not None: + _CUPTI_FILTER_PLAN_CACHE[full_cache_key] = plan + return plan + + +def _time_kernel_cuda_graph( + args, + run_fn, + reset_fn, + tag: str, + *, + expected_K: int, + pre_iter_fn=None, + pre_iter_group_factory=None, + iters_override: int | None = None, + cupti_plan_key: tuple | None = None, +) -> dict: + """CUDA-graph CUPTI timer (graph-per-iter design). + + Captures one CUDA graph holding a small group of logical iterations + (per-iter setup + reset + l2_flush + run_fn) and replays it enough + times to cover `warmup + iters`. + + Why graph-per-iter (vs the older "one giant graph holding all iters" + design): instantiating a CUDA graph is expensive — proportional to + graph size — so a single small graph instantiated once is much + cheaper than one big graph instantiated for each cell of a sweep. + Replays are cheap regardless. + + Mix cells use a per-replay device window: an outside-graph copy loads + the next group of PNAT/n_writes samples, then graph-captured per-iter + copies update kernel inputs before each reset + L2 flush + run. + + Pre-graph eager warmup: forces PyTorch's caching allocator + + Triton's autotune cache to settle before capture so the graph + doesn't bake in init-only allocations. + + ``iters_override`` (if not None) overrides ``args.iters`` for this + call. Used to give mix scenarios a higher iter count than pure + (more iters = more independent mix draws averaged in). + """ + host_timing = _HostTiming(bool(getattr(args, "host_timing", False))) + timer = CuptiKernelTimer.get() + warmup = args.warmup + iters = iters_override if iters_override is not None else args.iters + + # Pre-graph eager warmup: full per-iter chain once. This settles + # Triton/PyTorch setup and wrapper-side intermediate allocations; + # skipping it risks lazy work leaking into graph capture. + warmup_iters = _PRE_GRAPH_WARMUP_ITERS + host_timing.add("pre_graph_warmup_iters", warmup_iters) + host_timing.start() + for _ in range(warmup_iters): + reset_fn() + if pre_iter_fn is not None: + pre_iter_fn(0) + run_fn() + if warmup_iters > 0: + torch.cuda.synchronize() + host_timing.stop("pre_graph_warmup_ms") + + total_iters = warmup + iters + group_iters = _graph_group_iters(args, total_iters, pre_iter_fn, pre_iter_group_factory) + # Args are rounded at argparse-time so warmup+iters/mix_iters are already + # multiples of the relevant group_iters. Assert here to catch any caller + # bypassing argparse. + assert total_iters % group_iters == 0, ( + f"total_iters={total_iters} not a multiple of group_iters={group_iters}; " + f"args.warmup/iters/mix_iters should be rounded post-argparse." + ) + pre_replay_fn = None + graph_pre_iter_fn = None + if pre_iter_group_factory is not None and group_iters > 1: + pre_replay_fn, graph_pre_iter_fn = pre_iter_group_factory(group_iters) + + # Reset just before capture so warmup state changes don't bleed in. + host_timing.start() + reset_fn() + torch.cuda.synchronize() + host_timing.stop("pre_capture_reset_ms") + + # Capture a small group of identical logical iterations. Mix/pre_iter + # cells can group when they provide a graph-side pre-iter updater backed + # by a per-replay device window. + host_timing.start() + g = _capture_group_graph(args, run_fn, reset_fn, group_iters, graph_pre_iter_fn) + host_timing.stop("graph_capture_ms") + + if pre_replay_fn is not None: + host_timing.start() + pre_replay_fn(0) + torch.cuda.synchronize() + host_timing.stop("graph_preload_ms") + + plan_cache_key = None if cupti_plan_key is None else (cupti_plan_key, group_iters) + host_timing.add("cupti_plan_cached", ( + plan_cache_key is not None and plan_cache_key in _CUPTI_FILTER_PLAN_CACHE + )) + host_timing.start() + records_per_replay, ordinal_names = _get_cupti_filter_plan( + timer, + g, + cupti_plan_key, + group_iters, + ) + host_timing.stop("cupti_plan_ms") + target_count = sum(name is not None for name in ordinal_names) + expected_targets_per_replay = expected_K * group_iters + if target_count != expected_targets_per_replay: + print( + f"[WARN] CUPTI calibration mismatch for {tag!r}: expected " + f"{expected_targets_per_replay} target records in a {group_iters}-iter graph replay, " + f"got {target_count} target records out of {records_per_replay} total records.", + file=sys.stderr, + ) + + # Time: replay the grouped graph enough times to cover warmup+iters. + # Mix cells preload one device window per replay on the same stream. + # CUPTI records every kernel launch; _stats_from_cupti_records + # validates against expected_K and slices warmup off the front. + graph_replays = total_iters // group_iters + filter_plan = ((records_per_replay, ordinal_names),) * graph_replays + cupti_flush_period_ms = max(0, int(getattr(args, "cupti_flush_period_ms", 0))) + host_timing.start() + timer.start( + filter_plan, + flush_period_ms=cupti_flush_period_ms, + collect_timing=host_timing.enabled, + ) + host_timing.stop("cupti_start_ms") + for key, value in timer.last_start_timing().items(): + host_timing.add(f"cupti_start_{key}", value) + torch.cuda.nvtx.range_push(tag) + host_timing.start() + for i in range(graph_replays): + if pre_replay_fn is not None: + pre_replay_fn(i) + elif pre_iter_fn is not None: + pre_iter_fn(i) + g.replay() + host_timing.stop("graph_enqueue_ms") + host_timing.start() + torch.cuda.synchronize() + host_timing.stop("graph_sync_ms") + torch.cuda.nvtx.range_pop() + expected_raw_record_count = records_per_replay * graph_replays + host_timing.start() + if int(getattr(args, "cupti_defer_depth", 1)) > 1: + generation, stop_timing = timer.stop_async( + collect_timing=host_timing.enabled, + stats_request={ + "warmup": warmup, + "iters": iters, + "tag": tag, + "expected_K": expected_K, + "include_details": bool(getattr(args, "json_detailed", False)), + }, + ) + host_timing.stop("cupti_stop_ms") + for key, value in timer.last_stop_timing().items(): + host_timing.add(f"cupti_stop_{key}", value) + host_timing.stop_total() + host_timing.add("graph_group_iters", group_iters) + host_timing.add("graph_replays", graph_replays) + host_timing.add("cupti_records_per_replay", records_per_replay) + host_timing.add("cupti_target_records_per_replay", target_count) + host_timing.add("cupti_raw_records_expected", expected_raw_record_count) + host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) + return _PendingCuptiStats( + timer, + generation, + stop_timing, + host_timing, + warmup=warmup, + iters=iters, + tag=tag, + expected_K=expected_K, + expected_raw_record_count=expected_raw_record_count, + ) + + records, zero_ts_count, zero_ts_names, raw_record_count = timer.stop( + collect_timing=host_timing.enabled, + ) + host_timing.stop("cupti_stop_ms") + for key, value in timer.last_stop_timing().items(): + host_timing.add(f"cupti_stop_{key}", value) + if raw_record_count != expected_raw_record_count: + print( + f"[WARN] CUPTI raw-record mismatch for {tag!r}: expected " + f"{records_per_replay} total records/replay × {graph_replays} replays " + f"= {expected_raw_record_count}, got {raw_record_count}. SKIPPING cell.", + file=sys.stderr, + ) + return None + + host_timing.start() + stats = _stats_from_cupti_records( + records, + warmup, + iters, + tag, + expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names, + include_details=bool(getattr(args, "json_detailed", False)), + ) + host_timing.stop("stats_ms") + host_timing.stop_total() + host_timing.add("graph_group_iters", group_iters) + host_timing.add("graph_replays", graph_replays) + host_timing.add("cupti_records_per_replay", records_per_replay) + host_timing.add("cupti_target_records_per_replay", target_count) + host_timing.add("cupti_raw_records", raw_record_count) + host_timing.add("cupti_raw_records_expected", expected_raw_record_count) + host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) + host_timing.attach(stats) + return stats + + +def _time_kernel_eager( + args, + run_fn, + reset_fn, + tag: str, + *, + expected_K: int, + pre_iter_fn=None, + iters_override: int | None = None, + cupti_plan_key: tuple | None = None, +) -> dict: + """Non-graph CUPTI timer (for ncu wrapping, debugging, etc.). + + Each iter runs serially with sync between, but kernel start/end still + come from CUPTI — same accuracy as the graph path, just slower per-iter + (extra Python + sync overhead). + """ + host_timing = _HostTiming(bool(getattr(args, "host_timing", False))) + timer = CuptiKernelTimer.get() + warmup = args.warmup + iters = iters_override if iters_override is not None else args.iters + + del cupti_plan_key + + def _run_eager_loop(): + torch.cuda.nvtx.range_push(tag) + # Unified warmup+iters loop; CUPTI filters by warmup count internally. + for i in range(warmup + iters): + reset_fn() + if args.l2_flush: + _flush_l2() # includes synchronize + if pre_iter_fn is not None: + pre_iter_fn(i) + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + + host_timing.start() + records, zero_ts_count, zero_ts_names = timer.capture_names(_run_eager_loop) + host_timing.stop("timed_loop_and_cupti_parse_ms") + + host_timing.start() + stats = _stats_from_cupti_records( + records, + warmup, + iters, + tag, + expected_K, + zero_ts_count=zero_ts_count, + zero_ts_names=zero_ts_names, + include_details=bool(getattr(args, "json_detailed", False)), + ) + host_timing.stop("stats_ms") + host_timing.stop_total() + host_timing.attach(stats) + return stats + + +def _run_kernel_untimed(args, run_fn, reset_fn, tag: str) -> dict: + """No in-bench timing: just run the kernels for an external profiler + (nsys / ncu) to time externally. Returns a stats dict full of zeros so + downstream code (table, JSON) doesn't break. + + Note: pre_iter_fn / iters_override aren't plumbed here yet — mix-mode + benchmarking relies on CUPTI. Add when a use-case lands. + """ + warmup = args.warmup + iters = args.iters + + if args.cuda_graph: + # Eager warmup before capture (Triton autotune) + reset_fn(); run_fn(); torch.cuda.synchronize() + reset_fn(); torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + for _ in range(warmup + iters): + reset_fn() + if args.l2_flush: + _l2_flush.fill_(0.0) + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_push(tag) + g.replay() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + else: + torch.cuda.nvtx.range_push(tag) + for _ in range(warmup + iters): + reset_fn() + if args.l2_flush: + _flush_l2() + run_fn() + torch.cuda.synchronize() + torch.cuda.nvtx.range_pop() + + spans_us = [0.0] * iters + out = _stats_from_spans(spans_us) + out["iters_us"] = spans_us + out["per_kernel"] = {} + return out + + +def _time_kernel( + args, run_fn, reset_fn, tag: str, + *, + expected_K: int, + pre_iter_fn=None, + pre_iter_group_factory=None, + iters_override: int | None = None, + cupti_plan_key: tuple | None = None, +) -> dict: + """Dispatch to graph-CUPTI / eager-CUPTI / no-timer path. + + --cupti: in-process CUPTI Activity API timing (default). Use --no-cupti + when running under nsys (in-process CUPTI conflicts with nsys's own + subscriber); the bench then runs the kernels for nsys to time externally. + + `expected_K` is the kernels-per-iter count the caller declares + (computed via _kernels_per_iter_*). CUPTI paths validate against it + explicitly; the no-timer fallback ignores it (no records to validate). + """ + if not getattr(args, "cupti", True): + if pre_iter_fn is not None: + raise RuntimeError( + "_time_kernel: pre_iter_fn requires CUPTI (mix-mode); " + "got --no-cupti. Re-run with CUPTI on or plumb pre_iter_fn " + "through _run_kernel_untimed." + ) + return _run_kernel_untimed(args, run_fn, reset_fn, tag) + if args.cuda_graph: + return _time_kernel_cuda_graph( + args, run_fn, reset_fn, tag, + expected_K=expected_K, + pre_iter_fn=pre_iter_fn, + pre_iter_group_factory=pre_iter_group_factory, + iters_override=iters_override, + cupti_plan_key=cupti_plan_key, + ) + return _time_kernel_eager( + args, run_fn, reset_fn, tag, + expected_K=expected_K, + pre_iter_fn=pre_iter_fn, + iters_override=iters_override, + cupti_plan_key=cupti_plan_key, + ) + + +# Per-config benchmark (consolidated baseline + replay) + + +def _warm_one_config(args, cfg, baseline_fn) -> None: + """Module-level worker for the compile-warmup process pool. + + Module-level so ProcessPoolExecutor can pickle it (nested functions + aren't picklable). Each worker process holds its own GIL → no + serialization between concurrent compiles. + + ``cfg`` is a tuple of (outer_cfg, inner_overrides_or_list): + * outer_cfg = (batch, mtp_len, prev_ks, state_dtype, act_dtype, + sr_mode, rect, write_ckpt, mode, + sort_slots, reverse_nowrite, hardcode_sort) + * inner_overrides_or_list = dict of args attribute name -> value-string, + OR a list of such dicts. In the list form (CPS-grouped task) the + worker compiles each entry sequentially within the same process so + Triton's in-process kernel cache catches value-spec hits across + related entries (e.g. CPS={1,2} and {4,8} each form a `div_by_16` + spec bucket; the second compile in a bucket short-circuits). + + ``baseline_fn`` is optional — when ``None``, only the checkpointing + kernel is warmed (the baseline-selection kernel can be warmed once in + the parent if needed). This lets us avoid pickling C-extension + function references across processes. + """ + outer_cfg, inner_overrides_or_list = cfg + overrides_list = (inner_overrides_or_list + if isinstance(inner_overrides_or_list, list) + else [inner_overrides_or_list]) + (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, + rect, write_ckpt, mode, sort_slots, reverse_nowrite, hardcode_sort) = outer_cfg + import argparse as _ap + for inner_overrides in overrides_list: + # Fresh clone per entry: prevents knob-value leakage between + # consecutive cells in a CPS-grouped task (entries may set + # different non-CPS knobs in degenerate edge cases). + args_copy = _ap.Namespace(**vars(args)) + for k, v in inner_overrides.items(): + setattr(args_copy, k, v) + _bench_config( + args_copy, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, + sr_mode=sr_mode, rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, mode=mode, + sort_slots=sort_slots, reverse_nowrite=reverse_nowrite, + hardcode_sort=hardcode_sort, + warmup_only=True, + ) + + +def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, + baseline_fn, max_workers: int) -> None: + _cw_t0 = time.perf_counter() + def _cw(label: str) -> None: + dt = time.perf_counter() - _cw_t0 + print(f"[compile-warmup] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) + _cw("entered _compile_warmup_phase") + """Parallel compile-warmup using a ProcessPoolExecutor with `spawn` + start method. + + Each worker process holds its own GIL and its own CUDA context, so + Triton compiles (Python AST/codegen + LLVM/ptxas) run truly in + parallel. Previous ThreadPoolExecutor design hit GIL contention + in the Python codegen phase, capping throughput at ~1-2 cores even + with 28 threads (observed: 4 R threads vs 28 in pool). + + Compiled binaries land in Triton's on-disk cache (TRITON_CACHE_DIR + or default ~/.triton/cache). Workers share the cache via filesystem + — first to write any given (kernel_source × constexpr_set) hash + wins; concurrent writes to the SAME hash are wasteful but not + corrupting. + + spawn start method avoids inheriting parent CUDA state (which is + unsafe after fork on Linux with active CUDA contexts). Per-worker + import + CUDA init costs ~10s, amortized over each worker's many + compiles. baseline_fn is intentionally NOT passed to workers to + avoid pickling complications; the parent compiles the baseline + kernel itself before launching the pool when applicable. + """ + from concurrent.futures import ProcessPoolExecutor + import multiprocessing + + sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) + + rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) + write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) + modes_list = getattr(args, "modes_list", ["persistent_dynamic"]) + sort_list = getattr(args, "sort_slots_list", [False]) + rev_list = getattr(args, "reverse_nowrite_list", [False]) + hsort_list = getattr(args, "hardcode_sort_list", [False]) + + # Compile-warmup task enumeration: outer × inner cartesian. + # CRITICAL: only enumerate axes that change the kernel's COMPILE signature. + # Drop runtime axes (batch, prev_k) that produce identical kernel hashes — + # otherwise we'd pay ~50-100ms of bench setup per redundant cache-hit task. + # + # Batches collapsed to first only: batch is a runtime int passed to the + # kernel, not a constexpr; all batches share the same compiled kernel. + # prev_k is already a list passed into _bench_config (not enumerated here). + configs = [] + _compile_batches = batch_sizes[:1] # collapse runtime axis + for batch in _compile_batches: + for mtp_len in mtp_lengths: + prev_ks = _resolve_prev_ks(args, mtp_len) + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + for sr_mode in sr_modes_list: + for mode in modes_list: + # write_checkpoint is per-slot from PNAT in both + # persistent modes; fix to True for the compile + # signature (the kernel always handles both). + for write_ckpt in [True]: + for rect in rect_list: + # persistent_main / persistent_dynamic + # both consume slot_perm and benefit from + # write-first clustering when there's a + # mix scenario. + can_sort = args.mix_csv is not None + effective_sort_list = ( + sort_list if can_sort else [False] + ) + effective_hsort_list = ( + hsort_list if can_sort else [False] + ) + for sort_slots in effective_sort_list: + effective_rev_list = ( + rev_list if sort_slots else [False] + ) + for reverse_nowrite in effective_rev_list: + for hardcode_sort in effective_hsort_list: + if sort_slots and hardcode_sort: + continue + configs.append(( + batch, mtp_len, prev_ks, + state_dtype, act_dtype, + sr_mode, rect, write_ckpt, mode, + sort_slots, reverse_nowrite, + hardcode_sort, + )) + + # Enumerate inner-knob signatures. Two paths: + # (1) --cell-list mode (preferred when set): pull exactly the cells + # that will be timed from args._cell_list_set. No synthetic + # cartesian — we only pre-compile what will run. + # (2) Sweep-args mode: cartesian over split-aware axes that read + # BOTH unsplit (args.X) and per-half (args.X_write/_nowrite) + # knob settings. Older code read only args.X and silently + # enumerated 1 inner combo when callers set only the per-half + # versions (all cell-list usage, plus any --block-size-m-write/ + # _nowrite CLI invocation), causing massive in-process JIT + # compile tax for persistent_main especially. + # + # In BOTH paths we GROUP tasks by non-CPS signature so each worker + # process compiles all CPS values for its group sequentially. + # NUM_PERSISTENT = CPS * num_sms is a runtime int but Triton auto- + # specializes on `div_by_16`, partitioning {CPS=1,2} (132,264) from + # {CPS=4,8} (528,1056) into two distinct compiled variants. By + # keeping all CPS variants for one (M,W,S,LS,TMA,...) signature in + # the same worker, the second compile in each spec bucket hits the + # in-process Triton cache (no disk-cache round trip). + + def _ps(val): + if val is None or (isinstance(val, str) and not val): + return [None] + if isinstance(val, str): + return [v.strip() for v in val.split(",") if v.strip()] + return [val] + + def _split_or_pair(shared_attr, w_attr, nw_attr): + """Return list of (write_val, nowrite_val) strings. + + Reads shared (args.X), write-side (args.X_write), and nowrite-side + (args.X_nowrite) values. If both per-half attrs are None, emits + tied pairs (v,v) over the shared values. If either per-half is + set, cartesian-iterates per-half values, falling back to shared + for whichever side is None. + """ + w = _ps(getattr(args, w_attr, None)) + nw = _ps(getattr(args, nw_attr, None)) + s = _ps(getattr(args, shared_attr, None)) + if w == [None] and nw == [None]: + return [(v, v) for v in s] + if w == [None]: + w = s + if nw == [None]: + nw = s + return [(a, b) for a in w for b in nw] + + _cw(f"built {len(configs)} outer configs") + cell_set = getattr(args, "_cell_list_set", set()) + cell_keys = getattr(args, "_cell_list_keys", ()) + # CPS keys are runtime ints (kernel value-specializes on `div_by_16`); + # cells differing only on CPS values can SHARE a worker so the second + # CPS value in a div_by_16 bucket hits the in-process Triton cache. + _cps_keys = ("cta_per_sm_write", "cta_per_sm_nowrite", "cta_per_sm") + + if cell_set: + # ============== CELL-LIST PATH ============== + # Build tasks DIRECTLY from cells. Each cell carries its OWN + # outer-axis values (RECT, MODE, SR, WC, SORT, REVN, HSORT) so we + # pair each cell with its specific outer config — NOT the union- + # cartesian of all cells' outer values. Previously the OUTER × + # CELL cartesian doubled task count when a cell-list spanned both + # RECT=0 and RECT=1 (or any other outer-axis split); half the + # tasks then failed the cell-list filter inside the worker and + # wasted dispatch overhead. This path is O(|unique cell groups|). + from collections import defaultdict as _dd + cell_groups: dict = _dd(list) + for tup in cell_set: + d = dict(zip(cell_keys, tup)) + cell_outer = ( + "SR" if d.get("SR", 0) else "RN", # sr_mode + bool(d.get("RECT", 0)), # rect + bool(d.get("WC", 1)), # write_ckpt + d.get("MODE", "persistent_dynamic"), # mode + bool(d.get("SORT", 0)), # sort_slots + bool(d.get("REVN", 0)), # reverse_nowrite + bool(d.get("HSORT", 0)), # hardcode_sort + ) + inner = {} + for k, v in d.items(): + if k in _CELL_LIST_KEY_TO_ARG: + inner[_CELL_LIST_KEY_TO_ARG[k]] = str(v) + non_cps_sig = tuple(sorted((k, v) for k, v in inner.items() if k not in _cps_keys)) + cell_groups[(cell_outer, non_cps_sig)].append(inner) + + # CLI-runtime axes (batch/mtp/dtype) are NOT in cell-list — they + # come from CLI args and cartesian here (typically just 1 combo). + cli_outers = [] + for _b in _compile_batches: + for _m in mtp_lengths: + _pk = _resolve_prev_ks(args, _m) + for _sd in state_dtypes: + for _ad in act_dtypes: + cli_outers.append((_b, _m, _pk, _sd, _ad)) + + tasks = [] + for cli_outer in cli_outers: + for (cell_outer, _sig), inner_list in cell_groups.items(): + outer_cfg = (*cli_outer, *cell_outer) + tasks.append((outer_cfg, inner_list)) + n_groups = len(cell_groups) + n_total_cells = sum(len(g) for g in cell_groups.values()) + n_outer_used = len(cli_outers) + else: + # ============== SWEEP-ARGS PATH ============== + # Build inner_dicts via cartesian over knob axes, then cross with + # the `configs` outer cartesian. Existing behavior. + m_pairs = _split_or_pair("block_size_m", "block_size_m_write", "block_size_m_nowrite") + w_pairs = _split_or_pair("num_warps", "num_warps_write", "num_warps_nowrite") + ns_pairs = _split_or_pair("num_stages", "num_stages_write", "num_stages_nowrite") + cps_pairs = _split_or_pair("cta_per_sm", "cta_per_sm_write", "cta_per_sm_nowrite") + ls_pairs = _split_or_pair("num_loop_stages", "num_loop_stages_write", "num_loop_stages_nowrite") + pw_vals = _ps(args.precompute_num_warps) + ps_vals = _ps(args.precompute_num_stages) + h_vals = _ps(args.heads_per_block) + mr_vals = _ps(args.maxnreg) + ct_vals = _ps(args.num_ctas) + fl_vals = _ps(args.flatten) + wsp_vals = _ps(args.warp_specialize) + trl_vals = _ps(args.use_tma_rect_load) + twl_vals = _ps(args.use_tma_replay_write_load) + tnl_vals = _ps(args.use_tma_replay_nowrite_load) + tws_vals = _ps(args.use_tma_replay_write_store) + import itertools as _it + inner_dicts = [] + for ((mw, mnw), (ww, wnw), (sw, snw), (cw, cnw), (lw, lnw), + pw, ps_, h, mr, ct, fl, wsp, + trl, twl, tnl, tws) in _it.product( + m_pairs, w_pairs, ns_pairs, cps_pairs, ls_pairs, + pw_vals, ps_vals, h_vals, mr_vals, ct_vals, + fl_vals, wsp_vals, + trl_vals, twl_vals, tnl_vals, tws_vals): + d = {} + for k, v in ( + ("block_size_m_write", mw), + ("block_size_m_nowrite", mnw), + ("num_warps_write", ww), + ("num_warps_nowrite", wnw), + ("num_stages_write", sw), + ("num_stages_nowrite", snw), + ("cta_per_sm_write", cw), + ("cta_per_sm_nowrite", cnw), + ("num_loop_stages_write", lw), + ("num_loop_stages_nowrite", lnw), + ("precompute_num_warps", pw), + ("precompute_num_stages", ps_), + ("heads_per_block", h), + ("maxnreg", mr), + ("num_ctas", ct), + ("flatten", fl), + ("warp_specialize", wsp), + ("use_tma_rect_load", trl), + ("use_tma_replay_write_load", twl), + ("use_tma_replay_nowrite_load", tnl), + ("use_tma_replay_write_store", tws), + ): + if v is not None: + d[k] = str(v) + inner_dicts.append(d) + + groups: dict = {} + for d in inner_dicts: + sig = tuple(sorted((k, v) for k, v in d.items() if k not in _cps_keys)) + groups.setdefault(sig, []).append(d) + tasks = [] + for outer in configs: + for sig, group in groups.items(): + tasks.append((outer, group)) + n_groups = len(groups) + n_total_cells = sum(len(g) for g in groups.values()) + n_outer_used = len(configs) + + # Shuffle ACROSS tasks (preserve within-group CPS sequence for in-process + # cache adjacency — within-group order is intentional, not shuffled). + import random as _r + _r.shuffle(tasks) + + _cw(f"built {len(tasks)} tasks covering {n_total_cells} cells in {n_groups} groups") + print(f"[compile-warmup] {len(tasks)} compile tasks " + f"({n_outer_used} outer × {n_groups} cell-groups " + f"covering {n_total_cells} cells, CPS-grouped" + + (", per-cell outer" if cell_set else "") + + f") across {max_workers} processes (ProcessPoolExecutor, {_MP_START_METHOD} start)") + t0 = time.perf_counter() + + ctx = multiprocessing.get_context(_MP_START_METHOD) + errors = [] + _cw("about to create ProcessPoolExecutor") + with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as ex: + _cw("ProcessPoolExecutor created, about to submit tasks") + # baseline_fn=None: workers compile only the checkpointing kernel. + # Baseline kernels (if any) get compiled lazily in the parent during + # the timing phase — usually just one extra compile, negligible. + futures = { + ex.submit(_warm_one_config, args, task, None): task + for task in tasks + } + _cw(f"submitted {len(futures)} tasks, waiting for results") + _n_done = 0 + for fut in futures: + try: + fut.result() + except Exception as e: + errors.append((futures[fut], e)) + _n_done += 1 + # Progress beacons at 10/25/50/75/100% to gauge effective parallelism. + if _n_done in (max(1, len(futures)//10), + max(1, len(futures)//4), + max(1, len(futures)//2), + max(1, (3*len(futures))//4), + len(futures)): + _cw(f"{_n_done}/{len(futures)} tasks complete") + + if errors: + for cfg, e in errors: + print(f"[compile-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", + file=sys.stderr) + raise errors[0][1] + + print(f"[compile-warmup] done in {time.perf_counter() - t0:.1f}s") + + +def _bench_config( + args, + batch: int, + mtp_len: int, + prev_ks: list[int], + state_dtype: torch.dtype, + act_dtype: torch.dtype, + baseline_fn, + sr_mode: str = "RN", + rectangle_for_nowrite: bool = False, + write_checkpoint: bool = True, + mode: str = "persistent_dynamic", + mix_samples_cpu=None, + mix_label: str = "", + sort_slots: bool = False, + reverse_nowrite: bool = False, + perm_samples_cpu=None, + hardcode_sort: bool = False, + mix_samples_sorted_cpu=None, + warmup_only: bool = False, +) -> None: + """ + Benchmark one (batch, mtp_len, dtype) configuration. + + Runs the baseline kernel (if baseline_fn is not None) followed by the + replay kernel for each prev_k value. Tensors are built once and + shared across all runs in this config. + + When ``warmup_only`` is True, calls each kernel exactly once instead of + timing it. Used by the parallel-warmup phase to populate Triton's + persistent compile cache across all configs concurrently. No timing + output is produced. + """ + state_dtype_name = str(state_dtype).split(".")[-1] + act_dtype_name = str(act_dtype).split(".")[-1] + + ( + state0, + state_scales0, + old_x0, + old_B0, + old_dt0, + old_dA_cumsum0, + cache_buf_idx0, + x, + dt, + B, + C, + A, + dt_bias, + D, + prev_tokens, + slot_perm_buf, + out_incr, + out_base, + intermediate_states_buffer, + xbc_input0, + conv_state0, + conv_weight, + conv_bias, + d_inner, + conv_dim, + ) = _build_tensors( + batch, + mtp_len, + state_dtype, + act_dtype, + args.tp_nheads, + args.head_dim, + args.d_state, + args.tp_ngroups, + max_window=getattr(args, "max_window", None) or None, + ) + + nheads = args.tp_nheads + ngroups = args.tp_ngroups + head_dim = args.head_dim + d_state = args.d_state + with_conv1d = getattr(args, "with_conv1d", False) + use_philox = (sr_mode == "SR") + variant_fn = _VARIANT_FNS[args.variant]() + + # SR rounding: allow fp16 and the quantized dtypes (int8/int16/fp8). + # bf16/fp32 SR is not supported (no PTX path for bf16; fp32 doesn't need + # rounding). When sweeping --sr-modes RN,SR over a mixed dtype set, + # silently skip the SR cell for unsupported dtypes — the RN cell still + # prints, and other dtypes still get their SR row. + rand_seed = None + _SR_SUPPORTED = ( + torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn, + ) + if use_philox: + if state_dtype not in _SR_SUPPORTED: + return + rand_seed = torch.randint(0, 2**62, (1,), device="cuda", dtype=torch.int64) + + is_quantized = state_dtype in (torch.int8, torch.int16, torch.float8_e4m3fn) + + state_work = state0.clone() + state_scales_work = state_scales0.clone() if state_scales0 is not None else None + old_x_work = old_x0.clone() + old_B_work = old_B0.clone() + old_dt_work = old_dt0.clone() + old_dA_cumsum_work = old_dA_cumsum0.clone() + cache_buf_idx_work = cache_buf_idx0.clone() + xbc_input_work = xbc_input0.clone() + conv_state_work = conv_state0.clone() + + def _reset(): + state_work.copy_(state0) + if state_scales_work is not None: + state_scales_work.copy_(state_scales0) + old_x_work.copy_(old_x0) + old_B_work.copy_(old_B0) + old_dt_work.copy_(old_dt0) + old_dA_cumsum_work.copy_(old_dA_cumsum0) + cache_buf_idx_work.copy_(cache_buf_idx0) + if with_conv1d: + conv_state_work.copy_(conv_state0) + + def _reset_conv1d_realistic(): + """Realistic reset: cold cache, L2 flush, then hot in_proj output.""" + # 1. Reset cold state (cache tensors, SSM state) + state_work.copy_(state0) + if state_scales_work is not None: + state_scales_work.copy_(state_scales0) + old_x_work.copy_(old_x0) + old_B_work.copy_(old_B0) + old_dt_work.copy_(old_dt0) + old_dA_cumsum_work.copy_(old_dA_cumsum0) + cache_buf_idx_work.copy_(cache_buf_idx0) + conv_state_work.copy_(conv_state0) + # 2. L2 flush (evicts cold state from cache) + if _l2_flush is not None: + _l2_flush.fill_(0.0) + # 3. Write hot tensors (simulates in_proj output landing in L2) + xbc_input_work.copy_(xbc_input0) + + # Silently skip the baseline row for any (baseline, state_dtype, SR) + # combo it can't run. Better than erroring on a partial sweep — our + # kernel rows still print. Compatibility: + # * Quantized states (int8 / int16 / fp8): no baseline supports them. + # * Triton baseline (selective_state_update): no rand_seed kwarg. + # * flashinfer baseline: rand_seed only on fp16 state. + def _baseline_supports() -> bool: + if baseline_fn is None: + return False + if is_quantized: + return False + if use_philox: + if args.baseline == "triton": + return False + if args.baseline == "flashinfer" and state_dtype != torch.float16: + return False + return True + + if baseline_fn is not None and not _baseline_supports(): + if not warmup_only: + sr_tag = " + SR" if use_philox else "" + print( + f"# Skipping {args.baseline} baseline for " + f"state_dtype={state_dtype_name}{sr_tag} (unsupported)." + ) + baseline_fn = None + + show_kernel_col = baseline_fn is not None + + def _conv1d_split(xbc_in, conv_st, launch_dependent_kernels=False): + """Run conv1d update and split output into (x, B, C) views. + + The input tensor's strides are preserved through conv1d and the + transpose+view chain. With the production-matching layout + (contiguous (batch*T, conv_dim) viewed as (batch, conv_dim, T)), + the output after transpose+view has stride(-1)==1 and + stride(1)==dim, satisfying both our kernel and flashinfer. + """ + xbc_result = causal_conv1d_update( + xbc_in, + conv_st, + conv_weight, + conv_bias, + activation="silu", + launch_dependent_kernels=launch_dependent_kernels, + ) + xbc_flat = xbc_result.transpose(1, 2).view(batch * mtp_len, conv_dim) + x_flat, B_flat, C_flat = torch.split( + xbc_flat, [d_inner, ngroups * d_state, ngroups * d_state], dim=-1 + ) + x_conv = x_flat.view(batch, mtp_len, nheads, head_dim) + B_conv = B_flat.view(batch, mtp_len, ngroups, d_state) + C_conv = C_flat.view(batch, mtp_len, ngroups, d_state) + return x_conv, B_conv, C_conv + + # --- Baseline --- + if baseline_fn is not None: + tag = f"base_b{batch}_mtp{mtp_len}_s{state_dtype_name}_a{act_dtype_name}" + + philox_kwargs = {} + if rand_seed is not None and args.baseline == "flashinfer": + philox_kwargs = {"rand_seed": rand_seed, "philox_rounds": args.philox_rounds} + + if with_conv1d: + + def _run_baseline(): + x_conv, B_conv, C_conv = _conv1d_split(xbc_input_work, conv_state_work) + baseline_fn( + state_work, + x=x_conv, + dt=dt, + A=A, + B=B_conv, + C=C_conv, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + out=out_base, + disable_state_update=True, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=mtp_len, + **philox_kwargs, + ) + else: + + def _run_baseline(): + baseline_fn( + state_work, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + out=out_base, + disable_state_update=True, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=mtp_len, + **philox_kwargs, + ) + + reset_fn = _reset_conv1d_realistic if with_conv1d else _reset + if warmup_only: + reset_fn() + _run_baseline() + torch.cuda.synchronize() + else: + stats = _time_kernel( + args, _run_baseline, reset_fn, tag, + expected_K=_kernels_per_iter_baseline(with_conv1d), + cupti_plan_key=( + "baseline", + args.baseline, + batch, + mtp_len, + state_dtype_name, + act_dtype_name, + with_conv1d, + bool(args.l2_flush), + bool(args.external_pdl), + bool(use_philox), + _kernels_per_iter_baseline(with_conv1d), + ), + ) + + _submit_result_job( + args, + stats, + show_kernel_col=show_kernel_col, + kernel_name=args.baseline, + batch=batch, + mtp_len=mtp_len, + prev_k="N/A", + state_dtype_name=state_dtype_name, + act_dtype_name=act_dtype_name, + skipped_tag=tag, + ) + + # --- Sweep parameter parsing (invariant across prev_k) --- + def _parse_sweep(val): + if val is None: + return [None] + return [int(v) for v in val.split(",")] + + block_size_m_values = _parse_sweep(args.block_size_m) + num_warps_values = _parse_sweep(args.num_warps) + num_stages_values = _parse_sweep(args.num_stages) + precompute_num_warps_values = _parse_sweep(args.precompute_num_warps) + precompute_num_stages_values = _parse_sweep(args.precompute_num_stages) + heads_per_block_values = _parse_sweep(args.heads_per_block) + maxnreg_values = _parse_sweep(args.maxnreg) + num_ctas_values = _parse_sweep(args.num_ctas) + # Persistent-only sweep dims; ignored when the cell's mode != persistent_main. + cta_per_sm_values = _parse_sweep(args.cta_per_sm) + num_loop_stages_values = _parse_sweep(args.num_loop_stages) + flatten_values = _parse_sweep(args.flatten) + warp_specialize_values = _parse_sweep(args.warp_specialize) + # Per-main split-knob sweeps. Default = same as the shared sweep (so each + # combo is tied). When set independently, the inner loop sweeps the + # cross-product (write × nowrite); --skip-diagonal drops the tied subset. + def _split_or_share(split_csv, shared_values): + return _parse_sweep(split_csv) if split_csv else shared_values + block_size_m_write_values = _split_or_share(args.block_size_m_write, block_size_m_values) + block_size_m_nowrite_values = _split_or_share(args.block_size_m_nowrite, block_size_m_values) + num_warps_write_values = _split_or_share(args.num_warps_write, num_warps_values) + num_warps_nowrite_values = _split_or_share(args.num_warps_nowrite, num_warps_values) + num_stages_write_values = _split_or_share(args.num_stages_write, num_stages_values) + num_stages_nowrite_values = _split_or_share(args.num_stages_nowrite, num_stages_values) + cta_per_sm_write_values = _split_or_share(args.cta_per_sm_write, cta_per_sm_values) + cta_per_sm_nowrite_values = _split_or_share(args.cta_per_sm_nowrite, cta_per_sm_values) + num_loop_stages_write_values = _split_or_share(args.num_loop_stages_write, num_loop_stages_values) + num_loop_stages_nowrite_values = _split_or_share(args.num_loop_stages_nowrite, num_loop_stages_values) + # Whether any *_write / *_nowrite knob was independently set — used by + # --skip-diagonal to know if the cross-product is non-trivial. Without + # any split, the per-main values == shared values and skip-diagonal is + # a no-op (which is correct). + _any_split = any(getattr(args, name) for name in ( + "block_size_m_write", "block_size_m_nowrite", + "num_warps_write", "num_warps_nowrite", + "num_stages_write", "num_stages_nowrite", + "cta_per_sm_write", "cta_per_sm_nowrite", + "num_loop_stages_write", "num_loop_stages_nowrite", + )) + # TMA toggles — independent 0/1 sweep per path. The skip-dupe at the + # top of the inner loop body collapses cells where a flag's path is + # unreachable for the current rectangle_for_nowrite setting. + use_tma_rect_load_values = _parse_sweep(args.use_tma_rect_load) + use_tma_replay_write_load_values = _parse_sweep(args.use_tma_replay_write_load) + use_tma_replay_nowrite_load_values = _parse_sweep(args.use_tma_replay_nowrite_load) + use_tma_replay_write_store_values = _parse_sweep(args.use_tma_replay_write_store) + + # --- Replay kernel --- + # Cache T-axis capacity (for prev_k validity check on the nowrite path). + max_window = getattr(args, "max_window", 0) or mtp_len + + # Build the list of scenarios to time. A scenario is one cell in the + # output: pure-mode scenarios fill prev_tokens with one constant before + # the timing loop; mix-mode scenarios feed a pre-baked per-iter samples + # tensor, with the per-iter copy captured inside the CUDA graph. Pure + # and mix can coexist in one call so a single nsys trace covers both. + scenarios = [] + if not (getattr(args, "mix_only", False) and mix_samples_cpu is not None): + for prev_k in prev_ks: + # Persistent modes dispatch per-slot from PNAT, so any prev_k + # <= max_window is valid. + scenarios.append({ + "label": f"k{prev_k}", + "print_label": prev_k, + "fill": prev_k, + "pre_iter": None, + "iters": None, # use args.iters + }) + # Mix scenario: bench pre-bakes both a per-iter PNAT samples tensor + # and (for persistent_main) a per-iter n_writes samples tensor; + # grouped graph capture copies window rows into kernel-input tensors + # before each in-graph L2 flush, so the timed kernels read PNAT cold. + if mix_samples_cpu is not None: + device = state_work.device + # Hardcode-sort: per-iter prev_tokens are CPU-sorted write-first. + # Kernel runs USE_PERM=False but the EO gate sees clustered modes. + # Output is scrambled (we don't permute x/B/C/dt to match) but + # timing is meaningful — isolates clustering benefit from the + # per-program perm-load overhead in --sort-slots. + src = mix_samples_sorted_cpu if (hardcode_sort and mix_samples_sorted_cpu is not None) else mix_samples_cpu + samples_gpu = torch.from_numpy(src).to(device=device, dtype=torch.int32) + + # For persistent_main + mix: pre-compute the per-iter n_writes + # (count of slots needing the write path = PNAT+T > max_window) + # and the (1,) scratch the kernel reads from. Both halves of + # persistent_main always launch in mix scenarios (host can't + # cheaply read n_writes per iter without a sync), so the kernel's + # slot-range derivation must be correct from device n_writes. + # persistent_dynamic doesn't need n_writes (the kernel ignores + # n_writes_dev when IS_DYNAMIC=True via Triton DCE), but we still + # allocate a sentinel scratch so the wrapper API is uniform. + n_writes_samples_gpu = None + n_writes_dev_mix = None + # n_writes per iter = number of slots that overflow the window. + # Computed for ALL mix scenarios so the JSON output (--json-detailed) + # can pair each iter's span_us with its mix composition for downstream + # analysis (group iters by # writes → per-bucket median → analytic + # expectation under the steady-state PNAT distribution). + n_writes_per_iter_all = ((src + mtp_len) > max_window).sum(axis=1).astype(np.int32) + if mode in ("persistent_main", "persistent_dynamic"): + n_writes_samples_gpu = torch.from_numpy(n_writes_per_iter_all).to( + device=device, dtype=torch.int32 + ) + n_writes_dev_mix = torch.zeros(1, dtype=torch.int32, device=device) + + # Build _mix_pre_iter — the closure that runs OUTSIDE the captured + # graph between replays. Updates: prev_tokens (always), + # slot_perm_buf (when sort_slots), n_writes_dev_mix (persistent). + perm_samples_gpu = None + if sort_slots and perm_samples_cpu is not None: + perm_samples_gpu = torch.from_numpy(perm_samples_cpu).to( + device=device, dtype=torch.int32 + ) + if n_writes_samples_gpu is not None: + def _mix_pre_iter(i, _s=samples_gpu, _ps=perm_samples_gpu, + _ns=n_writes_samples_gpu, _pt=prev_tokens, + _pm=slot_perm_buf, _nw=n_writes_dev_mix): + _pt.copy_(_s[i]) + _pm.copy_(_ps[i]) + _nw.copy_(_ns[i:i+1]) + else: + def _mix_pre_iter(i, _s=samples_gpu, _ps=perm_samples_gpu, + _pt=prev_tokens, _pm=slot_perm_buf): + _pt.copy_(_s[i]) + _pm.copy_(_ps[i]) + else: + if n_writes_samples_gpu is not None: + def _mix_pre_iter(i, _s=samples_gpu, _ns=n_writes_samples_gpu, + _pt=prev_tokens, _nw=n_writes_dev_mix): + _pt.copy_(_s[i]) + _nw.copy_(_ns[i:i+1]) + else: + def _mix_pre_iter(i, _s=samples_gpu, _pt=prev_tokens): + _pt.copy_(_s[i]) + + def _mix_pre_iter_group_factory( + group_iters, + _s=samples_gpu, + _ps=perm_samples_gpu, + _ns=n_writes_samples_gpu, + _pt=prev_tokens, + _pm=slot_perm_buf, + _nw=n_writes_dev_mix, + ): + sample_window = torch.empty( + (group_iters, _s.shape[1]), device=_s.device, dtype=_s.dtype, + ) + perm_window = ( + torch.empty((group_iters, _ps.shape[1]), device=_ps.device, dtype=_ps.dtype) + if _ps is not None else None + ) + nw_window = ( + torch.empty((group_iters,), device=_ns.device, dtype=_ns.dtype) + if _ns is not None else None + ) + + def _pre_replay(replay_idx): + start = replay_idx * group_iters + end = start + group_iters + sample_window.copy_(_s[start:end]) + if perm_window is not None: + perm_window.copy_(_ps[start:end]) + if nw_window is not None: + nw_window.copy_(_ns[start:end]) + + def _graph_pre_iter(j): + _pt.copy_(sample_window[j]) + if perm_window is not None: + _pm.copy_(perm_window[j]) + if nw_window is not None: + _nw.copy_(nw_window[j:j + 1]) + + return _pre_replay, _graph_pre_iter + + # Mix iters override: if --mix-iters set, use it; else use args.iters. + mix_iters = getattr(args, "mix_iters", None) + scenarios.append({ + "label": f"mix{mix_label}", + "print_label": "mix", + "fill": None, + "pre_iter": _mix_pre_iter, + "pre_iter_group_factory": _mix_pre_iter_group_factory, + "iters": mix_iters, # None => use args.iters + # Pass through to _run_incr so the wrapper receives _n_writes_dev + # (mix scenarios) instead of _n_writes (pure scenarios). + "n_writes_dev": n_writes_dev_mix, + # Full per-iter n_writes array (size = warmup + iters). Used by + # the JSON-detailed output to pair each iter's span with its + # mix composition for post-hoc bucketing analysis. + "n_writes_per_iter": n_writes_per_iter_all, + }) + + # Pure scenarios don't pre-allocate n_writes_dev; mix scenarios do. + # Default empty-halves skip: True for pure (host knows n_writes, + # production-equivalent host-skip), False for mix (host can't read + # device n_writes per iter without sync, must always launch both). + for scn in scenarios: + scenario_n_writes_dev = scn.get("n_writes_dev") # None for pure + scenario_skip_empty = scenario_n_writes_dev is None + if scn["fill"] is not None: + prev_tokens.fill_(scn["fill"]) + prev_k_for_print = scn["print_label"] + scenario_pre_iter = scn["pre_iter"] + scenario_pre_iter_group_factory = scn.get("pre_iter_group_factory") + scenario_iters = scn.get("iters") # None => use args.iters + tag = f"incr_b{batch}_mtp{mtp_len}_{scn['label']}_s{state_dtype_name}_a{act_dtype_name}" + + # Iteration over per-cell knob combos. + # When NO per-main split is requested (_any_split=False), each row in + # the cross-product gives the same value to both write_main and + # nowrite_main (current behavior — backward-compat). When ANY split + # IS requested, we iterate the write and nowrite axes independently + # (cross-product blowup is the user's responsibility — they typically + # pair this with --skip-diagonal to drop the tied subset). + if _any_split: + _iter_axes = ( + block_size_m_write_values, block_size_m_nowrite_values, + num_warps_write_values, num_warps_nowrite_values, + num_stages_write_values, num_stages_nowrite_values, + precompute_num_warps_values, + precompute_num_stages_values, + heads_per_block_values, + maxnreg_values, num_ctas_values, + cta_per_sm_write_values, cta_per_sm_nowrite_values, + num_loop_stages_write_values, num_loop_stages_nowrite_values, + flatten_values, warp_specialize_values, + use_tma_rect_load_values, + use_tma_replay_write_load_values, + use_tma_replay_nowrite_load_values, + use_tma_replay_write_store_values, + ) + else: + # Tied: one value per shared knob. Wrap in single-element list for + # uniform iteration; the body sets w/nw both to the shared value. + _iter_axes = ( + block_size_m_values, [None], + num_warps_values, [None], + num_stages_values, [None], + precompute_num_warps_values, + precompute_num_stages_values, + heads_per_block_values, + maxnreg_values, num_ctas_values, + cta_per_sm_values, [None], + num_loop_stages_values, [None], + flatten_values, warp_specialize_values, + use_tma_rect_load_values, + use_tma_replay_write_load_values, + use_tma_replay_nowrite_load_values, + use_tma_replay_write_store_values, + ) + # Iteration source: when --cell-list is active AND this is the main + # timing path (not a compile-warmup worker), iterate the cell set + # DIRECTLY (one yield per cell). The earlier design iterated the + # full inner cartesian and filtered each iteration via membership in + # args._cell_list_set — that's O(cartesian) which blows up to + # billions of iterations when the cell-list spans wide split-knob + # values (CPS, LS, M, W, S each contributing a Wx*Wnw factor on top + # of TMA flags), producing 50+ min of CPU spin per bench call before + # any actual timing. Direct iteration is O(|cell_list|). + # + # IMPORTANT exception for workers (warmup_only=True): _warm_one_config + # clamps args.*_write/_nowrite via inner_overrides to single values, + # making the cartesian 1×1×...×1 = 1 iter, which is exactly the one + # cell that worker was given. If we used cell-list-direct iteration + # here, every worker would iterate ALL 2884 cells instead of just + # its assigned one — turning compile-warmup into 28-way duplication. + # (Observed: 256 tasks in 233s under that bug vs ~18s correct.) + if getattr(args, "_cell_list_set", None) and not warmup_only: + def _gen_from_cell_list(): + keys = args._cell_list_keys + for tup in args._cell_list_set: + d = dict(zip(keys, tup)) + yield ( + d.get("Mw"), d.get("Mnw"), + d.get("Ww"), d.get("Wnw"), + d.get("Sw"), d.get("Snw"), + d.get("pW"), d.get("pS"), + d.get("H"), + d.get("R"), d.get("CT"), + d.get("CPSw"), d.get("CPSnw"), + d.get("LSw"), d.get("LSnw"), + d.get("FL"), d.get("WS"), + d.get("TMARL"), d.get("TMAWL"), + d.get("TMANL"), d.get("TMAWS"), + ) + _iter_source = _gen_from_cell_list() + else: + _iter_source = itertools.product(*_iter_axes) + + for ( + block_size_m_w, + block_size_m_nw, + num_warps_w, + num_warps_nw, + num_stages_w, + num_stages_nw, + precompute_num_warps, + precompute_num_stages, + heads_per_block, + maxnreg, + num_ctas, + cta_per_sm_w, + cta_per_sm_nw, + num_loop_stages_w, + num_loop_stages_nw, + flatten, + warp_specialize, + use_tma_rect_load, + use_tma_replay_write_load, + use_tma_replay_nowrite_load, + use_tma_replay_write_store, + ) in _iter_source: + # When tied, _nw values were placeholder None; fill from _w (the + # shared value). When split, _w and _nw came from independent lists. + if not _any_split: + block_size_m_nw = block_size_m_w + num_warps_nw = num_warps_w + num_stages_nw = num_stages_w + cta_per_sm_nw = cta_per_sm_w + num_loop_stages_nw = num_loop_stages_w + # Skip-diagonal: when split is on, drop the tied subset (same as a + # prior shared-knob sweep would cover). + if _any_split and args.skip_diagonal and ( + block_size_m_w == block_size_m_nw and + num_warps_w == num_warps_nw and + num_stages_w == num_stages_nw and + cta_per_sm_w == cta_per_sm_nw and + num_loop_stages_w == num_loop_stages_nw + ): + continue + # Backward-compat aliases used by the existing body below. When + # tied, these are simply the shared value. When split, the + # _write copy is used for sweep_tag and grouping (a stable choice + # so the tag is unique per (write, nowrite) combo). + block_size_m = block_size_m_w + num_warps = num_warps_w + num_stages = num_stages_w + cta_per_sm = cta_per_sm_w + num_loop_stages = num_loop_stages_w + # Skip-dupe for TMA flag sweeps: a flag whose code path isn't + # reachable in this cell produces identical timing for value=0 + # and value=1. We canonicalize by skipping value=1 cells when + # the flag's path is unreachable. Path reachability rules: + # * write path (replay write-load + write-store): always true. + # * rect path (rect-load): rectangle_for_nowrite=True. + # * replay-nowrite path (nowrite-load): rect isn't taking it. + _write_path = True # both halves exist for persistent modes + _rect_path = rectangle_for_nowrite + _replay_nowrite_path = not rectangle_for_nowrite + def _set(v): # flag set to a non-zero sweep value + return v is not None and v != 0 + if (_set(use_tma_rect_load) and not _rect_path + or _set(use_tma_replay_write_load) and not _write_path + or _set(use_tma_replay_nowrite_load) and not _replay_nowrite_path + or _set(use_tma_replay_write_store) and not _write_path): + continue + + # Pre-allocate n_writes_dev tensor OUTSIDE the captured graph for + # persistent modes in pure scenarios. Mix scenarios already have + # `scenario_n_writes_dev` pre-allocated. The wrapper's fallback + # `torch.tensor([...], device=...)` allocation would invalidate + # the CUDA-graph capture stream — must allocate here, before the + # `_run_incr` lambda (which is what gets captured) is defined. + # For persistent_dynamic the kernel ignores the value (IS_DYNAMIC + # DCE's the load); we still need a valid pointer. For + # persistent_main pure, the value is constant per cell so we set + # it once here. + _n_writes_dev_pure: torch.Tensor | None = None + _host_n_writes_pure: int | None = None + if mode in ("persistent_main", "persistent_dynamic") and scenario_n_writes_dev is None: + _n_writes_dev_pure = torch.zeros(1, dtype=torch.int32, device=state_work.device) + if mode == "persistent_main": + scn_fill = scn["fill"] + is_write_scenario_local = (scn_fill + mtp_len) > max_window + _host_n_writes_pure = batch if is_write_scenario_local else 0 + _n_writes_dev_pure.fill_(_host_n_writes_pure) + + def _run_incr( + block_size_m=block_size_m, + num_warps=num_warps, + num_stages=num_stages, + precompute_num_warps=precompute_num_warps, + precompute_num_stages=precompute_num_stages, + heads_per_block=heads_per_block, + maxnreg=maxnreg, + num_ctas=num_ctas, + cta_per_sm=cta_per_sm, + num_loop_stages=num_loop_stages, + flatten=flatten, + warp_specialize=warp_specialize, + use_tma_rect_load=use_tma_rect_load, + use_tma_replay_write_load=use_tma_replay_write_load, + use_tma_replay_nowrite_load=use_tma_replay_nowrite_load, + use_tma_replay_write_store=use_tma_replay_write_store, + ): + if with_conv1d: + x_call, B_call, C_call = _conv1d_split( + xbc_input_work, conv_state_work, launch_dependent_kernels=args.external_pdl + ) + extra_kwargs = {"launch_with_pdl": args.external_pdl} + else: + x_call, B_call, C_call = x, B, C + extra_kwargs = {} + # write_checkpoint is only meaningful for the checkpointing + # variant; replay variant ignores the kwarg. state_scales + # is also checkpointing-only (replay kernel doesn't quantize). + if args.variant == "checkpointing": + extra_kwargs["write_checkpoint"] = write_checkpoint + extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite + extra_kwargs["mode"] = mode + if sort_slots: + extra_kwargs["slot_perm"] = slot_perm_buf + # reverse_nowrite kwarg dropped from the slim kernel + # wrapper (was a maindl/dlgrouped feature). The sweep + # axis is retained here only so cell-list cells emitted + # by the search driver before the drop still parse — + # they always set REVN=0. + if state_scales_work is not None: + extra_kwargs["state_scales"] = state_scales_work + if use_tma_rect_load: # 1 → True, 0/None → False + extra_kwargs["_use_tma_rect_load"] = True + if use_tma_replay_write_load: + extra_kwargs["_use_tma_replay_write_load"] = True + if use_tma_replay_nowrite_load: + extra_kwargs["_use_tma_replay_nowrite_load"] = True + if use_tma_replay_write_store: + extra_kwargs["_use_tma_replay_write_store"] = True + # persistent_main needs n_writes (count of write-mode + # slots in the pre-sorted batch) as a host-side int. + # Pure scenarios: every slot has the same PNAT, so + # n_writes is either 0 (all nowrite) or batch (all + # write) depending on whether PNAT+T overflows the + # window. Mix scenarios are skipped earlier. + if mode in ("persistent_main", "persistent_dynamic"): + # Per-cell sweep values for persistent-only knobs. + # Apply to both persistent variants. _parse_sweep + # returns [None] when the user didn't pass the flag, + # in which case we leave the wrapper's defaults. + if cta_per_sm is not None: + extra_kwargs["_cta_per_sm"] = cta_per_sm + if num_loop_stages is not None: + extra_kwargs["_num_loop_stages"] = num_loop_stages + if flatten is not None: + extra_kwargs["_flatten"] = bool(flatten) + if warp_specialize is not None: + extra_kwargs["_warp_specialize"] = bool(warp_specialize) + if mode in ("persistent_main", "persistent_dynamic"): + # persistent_main + mix REQUIRES sort: the kernel + # partitions slots [0, n_writes) = write half, + # [n_writes, batch) = nowrite half. This only holds + # if PNAT is monotone (writes first), which sort + # provides via either: + # sort_slots=1 → USE_PERM reads slot_perm to remap + # hardcode_sort=1 → PNAT itself is CPU-pre-sorted + # persistent_dynamic doesn't need sort (per-slot + # runtime dispatch); persistent_main pure scenarios + # are trivially sorted (homogeneous PNAT). + if (mode == "persistent_main" + and scenario_n_writes_dev is not None + and not (sort_slots or hardcode_sort)): + raise AssertionError( + "persistent_main + mix requires sort_slots=1 " + "or hardcode_sort=1 — kernel partitions slots " + "by index, which is only valid when PNAT is " + "monotone (writes first). Without sort, the " + "partition silently mismatches actual slot " + "modes. Re-run with --sort-slots 1 or " + "--hardcode-sort 1." + ) + # n_writes plumbing: pure scenarios pass an int + # (host knows the value, can host-skip empty halves); + # mix scenarios pass a (1,) device tensor updated + # per iter by the benchmark pre-iter path. + # _persistent_skip_empty_halves=False on mix so both + # halves always launch (kernel uses device n_writes + # to derive its slot range). + if scenario_n_writes_dev is not None: + # Mix path: caller-allocated tensor, updated per + # iter by scenario_pre_iter outside capture. + extra_kwargs["_n_writes_dev"] = scenario_n_writes_dev + extra_kwargs["_persistent_skip_empty_halves"] = False + elif mode == "persistent_main": + # Pure: caller pre-allocated `_n_writes_dev_pure` + # outside this lambda (so the alloc doesn't land + # inside the captured graph). Pass both the + # tensor and the host int so the wrapper can use + # host-skip when `_persistent_skip_empty_halves`. + extra_kwargs["_n_writes"] = _host_n_writes_pure + extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure + extra_kwargs["_persistent_skip_empty_halves"] = scenario_skip_empty + elif mode == "persistent_dynamic": + # persistent_dynamic pure: kernel ignores n_writes + # via IS_DYNAMIC DCE, but the wrapper needs a + # valid (1,) tensor pointer. Pass the pre-allocated + # zero tensor to avoid any in-capture alloc. + extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure + variant_fn( + state_work, + old_x_work, + old_B_work, + old_dt_work, + old_dA_cumsum_work, + cache_buf_idx_work, + prev_tokens, + x=x_call, + dt=dt, + A=A, + B=B_call, + C=C_call, + out=out_incr, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + rand_seed=rand_seed, + philox_rounds=args.philox_rounds, + use_internal_pdl=args.internal_pdl, + _block_size_m=block_size_m, + _num_warps=num_warps, + _num_stages=num_stages, + _precompute_num_warps=precompute_num_warps, + _precompute_num_stages=precompute_num_stages, + _heads_per_block=heads_per_block, + _maxnreg=maxnreg, + _num_ctas=num_ctas, + # Per-main overrides (None = tied to shared above; explicit + # only when the inner loop is iterating split axes). + _block_size_m_write=block_size_m_w if _any_split else None, + _block_size_m_nowrite=block_size_m_nw if _any_split else None, + _num_warps_write=num_warps_w if _any_split else None, + _num_warps_nowrite=num_warps_nw if _any_split else None, + _num_stages_write=num_stages_w if _any_split else None, + _num_stages_nowrite=num_stages_nw if _any_split else None, + _cta_per_sm_write=cta_per_sm_w if _any_split else None, + _cta_per_sm_nowrite=cta_per_sm_nw if _any_split else None, + _num_loop_stages_write=num_loop_stages_w if _any_split else None, + _num_loop_stages_nowrite=num_loop_stages_nw if _any_split else None, + **extra_kwargs, + ) + + parts = [] + # When tied (not _any_split), emit the shared single-value tag + # (M=8 etc). When split, emit explicit Mw / Mnw tags so cells + # with the same shared value but different per-main values get + # unique JSON keys. + def _emit_split(name_w, name_nw, val_w, val_nw): + if val_w is None and val_nw is None: + return + if not _any_split or val_w == val_nw: + parts.append(f"{name_w[:-1]}={val_w}") # strip the 'w' suffix + else: + parts.append(f"{name_w}={val_w}") + parts.append(f"{name_nw}={val_nw}") + _emit_split("Mw", "Mnw", block_size_m_w, block_size_m_nw) + _emit_split("Ww", "Wnw", num_warps_w, num_warps_nw) + _emit_split("Sw", "Snw", num_stages_w, num_stages_nw) + if precompute_num_warps is not None: + parts.append(f"pW={precompute_num_warps}") + if precompute_num_stages is not None: + parts.append(f"pS={precompute_num_stages}") + if heads_per_block is not None: + parts.append(f"H={heads_per_block}") + if maxnreg is not None: + parts.append(f"R={maxnreg}") + if num_ctas is not None: + parts.append(f"CT={num_ctas}") + # Persistent-only knobs (only meaningful when MODE=persistent_main; + # printed unconditionally so output rows are uniformly comparable + # across modes when the user passed these sweeps). + _emit_split("CPSw", "CPSnw", cta_per_sm_w, cta_per_sm_nw) + _emit_split("LSw", "LSnw", num_loop_stages_w, num_loop_stages_nw) + if flatten is not None: + parts.append(f"FL={flatten}") + if warp_specialize is not None: + parts.append(f"WS={warp_specialize}") + # TMA sweep tags. Four wrapper-level flags map to three + # kernel-level constexprs (rect-load and replay-nowrite-load + # share `USE_TMA_LOAD_NOWRITE`, picked by the wrapper based on + # RECTANGLE). TMARL specifically gates the rectangle path's + # state load; TMANL specifically gates the replay-style + # nowrite path's state load. Distinct because their measured + # perf profiles differ (see CHECKPOINTING_DESIGN.md item #17: + # rect TMA is "not a win" while replay-nowrite TMA is the + # biggest measured win at int8 b>=64). + if use_tma_rect_load is not None: + parts.append(f"TMARL={use_tma_rect_load}") # rect path load + if use_tma_replay_write_load is not None: + parts.append(f"TMAWL={use_tma_replay_write_load}") # replay-write load + if use_tma_replay_nowrite_load is not None: + parts.append(f"TMANL={use_tma_replay_nowrite_load}") # replay-NOWRITE load (NOT rect) + if use_tma_replay_write_store is not None: + parts.append(f"TMAWS={use_tma_replay_write_store}") # replay-write store + parts.append(f"SR={1 if use_philox else 0}") + parts.append(f"RECT={1 if rectangle_for_nowrite else 0}") + parts.append(f"WC={1 if write_checkpoint else 0}") + parts.append(f"MODE={mode}") + parts.append(f"SORT={1 if sort_slots else 0}") + parts.append(f"REVN={1 if reverse_nowrite else 0}") + parts.append(f"HSORT={1 if hardcode_sort else 0}") + sweep_suffix = (" " + ",".join(parts)) if parts else "" + sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") + + reset_fn = _reset_conv1d_realistic if with_conv1d else _reset + # --cell-list filter: only time cells whose (canonical-knob-values) + # tuple is in the loaded set. Robust to bench gaining new knobs + # (old cell-list files keep working: any keys they don't list + # become wildcards that retain CLI defaults). + if args._cell_list_keys: + _tup = _current_cell_tuple(args, locals()) + if _tup is None or _tup not in args._cell_list_set: + continue + # Resume from JSONL: skip cells already recorded. Built the same + # way _print_row builds JSON keys; must stay in sync. + done_keys = getattr(args, "_done_keys", None) + if done_keys: + # One key per scenario (k=6, k=11, mix) — skip the whole cell + # only if ALL of its scenarios are already done. We don't + # know which scenarios will be emitted here without + # re-evaluating the inner scenario loop; conservatively skip + # only when the prev_k_for_print's specific key is done. + _resume_key = _build_json_key( + args.variant, batch, mtp_len, prev_k_for_print, + state_dtype_name, sweep_suffix, args.tp_size, + ) + if _resume_key in done_keys: + continue + if warmup_only: + reset_fn() + if scenario_pre_iter is not None: + scenario_pre_iter(0) + _run_incr() + torch.cuda.synchronize() + else: + # Inline retry: CUPTI sometimes loses records under PDL + + # high cell count; retrying the SAME cell often catches it + # because the failure is transient at the kernel-launch level. + # Per --cupti-retry budget. On final failure, append tag to + # the skipped list for an external rerun in a fresh process. + defer_results = ( + args.cuda_graph + and getattr(args, "cupti", True) + and int(getattr(args, "cupti_defer_depth", 1)) > 1 + ) + retry_budget = 0 if defer_results else max(0, getattr(args, "cupti_retry", 1)) + stats = None + expected_K = _kernels_per_iter_incremental( + mode, with_conv1d=with_conv1d, + persistent_skip_empty=scenario_skip_empty, + ) + plan_key = ( + "incremental", + args.variant, + mode, + batch, + mtp_len, + state_dtype_name, + act_dtype_name, + with_conv1d, + bool(args.l2_flush), + bool(args.external_pdl), + bool(args.internal_pdl), + bool(use_philox), + bool(rectangle_for_nowrite), + bool(write_checkpoint), + bool(sort_slots), + bool(reverse_nowrite), + bool(hardcode_sort), + scenario_pre_iter is not None, + expected_K, + ) + for attempt in range(retry_budget + 1): + stats = _time_kernel( + args, _run_incr, reset_fn, sweep_tag, + expected_K=expected_K, + pre_iter_fn=scenario_pre_iter, + pre_iter_group_factory=scenario_pre_iter_group_factory, + iters_override=scenario_iters, + cupti_plan_key=plan_key, + ) + if stats is not None: + break + if attempt < retry_budget: + print( + f"[retry] CUPTI mismatch on {sweep_tag!r}; " + f"retrying ({attempt + 1}/{retry_budget})", + file=sys.stderr, + flush=True, + ) + if stats is None: + args._skipped_cells.append(sweep_tag) + + # Attach n_writes_per_iter when it is needed for scoring. + # For pure scenarios, n_writes is constant: 0 (nowrite) or + # batch (write), determined by scn["fill"] + mtp_len > max_window. + # For mix, scn carries the precomputed per-iter array. + per_iter_nw = None + if stats is not None and ( + getattr(args, "json_detailed", False) or scn["fill"] is None + ): + eff_iters = scenario_iters if scenario_iters is not None else args.iters + if scn["fill"] is not None: + if getattr(args, "json_detailed", False): + # Pure scenario: constant n_writes for every iter. + is_write = (scn["fill"] + mtp_len > max_window) + per_iter_nw = [batch if is_write else 0] * eff_iters + else: + # Mix scenario: slice off warmup, keep timed iters. + nw_full = scn.get("n_writes_per_iter") + if nw_full is not None: + per_iter_nw = nw_full[args.warmup:args.warmup + eff_iters].tolist() + else: + per_iter_nw = None + + if stats is not None: + _submit_result_job( + args, + stats, + show_kernel_col=show_kernel_col, + kernel_name=args.variant, + batch=batch, + mtp_len=mtp_len, + prev_k=prev_k_for_print, + state_dtype_name=state_dtype_name, + act_dtype_name=act_dtype_name, + sweep_suffix=sweep_suffix, + per_iter_nw=per_iter_nw, + skipped_tag=sweep_tag, + ) + + +# Map full torch dtype name → short tag used in JSON keys (matches collect.py). +_DTYPE_SHORT = { + "float32": "fp32", "bfloat16": "bf16", "float16": "fp16", + "int8": "int8", "int16": "int16", "float8_e4m3fn": "fp8", +} + + +def _build_json_key( + kernel_name, batch, mtp_len, prev_k, state_dtype_name, sweep_suffix, tp_size +): + """Build a key matching collect.py's kernel_data.json convention: + + incremental/{batch}/{mtp}/{sd}/k{k}/{sweep_parts}/tp{tp} + triton/{batch}/{mtp}/{sd}/tp{tp} + flashinfer/{batch}/{mtp}/{sd}/tp{tp} + + `kernel_name` is what _print_row receives: variant name for the timed + kernel (replay/checkpointing) or baseline name for the baseline row. + Variant rows collapse to "incremental" — the variant choice is captured + by the sweep flags collect.py would otherwise apply via --variant. + """ + if kernel_name in ("replay", "checkpointing"): + kind = "incremental" + else: + kind = kernel_name # "triton" / "flashinfer" + + sd = _DTYPE_SHORT.get(state_dtype_name, state_dtype_name) + parts = [kind, str(batch), str(mtp_len), sd] + if prev_k != "N/A": + parts.append(f"k{prev_k}") + if sweep_suffix: + # sweep_suffix format: " M=4,W=1,S=1,SR=0,RECT=0,WC=1" + # collect.py format: "M4_W1_S1_SR0_RECT0_WC0" + # Strip leading/trailing whitespace, drop '=', commas → underscores. + parts.append( + sweep_suffix.strip().replace("=", "").replace(",", "_") + ) + parts.append(f"tp{tp_size}") + return "/".join(parts) + + +def _print_row( + show_kernel_col, + kernel_name, + batch, + mtp_len, + prev_k, + state_dtype_name, + act_dtype_name, + stats, + sweep_suffix="", + tp_size=None, + json_detailed=False, + jsonl_path=None, + jsonl_host=None, + jsonl_gpu=None, +): + """Print one summary row and append the result to the JSONL sidecar. + + `stats` is a dict from _time_kernel: {median, p95, p99, n, iters_us, + [n_writes_per_iter], [per_kernel]}. The summary table only shows the + headline percentiles. JSONL captures the compact per-iter spans + + n_writes_per_iter by default; with json_detailed=True it also captures + per-kernel data. + + When `jsonl_path` is provided, appends one JSON line per row to the + JSONL sidecar (crash-safe incremental persistence; lets a killed sweep + resume from the last completed cell on rerun, even across hosts). Open + per-write because `args` is pickled to ProcessPoolExecutor workers and + file handles aren't picklable. JSONL is the canonical artifact — the + bench no longer writes a final `.json` summary; use `jsonl_to_json.py` + if a one-shot `.json` snapshot is needed. + """ + kernel_col = f"{kernel_name:>11} | " if show_kernel_col else "" + print( + f"| {kernel_col}{batch:>5} | {mtp_len:>7} | {str(prev_k):>6} | " + f"{state_dtype_name:>11} | {act_dtype_name:>9} | " + f"{stats['median']:>9.2f} | {stats['p95']:>7.2f} | {stats['p99']:>7.2f} |" + f"{sweep_suffix}" + ) + if jsonl_path is not None: + key = _build_json_key( + kernel_name, batch, mtp_len, prev_k, state_dtype_name, + sweep_suffix, tp_size, + ) + if json_detailed: + row_stats = stats + else: + row_stats = { + k: stats[k] + for k in ("median", "p95", "p99", "n", "iters_us", "n_writes_per_iter") + if k in stats + } + if "host_timing" in stats: + row_stats["host_timing"] = stats["host_timing"] + # Append to JSONL sidecar if a path is set (incremental persistence). + # Open per-write because args is pickled to ProcessPoolExecutor + # workers, and file handles aren't picklable. A clean SIGTERM or + # Python exception will leave the file consistent up to the last + # newline; catastrophic kills can leave a partial last line, which + # the resume reader tolerates via json.JSONDecodeError pass. + if jsonl_path is not None: + # Wall-clock timestamp (float seconds since UNIX epoch) at write + # time. Lets post-hoc analysis diff consecutive rows to derive + # per-cell wall budget and identify startup-bound vs steady-state + # segments (cells/sec, downtime between bench invocations) without + # needing to instrument the bench's outer loops separately. + import time as _time + rec = {"key": key, "stats": row_stats, "t": _time.time()} + if jsonl_host is not None: + rec["host"] = jsonl_host + if jsonl_gpu is not None: + rec["gpu"] = jsonl_gpu + with open(jsonl_path, "a") as f: + f.write(json.dumps(rec) + "\n") + + +def _finish_result_job(args, job: dict) -> None: + result = job["result"] + if isinstance(result, _PendingCuptiStats): + stats = result.resolve() + else: + stats = result + + if stats is None: + skipped_tag = job.get("skipped_tag") + if skipped_tag is not None: + args._skipped_cells.append(skipped_tag) + return + + per_iter_nw = job.get("per_iter_nw") + if per_iter_nw is not None: + stats["n_writes_per_iter"] = per_iter_nw + + _print_row( + job["show_kernel_col"], + job["kernel_name"], + job["batch"], + job["mtp_len"], + job["prev_k"], + job["state_dtype_name"], + job["act_dtype_name"], + stats, + job.get("sweep_suffix", ""), + tp_size=args.tp_size, + json_detailed=getattr(args, "json_detailed", False), + jsonl_path=getattr(args, "_jsonl_path", None), + jsonl_host=getattr(args, "_jsonl_host", None), + jsonl_gpu=getattr(args, "_jsonl_gpu", None), + ) + + +def _drain_pending_results(args, *, force: bool = False) -> None: + pending_results = getattr(args, "_pending_results", None) + if not pending_results: + return + + max_pending = max(1, int(getattr(args, "cupti_defer_depth", 1))) + while pending_results: + first_result = pending_results[0]["result"] + should_block = force or len(pending_results) >= max_pending + if ( + not should_block + and isinstance(first_result, _PendingCuptiStats) + and not first_result.is_ready() + ): + break + job = pending_results.pop(0) + _finish_result_job(args, job) + + +def _submit_result_job( + args, + result, + *, + show_kernel_col, + kernel_name, + batch, + mtp_len, + prev_k, + state_dtype_name, + act_dtype_name, + sweep_suffix="", + per_iter_nw=None, + skipped_tag=None, +) -> None: + job = { + "result": result, + "show_kernel_col": show_kernel_col, + "kernel_name": kernel_name, + "batch": batch, + "mtp_len": mtp_len, + "prev_k": prev_k, + "state_dtype_name": state_dtype_name, + "act_dtype_name": act_dtype_name, + "sweep_suffix": sweep_suffix, + "per_iter_nw": per_iter_nw, + "skipped_tag": skipped_tag, + } + if isinstance(result, _PendingCuptiStats): + args._pending_results.append(job) + _drain_pending_results(args) + else: + _finish_result_job(args, job) + + +# Cell-list mode — canonical knob-key mapping to argparse args + local +# loop variable. See _load_cell_list_into_args / inner-loop filter. +# +# Each entry: cell-key → (args attribute name, comma-separated string flag) +# For split (write/nowrite) knobs, we use Xw / Xnw keys. Tied forms (M, W, +# S, CPS, LS) accepted on load and expanded to their w/nw variants. +_CELL_LIST_KEY_TO_ARG = { + "Mw": "block_size_m_write", + "Mnw": "block_size_m_nowrite", + "Ww": "num_warps_write", + "Wnw": "num_warps_nowrite", + "Sw": "num_stages_write", + "Snw": "num_stages_nowrite", + "CPSw": "cta_per_sm_write", + "CPSnw": "cta_per_sm_nowrite", + "LSw": "num_loop_stages_write", + "LSnw": "num_loop_stages_nowrite", + "pW": "precompute_num_warps", + "pS": "precompute_num_stages", + "H": "heads_per_block", + "R": "maxnreg", + "CT": "num_ctas", + "FL": "flatten", + "WS": "warp_specialize", + "TMARL": "use_tma_rect_load", + "TMAWL": "use_tma_replay_write_load", + "TMANL": "use_tma_replay_nowrite_load", + "TMAWS": "use_tma_replay_write_store", + "RECT": "rectangle_for_nowrite", + "WC": "write_modes", + "SORT": "sort_slots", + "REVN": "reverse_nowrite", + "HSORT": "hardcode_sort", + # MODE and SR get special handling (string values): + # MODE → args.modes (single mode name) + # SR → args.sr_modes ("RN" if 0, "SR" if 1) +} + +# Split-knob tied form: "M" expands to both "Mw" and "Mnw". +_CELL_LIST_TIED_EXPANSIONS = { + "M": ("Mw", "Mnw"), + "W": ("Ww", "Wnw"), + "S": ("Sw", "Snw"), + "CPS": ("CPSw", "CPSnw"), + "LS": ("LSw", "LSnw"), +} + + +def _normalize_cell(cell: dict) -> dict: + """Expand tied-form keys (M, W, S, CPS, LS) to their w/nw variants. + Returns a new dict with only canonical split-or-plain keys. + """ + out = dict(cell) + for tied, (w_key, nw_key) in _CELL_LIST_TIED_EXPANSIONS.items(): + if tied in out: + v = out.pop(tied) + out.setdefault(w_key, v) + out.setdefault(nw_key, v) + return out + + +def _load_cell_list_into_args(args) -> None: + """Read --cell-list JSON, normalize, override args.* knob ranges, and + populate args._cell_list_keys + args._cell_list_set for the inner-loop + filter. Errors out if cells aren't uniform (different key sets). + """ + with open(args.cell_list) as f: + raw = json.load(f) + if not isinstance(raw, list): + sys.exit(f"--cell-list: expected JSON list, got {type(raw).__name__}") + cells = [_normalize_cell(c) for c in raw] + if not cells: + print("[cell-list] empty list — nothing to time", file=sys.stderr) + return + # All cells must share the same key set (uniform schema) + keys0 = frozenset(cells[0].keys()) + for i, c in enumerate(cells[1:], start=1): + if frozenset(c.keys()) != keys0: + sys.exit( + f"--cell-list: cells must have uniform key sets; cell[0] " + f"has {sorted(keys0)} but cell[{i}] has {sorted(c.keys())}" + ) + + # Auto-cover: collect per-knob value set across all cells + cover: dict = {} + for c in cells: + for k, v in c.items(): + cover.setdefault(k, set()).add(v) + # Apply overrides + for key, vals in cover.items(): + if key in _CELL_LIST_KEY_TO_ARG: + arg_name = _CELL_LIST_KEY_TO_ARG[key] + vals_str = ",".join(str(v) for v in sorted(vals)) + setattr(args, arg_name, vals_str) + elif key == "MODE": + args.modes = ",".join(sorted({str(v) for v in vals})) + elif key == "SR": + args.sr_modes = ",".join(sorted({"SR" if v else "RN" for v in vals})) + else: + print(f"[cell-list] WARNING: unknown key {key!r} in cells; " + f"will not override any args.* attribute (the value will " + f"still be matched in the filter if a matching local var " + f"is in scope)", file=sys.stderr) + + # Canonical key order (sorted) for tuple matching in the inner loop + args._cell_list_keys = tuple(sorted(keys0)) + args._cell_list_set = { + tuple(c[k] for k in args._cell_list_keys) for c in cells + } + print(f"[cell-list] loaded {len(cells)} cells with keys " + f"{list(args._cell_list_keys)}; overrode args.* to auto-cover", + file=sys.stderr) + + +# Maps cell-list key → name of the local variable in _bench_config's inner +# loop. Used to extract the "current cell" tuple for the filter check. +# Keep in sync with the loop-variable names; the filter is lenient about +# missing names (it picks them up from the inner scope at runtime). +_CELL_LIST_KEY_TO_LOCAL = { + "Mw": "block_size_m_w", + "Mnw": "block_size_m_nw", + "Ww": "num_warps_w", + "Wnw": "num_warps_nw", + "Sw": "num_stages_w", + "Snw": "num_stages_nw", + "CPSw": "cta_per_sm_w", + "CPSnw": "cta_per_sm_nw", + "LSw": "num_loop_stages_w", + "LSnw": "num_loop_stages_nw", + "pW": "precompute_num_warps", + "pS": "precompute_num_stages", + "H": "heads_per_block", + "R": "maxnreg", + "CT": "num_ctas", + "FL": "flatten", + "WS": "warp_specialize", + "TMARL": "use_tma_rect_load", + "TMAWL": "use_tma_replay_write_load", + "TMANL": "use_tma_replay_nowrite_load", + "TMAWS": "use_tma_replay_write_store", + "RECT": "rectangle_for_nowrite", + "WC": "write_checkpoint", + "MODE": "mode", + "SORT": "sort_slots", + "REVN": "reverse_nowrite", + "HSORT": "hardcode_sort", + "SR": "use_philox", +} + + +def _current_cell_tuple(args, locals_dict: dict) -> tuple | None: + """Build the (key1=val1, key2=val2, ...) tuple for the current inner-loop + iteration, matching args._cell_list_keys' order. Used by the inner-loop + filter to check membership in args._cell_list_set. Returns None if any + expected local is missing (the bench evolved a knob name — caller skips). + """ + if not args._cell_list_keys: + return None + vals = [] + for k in args._cell_list_keys: + local_name = _CELL_LIST_KEY_TO_LOCAL.get(k, k) + if local_name not in locals_dict: + return None + v = locals_dict[local_name] + # Coerce bools to ints to match cell-list JSON (1/0) + if isinstance(v, bool): + v = int(v) + vals.append(v) + return tuple(vals) + + +# Main benchmark loop + + +def _run_benchmark(args) -> None: + # Phase-timing markers — emit timestamped checkpoints so a captured-stdout + # run can later attribute wall time to setup vs compile-warmup vs prewarm + # vs timing. Single-line format makes log-grepping trivial. + _phase_t0 = time.perf_counter() + def _phase(label: str) -> None: + dt = time.perf_counter() - _phase_t0 + print(f"[phase] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) + _phase("enter _run_benchmark") + + # Pending-results FIFO for srxl's deferred CUPTI parsing pipeline. Each + # entry holds a _PendingCuptiStats handle; _drain_pending_results pulls + # ready entries and routes them to _print_row (which appends to JSONL). + args._pending_results = [] + + # JSONL incremental sidecar. Path = `.jsonl`. Each completed + # cell appends one line `{"key": , "stats": {...}, "host": }` + # to this file as it finishes timing. On startup we read this sidecar (if + # present) and populate _done_keys so a killed bench can resume without + # redoing already-timed cells. Crash-safe by construction: append-only + # writes survive SIGTERM/SIGKILL/reboot mid-sweep. + # + # Resume is host-blind: _done_keys includes records from any host, so a + # bench restarted on a different node fills in the missing cells without + # redoing cells already covered elsewhere. Cross-host *timings* aren't + # directly comparable, but each JSONL record carries its `host` stamp so + # the analyzer can group/compare per host. This bench no longer writes a + # final `.json` summary — the JSONL is the canonical artifact; use the + # `jsonl_to_json.py` helper if a one-shot `.json` snapshot is needed. + # + # Note: we store only paths/strings on `args` because args is pickled to + # ProcessPoolExecutor workers during compile-warmup, and file handles + # (TextIOWrapper) aren't picklable. _print_row open-appends per cell. + args._jsonl_path = None + args._done_keys: set[str] = set() + args._jsonl_host = None # hostname stamp for the current run + args._jsonl_gpu = None # GPU device id stamp (current process visibility) + if getattr(args, "json_output", None): + import socket + args._jsonl_host = socket.gethostname() + # Capture GPU id once at startup. Used by the oracle-cache layer in + # search_driver to attribute timings to a specific (host, gpu) pair + # for cross-process pruning. os.environ['CUDA_VISIBLE_DEVICES'] + # is the right source pre-torch-init (it's what the harness sets); + # post-init we could use torch.cuda.current_device() but we keep it + # to env to avoid forcing a CUDA init at this point in startup. + args._jsonl_gpu = os.environ.get("CUDA_VISIBLE_DEVICES", "") + args._jsonl_path = args.json_output + ".jsonl" + # Read existing JSONL if present: load every record's key into the + # skip set regardless of host (gap-fill on a new node). + if os.path.exists(args._jsonl_path): + n_loaded = 0 + host_counts: dict[str, int] = {} + with open(args._jsonl_path) as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except json.JSONDecodeError: + # Tolerate partial last line from a crash mid-write. + continue + k = rec.get("key") + if k is None: + continue + args._done_keys.add(k) + n_loaded += 1 + rec_host = rec.get("host") + if rec_host: + host_counts[rec_host] = host_counts.get(rec_host, 0) + 1 + if n_loaded: + host_summary = ", ".join( + f"{h}={n}" for h, n in sorted(host_counts.items()) + ) if host_counts else "(no host stamps)" + print( + f"[resume] {args._jsonl_path}: loaded {n_loaded} prior " + f"cell results across hosts [{host_summary}]; sweep will " + f"skip them. New cells stamp host={args._jsonl_host}.", + file=sys.stderr, + ) + + # Sidecar metadata: cmd, host, tp_size, variant, cupti, etc. Written + # once at startup; helps later analysis identify how this JSONL was + # produced even though there's no top-level .json wrapper anymore. + meta_path = args.json_output + ".meta.json" + meta_payload = { + "timestamp": datetime.now().isoformat(), + "host": args._jsonl_host, + "cmd": " ".join(sys.argv), + "tp_size": getattr(args, "tp_size", None), + "warmup": getattr(args, "warmup", None), + "iters": getattr(args, "iters", None), + "variant": getattr(args, "variant", None), + "cupti": getattr(args, "cupti", False), + } + # Append to a list so successive runs (gap-fill, retry) keep history. + existing_meta = [] + if os.path.exists(meta_path): + try: + with open(meta_path) as f: + existing_meta = json.load(f) + if not isinstance(existing_meta, list): + existing_meta = [existing_meta] + except (OSError, json.JSONDecodeError): + existing_meta = [] + existing_meta.append(meta_payload) + # Bench is sometimes invoked with --json-output pointing into a dir + # the caller hasn't created (subprocess driver, search loop, etc.). + # Ensure the dir exists before writing the meta sidecar OR the JSONL. + os.makedirs(os.path.dirname(os.path.abspath(meta_path)), exist_ok=True) + tmp = meta_path + ".tmp" + with open(tmp, "w") as f: + json.dump(existing_meta, f, indent=2) + os.replace(tmp, meta_path) + + # Skipped cells accumulator — populated by _bench_config when CUPTI capture + # mismatch causes a cell to be skipped. Written to args.skipped_output + # (or derived from json_output) at end of run. + args._skipped_cells = [] + + # Cell-list filter (replaces the old --retry-cells tag-string filter). + # When set, the sweep iterates ONLY the cells described in the list. + # + # Each entry in the JSON file is a dict of canonical knob keys → values, + # using the same names that appear in the sweep_tag (Mw/Mnw, Ww/Wnw, + # Sw/Snw, pW, pS, H, R, CT, CPSw/CPSnw, LSw/LSnw, FL, WS, TMARL, + # TMAWL, TMANL, TMAWS, SR, RECT, WC, MODE, SORT, REVN, HSORT). Each + # cell may also use the tied forms M / W / S / CPS / LS (single value + # applied to both write and nowrite halves). + # + # On load we: + # - Override the bench's CLI knob args (`args.block_size_m_write`, + # etc.) with the union of values present across all cells per knob, + # so the cartesian iteration auto-covers the list. + # - Build `args._cell_list_keys` (the canonical key order used by + # every cell — must be uniform across the list) and + # `args._cell_list_set` (frozen tuples for O(1) membership check + # inside the inner loop). + # + # In the inner loop, we build the current iteration's tuple and skip + # cells not in the set. Dict-matching is robust to bench gaining new + # knobs (old cell-list files keep working — newly-added knobs simply + # aren't matched on, so they retain CLI defaults). + # Cell-list state may already have been populated by main() (so that + # the args.*_list derivations downstream see the override). Default to + # empty if not. + _phase(f"done loading _done_keys ({len(args._done_keys)} entries)") + + if not hasattr(args, "_cell_list_keys"): + args._cell_list_keys: tuple = () + args._cell_list_set: set = set() + if getattr(args, "cell_list", None): + _load_cell_list_into_args(args) + _phase(f"done loading cell-list ({len(args._cell_list_set)} cells)") + + assert args.nheads % args.tp_size == 0, ( + f"nheads ({args.nheads}) must be divisible by tp_size ({args.tp_size})" + ) + assert args.ngroups % args.tp_size == 0, ( + f"ngroups ({args.ngroups}) must be divisible by tp_size ({args.tp_size})" + ) + args.tp_nheads = args.nheads // args.tp_size + args.tp_ngroups = args.ngroups // args.tp_size + + batch_sizes = [int(x) for x in args.batch_sizes.split(",")] + mtp_lengths = [int(x) for x in args.mtp_lengths.split(",")] + + dtype_map = { + "bf16": torch.bfloat16, + "fp32": torch.float32, + "fp16": torch.float16, + "int8": torch.int8, + "int16": torch.int16, + "fp8": torch.float8_e4m3fn, + } + state_dtypes = [dtype_map[s] for s in args.state_dtypes.split(",")] + act_dtypes = [dtype_map[s] for s in args.act_dtypes.split(",")] + + # Resolve baseline function + if args.baseline == "flashinfer": + from flashinfer.mamba import selective_state_update as baseline_fn + elif args.baseline == "triton": + baseline_fn = selective_state_update + else: + baseline_fn = None + + # --with-conv1d uses its own realistic L2 flush (cold cache flush then + # hot in_proj write). Override the generic l2_flush to avoid double-flushing. + if args.with_conv1d: + args.l2_flush = False + _init_l2_flush() # still needed for the realistic reset's flush step + elif args.l2_flush: + _init_l2_flush() + + _phase("about to enter compile-warmup") + if args.compile_threads > 0: + _compile_warmup_phase( + args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, + baseline_fn, max_workers=args.compile_threads, + ) + _phase("returned from compile-warmup") + + # Pre-warm the per-(state_dtype, act_dtype, mtp_len, ...) tensor cache at + # the largest requested batch size. Without this, the timing loop would + # progressively grow the cache as it encounters larger batches (e.g., + # iterate 1 -> 16 -> 64 -> 128 -> 512 = 5 separate growth allocations, + # each freeing the previous buffers). Pre-warming at max-batch up front + # makes every subsequent timing cell a view-slice (zero alloc cost). + _max_batch = max(batch_sizes) + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + for mtp_len in mtp_lengths: + _build_tensors( + _max_batch, mtp_len, state_dtype, act_dtype, + args.tp_nheads, args.head_dim, args.d_state, args.tp_ngroups, + max_window=getattr(args, "max_window", None) or None, + ) + _phase("done tensor prewarm — entering timing") + + if args.profile: + torch.cuda.cudart().cudaProfilerStart() + + # Print header + if baseline_fn is not None: + print( + f"| {'kernel':>11} | {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " + f"{'state_dtype':>11} | {'act_dtype':>9} | " + f"{'median_us':>9} | {'p95_us':>7} | {'p99_us':>7} |" + ) + print( + f"|{'-' * 13}|{'-' * 7}|{'-' * 9}|{'-' * 8}|" + f"{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" + ) + else: + print( + f"| {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " + f"{'state_dtype':>11} | {'act_dtype':>9} | " + f"{'median_us':>9} | {'p95_us':>7} | {'p99_us':>7} |" + ) + print( + f"|{'-' * 7}|{'-' * 9}|{'-' * 8}|{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" + ) + + sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) + rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) + write_modes_list = getattr(args, "write_modes_list", [args.write_checkpoint]) + modes_list = getattr(args, "modes_list", ["persistent_dynamic"]) + sort_list = getattr(args, "sort_slots_list", [False]) + rev_list = getattr(args, "reverse_nowrite_list", [False]) + hsort_list = getattr(args, "hardcode_sort_list", [False]) + + # Pre-load AL distribution for mix mode (if --mix-csv set). + mix_al = None + mix_label = "" + if args.mix_csv is not None: + from pathlib import Path as _Path + from checkpoint_mix_sim import load_al_distribution as _load_al + mix_label = _Path(args.mix_csv).stem + # T (= mtp_len) varies per cell; load once with the LARGEST mtp so + # we have enough columns; the loader normalizes the dist anyway. + mix_al = _load_al(_Path(args.mix_csv), T=max(mtp_lengths), column=args.mix_csv_column) + + for batch in batch_sizes: + for mtp_len in mtp_lengths: + # Resolve prev_k fractions → clamped integers in [0, mtp_len] + prev_ks = _resolve_prev_ks(args, mtp_len) + + # Pre-generate mix samples once per (batch, mtp_len) cell so all + # tuning configs see the same per-iter prev_tokens vectors — + # tuning differences become signal, mix-noise is shared. + # Size the sample buffer for the LARGER of args.iters and + # args.mix_iters since mix scenarios use mix_iters. + mix_samples_cpu = None + perm_samples_cpu = None # per-iter slot perm sorted write-first + mix_samples_sorted_cpu = None # per-iter prev_tokens, write-first + if mix_al is not None: + from checkpoint_mix_sim import sample_steady_state_pnat as _sample_pnat + _max_window = getattr(args, "max_window", 0) or mtp_len + _max_iters = max(args.iters, getattr(args, "mix_iters", None) or args.iters) + mix_samples_cpu = _sample_pnat( + mix_al, T=mtp_len, window=_max_window, batch=batch, + K=args.warmup + _max_iters, seed=args.mix_seed, + ) + if any(sort_list) or any(hsort_list): + # write-first stable argsort: kind='stable' preserves + # original-slot order within each mode group. + write_mask = ( + mix_samples_cpu + mtp_len > _max_window + ).astype(np.int8) # 1 = write, 0 = nowrite + perm_idx = np.argsort( + -write_mask, kind="stable", axis=-1 + ).astype(np.int32) + if any(sort_list): + perm_samples_cpu = perm_idx + if any(hsort_list): + # Apply the perm to the prev_tokens samples themselves. + # Result row i = mix_samples_cpu[i] reordered such + # that write-mode entries come first. + mix_samples_sorted_cpu = np.take_along_axis( + mix_samples_cpu, perm_idx, axis=-1 + ).astype(mix_samples_cpu.dtype) + + for state_dtype in state_dtypes: + for act_dtype in act_dtypes: + for sr_mode in sr_modes_list: + for mode in modes_list: + # Persistent modes ignore write_checkpoint + # (per-slot from PNAT); fix to True. + for write_ckpt in [True]: + for rect in rect_list: + # Sort/reverse only meaningful when a + # mix scenario is present (the actual + # sort experiment). Both persistent + # modes consume slot_perm. + can_sort = mix_samples_cpu is not None + effective_sort_list = ( + sort_list if can_sort else [False] + ) + effective_hsort_list = ( + hsort_list if can_sort else [False] + ) + for sort_slots in effective_sort_list: + for hardcode_sort in effective_hsort_list: + # sort_slots and hardcode_sort + # are alternative experiments + # for the same idea — skip the + # combined cell to avoid double + # interpretation. + if sort_slots and hardcode_sort: + continue + # rev=1 is meaningful with EITHER + # sort_slots=1 (perm-based) or + # hardcode_sort=1 (raw pid_b + # subtraction in unsorted-perm + # path). rev=1 with both 0 is + # a no-op. + effective_rev_list = ( + rev_list if (sort_slots or hardcode_sort) else [False] + ) + for reverse_nowrite in effective_rev_list: + _bench_config( + args, batch, mtp_len, + prev_ks, state_dtype, + act_dtype, baseline_fn, + sr_mode=sr_mode, + rectangle_for_nowrite=rect, + write_checkpoint=write_ckpt, + mode=mode, + mix_samples_cpu=mix_samples_cpu, + mix_label=mix_label, + sort_slots=sort_slots, + reverse_nowrite=reverse_nowrite, + perm_samples_cpu=perm_samples_cpu, + hardcode_sort=hardcode_sort, + mix_samples_sorted_cpu=mix_samples_sorted_cpu, + ) + + _drain_pending_results(args, force=True) + + if args.profile: + torch.cuda.cudart().cudaProfilerStop() + + # JSONL is the canonical artifact (written incrementally per cell with + # host stamps). No clean-exit `.json` write — use `jsonl_to_json.py` to + # materialize a snapshot when an analyzer wants one. + if args.json_output and args._jsonl_path is not None: + print(f"\nJSONL results: {args._jsonl_path} " + f"(meta sidecar: {args.json_output}.meta.json)") + + # Write the skipped-cells sidecar. Caller can convert this list to a + # --cell-list JSON (one dict per skipped cell) to drive a retry pass in + # a fresh process. + skipped_path = getattr(args, "skipped_output", None) + if skipped_path is None and args.json_output: + # Derive default: foo.json -> foo.skipped.json + skipped_path = args.json_output.rsplit(".", 1)[0] + ".skipped.json" + if skipped_path is not None and args._skipped_cells: + payload = { + "metadata": { + "timestamp": datetime.now().isoformat(), + "cmd": " ".join(sys.argv), + "skipped_count": len(args._skipped_cells), + }, + "skipped": args._skipped_cells, + } + tmp = skipped_path + ".tmp" + with open(tmp, "w") as f: + json.dump(payload, f, indent=2) + os.replace(tmp, skipped_path) + print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " + f"tags written to: {skipped_path}", file=sys.stderr) + elif args._skipped_cells: + # No output path but there are skipped cells — emit a stderr summary. + print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " + f"first 5: {args._skipped_cells[:5]}", file=sys.stderr) + + +# CLI + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Benchmark replay_selective_state_update Triton kernel", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--nheads", + type=int, + default=NHEADS, + help="Full-model nheads (divided by --tp-size for per-GPU slice)", + ) + parser.add_argument( + "--ngroups", + type=int, + default=NGROUPS, + help="Full-model ngroups (divided by --tp-size for per-GPU slice)", + ) + parser.add_argument( + "--head-dim", type=int, default=HEAD_DIM, help="Head dimension (not TP-split)" + ) + parser.add_argument( + "--d-state", type=int, default=D_STATE, help="SSM state dimension (not TP-split)" + ) + parser.add_argument( + "--tp-size", + type=int, + default=TP_SIZE, + help="Tensor parallel size; divides nheads and ngroups", + ) + parser.add_argument( + "--batch-sizes", default="1,2,4,8", help="Comma-separated decode batch sizes" + ) + parser.add_argument( + "--mtp-lengths", + default="1,2,4,8", + help="Comma-separated per-request sequence lengths (num_draft_tokens + 1 target)", + ) + parser.add_argument( + "--state-dtypes", + default="fp32", + help="Comma-separated state dtypes: fp16,bf16,fp32,int8,int16,fp8. " + "Quantized dtypes (int8/int16/fp8) require the checkpointing variant " + "and skip baselines (selective_state_update doesn't accept them).", + ) + parser.add_argument( + "--act-dtypes", + default="bf16", + help="Comma-separated activation dtypes for x/B/C/dt: fp32,bf16", + ) + parser.add_argument("--warmup", type=int, default=4, + help="Number of warmup iterations. Default aligns with " + "the graph group-iters (default 4 for mix scenarios) so " + "warmup + iters / mix-iters lands on a clean multiple " + "without per-args rounding overhead. Earlier default of " + "20 was overkill for steady-state warming.") + parser.add_argument("--iters", type=int, default=100, help="Number of timed iterations") + parser.add_argument( + "--compile-threads", + type=int, + default=64, + help="Number of THREADS used in the compile-warmup phase (one call " + "per (batch, mtp_len, prev_k, dtype, sweep) cell, parallelized over " + "N threads). Triton compile releases the GIL, so threads compile " + "in parallel and populate the persistent cache for free hits during " + "the sequential timed phase. 0 disables the phase. Default 64.", + ) + parser.add_argument( + "--mp-start-method", + choices=("spawn", "forkserver"), + default="spawn", + help="multiprocessing start method for compile-warmup workers AND " + "the CUPTI parser child process. 'spawn' (default) is robust but " + "each child re-imports the bench module (~15s torch+triton import " + "cost). 'forkserver' starts a server once, preloads the bench " + "module ONCE, then forks children cheaply (~1s each). When 4 " + "benches run concurrently with --compile-threads 26 each, spawn " + "still incurs 4*26=104 imports per round; forkserver cuts this to " + "4 (one per server).", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Wrap timed region in cudaProfilerStart/Stop (for ncu --target-processes all)", + ) + parser.add_argument( + "--l2-flush", + action=argparse.BooleanOptionalAction, + default=True, + help="L2 eviction between iterations", + ) + parser.add_argument( + "--cuda-graph", + action=argparse.BooleanOptionalAction, + default=True, + help="Capture all warmup + timed iterations in a " + "single CUDA graph with per-iteration events " + "inside the graph, eliminating all host overhead.", + ) + parser.add_argument( + "--cuda-graph-group-iters", + type=int, + default=None, + help="Capture this many logical benchmark iterations per graph " + "replay when warmup + iters is divisible by this value. Default " + f"auto-selects {_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE} for pure " + f"cells and {_DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX} for mix cells. " + "Mix cells use a per-replay device window so they can group " + "iterations too.", + ) + parser.add_argument( + "--cupti", + action=argparse.BooleanOptionalAction, + default=True, + help="Time kernels via CUPTI Activity API (1 ns from the GPU " + "profiling fabric); per-iter span = max(kernel_end) - " + "min(kernel_start). Default ON. --no-cupti disables in-bench " + "timing entirely (kernels still run, but median/p95/p99 are zero) " + "— use when wrapping the bench in nsys/ncu, where the external " + "profiler provides timings and our CUPTI subscriber would conflict.", + ) + parser.add_argument( + "--cupti-flush-period-ms", + type=int, + default=0, + help="If >0, ask CUPTI to periodically flush activity buffers during " + "the timed CUDA-graph region. This can overlap raw-buffer parsing with " + "long timed cells; 0 leaves flushing explicit at the end of each cell.", + ) + parser.add_argument( + "--cupti-defer-depth", + type=int, + default=4, + help="Maximum number of CUDA-graph CUPTI timing results that may be " + "left for the parser process while the main process starts later cells. " + "1 preserves synchronous per-cell parsing and inline retry behavior.", + ) + parser.add_argument( + "--json-output", + default=None, + help="If set, write per-cell results to this JSON file in the " + "shape consumed by collect.py / report.py. See the 'JSON output " + "schema' section at the top of this file.", + ) + parser.add_argument( + "--json-detailed", + action=argparse.BooleanOptionalAction, + default=False, + help="When --json-output is set, also include per_kernel " + "(per-iter relative start/end timestamps for each kernel). Compact " + "JSON always includes iters_us, and mix rows include " + "n_writes_per_iter. Default off keeps records compact.", + ) + parser.add_argument( + "--host-timing", + action=argparse.BooleanOptionalAction, + default=False, + help="Attach benchmark host-side phase timings to JSON/JSONL results. " + "Useful for diagnosing benchmark overhead, but it adds roughly 1 KB " + "per compact JSONL row and several perf_counter calls per cell.", + ) + parser.add_argument( + "--cupti-retry", + type=int, + default=1, + help="On CUPTI capture mismatch (kernel record count != expected), " + "retry the cell this many times in-process before giving up. CUPTI " + "gets racy after thousands of cells in one process (PDL + small " + "kernels occasionally lose records); a single retry usually catches " + "transient cases. Set 0 to disable and skip on first mismatch.", + ) + parser.add_argument( + "--skipped-output", + default=None, + help="Path to write the list of cells that failed CUPTI capture even " + "after --cupti-retry retries (JSON list of sweep_tag strings). " + "Default: derived from --json-output by replacing .json with " + ".skipped.json.", + ) + parser.add_argument( + "--cell-list", + default=None, + help="Path to a JSON list of cell dicts (one per cell to time). " + "Each dict has canonical knob keys → values: Mw, Mnw, Ww, Wnw, Sw, " + "Snw, pW, pS, H, R, CT, CPSw, CPSnw, LSw, LSnw, FL, WS, TMARL, " + "TMAWL, TMANL, TMAWS, SR, RECT, WC, MODE, SORT, REVN, HSORT (tied " + "forms M / W / S / CPS / LS are also accepted and auto-expanded). " + "When set, bench's CLI knob ranges are auto-overridden to the " + "per-knob union across all cells, and the inner-loop filter skips " + "any iteration whose knob-value tuple isn't in the list. All cells " + "must share the same key set (uniform schema).", + ) + parser.add_argument( + "--prev-tokens-fracs", + default="0,0.5,1.0", + type=lambda s: [float(x) for x in s.split(",")], + help="Fractions of mtp_len to use as prev_num_accepted_tokens " + "for the replay kernel sweep. Values are rounded " + "and clamped to [0, mtp_len].", + ) + parser.add_argument( + "--baseline", + default=None, + nargs="?", + const="triton", + choices=[None, "triton", "flashinfer"], + help="Baseline to benchmark alongside the replay kernel. " + "'triton': native Triton selective_state_update. " + "'flashinfer': flashinfer selective_state_update (same signature). " + "Pass --baseline alone for 'triton'. Default: no baseline.", + ) + parser.add_argument( + "--output", + default=None, + help="Path to save results (file or directory). " + "If a directory, writes benchmark_replay_.txt inside it.", + ) + parser.add_argument( + "--block-size-m", + type=str, + default=None, + help="Override BLOCK_SIZE_M: single value or comma-separated sweep (e.g. '4,8,16,32').", + ) + parser.add_argument( + "--num-warps", + type=str, + default=None, + help="Override num_warps: single value or comma-separated sweep (e.g. '1,2,4').", + ) + parser.add_argument( + "--internal-pdl", + action=argparse.BooleanOptionalAction, + default=True, + help="Internal PDL between precompute and main kernels (default: on).", + ) + parser.add_argument( + "--num-stages", + type=str, + default=None, + help="Override num_stages for the main kernel (comma-separated sweep).", + ) + parser.add_argument( + "--block-size-m-write", type=str, default=None, + help="Sweep BLOCK_SIZE_M for the WRITE main only (overrides --block-size-m " + "for the write half). Tied to --block-size-m if unset.", + ) + parser.add_argument( + "--block-size-m-nowrite", type=str, default=None, + help="Sweep BLOCK_SIZE_M for the NOWRITE main only. Tied to --block-size-m if unset.", + ) + parser.add_argument( + "--num-warps-write", type=str, default=None, + help="Sweep num_warps for the WRITE main only. Tied to --num-warps if unset.", + ) + parser.add_argument( + "--num-warps-nowrite", type=str, default=None, + help="Sweep num_warps for the NOWRITE main only. Tied to --num-warps if unset.", + ) + parser.add_argument( + "--num-stages-write", type=str, default=None, + help="Sweep num_stages for the WRITE main only. Tied to --num-stages if unset.", + ) + parser.add_argument( + "--num-stages-nowrite", type=str, default=None, + help="Sweep num_stages for the NOWRITE main only. Tied to --num-stages if unset.", + ) + parser.add_argument( + "--cta-per-sm-write", type=str, default=None, + help="Sweep cta_per_sm for the WRITE persistent_main only. Tied to --cta-per-sm if unset.", + ) + parser.add_argument( + "--cta-per-sm-nowrite", type=str, default=None, + help="Sweep cta_per_sm for the NOWRITE persistent_main only. Tied to --cta-per-sm if unset.", + ) + parser.add_argument( + "--num-loop-stages-write", type=str, default=None, + help="Sweep num_loop_stages for the WRITE persistent_main only. Tied to --num-loop-stages if unset.", + ) + parser.add_argument( + "--num-loop-stages-nowrite", type=str, default=None, + help="Sweep num_loop_stages for the NOWRITE persistent_main only. Tied to --num-loop-stages if unset.", + ) + parser.add_argument( + "--skip-diagonal", action=argparse.BooleanOptionalAction, default=False, + help="When sweeping any per-main *_write / *_nowrite knobs, skip cells " + "where ALL splittable knobs satisfy write_value == nowrite_value (i.e. " + "the 'diagonal' that's already covered by a prior shared-knob sweep). " + "Useful for incremental sweeps that extend earlier results without redoing " + "the tied-knob cells.", + ) + parser.add_argument( + "--precompute-num-warps", + type=str, + default=None, + help="Override num_warps for precompute kernel (comma-separated sweep).", + ) + parser.add_argument( + "--precompute-num-stages", + type=str, + default=None, + help="Override num_stages for precompute kernel (comma-separated sweep).", + ) + parser.add_argument( + "--max-window", + type=int, + default=16, + help="Cache T-axis capacity (max replay buffer length). Default 16 " + "matches Nemotron-3-Super-120B production. Pass 0 to fall back to " + "mtp_len (degenerate every-step-checkpoint case, mostly unused).", + ) + parser.add_argument( + "--prev-tokens-int", + type=lambda s: [int(x) for x in s.split(",")] if s else None, + default=None, + help="Absolute prev_num_accepted_tokens values to test, comma-separated " + "(e.g. '0,10,11,16'). Clamped to [0, max_window]. When set, " + "overrides --prev-tokens-fracs.", + ) + parser.add_argument( + "--write-checkpoint", + action=argparse.BooleanOptionalAction, + default=True, + help="Whether the checkpointing kernel should write the post-replay " + "state to HBM. True = checkpoint step (default). False = " + "non-checkpoint step (skip state HBM write + Philox). No effect on " + "the replay variant. Ignored if --write-modes is set.", + ) + parser.add_argument( + "--write-modes", + type=str, + default=None, + help="Comma-separated 0/1 values to sweep both write modes in a " + "single nsys process — for apples-to-apples comparison of write " + "vs nowrite (replay) vs nowrite (rectangle) within one timeline. " + "Skips silently for (write=False, prev_k+T>max_window) combos. " + "When set, overrides --write-checkpoint.", + ) + parser.add_argument( + "--with-conv1d", + action="store_true", + help="Include conv1d kernel before replay SSM. " + "Uses realistic L2 flush: cold caches flushed, hot in_proj output " + "kept warm. Measures conv1d → precompute → main span.", + ) + parser.add_argument( + "--external-pdl", + action=argparse.BooleanOptionalAction, + default=True, + help="External PDL: conv1d launches dependents, precompute waits. " + "Only relevant with --with-conv1d. --no-external-pdl disables.", + ) + parser.add_argument( + "--heads-per-block", + type=str, + default=None, + help="Override HEADS_PER_BLOCK for precompute kernel (comma-separated sweep).", + ) + parser.add_argument( + "--maxnreg", + type=str, + default=None, + help="Override maxnreg for the main kernel (comma-separated sweep).", + ) + parser.add_argument( + "--num-ctas", + type=str, + default=None, + help="Override num_ctas for the main kernel (comma-separated sweep).", + ) + parser.add_argument( + "--cta-per-sm", + type=str, + default=None, + help="CTAs per SM in the 1D persistent grid for mode=persistent_main " + "(comma-separated sweep). num_persistent = cta_per_sm × NUM_SMS. " + "Default = 1 (one CTA per SM). Replaces the old --num-persistent. " + "Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--num-loop-stages", + type=str, + default=None, + help="num_stages on the inner tl.range(...) persistent loop for " + "mode=persistent_main (comma-separated sweep). Default = 2. Note: " + "this is loop-level, NOT the kernel-arg num_stages (which only " + "pipelines dot-feeding loads). Watch Triton issue #8259 — " + "num_stages>1 + flatten=True can corrupt stores in non-dot kernels. " + "Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--flatten", + type=str, + default=None, + help="`flatten` arg on tl.range(...) for mode=persistent_main " + "(comma-separated 0/1 sweep). Default = 1. Ignored for " + "non-persistent_main modes.", + ) + parser.add_argument( + "--warp-specialize", + type=str, + default=None, + help="`warp_specialize` arg on tl.range(...) for mode=persistent_main " + "(comma-separated 0/1 sweep). Default = 0. Triton 3.6 only " + "supports it on simple matmul loops; our scan loop probably won't " + "pattern-match — exposed as a knob for sweep experiments. Requires " + "num_warps >= 4 if 1. Ignored for non-persistent_main modes.", + ) + parser.add_argument( + "--sr-modes", + type=str, + default="RN", + help="Comma-separated rounding modes to sweep: any combination of " + "{RN, SR}. SR (stochastic rounding) is silently skipped for state " + "dtypes that don't support it (bf16, fp32). Default 'RN' matches " + "legacy --philox-rounding=False behavior.", + ) + parser.add_argument( + "--rectangle-for-nowrite", + type=str, + default="0", + help="Comma-separated 0/1 values: 0 = replay-style nowrite kernel, " + "1 = dedicated rectangle nowrite kernel. Sweep both with '0,1' to " + "compare in one invocation. Silently no-op for write cells (the " + "write path always uses replay-style). Only applies to the " + "checkpointing variant.", + ) + parser.add_argument( + "--use-tma-rect-load", + type=str, + default=None, + help="Comma-separated 0/1 sweep. Use TMA (host-built tensor " + "descriptor) for state load in the rectangle nowrite path. " + "Cells where the rect path isn't reachable (rectangle_for_nowrite=False) " + "skip the value=1 case as a dupe.", + ) + parser.add_argument( + "--use-tma-replay-write-load", + type=str, + default=None, + help="Comma-separated 0/1 sweep. TMA state LOAD in replay main " + "when WC=True. Independent from nowrite-load and rect TMA — see " + "CHECKPOINTING_DESIGN.md item #17 for measured perf.", + ) + parser.add_argument( + "--use-tma-replay-nowrite-load", + type=str, + default=None, + help="Comma-separated 0/1 sweep. TMA state LOAD in replay main " + "when WC=False. Design doc reports the largest win on this path " + "(int8 b>=64: -8 to -12%%).", + ) + parser.add_argument( + "--use-tma-replay-write-store", + type=str, + default=None, + help="Comma-separated 0/1 sweep. TMA state STORE in replay main " + "(WC=True path only — no-op for WC=False). Independent from all " + "load TMA flags.", + ) + parser.add_argument( + "--modes", + type=str, + default="persistent_dynamic", + help="Comma-separated dispatch modes to sweep, any of " + "{persistent_dynamic, persistent_main}. " + "persistent_dynamic = single persistent-CTA kernel that dispatches " + "per-slot at runtime based on PNAT. " + "persistent_main = persistent-CTA kernel with two halves (write + " + "nowrite), requires caller-provided _n_writes and a write-first " + "sorted slot_perm. Both modes ignore --write-modes (per-slot from PNAT).", + ) + parser.add_argument( + "--mix-csv", + type=str, + default=None, + help="Path to AL histogram CSV (cols: AL, count). When set, an " + "additional 'mix' cell is emitted per (batch, mtp, dtype, sr, " + "mode, RECT, M, W, ...) combo where prev_tokens varies per iter, " + "drawn from the steady-state PNAT distribution induced by the " + "AL histogram. Both persistent modes support mix scenarios. " + "Each iteration of the captured CUDA graph has a different " + "pre-baked prev_tokens vector; warmup iters use distinct samples " + "from the timed iters so nsys-included warmup leaks don't bias.", + ) + parser.add_argument( + "--mix-csv-column", + type=int, + default=1, + help="Column index (0-based) in the AL histogram CSV for the " + "count/probability column. Default 1 (second column).", + ) + parser.add_argument( + "--mix-seed", + type=int, + default=42, + help="RNG seed for the steady-state PNAT sampler. Same seed " + "across runs => same per-slot samples for reproducible " + "comparisons.", + ) + parser.add_argument( + "--sort-slots", + type=str, + default="0", + help="Comma-separated 0/1. When 1, mix scenarios pre-sort slots " + "write-first (write slots at the head of slot_perm, nowrite at the " + "tail) and the persistent kernels read pid_b through that perm — " + "clusters early-outs at one end of the grid. Pure-batch cells " + "skip sort=1 (no mix to sort).", + ) + parser.add_argument( + "--reverse-nowrite", + type=str, + default="0", + help="Comma-separated 0/1. When 1 (and --sort-slots 1), the " + "nowrite-side persistent kernel walks the perm in reverse so the " + "two halves front-load real work from both ends. reverse=1 with " + "sort=0 is skipped (no perm to reverse).", + ) + parser.add_argument( + "--hardcode-sort", + type=str, + default="0", + help="Comma-separated 0/1. When 1, the per-iter prev_tokens " + "samples are pre-sorted write-first OFFLINE (CPU-side) before " + "the timed region — kernel runs unchanged (USE_PERM=False) but " + "the EO gate sees sorted PNAT so early-outs cluster naturally. " + "Zero per-program load cost vs --sort-slots; output is " + "scrambled (we don't permute x/B/C/dt) but timing is meaningful. " + "Used to isolate whether clustering helps independent of the " + "perm-load overhead in the sort-slots path.", + ) + parser.add_argument( + "--mix-iters", + type=int, + default=None, + help="Iteration count override for mix scenarios (each iter is a " + "different per-slot prev_tokens draw). Default (None) uses " + "--iters. Mix scenarios benefit from more iters since each " + "iter samples a different mix; pure scenarios don't.", + ) + parser.add_argument( + "--mix-only", + action=argparse.BooleanOptionalAction, + default=False, + help="When --mix-csv is set, emit only mix scenarios and skip the " + "pure prev_k sibling scenarios. Default: false.", + ) + parser.add_argument( + "--philox-rounding", + action="store_true", + help="DEPRECATED — equivalent to --sr-modes SR. Retained for " + "backward compatibility; use --sr-modes for new scripts. fp16 SR " + "and fp8 SR require sm_100a (Blackwell B200+).", + ) + parser.add_argument( + "--philox-rounds", + type=int, + default=5, + help="Number of Philox PRNG rounds. Default 5 matches the " + "Nemotron-3-Super-120B production config (mamba_ssm_philox_rounds=5 " + "in examples/configs and tests/integration/perf configs). The " + "wrapper's generic fallback default is 10; callers without explicit " + "config see 10. Only consulted when --philox-rounding is enabled.", + ) + parser.add_argument( + "--variant", + choices=["replay", "checkpointing"], + default="replay", + help="Which kernel to time as the 'replay' row. 'replay' = today's " + "kernel (selective_state_update.py:replay). 'checkpointing' = " + "checkpointing_state_update.py. Both share the same wrapper signature.", + ) + parser.add_argument( + "--full-import", + action="store_true", + help="Use standard tensorrt_llm import path instead of fast direct " + "module loading. Slower (~40s startup) but guaranteed correct " + "if the fast path breaks due to package changes.", + ) + args = parser.parse_args() + if args.mix_only and args.mix_csv is None: + parser.error("--mix-only requires --mix-csv") + + # Round iter counts up so warmup + iters (and warmup + mix_iters) are clean + # multiples of the graph group-iters used downstream. Default mix group is + # 4, default pure group is 2. An explicit --cuda-graph-group-iters can + # request a larger group. We round to the max of the two so all scenarios + # in a single run (pure + mix) share a clean total_iters. The cost is at + # most (group-1) extra iters per scenario — negligible — and the win is + # that graph_group_iters never falls back to 1 (which caused ~5x slowdown + # in observed benchmark walls). + _group_for_rounding = max( + _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX, + _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE, + getattr(args, "cuda_graph_group_iters", None) or 0, + ) + def _round_iters_to_group(name, val): + total = args.warmup + val + if total % _group_for_rounding == 0: + return val + new_total = ((total + _group_for_rounding - 1) // _group_for_rounding) * _group_for_rounding + new_val = new_total - args.warmup + print(f"[bench] rounding --{name} {val} → {new_val} so warmup+{name} " + f"({new_total}) is a multiple of graph group_iters={_group_for_rounding}", + file=sys.stderr) + return new_val + args.iters = _round_iters_to_group("iters", args.iters) + if getattr(args, "mix_iters", None): + args.mix_iters = _round_iters_to_group("mix-iters", args.mix_iters) + + # Cell-list (if any) must be applied BEFORE the post-argparse string→list + # derivations below — those build args.*_list from args.* strings, so a + # cell-list override of e.g. args.modes='persistent_main' needs to land + # before args.modes_list is computed. The function populates args._cell_list_keys + # and args._cell_list_set, plus overrides args.* knob strings to the + # per-knob union of values across the listed cells. + if getattr(args, "cell_list", None): + _load_cell_list_into_args(args) + + # Backward-compat: --philox-rounding implies --sr-modes SR if --sr-modes + # was left at the default. If both are set explicitly, error. + sr_modes_default = (args.sr_modes == "RN") + if args.philox_rounding: + if not sr_modes_default and args.sr_modes != "SR": + parser.error( + "--philox-rounding (deprecated) is incompatible with explicit " + f"--sr-modes={args.sr_modes!r}. Use --sr-modes SR (or " + "RN,SR) instead and drop --philox-rounding." + ) + args.sr_modes = "SR" + + sr_modes = [m.strip() for m in args.sr_modes.split(",") if m.strip()] + for m in sr_modes: + if m not in ("RN", "SR"): + parser.error(f"--sr-modes value must be RN or SR, got {m!r}") + args.sr_modes_list = sr_modes + + rect_modes = [v.strip() for v in args.rectangle_for_nowrite.split(",") if v.strip()] + rect_list = [] + for v in rect_modes: + if v not in ("0", "1"): + parser.error(f"--rectangle-for-nowrite value must be 0 or 1, got {v!r}") + rect_list.append(v == "1") + args.rectangle_for_nowrite_list = rect_list + + sort_modes = [v.strip() for v in (args.sort_slots or "0").split(",") if v.strip()] + sort_list = [] + for v in sort_modes: + if v not in ("0", "1"): + parser.error(f"--sort-slots value must be 0 or 1, got {v!r}") + sort_list.append(v == "1") + args.sort_slots_list = sort_list + + rev_modes = [v.strip() for v in (args.reverse_nowrite or "0").split(",") if v.strip()] + rev_list = [] + for v in rev_modes: + if v not in ("0", "1"): + parser.error(f"--reverse-nowrite value must be 0 or 1, got {v!r}") + rev_list.append(v == "1") + args.reverse_nowrite_list = rev_list + + hsort_modes = [v.strip() for v in (args.hardcode_sort or "0").split(",") if v.strip()] + hsort_list = [] + for v in hsort_modes: + if v not in ("0", "1"): + parser.error(f"--hardcode-sort value must be 0 or 1, got {v!r}") + hsort_list.append(v == "1") + args.hardcode_sort_list = hsort_list + + if args.write_modes is not None: + wm = [v.strip() for v in args.write_modes.split(",") if v.strip()] + write_list = [] + for v in wm: + if v not in ("0", "1"): + parser.error(f"--write-modes value must be 0 or 1, got {v!r}") + write_list.append(v == "1") + args.write_modes_list = write_list + else: + args.write_modes_list = [args.write_checkpoint] + + modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] + valid_modes = { + "persistent_main", "persistent_dynamic", + } + for m in modes_raw: + if m not in valid_modes: + parser.error( + f"--modes value must be one of {sorted(valid_modes)}, got {m!r}" + ) + args.modes_list = modes_raw or ["persistent_dynamic"] + return args + + +class _Tee: + """Write to both stdout and a file simultaneously.""" + + def __init__(self, path: str): + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + self._file = open(path, "w") # noqa: SIM115 + self._stdout = sys.stdout + + def write(self, data): + self._stdout.write(data) + self._file.write(data) + + def flush(self): + self._stdout.flush() + self._file.flush() + + def close(self): + self._file.close() + + +if __name__ == "__main__": + _args = _parse_args() + + # Configure multiprocessing start method early — must be before any + # mp.get_context() that uses the chosen method. For forkserver, also + # add this file's dir to sys.path so the forkserver can import this + # module by basename for preload (otherwise it tries to import + # __main__, which is a different beast across processes). + if _args.mp_start_method == "forkserver": + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + mp.set_start_method("forkserver", force=True) + try: + mp.set_forkserver_preload([ + "benchmark_replay_selective_state_update", + ]) + except Exception as _e: + print(f"[warn] set_forkserver_preload failed: {_e!r}; " + f"forks will still work but pay full import cost", + file=sys.stderr) + _MP_START_METHOD = _args.mp_start_method + + _out_path = None + if _args.output != "-": + _ts = datetime.now().strftime("%Y%m%d_%H%M%S") + _fname = f"benchmark_replay_{_ts}.txt" + if _args.output is None: + _out_path = os.path.expanduser(f"~/nemo_logs/{_fname}") + elif os.path.isdir(_args.output) or _args.output.endswith("/"): + _out_path = os.path.join(_args.output, _fname) + else: + _out_path = _args.output + + if _out_path: + _tee = _Tee(_out_path) + sys.stdout = _tee + print(f"# benchmark_replay_selective_state_update {datetime.now().isoformat()}") + print(f"# cmd: {' '.join(sys.argv)}") + + try: + _run_benchmark(_args) + finally: + if _out_path: + sys.stdout = _tee._stdout + _tee.close() + print(f"\nResults saved to: {_out_path}") diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_slim.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_slim.py new file mode 100644 index 000000000000..a8d11a254331 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_slim.py @@ -0,0 +1,1915 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +import math + +import pytest +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import repeat + +from tensorrt_llm._torch.modules.mamba.checkpointing_state_update_slim import ( + _stochastic_round_int8_packed, + _stochastic_round_int16_packed, + checkpointing_state_update, +) +from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update +from tensorrt_llm._utils import get_sm_version + +# Philox stochastic rounding uses PTX cvt.rs.f16x2.f32 which requires sm >= 100. +_skip_pre_sm100 = pytest.mark.skipif( + get_sm_version() < 100, reason="Philox stochastic rounding needs sm >= 100" +) + +# Configs derived from NVIDIA-Nemotron-3-Super-120B-A12B Mamba2 parameters +# (nheads=128, headdim=64, d_state=128, ngroups=8) with TP split applied: +# TP=8: nheads=16, ngroups=1 — primary production config +# TP=4: nheads=32, ngroups=2 — exercises ngroups>1 (grouped B/C path) +_CONFIGS = [ + # (nheads, head_dim, d_state, ngroups) + (16, 64, 128, 1), # TP=8 production config + (32, 64, 128, 2), # TP=4, ngroups>1 (more heads than B/C groups) +] + +# Quantized state dtypes and their representable-magnitude limits (== QUANT_MAX +# in the kernel). fp8_e4m3fn cells require SM 89+ for the fp32↔fp8 cvt PTX +# instructions; SR variants of fp16/fp8 additionally need SM 100+. +_QUANT_MAX_BY_DTYPE = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, +} + + +def _quantize_state(state_fp32: torch.Tensor, state_dtype: torch.dtype, quant_max: float): + """Quantize fp32 state to (state_quant, decode_scale) using the same + per-(head, dim) channel scheme the kernel does on store. decode_scale = + max_abs_per_channel / quant_max (= 1/encode_scale). + """ + amax = state_fp32.abs().amax(dim=-1) # (cache, nheads, head_dim) + encode_scale = quant_max / amax.clamp(min=1e-30) + decode_scale = 1.0 / encode_scale + scaled = state_fp32 * encode_scale.unsqueeze(-1) + if state_dtype == torch.float8_e4m3fn: + # Native cast does RN at the fp8 grid; explicit round() would destroy + # sub-integer precision (matches the kernel's fp8 RN path). + state_quant = scaled.clamp(-quant_max, quant_max).to(state_dtype) + else: + state_quant = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) + return state_quant, decode_scale + + +def _dequantize_state(state_quant: torch.Tensor, decode_scale: torch.Tensor): + return state_quant.to(torch.float32) * decode_scale.unsqueeze(-1) + + +def _maybe_skip_dtype(state_dtype, use_sr): + """Skip on insufficient SM. fp8 e4m3fn (any) needs SM 89+; fp16/fp8 SR + needs SM 100+; int8/int16 (RN or SR) runs anywhere.""" + if state_dtype == torch.float8_e4m3fn and get_sm_version() < 89: + pytest.skip("fp8_e4m3fn requires SM 89+ (Ada Lovelace / Hopper / Blackwell)") + if use_sr and state_dtype in (torch.float16, torch.float8_e4m3fn) and get_sm_version() < 100: + pytest.skip(f"{state_dtype} stochastic rounding requires SM 100+ (Blackwell B200+)") + + +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize( + "state_dtype", + [ + torch.float16, + torch.bfloat16, + torch.float32, + torch.int8, + torch.int16, + torch.float8_e4m3fn, + ], + ids=["fp16", "bf16", "fp32", "int8", "int16", "fp8"], +) +@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) +@pytest.mark.parametrize( + "T", [6, 10, 16, 27, 32, 55], ids=["T6", "T10", "T16", "T27", "T32", "T55"] +) +@pytest.mark.parametrize( + "write_checkpoint,rectangle_for_nowrite", + [ + (True, False), # write path (rectangle_for_nowrite is ignored) + (False, False), # nowrite path via replay-style kernels + (False, True), # nowrite path via dedicated rectangle kernels + ], + ids=["write", "no_write_replay", "no_write_rectangle"], +) +@pytest.mark.parametrize( + "mode", + ["persistent_dynamic", "persistent_main"], + ids=["persistent_dynamic", "persistent_main"], +) +def test_checkpointing_state_update( + nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T, + write_checkpoint, rectangle_for_nowrite, mode, +): + """ + Verify that: + checkpointing_state_update(state0, old_caches, k, new_x, ...) + produces the same output as: + selective_state_update(state_after_k_old_tokens, new_x, ...) + and writes state_after_k_old_tokens back to the state tensor. + + Quantized state dtypes (int8/int16/fp8) follow the same flow with + a per-(head, dim) channel decode-scale tensor; comparison is done + via dequant(state, scales) against the fp32 reference. + """ + _maybe_skip_dtype(state_dtype, use_sr=False) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + + batch = 2 + device = "cuda" + dtype = torch.bfloat16 # input activations are bf16 + assert nheads % ngroups == 0 + + # Cache T-axis size (max_window). Use the kernel's BLOCK_SIZE_T as the + # ceiling — this is what the wrapper allows and enables PNAT-aware writes + # at [PNAT, PNAT+T) for no-checkpoint mode. For T=6 that's 16 (production + # max_window); for larger T it scales with np2(T). + max_window = max(triton.next_power_of_2(T), 16) + + if paged_cache: + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + else: + cache_size = batch + state_batch_indices = None + + torch.manual_seed(42) + + # A: (nheads, head_dim, d_state) with stride(-2)=0, stride(-1)=0 [tie_hdim] + A_base = -torch.rand(nheads, device=device) - 0.5 # float32, negative + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + + # dt_bias: (nheads, head_dim) with stride(-1)=0 [tie_hdim] + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + + # D: (nheads, head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + # Initial SSM state (cache_size slots). Quantized dtypes need a separate + # init: derive scales from a fp32 source so the quantized state isn't + # garbage on dequant. ref_input_state is what the fp32 reference run + # sees — for non-quant it's state0 (cast to fp32 inside reference); for + # quant it's the lossy dequant of state0 (matches what the kernel sees + # internally on load). + if is_quantized: + state0_fp32 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + state0, state0_scales = _quantize_state(state0_fp32, state_dtype, quant_max) + ref_input_state = _dequantize_state(state0, state0_scales) + else: + state0 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state0_scales = None + ref_input_state = state0.float() + + # Old inputs: up to `max_window` tokens per batch request, so the test + # loop can probe PNAT > T-1 (which the prior T-token setup couldn't + # reach). step1_T = max_window covers the full PNAT range we sweep. + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) # stride(-1)=0 + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + # Capture intermediate SSM states using selective_state_update across + # all step1_T positions — gives us reference states for k ∈ [0, step1_T]. + states_buffer_f32 = torch.zeros( + cache_size, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = ( + state_batch_indices + if paged_cache + else torch.arange(batch, device=device, dtype=torch.int32) + ) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, + dt1, + A, + B1, + C1, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + # Build cache tensors for the replay kernel. + # old_x: (cache, max_window, nheads, dim) bf16 — single-buffered + # old_B: (cache, 2, max_window, ngroups, dstate) bf16 — double-buffered + # old_dt: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous + # old_dA_cumsum: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous + # cache_buf_idx: random 0s and 1s to verify indexing correctness + old_x = torch.zeros(cache_size, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) + + # Fill each slot's active buffer (= cache_buf_idx) with step 1's data at + # positions [0:step1_T) = [0:max_window). Whole buffer covered so PNAT + # values up to max_window are exercised. Inactive buffer has random + # garbage to catch indexing bugs. + slots = state_batch_indices if paged_cache else slice(None) + old_x[slots, :step1_T] = x1 + + # Compute processed dt and dA_cumsum for step 1 + dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) + + # Write to each slot's active buffer based on its cache_buf_idx + slot_indices = state_batch_indices.tolist() if paged_cache else list(range(cache_size)) + for i, slot in enumerate(slot_indices): + buf = cache_buf_idx[slot].item() + batch_idx = i # maps slot back to the batch index + old_B[slot, buf, :step1_T] = B1[batch_idx] + old_dt[slot, buf, :, :step1_T] = dt1[batch_idx].T # (step1_T, nheads) → (nheads, step1_T) + old_dA_cumsum[slot, buf, :, :step1_T] = dA_cumsum1[batch_idx].T + + # Main loop: test each k (number of old tokens replayed). + # write_checkpoint=False (nowrite): k ∈ [0, max_window-T] — new tokens + # append at [k, k+T) of the active buffer; need k+T ≤ max_window. + # write_checkpoint=True (write): k ∈ [max_window-T+1, max_window] — + # new tokens land in the staging buffer at [0, T); k > max_window-T + # captures the overflow case that triggers a checkpoint in production. + # Combined sweep covers the full k ∈ [0, max_window] with the + # appropriate boundary handling per mode. + if write_checkpoint: + k_lo = max(0, max_window - T + 1) + k_hi = max_window + 1 # exclusive + else: + k_lo = 0 + k_hi = max_window - T + 1 # exclusive + for k in range(k_lo, k_hi): + torch.manual_seed(k + 100) + + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + # Reference (fp32, starting from the same lossy-or-not state the + # kernel sees). + ref_state_f32 = ref_input_state.clone() + if k > 0: + ref_state_f32[slots] = states_buffer_f32[slots, k - 1] + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=(state_batch_indices if paged_cache else None), + out=ref_out, + ) + + # Replay kernel — clone caches into mutable working copies that we + # can inspect AFTER the call to verify cache postconditions. + test_state = state0.clone() + test_scales = state0_scales.clone() if is_quantized else None + prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_x_w = old_x.clone() + old_B_w = old_B.clone() + old_dt_w = old_dt.clone() + old_dA_cumsum_w = old_dA_cumsum.clone() + # cache_buf_idx stays at its random values — each slot reads from its own buffer + + # Persistent_main needs _n_writes (number of write slots) and a sort + # perm. For pure-write or pure-nowrite cases here, all slots have the + # same status, so _n_writes is batch or 0 and the perm is identity. + # Persistent_dynamic ignores both (kernel uses runtime PNAT dispatch). + pm_kwargs = ( + {"_n_writes": batch if write_checkpoint else 0} + if mode == "persistent_main" else {} + ) + checkpointing_state_update( + test_state, + old_x_w, + old_B_w, + old_dt_w, + old_dA_cumsum_w, + cache_buf_idx.clone(), + prev_tokens, + x=x2, + dt=dt2, + A=A, + B=B2, + C=C2, + out=test_out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + state_scales=test_scales, + write_checkpoint=write_checkpoint, + rectangle_for_nowrite=rectangle_for_nowrite, + mode=mode, + **pm_kwargs, + ) + + # Tolerance rationale: the replay kernel uses bf16 tl.dot for four + # matmuls (dB_scaled @ old_x, C @ state, CB_scaled @ x, and C @ B in + # precompute). The reference (selective_state_update) and flashinfer + # baseline use fp32 element-wise MACs. The bf16 input casts lose the + # dt_bias/A-derived bits that the baselines keep — per-element rounding, + # not accumulating. Prefill (ssd_chunk_scan) does identical bf16 tl.dot + # casts, so we match prefill precision exactly. Empirical: max ~1.0 at + # T<=16, ~2.0 at T=32-55; mean ~0.014; <0.02% of elements exceed 0.5. + # State dtype (fp16/bf16/fp32) doesn't shift the error — bf16 dot + # inputs dominate, not state storage. + # + # Quantized states add a per-element state quant error eps that + # propagates through C @ state in the output dot. With dstate=128 + # and C ~ N(0,1), the output channel std from this noise is roughly + # eps * sqrt(128/3) ≈ 6.5 * eps. Stack with the bf16 baseline: + # out_atol = bf16_atol + 6.5 * eps_max + # where eps_max is the worst-case per-element error at the + # post-replay state magnitude (T=55 → amax ≈ 23). + # + # Per-element error (eps_max for T=55): + # int8 (uniform grid): amax/(2*127) ≈ 0.091 + # int16 (uniform grid): amax/(2*32767) ≈ 3.5e-4 + # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 (worst-case + # cell at top of channel; smaller for + # smaller-magnitude elements) + out_atol = ( + {torch.int8: 1.6, torch.int16: 1.05, torch.float8_e4m3fn: 4.0}[state_dtype] + if is_quantized else 1.0 + ) + out_rtol = ( + {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] + if is_quantized else 2e-2 + ) + out_diff = (test_out.float() - ref_out.float()).abs() + out_max = out_diff.max().item() + out_mean = out_diff.mean().item() + try: + torch.testing.assert_close( + test_out, ref_out, rtol=out_rtol, atol=out_atol, + msg=f"Output mismatch at k={k}", + ) + except AssertionError: + print( + f"k={k} out: max={out_max:.4f} mean={out_mean:.4f} " + f"nan={torch.isnan(test_out).any().item()} " + f"inf={torch.isinf(test_out).any().item()}" + ) + raise + + # State expectation depends on write_checkpoint: + # True → kernel writes the post-replay state; expect the + # selective_state_update reference's state at step k-1. + # False → kernel skips the HBM store; state must be UNCHANGED + # from the input (state0; for quant, scales also unchanged). + if is_quantized: + if write_checkpoint: + # Compare via dequant against the fp32 reference state. + expected_fp32 = ( + ref_input_state[slots] if k == 0 else states_buffer_f32[slots, k - 1] + ) + actual_fp32 = _dequantize_state(test_state[slots], test_scales[slots]) + # State diff = bf16_replay_error + quant_error (per element). + # The bf16 component is the SAME error source the non-quant + # test absorbs in its atol=1.0 baseline (replay's tl.dot is + # bf16-input fp32-accum; per-element error ~ 2^-7 * amax, + # empirically ≤ ~0.2 at T=55 amax≈23). Quant adds: + # int8: amax/(2*127) ≈ 0.091 worst-case + # int16: amax/(2*32767) ≈ 3.5e-4 (negligible vs bf16) + # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 worst-case + # Atol = bf16_baseline (1.0) + quant_eps_max. + state_atol = { + torch.int8: 1.1, torch.int16: 1.0, torch.float8_e4m3fn: 2.5, + }[state_dtype] + state_rtol = { + torch.int8: 5e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 1e-1, + }[state_dtype] + try: + torch.testing.assert_close( + actual_fp32, expected_fp32, + rtol=state_rtol, atol=state_atol, + msg=f"State mismatch at k={k} dtype={state_dtype}", + ) + except AssertionError: + diff = (actual_fp32 - expected_fp32).abs() + print( + f"k={k} state(dequant): max={diff.max().item():.4f} " + f"mean={diff.mean().item():.4f}" + ) + raise + # Scales sanity (fp32, finite, positive). + assert test_scales.dtype == torch.float32 + assert torch.isfinite(test_scales[slots]).all(), ( + f"state_scales has non-finite values at k={k}" + ) + assert (test_scales[slots] > 0).all(), ( + f"state_scales has non-positive values at k={k}" + ) + else: + # No write: raw quant state and scales unchanged. Use + # torch.equal for byte-level equality (dtype-agnostic; works + # for int8 / int16 / fp8 alike). + assert torch.equal(test_state[slots], state0[slots]), ( + f"Quant state changed at k={k} write_checkpoint=False" + ) + assert torch.equal(test_scales[slots], state0_scales[slots]), ( + f"State scales changed at k={k} write_checkpoint=False" + ) + else: + if write_checkpoint: + expected_state = ( + state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) + ) + else: + expected_state = state0[slots] + state_diff = (test_state[slots].float() - expected_state.float()).abs() + state_max = state_diff.max().item() + state_mean = state_diff.mean().item() + try: + torch.testing.assert_close( + test_state[slots], + expected_state, + rtol=2e-2, + atol=1.0 if write_checkpoint else 0.0, + msg=f"State mismatch at k={k} (write_checkpoint={write_checkpoint})", + ) + except AssertionError: + print( + f"k={k} state: max={state_max:.4f} mean={state_mean:.4f} " + f"nan={torch.isnan(test_state).any().item()} " + f"inf={torch.isinf(test_state).any().item()}" + ) + raise + + # --- Cache postconditions --- + # Compute step 2's processed values (what the kernel should have + # stored at [write_offset : write_offset+T) of write_buf): + # write_buf = (1 - active_buf) if write_checkpoint else active_buf + # write_offset = 0 if write_checkpoint else k + # Untouched cache regions must equal their pre-call snapshots + # (old_x / old_B / old_dt / old_dA_cumsum captured before the call). + dt2_proc = F.softplus(dt2_base.float() + dt_bias_base.float()[None, None, :]) # (B,T,H) + dA_cumsum2 = torch.cumsum(A_base.float()[None, None, :] * dt2_proc, dim=1) + write_offset = 0 if write_checkpoint else k + + for batch_idx, slot in enumerate(slot_indices): + active = cache_buf_idx[slot].item() + wb = (1 - active) if write_checkpoint else active + + # --- old_x (single-buffered): write at [write_offset : +T) of slot --- + written_x = old_x_w[slot, write_offset : write_offset + T] + torch.testing.assert_close( + written_x, x2[batch_idx], rtol=0, atol=0, + msg=f"old_x written region wrong at k={k} write={write_checkpoint}", + ) + # Untouched ranges of old_x[slot] + if write_offset > 0: + torch.testing.assert_close( + old_x_w[slot, :write_offset], old_x[slot, :write_offset], + rtol=0, atol=0, + msg=f"old_x [0:{write_offset}) modified at k={k} write={write_checkpoint}", + ) + if write_offset + T < max_window: + torch.testing.assert_close( + old_x_w[slot, write_offset + T:], old_x[slot, write_offset + T:], + rtol=0, atol=0, + msg=f"old_x [{write_offset+T}:) modified at k={k} write={write_checkpoint}", + ) + + # --- old_B (double-buffered): write at write_buf, [write_offset:+T) --- + torch.testing.assert_close( + old_B_w[slot, wb, write_offset : write_offset + T], + B2[batch_idx], rtol=0, atol=0, + msg=f"old_B written region wrong at k={k} write={write_checkpoint}", + ) + # Other-buffer (= 1-wb) untouched + torch.testing.assert_close( + old_B_w[slot, 1 - wb], old_B[slot, 1 - wb], + rtol=0, atol=0, + msg=f"old_B inactive buffer modified at k={k} write={write_checkpoint}", + ) + + # --- old_dt (double-buffered, fp32, layout (heads, T)): --- + torch.testing.assert_close( + old_dt_w[slot, wb, :, write_offset : write_offset + T], + dt2_proc[batch_idx].T, + rtol=1e-4, atol=1e-4, + msg=f"old_dt written region wrong at k={k} write={write_checkpoint}", + ) + torch.testing.assert_close( + old_dt_w[slot, 1 - wb], old_dt[slot, 1 - wb], + rtol=0, atol=0, + msg=f"old_dt inactive buffer modified at k={k} write={write_checkpoint}", + ) + + # --- old_dA_cumsum (double-buffered, fp32, layout (heads, T)): --- + torch.testing.assert_close( + old_dA_cumsum_w[slot, wb, :, write_offset : write_offset + T], + dA_cumsum2[batch_idx].T, + rtol=1e-4, atol=1e-4, + msg=f"old_dA_cumsum written region wrong at k={k} write={write_checkpoint}", + ) + torch.testing.assert_close( + old_dA_cumsum_w[slot, 1 - wb], old_dA_cumsum[slot, 1 - wb], + rtol=0, atol=0, + msg=f"old_dA_cumsum inactive buf modified at k={k} write={write_checkpoint}", + ) + + +@pytest.mark.parametrize( + "scenario,pnat_per_slot_list,n_writes_expected,slot_perm_list,rectangle_for_nowrite", + [ + # All-write: every slot has PNAT triggering write + # (PNAT + T > max_window). No permutation needed. + ("all_write", [12, 13, 14, 15], 4, [0, 1, 2, 3], False), + # All-nowrite: every slot fits in the window. n_writes = 0. + ("all_nowrite", [3, 4, 5, 6], 0, [0, 1, 2, 3], False), + # Mixed (write-first sorted via slot_perm): physical slots 2, 3 + # are writes; physical slots 0, 1 are nowrites. slot_perm + # remaps grid pid_b 0..3 to physical slots 2, 3, 0, 1 — so the + # first n_writes=2 grid programs hit write slots and the rest + # hit nowrite slots. + ("mixed_sorted", [3, 10, 12, 16], 2, [2, 3, 0, 1], False), + # Same as mixed_sorted but with rectangle_for_nowrite=True — the + # nowrite half of the persistent loop dispatches to the rectangle + # impl instead of replay-nowrite. Covers the rect-path correctness + # when mixed with the write half in a single persistent kernel. + ("mixed_sorted_rect", [3, 10, 12, 16], 2, [2, 3, 0, 1], True), + ], + ids=["all_write", "all_nowrite", "mixed_sorted", "mixed_sorted_rect"], +) +def test_checkpointing_state_update_persistent_main( + scenario, pnat_per_slot_list, n_writes_expected, slot_perm_list, + rectangle_for_nowrite, +): + """ + Persistent-CTA main kernel: 1D-grid kernel that loops over + (slot, M-tile, head) work units via tl.range. Caller pre-sorts + slots write-first and passes _n_writes (count of write slots) so + the kernel can split the persistent loop into write and nowrite + halves with the right WRITE_CHECKPOINT constexpr each time. + + Setup mirrors test_checkpointing_state_update_sorted_dispatch + (same fixed seeds, same input shapes) so the reference state + evolution is identical and we can compare per-slot output and + HBM-state postconditions to the same reference. + + Cases: + - all_write (n_writes=B): every slot exercises the + WRITE_CHECKPOINT=True branch of the persistent loop. + - all_nowrite (n_writes=0): every slot exercises the + WRITE_CHECKPOINT=False branch. Verifies the kernel handles + the "write half is empty" launch (n_slots=0 → early return). + - mixed_sorted: slots [2, 3] are writes, slots [0, 1] are + nowrites. slot_perm = [2, 3, 0, 1]. Persistent kernel + should call its impl with pid_b ∈ {2, 3} for the write half + and pid_b ∈ {0, 1} for the nowrite half, even though the + grid pid_b_grid is 0..n_slots-1 in each. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + + pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() + slot_perm = torch.tensor(slot_perm_list, device=device, dtype=torch.int32) + # Sanity: caller-supplied n_writes must match the actual count of + # write slots in the post-perm order. + write_count = sum(pnat_means_write) + assert write_count == n_writes_expected, ( + f"test setup error: expected {n_writes_expected} writes, got {write_count}" + ) + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=dtype + ) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + batch, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + test_state, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + mode="persistent_main", + rectangle_for_nowrite=rectangle_for_nowrite, + slot_perm=slot_perm, + _n_writes=n_writes_expected, + ) + + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (scenario={scenario})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (scenario={scenario})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", + ) + + +@pytest.mark.parametrize( + "scenario,pnat_per_slot_list,rectangle_for_nowrite", + [ + # All-write: every slot has PNAT triggering write (PNAT + T > max_window). + ("all_write", [12, 13, 14, 15], False), + # All-nowrite: every slot fits in the window. + ("all_nowrite", [3, 4, 5, 6], False), + # Mixed: some slots write, some nowrite. No pre-sort needed; the + # dynamic kernel dispatches per-slot at runtime via PNAT load. + ("mixed_unsorted", [3, 12, 10, 15], False), + # Mixed with rectangle_for_nowrite=True — the per-slot runtime + # dispatch picks the rect impl for nowrite slots. Covers the + # rect-path under pd's runtime branch (no pre-sort). + ("mixed_unsorted_rect", [3, 12, 10, 15], True), + ], + ids=["all_write", "all_nowrite", "mixed_unsorted", "mixed_unsorted_rect"], +) +def test_checkpointing_state_update_persistent_dynamic( + scenario, pnat_per_slot_list, rectangle_for_nowrite, +): + """ + Persistent-dynamic kernel: 1D persistent-CTA grid covering the full + batch, with runtime per-slot WRITE_CHECKPOINT branch derived from + each slot's PNAT. Single launch, no half-split, no n_writes needed, + no slot_perm needed (handles unsorted batches natively). + + Same setup as test_checkpointing_state_update_persistent_main; we + verify all three scenarios — including a mixed-unsorted batch the + persistent_main kernel can't handle without pre-sorting. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + + pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=dtype + ) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + batch, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + test_state, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + mode="persistent_dynamic", + rectangle_for_nowrite=rectangle_for_nowrite, + ) + + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (scenario={scenario})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (scenario={scenario})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", + ) + + +@pytest.mark.parametrize( + "scenario,pnat_per_slot_list,n_writes_expected,slot_perm_list", + [ + # Same input shape as the mixed_sorted case in + # test_checkpointing_state_update_persistent_main, but exercises + # the device-tensor n_writes plumbing + skip_empty_halves=False + # path that the bench's mix-mode benchmarking depends on. + ("mixed_sorted", [3, 10, 12, 16], 2, [2, 3, 0, 1]), + # Boundary: n_writes=batch — when skip_empty_halves=False, the + # nowrite half launches with an empty slot range and the kernel + # must do nothing useful (tl.range covers 0 iterations). + ("all_write_noskip", [12, 13, 14, 15], 4, [0, 1, 2, 3]), + # Boundary: n_writes=0 — write half launches with empty range. + ("all_nowrite_noskip", [3, 4, 5, 6], 0, [0, 1, 2, 3]), + ], + ids=["mixed_sorted", "all_write_noskip", "all_nowrite_noskip"], +) +def test_checkpointing_state_update_persistent_main_device_n_writes( + scenario, pnat_per_slot_list, n_writes_expected, slot_perm_list, +): + """ + Persistent_main with the device-tensor n_writes plumbing. + + The bench's mix-mode benchmarking captures a single CUDA graph and + replays it many times with varying per-iter n_writes. To do that the + wrapper takes `_n_writes_dev` (a (1,) int32 device tensor) instead of + `_n_writes` (host int), and the kernel reads `n_writes` from device + memory at entry — same captured pointer across replays, value can + change between replays via an outside-graph copy. Also exercises + `_persistent_skip_empty_halves=False`: both halves of persistent_main + always launch, even when one half has no work (mix scenarios can't + cheaply know n_writes host-side per iter to skip). + + Verifies: + 1. Kernel reads device n_writes correctly (output matches reference). + 2. Empty-half launches don't corrupt state (skip_empty_halves=False + + n_writes=0 / =batch boundary cases). + 3. slot_perm + USE_PERM still works through the device-n_writes path. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + + pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() + slot_perm = torch.tensor(slot_perm_list, device=device, dtype=torch.int32) + write_count = sum(pnat_means_write) + assert write_count == n_writes_expected, ( + f"test setup error: expected {n_writes_expected} writes, got {write_count}" + ) + + # Device-tensor n_writes (the new path). Caller mutates between iters + # in mix-mode benchmarking; we only run one iter here so a single fill. + n_writes_dev = torch.tensor([n_writes_expected], device=device, dtype=torch.int32) + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=dtype + ) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, dt1_input, A, B1, C1, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + batch, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + old_x[:, :step1_T] = x1 + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + test_state, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, dt=dt2, A=A, B=B2, C=C2, + out=test_out, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, + mode="persistent_main", + slot_perm=slot_perm, + # NEW PATHS: + _n_writes_dev=n_writes_dev, # device tensor (not host int) + _persistent_skip_empty_halves=False, # both halves always launch + ) + + torch.testing.assert_close( + test_out.float(), ref_out.float(), + atol=1.0, rtol=0.05, + msg=f"Output mismatch (scenario={scenario})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), ref_state_after_replay[i].float(), + atol=1.0, rtol=0.05, + msg=f"Write slot {i}: state mismatch (scenario={scenario})", + ) + else: + torch.testing.assert_close( + test_state[i], state0[i], rtol=0, atol=0, + msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", + ) + + +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize( + "state_dtype", + [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], + ids=["fp16", "int8", "int16", "fp8"], +) +@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) +@pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) +def test_checkpointing_state_update_philox(state_dtype, nheads, head_dim, d_state, ngroups, paged_cache, T): + """ + Verify that Philox stochastic rounding produces correct results across + all SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). + + Runs our kernel twice with identical inputs — once without rand_seed + (deterministic RN), once with rand_seed (Philox SR) — and confirms: + - Outputs are within bf16-dot tolerance (state perturbation ≤ 1 ULP). + - State dtype is preserved. + - State difference is bounded by ~1 ULP of the chosen grid. + """ + _maybe_skip_dtype(state_dtype, use_sr=True) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + + batch = 2 + device = "cuda" + dtype = torch.bfloat16 + assert nheads % ngroups == 0 + + if paged_cache: + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + else: + cache_size = batch + state_batch_indices = None + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + if is_quantized: + state0_fp32 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + state0, state0_scales = _quantize_state(state0_fp32, state_dtype, quant_max) + else: + state0 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state0_scales = None + + # Cache tensors + old_x = torch.randn(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) + + # New token inputs + x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) + B = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + prev_tokens = torch.full((cache_size,), T // 2, device=device, dtype=torch.int32) + + common_kwargs = dict( + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + ) + + # --- Run without rounding (deterministic RN store) --- + state_no_round = state0.clone() + scales_no_round = state0_scales.clone() if is_quantized else None + out_no_round = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_no_round, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + out=out_no_round, + state_scales=scales_no_round, + **common_kwargs, + ) + + # --- Run with Philox rounding --- + rand_seed = torch.tensor([12345], device=device, dtype=torch.int64) + state_rounded = state0.clone() + scales_rounded = state0_scales.clone() if is_quantized else None + out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_rounded, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + out=out_rounded, + rand_seed=rand_seed, + philox_rounds=10, + state_scales=scales_rounded, + **common_kwargs, + ) + + # Outputs should be nearly identical — rounding only perturbs the + # post-replay state by ±1 ULP before the output phase reads it. + # Out_atol = bf16_baseline + 6.5 * per_elem_ULP_after_dequant: + # non-quant fp16: fp16 ULP at typical magnitude is tiny → 1.0 + # int8: amax/127 ≈ 23/127 → 6.5*0.18 ≈ 1.2 + bf16_baseline + # int16: amax/32767 ≈ 7e-4 → ~bf16_baseline only + # fp8: amax/14 ≈ 23/14 → 6.5*1.6 ≈ 10.7 + bf16_baseline + out_atol = ( + {torch.int8: 1.5, torch.int16: 1.0, torch.float8_e4m3fn: 6.0}[state_dtype] + if is_quantized else 1.0 + ) + out_rtol = ( + {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] + if is_quantized else 2e-2 + ) + torch.testing.assert_close( + out_rounded, out_no_round, rtol=out_rtol, atol=out_atol, + msg=f"Output diverged with Philox rounding ({state_dtype})", + ) + + # State dtype preserved. + assert state_rounded.dtype == state_dtype + + # State diff between RN and SR is bounded by 1 quant cell per element. + # Per-channel decode_scale varies by 10x+ across channels (amax depends + # on randn extremes), so a single flat atol can't bound it accurately — + # use per-channel ULP-aware comparison. + slots = state_batch_indices if paged_cache else slice(None) + if is_quantized: + rounded_fp32 = _dequantize_state(state_rounded[slots], scales_rounded[slots]) + no_round_fp32 = _dequantize_state(state_no_round[slots], scales_no_round[slots]) + diff = (rounded_fp32 - no_round_fp32).abs() + # Per-element bound = max(decode_scale_no_round, decode_scale_rounded). + # decode_scale is shape (cache, nheads, dim); broadcast over dstate. + scale_bound = torch.maximum( + scales_no_round[slots], scales_rounded[slots] + ).unsqueeze(-1) + # int8 / int16: 1 cell after dequant = decode_scale exactly. + # fp8_e4m3: variable grid; the largest cell within a channel scaled + # to fit ±448 is at the channel's max-magnitude element, where the + # cell is ~32x larger than the average. Bound = decode_scale * 32. + # Apply a 1.5x slack pad for floating-point compare quirks at the + # exact-cell boundary. + cell_pad = ( + 32.0 if state_dtype == torch.float8_e4m3fn else 1.0 + ) + bound = scale_bound * (cell_pad * 1.5) + if not (diff <= bound).all(): + offenders = (diff > bound).sum().item() + n_total = diff.numel() + pytest.fail( + f"State RN-SR diff exceeds 1 cell per element for " + f"{offenders}/{n_total} elements ({state_dtype}). " + f"max_diff={diff.max().item():.4g}, " + f"max_bound={bound.max().item():.4g}." + ) + else: + # fp16 ULP depends on magnitude — rtol absorbs that. + torch.testing.assert_close( + state_rounded[slots], + state_no_round[slots], + rtol=2e-3, + atol=0.2, + msg=f"State diverged with Philox rounding ({state_dtype})", + ) + + +@pytest.mark.parametrize( + "state_dtype", + [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], + ids=["fp16", "int8", "int16", "fp8"], +) +def test_philox_rounding_unbiased(state_dtype): + """ + Verify that Philox stochastic rounding is unbiased across all + SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). + + Captures the true fp32 post-replay state by running with fp32 storage, + then runs the kernel with the target dtype + Philox SR. Compares the + SR rounding residual against the deterministic-RN residual: SR should + have mean residual closer to zero than RN, since RN has a systematic + round-to-nearest-even bias and SR is unbiased by construction. + + Uses a large batch (16) for ~2M state elements — plenty of statistics. + """ + _maybe_skip_dtype(state_dtype, use_sr=True) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + batch, T = 16, 6 + device = "cuda" + dtype = torch.bfloat16 + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + # fp32 reference state — replay produces values that don't fit cleanly + # in the target dtype's grid, exposing the rounding bias. + state0_fp32 = torch.randn( + batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + + old_x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) + + x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt_val = repeat(dt_base, "b t h -> b t h p", p=head_dim) + B = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + prev_tokens = torch.full((batch,), T, device=device, dtype=torch.int32) + + common_kwargs = dict( + x=x, dt=dt_val, A=A, B=B, C=C, D=D, dt_bias=dt_bias, dt_softplus=True, + ) + + # 1. fp32 state — captures true post-replay fp32 state. + state_fp32 = state0_fp32.clone() + out_fp32 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_fp32, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), prev_tokens, out=out_fp32, **common_kwargs, + ) + + # 2. Target dtype + Philox SR. For quant we also need scales (derived + # from the same per-channel amax used by the kernel on store). + rand_seed = torch.tensor([99999], device=device, dtype=torch.int64) + if is_quantized: + state_rounded, scales_rounded = _quantize_state(state0_fp32, state_dtype, quant_max) + else: + state_rounded = state0_fp32.to(state_dtype) + scales_rounded = None + out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + checkpointing_state_update( + state_rounded, + old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + cache_buf_idx.clone(), prev_tokens, out=out_rounded, + rand_seed=rand_seed, philox_rounds=10, + state_scales=scales_rounded, + **common_kwargs, + ) + + # Compute residuals. For non-quant: stochastic_residual = SR(fp32) - + # fp32, deterministic_residual = RN(fp32) - fp32. For quant: dequant + # both, comparing in fp32. + if is_quantized: + fp32_vals = state_fp32.flatten() + stochastic_residual = ( + _dequantize_state(state_rounded, scales_rounded).flatten() - fp32_vals + ) + # Deterministic reference: do the same per-channel quant on the + # captured fp32 state, then dequant. This is what the kernel would + # have produced with rand_seed=None. + det_quant, det_scales = _quantize_state(state_fp32, state_dtype, quant_max) + deterministic_residual = ( + _dequantize_state(det_quant, det_scales).flatten() - fp32_vals + ) + else: + fp32_vals = state_fp32.flatten() + stochastic_residual = state_rounded.float().flatten() - fp32_vals + deterministic_residual = fp32_vals.to(state_dtype).float() - fp32_vals + + # Only consider elements where rounding matters (non-zero residual possible). + nonzero_mask = deterministic_residual.abs() > 0 + num_nonzero = nonzero_mask.sum().item() + assert num_nonzero > 1000, f"Too few roundable elements: {num_nonzero}" + + stochastic_mean = stochastic_residual[nonzero_mask].mean().item() + stochastic_std = stochastic_residual[nonzero_mask].std().item() + deterministic_mean = deterministic_residual[nonzero_mask].mean().item() + + # SE-based bias check. An unbiased estimator's sample mean has standard + # error SE = std / sqrt(n). We require |sr_mean| < K*SE (K=4 ≈ ~3.2e-5 + # one-sided false-positive rate). This auto-calibrates per dtype: + # * int16: residual std ~1e-4 → SE ~9e-8 (very tight bound) + # * int8: residual std ~3e-2 → SE ~2e-5 + # * fp8: residual std ~1e-1 → SE ~9e-5 (loosest, magnitude-driven) + # The previous fixed-1e-5 threshold was below SE for int8/fp8 and would + # always fail by chance. Note the |sr|<|det| fallback was also dropped: + # on Gaussian (symmetric) inputs RN's bias is ~0 by symmetry, so SR vs RN + # is just two unbiased estimators racing — unreliable as a unbias test. + se_sr = stochastic_std / (num_nonzero ** 0.5) + K = 4 + assert abs(stochastic_mean) < K * se_sr, ( + f"SR mean exceeds {K}*SE (likely biased) ({state_dtype}): " + f"stochastic_mean={stochastic_mean:.3e}, " + f"SE={se_sr:.3e} (K*SE={K * se_sr:.3e}), " + f"deterministic_mean={deterministic_mean:.3e} (for reference), " + f"n_elements={num_nonzero}" + ) + + +# HEADS_PER_BLOCK > 1 test. The default heuristic only picks HPB > 1 at large +# total_heads (>= 256-512), which the main test with batch=2 never reaches. +# This test overrides _heads_per_block to exercise the two-loop structure in +# the precompute kernel (store-then-reload of per-head dt/dA_cumsum). +# Configs: (nheads=16, ngroups=1) and (nheads=32, ngroups=2) both have +# heads_per_group=16. The heuristic caps HPB at min(2|4, hpg), so HPB=2, 4. +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) +@pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) +def test_checkpointing_heads_per_block( + nheads, + head_dim, + d_state, + ngroups, + state_dtype, + T, + heads_per_block, +): + # PDL flags use wrapper defaults; trimming the parametrize keeps this + # suite fast. Coverage of {launch_with_pdl, use_internal_pdl} variations + # lives in the dedicated correctness tests above (test_checkpointing_state_update). + batch = 8 + """ + Verify checkpointing_state_update produces correct results when + _heads_per_block > 1, exercising the precompute kernel's two-loop + structure (store per-head dt/dA_cumsum in loop 1, reload in loop 2). + """ + device = "cuda" + dtype = torch.bfloat16 + + if nheads % heads_per_block != 0: + pytest.skip(f"nheads ({nheads}) not divisible by heads_per_block ({heads_per_block})") + if heads_per_block > nheads // ngroups: + pytest.skip( + f"heads_per_block ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})" + ) + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + cache_size = batch + state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + + x1 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + cache_size, T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + state0.clone(), + x1, + dt1, + A, + B1, + C1, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) + + old_x[:] = x1 + dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) + + for slot in range(cache_size): + buf = cache_buf_idx[slot].item() + old_B[slot, buf] = B1[slot] + old_dt[slot, buf] = dt1[slot].T + old_dA_cumsum[slot, buf] = dA_cumsum1[slot].T + + k = T + torch.manual_seed(123) + + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = state0.float().clone() + ref_state_f32[:] = states_buffer_f32[:, k - 1] + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + out=ref_out, + ) + + test_state = state0.clone() + prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + + checkpointing_state_update( + test_state, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + x=x2, + dt=dt2, + A=A, + B=B2, + C=C2, + out=test_out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + _heads_per_block=heads_per_block, + ) + + torch.testing.assert_close( + test_out, + ref_out, + rtol=2e-2, + atol=1.0, + msg=f"Output mismatch with HPB={heads_per_block}, T={T}, " + f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", + ) + + expected_state = states_buffer_f32[:, k - 1].to(state_dtype) + torch.testing.assert_close( + test_state, + expected_state, + rtol=2e-2, + atol=1.0, + msg=f"State mismatch with HPB={heads_per_block}, T={T}, " + f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", + ) + + +# HPB > 1 multi-step test. Production chains decode steps; bugs in +# buffer ordering or stale cache values accumulate across steps and can +# be invisible in a single-step test. +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("T", [6, 16], ids=["T6", "T16"]) +@pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) +@pytest.mark.parametrize("paged_cache", [False, True], ids=["contig", "paged"]) +def test_checkpointing_heads_per_block_multistep( + nheads, head_dim, d_state, ngroups, state_dtype, T, heads_per_block, paged_cache +): + """ + Chain N decode steps with HPB > 1 and verify each step's output matches + a fresh reference. A bug that mixes up WRITE/READ buffers, writes wrong + data to cache, or races in the two-loop structure would accumulate + across steps. + """ + batch = 2 + device = "cuda" + dtype = torch.bfloat16 + n_steps = 8 + + if nheads % heads_per_block != 0: + pytest.skip(f"nheads ({nheads}) not divisible by HPB ({heads_per_block})") + if heads_per_block > nheads // ngroups: + pytest.skip(f"HPB ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})") + + torch.manual_seed(42) + + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + if paged_cache: + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + slots = state_batch_indices + else: + cache_size = batch + state_batch_indices = None + slots = slice(None) + + all_x = [] + all_dt = [] + all_B = [] + all_C = [] + for step in range(n_steps): + torch.manual_seed(1000 + step) + all_x.append(torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype)) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + all_dt.append(repeat(dt_base, "b t h -> b t h p", p=head_dim)) + all_B.append(torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype)) + all_C.append(torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype)) + + torch.manual_seed(999) + state_init = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + + ref_state = state_init.float().clone() + ref_outs = [] + ref_slots = ( + state_batch_indices + if paged_cache + else torch.arange(batch, device=device, dtype=torch.int32) + ) + for step in range(n_steps): + out_step = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state, + all_x[step], + all_dt[step], + A, + all_B[step], + all_C[step], + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=ref_slots, + out=out_step, + ) + ref_outs.append(out_step) + + test_state = state_init.clone() + old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.zeros(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_dA_cumsum = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) + + for step in range(n_steps): + k = T if step > 0 else 0 + prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + + checkpointing_state_update( + test_state, + old_x, + old_B, + old_dt, + old_dA_cumsum, + cache_buf_idx, + prev_tokens, + x=all_x[step], + dt=all_dt[step], + A=A, + B=all_B[step], + C=all_C[step], + out=test_out, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + _heads_per_block=heads_per_block, + ) + + if paged_cache: + cache_buf_idx[slots] = 1 - cache_buf_idx[slots] + else: + cache_buf_idx[:] = 1 - cache_buf_idx + + torch.testing.assert_close( + test_out, + ref_outs[step], + rtol=2e-2, + atol=2.0, + msg=f"Output mismatch at step {step} with HPB={heads_per_block}, " + f"T={T}, nheads={nheads}, ngroups={ngroups}, " + f"state_dtype={state_dtype}, paged_cache={paged_cache}", + ) + + +# ----- SR grid-bracket tests (fp8 and fp16) ----- +# +# Verify that each PTX SR output lands on the destination dtype's grid as +# a bracket neighbour of the fp32 input. Catches byte-order traps in the +# inline-asm source-register specifier: +# * fp8: cvt.rs.satfinite.e4m3x4.f32 with pack=4, asm "{$4,$3,$2,$1}" +# * fp16: cvt.rs.f16x2.f32 with pack=2, asm "$0, $2, $1, $3" +# The unbiased test (test_philox_rounding_unbiased) wouldn't catch a +# shuffle: outputs that are still on-grid but swapped within a pack still +# average correctly. Only the per-element bracket check exposes it. +# +# Both kernels are inline copies of the production helpers — kept here so +# the test exercises the exact PTX form independent of wrapper changes. + + +@triton.jit +def _packed_int8_sr_kernel(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + (offs // 4)) + y = _stochastic_round_int8_packed(x, rand, offs) + tl.store(out_ptr + offs, y.to(tl.int8)) + + +@triton.jit +def _packed_int16_sr_kernel(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + (offs // 2)) + y = _stochastic_round_int16_packed(x, rand, offs) + tl.store(out_ptr + offs, y.to(tl.int16)) + + +def _bitrev_int(x: int, bits: int) -> int: + out = 0 + for _ in range(bits): + out = (out << 1) | (x & 1) + x >>= 1 + return out + + +def _rand_words(rand: torch.Tensor) -> list[int]: + return [int(v) & 0xFFFFFFFF for v in rand.cpu().tolist()] + + +def test_packed_int_sr_matches_reference(): + device = "cuda" + n = 1024 + offs = torch.arange(n, device=device, dtype=torch.float32) + x = ((offs % 37) - 18.0) + (((offs * 13.0) % 97.0) + 0.3) / 128.0 + + torch.manual_seed(42) + rand_i8 = torch.randint(-(2**31), 2**31, (n // 4,), device=device, dtype=torch.int32) + out_i8 = torch.empty(n, device=device, dtype=torch.int8) + _packed_int8_sr_kernel[(1,)](x, rand_i8, out_i8, BLOCK=n) + + x_cpu = x.cpu().tolist() + rand_i8_words = _rand_words(rand_i8) + ref_i8 = [] + for i, value in enumerate(x_cpu): + word = rand_i8_words[i // 4] + low = word & 0x0000FFFF + high = (word >> 16) & 0x0000FFFF + pos = i & 3 + if pos == 0: + rand16 = low + elif pos == 1: + rand16 = _bitrev_int(low, 16) + elif pos == 2: + rand16 = high + else: + rand16 = _bitrev_int(high, 16) + ref_i8.append(math.floor(value + rand16 / float(1 << 16))) + + torch.testing.assert_close( + out_i8.cpu().to(torch.int16), + torch.tensor(ref_i8, dtype=torch.int16), + rtol=0, + atol=0, + ) + + rand_i16 = torch.randint(-(2**31), 2**31, (n // 2,), device=device, dtype=torch.int32) + out_i16 = torch.empty(n, device=device, dtype=torch.int16) + _packed_int16_sr_kernel[(1,)](x, rand_i16, out_i16, BLOCK=n) + + rand_i16_words = _rand_words(rand_i16) + ref_i16 = [] + for i, value in enumerate(x_cpu): + word = rand_i16_words[i // 2] + rand_bits = word if (i & 1) == 0 else _bitrev_int(word, 32) + rand24 = rand_bits & 0x00FFFFFF + ref_i16.append(math.floor(value + rand24 / float(1 << 24))) + + torch.testing.assert_close( + out_i16.cpu(), + torch.tensor(ref_i16, dtype=torch.int16), + rtol=0, + atol=0, + ) + + +@triton.jit +def _bracket_kernel_fp8(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + offs) + y = tl.inline_asm_elementwise( + asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", + constraints="=r,r,r,r,r,r,r,r,r", + args=(x, rand), + dtype=tl.float8e4nv, + is_pure=True, + pack=4, + ) + tl.store(out_ptr + offs, y) + + +@triton.jit +def _bracket_kernel_fp16(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + offs) + y = tl.inline_asm_elementwise( + asm="""{ + cvt.rs.f16x2.f32 $0, $2, $1, $3; + }""", + constraints=("=r,r,r,r,r"), + args=(x, rand), + dtype=tl.float16, + is_pure=True, + pack=2, + ) + tl.store(out_ptr + offs, y) + + +_BRACKET_KERNEL = { + torch.float8_e4m3fn: _bracket_kernel_fp8, + torch.float16: _bracket_kernel_fp16, +} + + +def _build_finite_grid(dtype: torch.dtype, device: str) -> torch.Tensor: + """Reinterpret all bit patterns of ``dtype`` as floats; return sorted + unique finite values (drops ±inf, NaNs).""" + if dtype == torch.float8_e4m3fn: + ints = torch.arange(256, dtype=torch.uint8, device=device) + full = ints.view(torch.float8_e4m3fn).to(torch.float32) + elif dtype == torch.float16: + # int16 view of all 65536 patterns (covers fp16 normals + subnormals + # + ±inf + NaN; we filter to finite below). + ints = torch.arange(65536, dtype=torch.int32, device=device).to(torch.int16) + full = ints.view(torch.float16).to(torch.float32) + else: + raise ValueError(f"Unsupported bracket-test dtype: {dtype}") + return full[torch.isfinite(full)].sort()[0].unique() + + +def _build_bracket_inputs(dtype: torch.dtype, n: int, device: str) -> torch.Tensor: + """Test inputs spanning the dtype's grid range. Includes on-grid points + so we exercise the no-rounding case; for fp8 also includes overflow to + test saturation (PTX `cvt.rs.satfinite.e4m3x4.f32` clamps in-op). + + fp16 inputs are kept inside the finite range — `cvt.rs.f16x2.f32` does + NOT have a `satfinite` modifier and produces ±inf for OOR inputs (not + a saturate-to-±max). The kernel only ever sees in-range fp32 state in + practice (state_amax is always ≪ fp16_max), so the test mirrors that. + """ + grid = _build_finite_grid(dtype, device) + g_min, g_max = grid[0].item(), grid[-1].item() + x = torch.empty(n, device=device, dtype=torch.float32) + if dtype == torch.float8_e4m3fn: + # 1.5x range exercises saturation; satfinite handles it in-op. + x.uniform_(g_min * 1.5, g_max * 1.5) + else: # fp16: four magnitude bands, all within finite range. + x[: n // 4].uniform_(-1.0, 1.0) + x[n // 4 : n // 2].uniform_(-100, 100) + x[n // 2 : 3 * n // 4].uniform_(-1000, 1000) + x[3 * n // 4 :].uniform_(g_min * 0.99, g_max * 0.99) + return x, grid + + +@_skip_pre_sm100 +@pytest.mark.parametrize( + "state_dtype", + [torch.float8_e4m3fn, torch.float16], + ids=["fp8", "fp16"], +) +def test_sr_grid_bracket(state_dtype): + """Verify SR PTX outputs each lie on the destination grid as a bracket + neighbour of the fp32 input.""" + device = "cuda" + n = 1024 # multiple of both pack=4 (fp8) and pack=2 (fp16) + + torch.manual_seed(42) + x, grid_finite = _build_bracket_inputs(state_dtype, n, device) + g_min, g_max = grid_finite[0].item(), grid_finite[-1].item() + + # Bracket [lo, hi] in the destination grid for each input. For + # out-of-range inputs the bracket is the saturating endpoint pair. + x_clamped = x.clamp(g_min, g_max) + idx = torch.searchsorted(grid_finite, x_clamped, right=False).clamp( + min=1, max=len(grid_finite) - 1 + ) + lo = grid_finite[idx - 1] + hi = grid_finite[idx] + # For x exactly on grid, idx points at it; lo = grid[i-1], hi = x — the + # bracket allows out==hi (=x) which is what RN-on-grid produces. + + kernel = _BRACKET_KERNEL[state_dtype] + + for seed in range(4): + torch.manual_seed(seed) + # int32 for raw random bits — PTX takes the bit pattern, sign + # interpretation doesn't matter. + rand = torch.randint(-(2**31), 2**31, (n,), device=device, dtype=torch.int32) + out = torch.empty(n, device=device, dtype=state_dtype) + kernel[(1,)](x, rand, out, BLOCK=n) + out_fp32 = out.to(torch.float32) + + on_grid = (out_fp32 == lo) | (out_fp32 == hi) + if not on_grid.all(): + offenders = ~on_grid + n_off = offenders.sum().item() + sample = ( + x[offenders][:5].tolist(), + lo[offenders][:5].tolist(), + hi[offenders][:5].tolist(), + out_fp32[offenders][:5].tolist(), + ) + pytest.fail( + f"{state_dtype} SR output not on grid bracket for {n_off}/{n} " + f"elements (seed={seed}). x={sample[0]} lo={sample[1]} " + f"hi={sample[2]} out={sample[3]}. Likely the PTX byte-order " + "bug (cvt.rs source-register order)." + ) + From 051511456e71a818e2def706cf507ae603843f16 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 18 May 2026 17:13:19 -0700 Subject: [PATCH 54/89] mamba_checkpointing: fix dA_cumsum cross-step continuity in slim precomputes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The slim kernel stored old_dA_cumsum as per-step cumsums each call, but the write replay assumed continuous cumsum across the buffer. After 2+ back-to-back nowrites the buffer tail at PNAT-1 reflected only the last step's partial total, so total_decay = exp(load) was off by exp(prefix) — producing diffs ~1000x on the next write step. Fix in both precompute kernels (_replay_precompute_impl, _rectangle_precompute_impl): load prefix at old_dA_cumsum[buf_active, head, PNAT-1] and add to the new-T cumsum before storing. Gated by PNAT > 0 and (replay) not write_checkpoint (write path starts fresh in staging buf). With continuous buffer values, the rectangle precompute formula simplifies: s_k = -old_dA_cumsum_all[k] (was total - old_dA_cumsum_all) and decay_vec_full[t] = exp(dA_cumsum_new[t]) (was total_decay * exp(...)). Drop the now-unused total_dA_cumsum load. Same one-exp-on- sum form — no numerical regression. Relax wrapper assert max_window <= BLOCK_SIZE_T: kernels already have separate BLOCK_SIZE_WINDOW/BLOCK_SIZE_K derived from MAX_REPLAY_BUFFER_LENGTH, so max_window can exceed BLOCK_SIZE_T freely. Unblocks max_window > T configurations in tests. Tests: - test_checkpointing_state_update: update expected old_dA_cumsum on the nowrite paths to include the prefix. - test_checkpointing_heads_per_block (single-step): full PNAT sweep (0, 1, T, threshold-1, threshold, threshold+1, max_window-1) and assert full write contract on old_x, old_B, old_dt, old_dA_cumsum plus untouched-region invariants. - test_checkpointing_heads_per_block_multistep: divergent per-slot acceptance — n_writes hits {0, 1, 2} and slot_perm exercises both identity and non-identity orderings across n_steps. - Merge test_..._persistent_main + test_..._persistent_dynamic into test_..._scenarios parametrized over mode (pm/pd) and explicit-vs-auto slot_perm. - Parametrize mode and rectangle_for_nowrite across all kernel- calling tests (rect skipped on write-only flows). Bench: rename n_writes_dev -> n_writes and drop _persistent_skip_empty_halves to match the slim wrapper API. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/checkpointing_state_update_slim.py | 578 ++++++++----- ...mark_replay_selective_state_update_slim.py | 23 +- .../test_checkpointing_state_update_slim.py | 786 +++++++++++------- 3 files changed, 874 insertions(+), 513 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py index 73047aec3240..9dad703a8cde 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py @@ -303,6 +303,26 @@ def _replay_precompute_impl( dA_cumsum = tl.cumsum(A[:, None] * dt, axis=1) # (H, T) decay_vec = tl.exp(dA_cumsum) # (H, T) + # Cross-step continuity for old_dA_cumsum: when appending to active_buf at + # offset PNAT > 0, the previous step left a running cumsum at [0, PNAT) + # whose tail value lives at active_buf[head, PNAT-1]. Add that tail to + # this step's per-step-restarted cumsum before storing so the buffer + # holds one continuous cumsum across N back-to-back nowrites. Write path + # (write_buf = 1 - buf_active, write_offset = 0) starts fresh, no prefix. + # Both branches are on scalar runtime values (write_checkpoint and PNAT), + # uniform across the block — use scalar if to short-circuit the load. + if write_checkpoint or prev_num_accepted_tokens == 0: + prev_total = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) + else: + last_cumsum_ptrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T + ) + prev_total = tl.load(last_cumsum_ptrs).to(tl.float32) + # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of write_buf. old_dt_addrs = ( old_dt_ptr @@ -320,7 +340,7 @@ def _replay_precompute_impl( + heads_block[:, None] * stride_old_dA_cumsum_head + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T ) - tl.store(old_dA_cumsum_addrs, dA_cumsum, mask=t_mask[None, :]) + tl.store(old_dA_cumsum_addrs, dA_cumsum + prev_total[:, None], mask=t_mask[None, :]) # decay_vec scratch — always at offs_t. decay_vec_addrs = ( @@ -530,6 +550,23 @@ def _rectangle_precompute_impl( A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) dA_cumsum = tl.cumsum(A * dt, axis=0) + # Cross-step continuity for old_dA_cumsum: rectangle precompute runs + # only on the nowrite path (write_buf == buf_active, write_offset == PNAT). + # Add the running tail from buf_active[head_idx, PNAT-1] so the buffer + # holds one continuous cumsum across back-to-back nowrites. PNAT is + # scalar/uniform, use scalar if to short-circuit the load at PNAT=0. + if prev_num_accepted_tokens == 0: + prev_total = 0.0 + else: + last_cumsum_ptr = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T + ) + prev_total = tl.load(last_cumsum_ptr).to(tl.float32) + # Store dt and dA_cumsum to write_buf at [write_offset, write_offset+T) # for next step's replay/rectangle use. old_dt_base = ( @@ -552,7 +589,7 @@ def _rectangle_precompute_impl( ) tl.store( old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, - dA_cumsum, + dA_cumsum + prev_total, mask=t_mask, ) @@ -585,9 +622,6 @@ def _rectangle_precompute_impl( # (H, T) and combo = factor_dt * exp_diff (H, T, K). Store decay_vec_full; # combo_block stays in registers across gdc_wait — used directly post-wait # to compute rect_CB_scaled without a global memory roundtrip. - prev_k_idx = tl.minimum( - tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 - ) offs_h = tl.arange(0, HEADS_PER_BLOCK) heads_block = first_head + offs_h # (H,) @@ -627,11 +661,10 @@ def _rectangle_precompute_impl( old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, mask=hk_mask, other=0.0, ).to(tl.float32) - # (H,) scalar-per-head: total_dA_cumsum at prev_k_idx. - total_dA_cumsum = tl.load( - old_dA_cumsum_read_h + prev_k_idx * stride_old_dA_cumsum_T - ).to(tl.float32) # (H, T) loads at [PNAT, PNAT+T) — this step's dA_cumsum_new from loop 1. + # With the cross-step continuity fix in loop 1, the values stored at + # [PNAT, PNAT+T) are the continuous cumsum (prefix + per-step new + # cumsum) — i.e., continuous_cumsum[PNAT..PNAT+T-1] in global indexing. ht_mask = t_mask[None, :] # (1, T) dA_cumsum_new = tl.load( old_dA_cumsum_write_h[:, None] @@ -651,11 +684,12 @@ def _rectangle_precompute_impl( mask=hkn_mask, other=0.0, ).to(tl.float32) - # decay_vec_full = total_decay * exp(cumAdt_new). (H, T). - total_decay = tl.where( - prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0 - ) # (H,) - decay_vec_full_block = total_decay[:, None] * tl.exp(dA_cumsum_new) # (H, T) + # decay_vec_full[t] = exp(continuous_cumsum[PNAT+t]) — directly the + # continuous value now stored at buffer position write_offset+t. Was + # decomposed as total_decay * exp(per_step_new[t]) when the buffer held + # per-step (non-continuous) cumsum; with the continuity fix the value + # IS continuous_cumsum[PNAT+t] so no decomposition is needed. + decay_vec_full_block = tl.exp(dA_cumsum_new) # (H, T) decay_vec_addrs = ( decay_vec_ptr + pid_b * stride_dv_batch @@ -665,11 +699,25 @@ def _rectangle_precompute_impl( tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) # combo_block = factor_dt * exp_diff — (H, T, K). Stays in registers - # across gdc_wait. + # across gdc_wait. With continuous cumsum in the buffer, s_k for any k + # (old or new) is simply -continuous_cumsum[k]; exp_diff[t, k] then + # equals exp(continuous_cumsum[PNAT+t] - continuous_cumsum[k]) — the + # decay weight for token k's contribution to the output at position + # PNAT+t. No need to subtract any "total" — the dA_cumsum_new[t] term + # already carries the full prefix. + # + # Numerical note: pre-fix this kernel computed `total - old_dA[k]` + # (small-minus-small) then summed `+ dA_cumsum_new[t]` (also small, + # per-step). Post-fix `s_k = -old_dA[k]` is large-magnitude positive + # and `dA_cumsum_new[t]` is large-magnitude negative; their sum + # cancels back to the same small value. Cancellation error is bounded + # by ulp(max_magnitude) ≈ 2^-23 · |continuous_cumsum| — negligible + # for max_window ≤ ~1024. Still one exp on the sum (not two muls of + # exps), so no overflow regression vs the original formulation. factor_dt = tl.where(is_old_k[None, :], old_dt_all, dt_at_kn) # (H, K) s_k = tl.where( is_old_k[None, :], - total_dA_cumsum[:, None] - old_dA_cumsum_all, + -old_dA_cumsum_all, -dA_cumsum_at_kn, ) # (H, K) # exp_diff (H, T, K) = exp(s_k (H, 1, K) + dA_cumsum_new (H, T, 1)). @@ -2098,6 +2146,161 @@ def _persistent_main_kernel( } +# --------------------------------------------------------------------------- +# Default tunings — looked up by (effective_batch, dtype, sr) when the caller +# leaves mode/knobs as None. +# +# Effective batch = raw_batch × nheads_per_rank. Our sweep was at TP=8 with +# the standard Mamba2 nheads; at call time we compute it from the input +# tensor shape so callers at other TP / nheads pick up the right cell. +# +# Schema: dict[(dtype_str, sr_str)] → list[(eff_batch_threshold, mode, knobs)] +# sorted by threshold ascending. Lookup finds the first threshold ≥ eff_b +# (so missing intermediate batches fall up to the next tuned cell). If +# eff_b exceeds the largest threshold, use the largest entry. +# +# Each `knobs` dict only contains keys for the chosen mode; the wrapper +# unpacks them with the same name as the matching kwargs. Caller-provided +# kwargs always win over table values. +# +# This table is intentionally NOT parameterized by T or max_window. Our +# sweep was T=6, max_window=16. Callers outside that regime silently get +# the same numbers — they may be suboptimal but they're correct. +# +# Source: audit_v2.py --emit-tuning. Auto-generated from per-cell search +# winners (best of pd / pm by bucket_expected_renorm). Sweep was TP=8 with +# NHEADS=128 → nheads_per_rank=16; thresholds are in effective_batch units. +# Missing dtype/SR combos (fp16/RN, int8/RN, fp8/*) fall back via the +# _resolve_tuning chain — RN→SR for same dtype, then fp8→int8/SR. +_DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { + ("fp32", "RN"): [ + ( 16, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 8, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.22us + ( 32, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.17us + ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.08us + ( 128, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=8, score=9.00us + ( 256, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=16, score=10.92us + ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=13.53us + ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=19.50us + ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=30.28us + ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=50.32us + ( 8192, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=90.99us + (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=171.69us + ], + ("fp16", "SR"): [ + ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.16us + ( 32, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.01us + ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=7.95us + ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 4, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=8.87us + ( 256, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 2, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.28us + ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=12.90us + ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=16.71us + ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=25.71us + ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=39.80us + ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=71.34us + (16384, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=133.51us + ], + ("int8", "SR"): [ + ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.34us + ( 32, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.36us + ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 2, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.40us + ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 10, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=9.37us + ( 256, "persistent_dynamic", {'_block_size_m': 16, '_cta_per_sm': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages': 1, '_num_stages': 4, '_num_warps': 1, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.02us + ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=32, score=13.15us + ( 1024, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=17.82us + ( 2048, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=27.01us + ( 4096, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=43.23us + ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=77.01us + (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=140.43us + ], +} + + +# Knob names that map between the modes' single-value (pd) and split-value +# (pm) namespaces. Used by `_bridge_tuning_knobs` when caller forces a mode +# different from the table's recommendation. +_PD_TO_PM_SPLIT_MAP = { # pd unsplit knob → (pm_write_knob, pm_nowrite_knob) + "_block_size_m": ("_block_size_m_write", "_block_size_m_nowrite"), + "_num_warps": ("_num_warps_write", "_num_warps_nowrite"), + "_num_stages": ("_num_stages_write", "_num_stages_nowrite"), + # CPS / LS are persistent-loop knobs; pd uses _cta_per_sm + _num_loop_stages + # as unsplit, pm uses _cta_per_sm_write/_nowrite + _num_loop_stages_write/_nowrite. + "_cta_per_sm": ("_cta_per_sm_write", "_cta_per_sm_nowrite"), + "_num_loop_stages": ("_num_loop_stages_write", "_num_loop_stages_nowrite"), +} + + +def _bridge_tuning_knobs(knobs: dict, from_mode: str, to_mode: str) -> dict: + """Convert a tuning dict between pd ↔ pm knob namespaces. + + pd → pm: copy each unsplit value to both write/nowrite split knobs; drop + the unsplit form (pm doesn't read it). + pm → pd: take the nowrite split value as the unsplit knob; drop the + write/nowrite split forms (pd doesn't read them). + Shape knobs that exist in both modes (_heads_per_block, _flatten, + _warp_specialize, TMA flags, rectangle_for_nowrite) carry over unchanged. + """ + out = dict(knobs) + if from_mode == "persistent_dynamic" and to_mode == "persistent_main": + for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): + if unsplit in out: + out.setdefault(pm_w, out[unsplit]) + out.setdefault(pm_nw, out[unsplit]) + del out[unsplit] + elif from_mode == "persistent_main" and to_mode == "persistent_dynamic": + for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): + if pm_nw in out: + out.setdefault(unsplit, out[pm_nw]) + out.pop(pm_w, None) + out.pop(pm_nw, None) + return out + + +def _resolve_tuning( + batch: int, nheads_per_rank: int, dt_str: str, sr_str: str, +) -> tuple[str, dict] | None: + """Look up the default mode + knobs for this (eff_batch, dt, sr) cell. + + Returns (mode, knobs_dict) or None if the table has no entry covering + this dtype/sr (including the fp8→int8/SR and dtype/RN→dtype/SR fallbacks). + Returning None lets the wrapper fall back to caller-provided kwargs or + kernel-side defaults. + """ + eff_b = batch * max(1, nheads_per_rank) + # Lookup chain. Order: + # 1. Exact (dt, sr). + # 2. (dt, SR) if RN missing for that dtype. + # 3. Cross-dtype fallback for dtypes we haven't tuned: + # bf16 / int16 → fp16/SR + # fp8 → int8/SR + # Unknown dtype → raise. + valid_dtypes = {"fp32", "fp16", "bf16", "int8", "int16", "fp8"} + if dt_str not in valid_dtypes: + raise ValueError( + f"checkpointing_state_update: unsupported state dtype {dt_str!r}; " + f"expected one of {sorted(valid_dtypes)}" + ) + keys_to_try = [(dt_str, sr_str)] + if sr_str == "RN": + keys_to_try.append((dt_str, "SR")) + if dt_str in ("bf16", "int16"): + keys_to_try.append(("fp16", "SR")) + elif dt_str == "fp8": + keys_to_try.append(("int8", "SR")) + entries = None + for k in keys_to_try: + if k in _DEFAULT_TUNING: + entries = _DEFAULT_TUNING[k] + break + if entries is None: + return None + # Find first threshold ≥ eff_b; if none, use largest entry. + for thresh, mode, knobs in entries: + if eff_b <= thresh: + return mode, dict(knobs) + thresh, mode, knobs = entries[-1] + return mode, dict(knobs) + + def checkpointing_state_update( state: torch.Tensor, old_x: torch.Tensor, @@ -2112,6 +2315,17 @@ def checkpointing_state_update( B: torch.Tensor, C: torch.Tensor, out: torch.Tensor, + # Required persistent-mode plumbing (REQUIRED for both pd and pm; pd + # ignores both internally but the wrapper still demands them): + # n_writes : (1,) int32 device tensor with the count of write-mode + # slots in the batch. pm uses it to size the two halves; + # pd ignores it (per-slot runtime PNAT check). + # slot_perm : (batch,) int32 device tensor remapping grid pid → slot. + # pm uses it to cluster writes first (kernel grid step is + # write_half then nowrite_half); pd ignores it. Callers + # that don't care about ordering should pass arange(batch). + n_writes: torch.Tensor, + slot_perm: torch.Tensor, D: torch.Tensor | None = None, z: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, @@ -2124,14 +2338,8 @@ def checkpointing_state_update( launch_with_pdl=False, use_internal_pdl=True, write_checkpoint: bool = True, - rectangle_for_nowrite: bool = False, - mode: str = "persistent_dynamic", - # Slot permutation: int32 (batch,) tensor mapping grid program_id -> - # original slot index. When provided, the persistent kernels read pid_b - # through this perm so callers can pre-sort slots write-first (required - # for `persistent_main`, optional perf hint for `persistent_dynamic`). - # None => identity perm. - slot_perm: torch.Tensor | None = None, + rectangle_for_nowrite: bool | None = None, + mode: str | None = None, _block_size_m: int | None = None, _num_warps: int | None = None, _num_stages: int | None = None, @@ -2160,38 +2368,19 @@ def checkpointing_state_update( # TMA state-tensor toggles — 4 independent paths (see CHECKPOINTING_DESIGN.md # item #17 for measured perf profiles). Each is False=raw load/store, True= # use a host-built TMA tensor_descriptor for that path. - _use_tma_rect_load: bool = False, # rect kernel's state load (nowrite-only) - _use_tma_replay_write_load: bool = False, # replay-style state load when WC=True - _use_tma_replay_write_store: bool = False, # replay-style state store when WC=True - _use_tma_replay_nowrite_load: bool = False, # replay-style state load when WC=False - # Persistent-mode bench kwargs (only consulted when mode == "persistent_main"): - # _n_writes : int — count of write-mode slots in the (pre-sorted) batch. - # Required when mode == "persistent_main"; the persistent kernel uses - # it as a runtime int32 to compute total_work for write/nowrite halves. + _use_tma_rect_load: bool | None = None, # rect kernel's state load (nowrite-only) + _use_tma_replay_write_load: bool | None = None, # replay-style state load when WC=True + _use_tma_replay_write_store: bool | None = None, # replay-style state store when WC=True + _use_tma_replay_nowrite_load: bool | None = None, # replay-style state load when WC=False + # Persistent-mode tuning kwargs (consulted for both pd and pm; pd uses + # _cta_per_sm / _num_loop_stages, pm uses the _write/_nowrite splits): # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally - # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. Default = 1. + # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. # _num_loop_stages : int — `num_stages` arg on the inner `tl.range(...)` # persistent loop. Note: this is loop-level, NOT the kernel-arg - # `num_stages` (which only pipelines dot-feeding loads). Default 2. - # _flatten : bool — `flatten` arg on `tl.range(...)`. Default True - # (the canonical Triton 3.6 persistent idiom). + # `num_stages` (which only pipelines dot-feeding loads). + # _flatten : bool — `flatten` arg on `tl.range(...)`. # _warp_specialize : bool — `warp_specialize` arg on `tl.range(...)`. - # Default False. Triton 3.6 only supports it on simple matmul loops; - # our scan loop probably won't pattern-match — but exposed as a knob - # for sweep experiments. Requires num_warps >= 4 if True. - _n_writes: int | None = None, - # Optional pre-allocated (1,) int32 device tensor for the persistent - # kernel's n_writes input. Bench passes this in mix scenarios so the - # captured CUDA graph can read varying n_writes per iter without - # re-capture. When None and `_n_writes` is provided, we allocate a - # scratch tensor and fill from `_n_writes` (pure scenarios). - _n_writes_dev: torch.Tensor | None = None, - # When True, persistent_main host-skips empty-half launches (n_writes=0 - # or =batch in pure scenarios). Default True preserves today's behavior. - # Set False to always launch both halves — used by mix scenarios (where - # host can't cheaply read n_writes per iter) and for fair K-consistent - # comparisons. - _persistent_skip_empty_halves: bool = True, _cta_per_sm: int | None = None, _num_loop_stages: int | None = None, _flatten: bool | None = None, @@ -2275,21 +2464,20 @@ def checkpointing_state_update( use_internal_pdl = False # Mode selection: - # mode="persistent_dynamic" (default): single persistent-CTA kernel - # covering the full batch. Each work-item dispatches via runtime - # PNAT check (is_write = (pnat + T) > MAX). No write/nowrite split. - # slot_perm is honored but optional. write_checkpoint is ignored - # (per-slot from PNAT). + # mode=None (default): look up the table-tuned mode + knobs for this + # (effective_batch, dtype, sr) cell. See `_resolve_tuning` above. + # mode="persistent_dynamic": single persistent-CTA kernel covering the + # full batch. Each work-item dispatches via runtime PNAT check + # (is_write = (pnat + T) > MAX). No write/nowrite split. + # slot_perm is honored but optional. write_checkpoint is ignored. # mode="persistent_main": persistent-CTA kernel with two launches - # (write half + nowrite half). Caller must pre-sort slots - # write-first via slot_perm and pass _n_writes / _n_writes_dev so - # the kernel can split the persistent loop into the two halves - # with the right WRITE_CHECKPOINT constexpr each time. RECTANGLE - # constexpr (= rectangle_for_nowrite) picks rect vs replay for the - # nowrite half. write_checkpoint is ignored (per-slot from PNAT). - assert mode in ("persistent_dynamic", "persistent_main"), ( - f"unknown mode {mode!r}; expected 'persistent_dynamic' or 'persistent_main'" - ) + # (write half + nowrite half). Caller MUST pre-sort slot_perm + # write-first; the n_writes tensor partitions the persistent loop + # into the two halves with the right WRITE_CHECKPOINT constexpr + # each time. RECTANGLE constexpr (= rectangle_for_nowrite) picks + # rect vs replay for the nowrite half. write_checkpoint is ignored. + # Note: mode-and-knob resolution from the default-tuning table happens + # below, after we have `batch` and `nheads`. # --- Hardware support gates --- # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX @@ -2356,11 +2544,72 @@ def checkpointing_state_update( ngroups = B.shape[2] assert nheads % ngroups == 0 - # --- Quantization plumbing --- + # --- Quantization plumbing (needed for SR/RN classification below) --- # QUANT_MAX > 0 ⇔ state is int8 / int16 / fp8_e4m3fn. Kernel-entry # static_assert on the Triton side mirrors this invariant. quant_max = _QUANT_MAX_BY_DTYPE.get(state.dtype, 0.0) is_quantized = quant_max > 0.0 + + # --- Default-tuning lookup --- + # Resolve (mode, knobs) from the table when caller leaves them None. + # Caller-provided kwargs always win. If the caller forces a mode that + # differs from the table's recommendation for this cell, we BRIDGE the + # table's knobs into the forced mode's knob namespace rather than fall + # back to (likely-terrible) kernel defaults: + # table pd → forced pm: copy each unsplit pd knob (M, W, S, CPS, LS) + # to both write and nowrite split knobs. + # table pm → forced pd: take the nowrite split values (Mnw, Wnw, Snw, + # CPSnw, LSnw) as the unsplit knobs. + # Empty table → no-op (caller passes whatever, mode falls back to pd). + _dt_str = { + torch.float32: "fp32", + torch.float16: "fp16", + torch.bfloat16: "bf16", + torch.int8: "int8", + torch.int16: "int16", + torch.float8_e4m3fn: "fp8", + }.get(state.dtype, str(state.dtype)) + _sr_str = "SR" if (rand_seed is not None and is_quantized) else "RN" + _table_entry = _resolve_tuning(batch, nheads, _dt_str, _sr_str) + if _table_entry is not None: + _table_mode, _table_knobs = _table_entry + if mode is None: + mode = _table_mode + if mode != _table_mode: + # Bridge across modes — see header comment above. + _table_knobs = _bridge_tuning_knobs(_table_knobs, _table_mode, mode) + # Fill None-valued kwargs from table. We can't reliably mutate + # locals() for re-read, so re-bind each kwarg explicitly. + if rectangle_for_nowrite is None and "rectangle_for_nowrite" in _table_knobs: + rectangle_for_nowrite = bool(_table_knobs["rectangle_for_nowrite"]) + _block_size_m = _block_size_m if _block_size_m is not None else _table_knobs.get("_block_size_m") + _num_warps = _num_warps if _num_warps is not None else _table_knobs.get("_num_warps") + _num_stages = _num_stages if _num_stages is not None else _table_knobs.get("_num_stages") + _heads_per_block = _heads_per_block if _heads_per_block is not None else _table_knobs.get("_heads_per_block") + _precompute_num_warps = _precompute_num_warps if _precompute_num_warps is not None else _table_knobs.get("_precompute_num_warps") + _precompute_num_stages = _precompute_num_stages if _precompute_num_stages is not None else _table_knobs.get("_precompute_num_stages") + _block_size_m_write = _block_size_m_write if _block_size_m_write is not None else _table_knobs.get("_block_size_m_write") + _block_size_m_nowrite = _block_size_m_nowrite if _block_size_m_nowrite is not None else _table_knobs.get("_block_size_m_nowrite") + _num_warps_write = _num_warps_write if _num_warps_write is not None else _table_knobs.get("_num_warps_write") + _num_warps_nowrite = _num_warps_nowrite if _num_warps_nowrite is not None else _table_knobs.get("_num_warps_nowrite") + _num_stages_write = _num_stages_write if _num_stages_write is not None else _table_knobs.get("_num_stages_write") + _num_stages_nowrite = _num_stages_nowrite if _num_stages_nowrite is not None else _table_knobs.get("_num_stages_nowrite") + _cta_per_sm = _cta_per_sm if _cta_per_sm is not None else _table_knobs.get("_cta_per_sm") + _num_loop_stages = _num_loop_stages if _num_loop_stages is not None else _table_knobs.get("_num_loop_stages") + _flatten = _flatten if _flatten is not None else _table_knobs.get("_flatten") + _warp_specialize = _warp_specialize if _warp_specialize is not None else _table_knobs.get("_warp_specialize") + _use_tma_rect_load = _use_tma_rect_load or bool(_table_knobs.get("_use_tma_rect_load", False)) + _use_tma_replay_write_load = _use_tma_replay_write_load or bool(_table_knobs.get("_use_tma_replay_write_load", False)) + _use_tma_replay_write_store = _use_tma_replay_write_store or bool(_table_knobs.get("_use_tma_replay_write_store", False)) + _use_tma_replay_nowrite_load = _use_tma_replay_nowrite_load or bool(_table_knobs.get("_use_tma_replay_nowrite_load", False)) + # Final defaults if neither caller nor table set them (empty table case). + if mode is None: + mode = "persistent_dynamic" + if rectangle_for_nowrite is None: + rectangle_for_nowrite = False + assert mode in ("persistent_dynamic", "persistent_main"), ( + f"unknown mode {mode!r}; expected 'persistent_dynamic' or 'persistent_main'" + ) if is_quantized: assert state_scales is not None, ( f"state.dtype={state.dtype} requires state_scales tensor " @@ -2378,18 +2627,11 @@ def checkpointing_state_update( # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the # placeholder degenerate case max_window = T (every step is a checkpoint # step). For real replay-style checkpointing, max_window > T and - # `prev_num_accepted_tokens` can be 0..max_window. + # `prev_num_accepted_tokens` can be 0..max_window. Window-axis kernel + # tiles (BLOCK_SIZE_WINDOW, BLOCK_SIZE_K) are derived independently from + # MAX_REPLAY_BUFFER_LENGTH so max_window can exceed BLOCK_SIZE_T freely. max_window = old_x.shape[1] assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" - # Replay-style code path uses BLOCK_SIZE_T = max(np2(T), 16) for the - # combined T-axis (T_new tile size) and reuses it for window loads. Until - # the heuristic is generalized to track max_window separately, require - # max_window to fit within that tile. - block_size_t = max(triton.next_power_of_2(T), 16) - assert max_window <= block_size_t, ( - f"max_window={max_window} exceeds BLOCK_SIZE_T={block_size_t} " - f"derived from T={T}; extend the heuristic to include max_window." - ) assert x.shape == (batch, T, nheads, dim) assert dt.shape == x.shape @@ -2539,7 +2781,11 @@ def checkpointing_state_update( if _num_warps is not None: num_warps = _num_warps if _heads_per_block is not None: - heads_per_block = _heads_per_block + # Cap at heads_per_group: HEADS_PER_BLOCK divides the kernel's head + # axis, so a table value larger than the model's heads-per-group + # would overshoot. Protects callers running smaller models than + # the one we tuned against. + heads_per_block = min(_heads_per_block, heads_per_group) if _precompute_num_warps is not None: precompute_num_warps = _precompute_num_warps @@ -2615,23 +2861,37 @@ def checkpointing_state_update( state_tma_descriptor_write = state # dummy; all consuming constexprs False state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False - # Slot permutation — pointer + USE_PERM gate. When the caller provides - # a perm tensor the dl-family launches read pid_b through it; otherwise - # we pass any valid pointer (state_batch_indices) and USE_PERM=False so - # the kernel falls back to pid_grid. - if slot_perm is not None: - assert slot_perm.dtype in (torch.int32, torch.int64), ( - f"slot_perm must be int32/int64, got {slot_perm.dtype}" - ) - assert slot_perm.numel() >= batch, ( - f"slot_perm has {slot_perm.numel()} entries; need >= batch ({batch})" - ) - slot_perm_arg = slot_perm - use_perm = True - else: - # Any valid ptr — gated by USE_PERM=False at compile time. - slot_perm_arg = state_batch_indices if state_batch_indices is not None else state - use_perm = False + # Slot permutation — pointer + USE_PERM gate. Always required: the + # persistent_main kernel reads pid_b through slot_perm to walk the + # write-first sorted batch. persistent_dynamic forces USE_PERM=False + # at the call site (see launch_persistent_dynamic_main below) so the + # perm value doesn't matter for pd, but the tensor must still be valid. + assert isinstance(slot_perm, torch.Tensor), ( + f"slot_perm must be a torch.Tensor, got {type(slot_perm).__name__}" + ) + assert slot_perm.device == device, ( + f"slot_perm must be on device {device}, got {slot_perm.device}" + ) + assert slot_perm.dtype in (torch.int32, torch.int64), ( + f"slot_perm must be int32/int64, got {slot_perm.dtype}" + ) + assert slot_perm.shape == (batch,), ( + f"slot_perm must have shape (batch={batch},), got {tuple(slot_perm.shape)}" + ) + assert isinstance(n_writes, torch.Tensor), ( + f"n_writes must be a torch.Tensor, got {type(n_writes).__name__}" + ) + assert n_writes.device == device, ( + f"n_writes must be on device {device}, got {n_writes.device}" + ) + assert n_writes.dtype == torch.int32, ( + f"n_writes must be int32, got {n_writes.dtype}" + ) + assert n_writes.shape == (1,), ( + f"n_writes must have shape (1,), got {tuple(n_writes.shape)}" + ) + slot_perm_arg = slot_perm + use_perm = True precomp_grid = (batch, nheads // heads_per_block) d_strides = (D.stride(0), D.stride(1)) if D is not None else (0, 0) @@ -2708,28 +2968,12 @@ def launch_dynamic_precompute(rectangle: bool): _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M def launch_persistent_main(write_checkpoint: bool, - n_writes_dev: torch.Tensor, *, - host_n_writes: int | None = None, - skip_empty_halves: bool = True, launch_dependent_kernels: bool = False, rectangle: bool = False): - # `n_writes_dev` is a (1,) int32 device tensor; the kernel reads - # the count from device memory. `host_n_writes` is the same value - # known host-side (when available — pure scenarios) and lets us - # skip the launch entirely if its half is empty. In mix scenarios - # the host doesn't know n_writes per iter without a sync, so - # `host_n_writes is None` and `skip_empty_halves` is forced False - # — both halves always launch and the kernel processes whatever - # range device-n_writes implies. - if skip_empty_halves and host_n_writes is not None: - n_slots_for_kernel = host_n_writes if write_checkpoint else (batch - host_n_writes) - if n_slots_for_kernel <= 0: - return - # Per-main knob selection. The two persistent_main launches (write - # half vs nowrite half) get independent BLOCK_SIZE_M / num_warps / - # num_stages / cta_per_sm / num_loop_stages. See the per-main args - # block in the wrapper signature. + # `n_writes` (wrapper-level) is the (1,) int32 device tensor with the + # write count. Both halves always launch; the kernel's runtime PNAT + # check iterates only the slots that belong to its half. _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE @@ -2739,18 +2983,11 @@ def launch_persistent_main(write_checkpoint: bool, _nls = _nls if _nls else 2 _num_persistent = _cps * _num_sms _num_pid_m_local = (dim + _bsm - 1) // _bsm - # Grid sizing: cap at min(full persistent grid, actual total_work). - # `n_slots` for this launch is `host_n_writes` (write half) / `batch - - # host_n_writes` (nowrite half) when host knows it (pure); else upper - # bound `batch` for mix scenarios where host can't read n_writes_dev - # without a sync. Upper-bound is fine — the kernel's runtime check - # only iterates actual work; the only cost of overcounting is a few - # extra CTAs. - if host_n_writes is not None: - _n_slots_for_launch = host_n_writes if write_checkpoint else (batch - host_n_writes) - else: - _n_slots_for_launch = batch - _total_work_launch = max(1, _n_slots_for_launch * _num_pid_m_local * nheads) + # Grid sizing: cap at min(full persistent grid, upper-bound total work). + # We use `batch` as the upper bound on slots-per-half — overcounting + # by a few CTAs is fine since the kernel's runtime check only + # iterates the slots that actually belong to its half. + _total_work_launch = max(1, batch * _num_pid_m_local * nheads) grid = (min(_num_persistent, _total_work_launch),) # Per-path TMA descriptor — block_shape[0] must match _bsm. _desc = (state_tma_descriptor_write if write_checkpoint @@ -2762,7 +2999,7 @@ def launch_persistent_main(write_checkpoint: bool, x, C, D, z, out, cb_scaled, decay_vec, state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, - n_writes_dev, batch, nheads, + n_writes, batch, nheads, T, max_window, dim, dstate, nheads // ngroups, state.stride(0), state.stride(1), state.stride(2), state.stride(3), state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], @@ -2839,7 +3076,7 @@ def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, x, C, D, z, out, cb_scaled, decay_vec, state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, - n_writes_dev, batch, nheads, + n_writes, batch, nheads, T, max_window, dim, dstate, nheads // ngroups, state.stride(0), state.stride(1), state.stride(2), state.stride(3), state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], @@ -2896,93 +3133,36 @@ def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, # ---- Mode dispatch ---------------------------------------------------- with torch.cuda.device(device.index): if mode == "persistent_dynamic": - # Single-launch persistent kernel covering the full batch. - # Each work-item dispatches via runtime PNAT check (is_write = - # (pnat + T) > MAX). No n_writes/half-split — kernel ignores - # n_writes_dev when IS_DYNAMIC=True (Triton DCEs the load). - # We still need a valid pointer to satisfy the kernel arg - # signature; allocate or reuse `_n_writes_dev`. - n_writes_dev_local = ( - _n_writes_dev if _n_writes_dev is not None - else torch.zeros(1, dtype=torch.int32, device=device) - ) + # Single-launch persistent kernel covering the full batch. Each + # work-item dispatches via runtime PNAT check. Kernel ignores + # n_writes (Triton DCEs the load) when IS_DYNAMIC=True; we still + # pass the wrapper-provided tensor as required by the signature. launch_dynamic_precompute(rectangle=rectangle_for_nowrite) launch_persistent_dynamic_main( - n_writes_dev_local, + n_writes, launch_dependent_kernels=False, rectangle=rectangle_for_nowrite, ) elif mode == "persistent_main": # Persistent-CTA main kernel. One shared dynamic_precompute - # (per-slot dispatch at runtime via PNAT) feeds two - # persistent_main launches (write half + nowrite half). + # (per-slot dispatch via PNAT) feeds two persistent_main + # launches (write half + nowrite half). Both halves ALWAYS + # launch; the kernel's runtime check iterates only the slots + # belonging to its half (write: [0, n_writes), nowrite: + # [n_writes, batch)). # - # Hard-sort contract: caller has pre-sorted slots host-side so - # PNAT is monotone (writes first). Pass the perm via - # slot_perm + USE_PERM. - # - # n_writes is read by the kernel from a (1,) int32 device - # tensor. The caller can provide: - # * _n_writes_dev only (mix): a pre-filled (1,) int32 tensor - # it updates per iter via pre_iter_fn outside the captured - # graph. host can't cheaply read it without a sync, so both - # halves always launch. - # * _n_writes only (non-graph callers, e.g. unit tests): host - # int. We allocate the scratch tensor on the fly. CANNOT - # be used inside CUDA-graph capture — alloc inside capture - # invalidates the stream. - # * Both (pure under graph capture): caller pre-allocates the - # tensor outside capture and tells us the host value too. - # We skip the internal allocation and apply host-skip when - # _persistent_skip_empty_halves=True. This is the - # production-equivalent path the bench's pure cells take. - if _n_writes_dev is not None: - n_writes_dev_local = _n_writes_dev # no allocation - if _n_writes is not None: - # Caller provided both: pure scenario with pre-allocated - # tensor. Use host_n_writes for the skip-empty fast path. - assert 0 <= _n_writes <= batch, ( - f"_n_writes={_n_writes} must be in [0, batch={batch}]" - ) - host_n_writes_local = _n_writes - skip_empty_local = _persistent_skip_empty_halves - else: - # Mix: host doesn't know n_writes without a sync. - host_n_writes_local = None - skip_empty_local = False - else: - # No pre-allocated tensor. Fall back to on-the-fly alloc - # from _n_writes (host int). NOT graph-capture-safe. - assert _n_writes is not None, ( - "mode='persistent_main' requires either _n_writes " - "(host int, non-graph callers) or _n_writes_dev (device " - "tensor, recommended for graph-capture callers)." - ) - assert 0 <= _n_writes <= batch, ( - f"_n_writes={_n_writes} must be in [0, batch={batch}]" - ) - n_writes_dev_local = torch.tensor( - [_n_writes], dtype=torch.int32, device=device, - ) - host_n_writes_local = _n_writes - skip_empty_local = _persistent_skip_empty_halves - # rectangle_for_nowrite=True: precompute populates cb_scaled - # for the rect path; nowrite half uses the rectangle impl; - # write half always replay-style (rect doesn't apply). + # Caller-provided contract: `n_writes` is a (1,) int32 device + # tensor (the kernel reads it at runtime, after the precompute); + # `slot_perm` is a (batch,) int32 device tensor pre-sorted + # write-first. launch_dynamic_precompute(rectangle=rectangle_for_nowrite) launch_persistent_main( write_checkpoint=True, - n_writes_dev=n_writes_dev_local, - host_n_writes=host_n_writes_local, - skip_empty_halves=skip_empty_local, launch_dependent_kernels=True, rectangle=False, # write always replay-style ) launch_persistent_main( write_checkpoint=False, - n_writes_dev=n_writes_dev_local, - host_n_writes=host_n_writes_local, - skip_empty_halves=skip_empty_local, launch_dependent_kernels=False, rectangle=rectangle_for_nowrite, ) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py index 793cbb22c95d..84d249ddfbb3 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py @@ -3044,27 +3044,18 @@ def _run_incr( "modes. Re-run with --sort-slots 1 or " "--hardcode-sort 1." ) - # n_writes plumbing: pure scenarios pass an int - # (host knows the value, can host-skip empty halves); - # mix scenarios pass a (1,) device tensor updated - # per iter by the benchmark pre-iter path. - # _persistent_skip_empty_halves=False on mix so both - # halves always launch (kernel uses device n_writes - # to derive its slot range). + # n_writes plumbing: pure scenarios pass an int (host + # knows the value, wrapper host-skips empty halves); + # mix scenarios pass only a (1,) device tensor updated + # per iter by the benchmark pre-iter path (wrapper + # cannot host-skip since host_n_writes is unknown). if scenario_n_writes_dev is not None: - # Mix path: caller-allocated tensor, updated per - # iter by scenario_pre_iter outside capture. + # Mix: caller-allocated tensor, updated per iter. extra_kwargs["_n_writes_dev"] = scenario_n_writes_dev - extra_kwargs["_persistent_skip_empty_halves"] = False elif mode == "persistent_main": - # Pure: caller pre-allocated `_n_writes_dev_pure` - # outside this lambda (so the alloc doesn't land - # inside the captured graph). Pass both the - # tensor and the host int so the wrapper can use - # host-skip when `_persistent_skip_empty_halves`. + # Pure pm: pre-allocated tensor + host int. extra_kwargs["_n_writes"] = _host_n_writes_pure extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure - extra_kwargs["_persistent_skip_empty_halves"] = scenario_skip_empty elif mode == "persistent_dynamic": # persistent_dynamic pure: kernel ignores n_writes # via IS_DYNAMIC DCE, but the wrapper needs a diff --git a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_slim.py b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_slim.py index a8d11a254331..7f9029dab56b 100644 --- a/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_slim.py +++ b/tests/unittest/_torch/modules/mamba/test_checkpointing_state_update_slim.py @@ -30,6 +30,35 @@ from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update from tensorrt_llm._utils import get_sm_version + +def _make_persistent_inputs(prev_tokens, T, max_window, batch, + state_batch_indices, device): + """Build the (n_writes, slot_perm) tensors the wrapper now REQUIRES. + + persistent_main reads n_writes (a (1,) int32 device tensor) to split + its persistent loop into write/nowrite halves, and slot_perm (a + (batch,) int32 device tensor) to walk the batch write-first. + persistent_dynamic ignores both internally but the wrapper still + requires them — call this helper everywhere the test calls the + wrapper. + + prev_tokens may be cache_size-shaped (paged cache) or batch-shaped. + state_batch_indices, if non-None, selects the active batch slots out + of a cache_size-sized prev_tokens. + """ + if state_batch_indices is not None: + active_pnat = prev_tokens[state_batch_indices.long()] + else: + active_pnat = prev_tokens[:batch] + write_mask = (active_pnat + T) > max_window + n_writes = write_mask.sum().to(torch.int32).reshape(1) + # writes-first stable argsort: True (write) sorts before False (nowrite). + slot_perm = torch.argsort( + (~write_mask).to(torch.int32), stable=True, + ).to(torch.int32) + return n_writes, slot_perm + + # Philox stochastic rounding uses PTX cvt.rs.f16x2.f32 which requires sm >= 100. _skip_pre_sm100 = pytest.mark.skipif( get_sm_version() < 100, reason="Philox stochastic rounding needs sm >= 100" @@ -318,9 +347,8 @@ def test_checkpointing_state_update( # perm. For pure-write or pure-nowrite cases here, all slots have the # same status, so _n_writes is batch or 0 and the perm is identity. # Persistent_dynamic ignores both (kernel uses runtime PNAT dispatch). - pm_kwargs = ( - {"_n_writes": batch if write_checkpoint else 0} - if mode == "persistent_main" else {} + n_writes_t, slot_perm_t = _make_persistent_inputs( + prev_tokens, T, max_window, batch, state_batch_indices, device, ) checkpointing_state_update( test_state, @@ -336,6 +364,8 @@ def test_checkpointing_state_update( B=B2, C=C2, out=test_out, + n_writes=n_writes_t, + slot_perm=slot_perm_t, D=D, dt_bias=dt_bias, dt_softplus=True, @@ -344,7 +374,6 @@ def test_checkpointing_state_update( write_checkpoint=write_checkpoint, rectangle_for_nowrite=rectangle_for_nowrite, mode=mode, - **pm_kwargs, ) # Tolerance rationale: the replay kernel uses bf16 tl.dot for four @@ -542,9 +571,18 @@ def test_checkpointing_state_update( ) # --- old_dA_cumsum (double-buffered, fp32, layout (heads, T)): --- + # WRITE: fresh staging buf starts from 0, store per-step cumsum. + # NOWRITE: append at offset k of active buf — values are continuous + # from the start of the buffer, so add the prefix at position k-1 + # (matches the kernel's cross-step continuity fix). + if write_checkpoint or k == 0: + expected_dAcs = dA_cumsum2[batch_idx].T + else: + prefix = old_dA_cumsum[slot, wb, :, k - 1] # (heads,) + expected_dAcs = dA_cumsum2[batch_idx].T + prefix[:, None] torch.testing.assert_close( old_dA_cumsum_w[slot, wb, :, write_offset : write_offset + T], - dA_cumsum2[batch_idx].T, + expected_dAcs, rtol=1e-4, atol=1e-4, msg=f"old_dA_cumsum written region wrong at k={k} write={write_checkpoint}", ) @@ -556,54 +594,59 @@ def test_checkpointing_state_update( @pytest.mark.parametrize( - "scenario,pnat_per_slot_list,n_writes_expected,slot_perm_list,rectangle_for_nowrite", + "scenario,pnat_per_slot_list,explicit_slot_perm,rectangle_for_nowrite", [ # All-write: every slot has PNAT triggering write - # (PNAT + T > max_window). No permutation needed. - ("all_write", [12, 13, 14, 15], 4, [0, 1, 2, 3], False), - # All-nowrite: every slot fits in the window. n_writes = 0. - ("all_nowrite", [3, 4, 5, 6], 0, [0, 1, 2, 3], False), - # Mixed (write-first sorted via slot_perm): physical slots 2, 3 - # are writes; physical slots 0, 1 are nowrites. slot_perm - # remaps grid pid_b 0..3 to physical slots 2, 3, 0, 1 — so the - # first n_writes=2 grid programs hit write slots and the rest - # hit nowrite slots. - ("mixed_sorted", [3, 10, 12, 16], 2, [2, 3, 0, 1], False), - # Same as mixed_sorted but with rectangle_for_nowrite=True — the - # nowrite half of the persistent loop dispatches to the rectangle - # impl instead of replay-nowrite. Covers the rect-path correctness - # when mixed with the write half in a single persistent kernel. - ("mixed_sorted_rect", [3, 10, 12, 16], 2, [2, 3, 0, 1], True), + # (PNAT + T > max_window). No permutation needed. Auto and + # explicit slot_perm both yield identity here. + ("all_write", [12, 13, 14, 15], None, False), + # All-nowrite: every slot fits in the window. + ("all_nowrite", [3, 4, 5, 6], None, False), + # Mixed PNATs with HAND-CODED slot_perm. The auto-computed perm + # for these PNATs would be [2, 3, 0, 1] — same as the explicit + # value — so this scenario is functionally redundant with the + # auto variant ONLY if `_make_persistent_inputs` (the test + # helper) is itself correct. Keeping a hand-coded copy guards + # against a buggy helper: the kernel still gets a known-good perm + # and the scenario would still pass even if the helper regressed. + ("mixed_explicit", [3, 10, 12, 16], [2, 3, 0, 1], False), + ("mixed_explicit_rect", [3, 10, 12, 16], [2, 3, 0, 1], True), + # Mixed PNATs with AUTO-COMPUTED slot_perm via `_make_persistent_inputs` + # — production-shaped flow. Different PNAT layout from the explicit + # case ([1, 3, 0, 2] perm) so the kernel sees a distinct permutation. + ("mixed_auto", [3, 12, 10, 15], None, False), + ("mixed_auto_rect", [3, 12, 10, 15], None, True), + ], + ids=[ + "all_write", "all_nowrite", + "mixed_explicit", "mixed_explicit_rect", + "mixed_auto", "mixed_auto_rect", ], - ids=["all_write", "all_nowrite", "mixed_sorted", "mixed_sorted_rect"], ) -def test_checkpointing_state_update_persistent_main( - scenario, pnat_per_slot_list, n_writes_expected, slot_perm_list, - rectangle_for_nowrite, +@pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) +def test_checkpointing_state_update_scenarios( + scenario, pnat_per_slot_list, explicit_slot_perm, rectangle_for_nowrite, mode, ): """ - Persistent-CTA main kernel: 1D-grid kernel that loops over - (slot, M-tile, head) work units via tl.range. Caller pre-sorts - slots write-first and passes _n_writes (count of write slots) so - the kernel can split the persistent loop into write and nowrite - halves with the right WRITE_CHECKPOINT constexpr each time. + Combined scenarios test covering both kernel modes (persistent_main, + persistent_dynamic) across a representative mix of write/nowrite + layouts: + + - all_write / all_nowrite: every slot on one branch — verifies the + empty-half early-return on pm and the all-uniform per-slot dispatch + on pd. + - mixed_explicit: hand-coded `slot_perm`, bypassing the test's + `_make_persistent_inputs` helper. Guards against a buggy helper: + if the auto-computation regressed, the auto scenarios would still + pass with the broken value, but this one runs against a known-good + perm and would still detect the kernel-side issue. + - mixed_auto: unsorted PNATs, `slot_perm` auto-computed via the + helper — the production-shaped flow. Setup mirrors test_checkpointing_state_update_sorted_dispatch - (same fixed seeds, same input shapes) so the reference state - evolution is identical and we can compare per-slot output and - HBM-state postconditions to the same reference. - - Cases: - - all_write (n_writes=B): every slot exercises the - WRITE_CHECKPOINT=True branch of the persistent loop. - - all_nowrite (n_writes=0): every slot exercises the - WRITE_CHECKPOINT=False branch. Verifies the kernel handles - the "write half is empty" launch (n_slots=0 → early return). - - mixed_sorted: slots [2, 3] are writes, slots [0, 1] are - nowrites. slot_perm = [2, 3, 0, 1]. Persistent kernel - should call its impl with pid_b ∈ {2, 3} for the write half - and pid_b ∈ {0, 1} for the nowrite half, even though the - grid pid_b_grid is 0..n_slots-1 in each. + (same fixed seeds, same input shapes) so the reference state evolution + is identical and we can compare per-slot output and HBM-state + postconditions to the same reference. """ nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 T = 6 @@ -614,13 +657,6 @@ def test_checkpointing_state_update_persistent_main( pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) pnat_means_write = (pnat_per_slot + T > max_window).tolist() - slot_perm = torch.tensor(slot_perm_list, device=device, dtype=torch.int32) - # Sanity: caller-supplied n_writes must match the actual count of - # write slots in the post-perm order. - write_count = sum(pnat_means_write) - assert write_count == n_writes_expected, ( - f"test setup error: expected {n_writes_expected} writes, got {write_count}" - ) torch.manual_seed(42) A_base = -torch.rand(nheads, device=device) - 0.5 @@ -698,6 +734,17 @@ def test_checkpointing_state_update_persistent_main( test_state = state0.clone() test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + # slot_perm: hand-coded if the scenario provides one, else auto. + # n_writes is always computed from PNATs (auto-computed `n_writes` is + # identical to the hand-coded value when the explicit perm is valid). + if explicit_slot_perm is not None: + slot_perm = torch.tensor(explicit_slot_perm, device=device, dtype=torch.int32) + n_writes_count = sum(pnat_means_write) + n_writes_t = torch.tensor([n_writes_count], dtype=torch.int32, device=device) + else: + n_writes_t, slot_perm = _make_persistent_inputs( + pnat_per_slot, T, max_window, batch, None, device, + ) checkpointing_state_update( test_state, old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), @@ -705,160 +752,11 @@ def test_checkpointing_state_update_persistent_main( pnat_per_slot, x=x2, dt=dt2, A=A, B=B2, C=C2, out=test_out, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, - mode="persistent_main", - rectangle_for_nowrite=rectangle_for_nowrite, + n_writes=n_writes_t, slot_perm=slot_perm, - _n_writes=n_writes_expected, - ) - - torch.testing.assert_close( - test_out.float(), ref_out.float(), - atol=1.0, rtol=0.05, - msg=f"Output mismatch (scenario={scenario})", - ) - - for i in range(batch): - if pnat_means_write[i]: - torch.testing.assert_close( - test_state[i].float(), ref_state_after_replay[i].float(), - atol=1.0, rtol=0.05, - msg=f"Write slot {i}: state mismatch (scenario={scenario})", - ) - else: - torch.testing.assert_close( - test_state[i], state0[i], rtol=0, atol=0, - msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", - ) - - -@pytest.mark.parametrize( - "scenario,pnat_per_slot_list,rectangle_for_nowrite", - [ - # All-write: every slot has PNAT triggering write (PNAT + T > max_window). - ("all_write", [12, 13, 14, 15], False), - # All-nowrite: every slot fits in the window. - ("all_nowrite", [3, 4, 5, 6], False), - # Mixed: some slots write, some nowrite. No pre-sort needed; the - # dynamic kernel dispatches per-slot at runtime via PNAT load. - ("mixed_unsorted", [3, 12, 10, 15], False), - # Mixed with rectangle_for_nowrite=True — the per-slot runtime - # dispatch picks the rect impl for nowrite slots. Covers the - # rect-path under pd's runtime branch (no pre-sort). - ("mixed_unsorted_rect", [3, 12, 10, 15], True), - ], - ids=["all_write", "all_nowrite", "mixed_unsorted", "mixed_unsorted_rect"], -) -def test_checkpointing_state_update_persistent_dynamic( - scenario, pnat_per_slot_list, rectangle_for_nowrite, -): - """ - Persistent-dynamic kernel: 1D persistent-CTA grid covering the full - batch, with runtime per-slot WRITE_CHECKPOINT branch derived from - each slot's PNAT. Single launch, no half-split, no n_writes needed, - no slot_perm needed (handles unsorted batches natively). - - Same setup as test_checkpointing_state_update_persistent_main; we - verify all three scenarios — including a mixed-unsorted batch the - persistent_main kernel can't handle without pre-sorting. - """ - nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 - T = 6 - max_window = 16 - batch = 4 - device = "cuda" - dtype = torch.bfloat16 - - pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) - pnat_means_write = (pnat_per_slot + T > max_window).tolist() - - torch.manual_seed(42) - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) - dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) - D_base = torch.randn(nheads, device=device, dtype=dtype) - D = repeat(D_base, "h -> h p", p=head_dim) - - state0 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=dtype - ) - ref_input_state = state0.float() - - step1_T = max_window - x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) - dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) - B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - - states_buffer_f32 = torch.zeros( - batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) - out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_input_state.clone(), - x1, dt1_input, A, B1, C1, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=cache_idx_for_capture, - intermediate_states_buffer=states_buffer_f32, - cache_steps=step1_T, - out=out1, - disable_state_update=True, - ) - - old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn( - batch, 2, nheads, max_window, device=device, dtype=torch.float32 - ) - cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) - - old_x[:, :step1_T] = x1 - dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) - dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) - for i in range(batch): - buf = cache_buf_idx[i].item() - old_B[i, buf, :step1_T] = B1[i] - old_dt[i, buf, :, :step1_T] = dt1_processed[i].T - old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T - - torch.manual_seed(123) - x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) - dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) - B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - - ref_state_f32 = ref_input_state.clone() - for i in range(batch): - k_i = pnat_per_slot[i].item() - if k_i > 0: - ref_state_f32[i] = states_buffer_f32[i, k_i - 1] - ref_state_after_replay = ref_state_f32.clone() - - ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state_f32, x2, dt2, A, B2, C2, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, out=ref_out, - ) - - test_state = state0.clone() - test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - checkpointing_state_update( - test_state, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), - cache_buf_idx.clone(), - pnat_per_slot, - x=x2, dt=dt2, A=A, B=B2, C=C2, - out=test_out, D=D, dt_bias=dt_bias, dt_softplus=True, state_batch_indices=None, - mode="persistent_dynamic", + mode=mode, rectangle_for_nowrite=rectangle_for_nowrite, ) @@ -899,8 +797,11 @@ def test_checkpointing_state_update_persistent_dynamic( ], ids=["mixed_sorted", "all_write_noskip", "all_nowrite_noskip"], ) +@pytest.mark.parametrize("rectangle_for_nowrite", [True, False], ids=["rect", "norect"]) +@pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_checkpointing_state_update_persistent_main_device_n_writes( scenario, pnat_per_slot_list, n_writes_expected, slot_perm_list, + rectangle_for_nowrite, mode, ): """ Persistent_main with the device-tensor n_writes plumbing. @@ -1023,13 +924,12 @@ def test_checkpointing_state_update_persistent_main_device_n_writes( pnat_per_slot, x=x2, dt=dt2, A=A, B=B2, C=C2, out=test_out, + n_writes=n_writes_dev, # device tensor (was named n_writes_dev in old API) + slot_perm=slot_perm, D=D, dt_bias=dt_bias, dt_softplus=True, state_batch_indices=None, - mode="persistent_main", - slot_perm=slot_perm, - # NEW PATHS: - _n_writes_dev=n_writes_dev, # device tensor (not host int) - _persistent_skip_empty_halves=False, # both halves always launch + mode=mode, + rectangle_for_nowrite=rectangle_for_nowrite, ) torch.testing.assert_close( @@ -1060,7 +960,10 @@ def test_checkpointing_state_update_persistent_main_device_n_writes( ) @pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) @pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) -def test_checkpointing_state_update_philox(state_dtype, nheads, head_dim, d_state, ngroups, paged_cache, T): +@pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) +def test_checkpointing_state_update_philox( + state_dtype, nheads, head_dim, d_state, ngroups, paged_cache, T, mode, +): """ Verify that Philox stochastic rounding produces correct results across all SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). @@ -1124,6 +1027,13 @@ def test_checkpointing_state_update_philox(state_dtype, nheads, head_dim, d_stat prev_tokens = torch.full((cache_size,), T // 2, device=device, dtype=torch.int32) + # max_window is old_x.shape[1] per the wrapper convention; the philox + # test sets old_x = (cache_size, T, ...) so max_window = T here. + _max_window_philox = T + _n_writes_philox, _slot_perm_philox = _make_persistent_inputs( + prev_tokens, T, _max_window_philox, batch, state_batch_indices, device, + ) + common_kwargs = dict( x=x, dt=dt, @@ -1134,6 +1044,9 @@ def test_checkpointing_state_update_philox(state_dtype, nheads, head_dim, d_stat dt_bias=dt_bias, dt_softplus=True, state_batch_indices=state_batch_indices, + n_writes=_n_writes_philox, + slot_perm=_slot_perm_philox, + mode=mode, ) # --- Run without rounding (deterministic RN store) --- @@ -1296,8 +1209,13 @@ def test_philox_rounding_unbiased(state_dtype): prev_tokens = torch.full((batch,), T, device=device, dtype=torch.int32) + # max_window = old_x.shape[1] = T + _n_writes_unb, _slot_perm_unb = _make_persistent_inputs( + prev_tokens, T, T, batch, None, device, + ) common_kwargs = dict( x=x, dt=dt_val, A=A, B=B, C=C, D=D, dt_bias=dt_bias, dt_softplus=True, + n_writes=_n_writes_unb, slot_perm=_slot_perm_unb, ) # 1. fp32 state — captures true post-replay fp32 state. @@ -1381,12 +1299,24 @@ def test_philox_rounding_unbiased(state_dtype): # total_heads (>= 256-512), which the main test with batch=2 never reaches. # This test overrides _heads_per_block to exercise the two-loop structure in # the precompute kernel (store-then-reload of per-head dt/dA_cumsum). +# +# Beyond OUTPUT and STATE checks, this also asserts the kernel's WRITE +# CONTRACT on every cache buffer (old_x, old_B, old_dt, old_dA_cumsum) +# across a sweep of PNATs covering nowrite (PNAT=0,1,T,max_window-T-1, +# max_window-T) and write (PNAT=max_window-T+1, max_window-1) paths. The +# old_dA_cumsum check is the one that originally hid the +# continuous-across-nowrite bug — direct verification prevents regression. +# Both `rectangle_for_nowrite` arms are exercised explicitly (don't rely on +# tuning). +# # Configs: (nheads=16, ngroups=1) and (nheads=32, ngroups=2) both have # heads_per_group=16. The heuristic caps HPB at min(2|4, hpg), so HPB=2, 4. @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) @pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) +@pytest.mark.parametrize("rectangle_nowrite", [True, False], ids=["rect", "norect"]) +@pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_checkpointing_heads_per_block( nheads, head_dim, @@ -1395,6 +1325,8 @@ def test_checkpointing_heads_per_block( state_dtype, T, heads_per_block, + rectangle_nowrite, + mode, ): # PDL flags use wrapper defaults; trimming the parametrize keeps this # suite fast. Coverage of {launch_with_pdl, use_internal_pdl} variations @@ -1404,6 +1336,13 @@ def test_checkpointing_heads_per_block( Verify checkpointing_state_update produces correct results when _heads_per_block > 1, exercising the precompute kernel's two-loop structure (store per-head dt/dA_cumsum in loop 1, reload in loop 2). + + In addition to the output + state checks, this verifies the full write + contract: per-slot the kernel touches the correct staging/active buffer + at the correct offset for old_x, old_B, old_dt, old_dA_cumsum; leaves + untouched regions and the other buffer byte-identical to pre-call; and + (for nowrite) preserves the dA_cumsum prefix continuity by adding the + pre-call old_dA_cumsum[slot, active_buf, head, PNAT-1] value. """ device = "cuda" dtype = torch.bfloat16 @@ -1424,20 +1363,27 @@ def test_checkpointing_heads_per_block( D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) + # max_window = 2*next_pow2(T) so the PNAT=T nowrite case is always + # valid (requires max_window >= 2T) and we have headroom for the + # full PNAT sweep below. + max_window = max(2 * triton.next_power_of_2(T), 16) + cache_size = batch state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) - x1 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + # Generate enough fill data to cover max_window steps (needed for the + # write-slot replay reference, which walks up to max_window-1 steps). + x1 = torch.randn(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, max_window, nheads, device=device, dtype=dtype) dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) - B1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + B1 = torch.randn(batch, max_window, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, max_window, ngroups, d_state, device=device, dtype=dtype) states_buffer_f32 = torch.zeros( - cache_size, T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + cache_size, max_window, nheads, head_dim, d_state, device=device, dtype=torch.float32 ) cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) - out1 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + out1 = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) selective_state_update( state0.clone(), x1, @@ -1450,64 +1396,129 @@ def test_checkpointing_heads_per_block( dt_softplus=True, state_batch_indices=cache_idx_for_capture, intermediate_states_buffer=states_buffer_f32, - cache_steps=T, + cache_steps=max_window, out=out1, disable_state_update=True, ) - old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + # Pre-fill cache buffers. old_x is single-buffer (no dbuf dim); old_B, + # old_dt, old_dA_cumsum are double-buffered. Initialize BOTH buffers + # with controlled random data so "outside write range / other buffer + # unchanged" assertions have well-defined expected values for both. + old_x_init = torch.randn( + cache_size, max_window, nheads, head_dim, device=device, dtype=dtype + ) + old_B_init = torch.randn( + cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype + ) + old_dt_init = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + old_dA_cumsum_init = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) - old_x[:] = x1 - dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) - dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) + # Capture-step data populates ONE buffer per slot (the active one selected + # by cache_buf_idx). Mirrors production: previous step wrote into the + # now-active buffer; the other buffer holds stale data the kernel must + # not touch on a nowrite call. + dt1_proc = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_proc, dim=1) + old_x_init[:] = x1 # single buffer for slot in range(cache_size): - buf = cache_buf_idx[slot].item() - old_B[slot, buf] = B1[slot] - old_dt[slot, buf] = dt1[slot].T - old_dA_cumsum[slot, buf] = dA_cumsum1[slot].T + buf = int(cache_buf_idx[slot].item()) + old_B_init[slot, buf] = B1[slot] + old_dt_init[slot, buf] = dt1_proc[slot].T # (nheads, max_window) + old_dA_cumsum_init[slot, buf] = dA_cumsum1[slot].T # (nheads, max_window) + + # --- PNAT sweep -------------------------------------------------------- + # Cover nowrite (PNAT+T <= max_window) and write (PNAT+T > max_window) + # paths plus the boundary, with both PNAT=0 (no prefix) and PNAT=T (the + # smallest prefix-load case the kernel cares about). + candidate_pnats = [ + 0, # nowrite, no prefix + 1, # nowrite, smallest nontrivial prefix + T, # nowrite, prefix length one stored step + max_window - T - 1, # nowrite, largest PNAT just below threshold + max_window - T, # nowrite, exactly at threshold + max_window - T + 1, # write, smallest PNAT above threshold + max_window - 1, # write, maximum + ] + seen = set() + pnat_list = [] + for p in candidate_pnats: + if 0 <= p < max_window and p not in seen: + seen.add(p) + pnat_list.append(p) + while len(pnat_list) < batch: + pnat_list.append(0) + pnat_list = pnat_list[:batch] + + has_write = any((p + T) > max_window for p in pnat_list) + has_nowrite = any((p + T) <= max_window for p in pnat_list) + assert has_write and has_nowrite, ( + f"PNAT sweep must cover both write and nowrite: {pnat_list}, " + f"T={T}, max_window={max_window}" + ) - k = T - torch.manual_seed(123) + prev_tokens = torch.tensor(pnat_list, device=device, dtype=torch.int32) + pnat_means_write = [(pnat_list[i] + T) > max_window for i in range(batch)] + torch.manual_seed(123) x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + # Per-step processed dt and per-step cumsum — what the kernel writes to + # old_dt and the cumsum-from-zero portion of old_dA_cumsum. + dt2_proc = F.softplus(dt2_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum2_step = torch.cumsum(A_base.float()[None, None, :] * dt2_proc, dim=1) + + # Reference state walk identical to the original test. ref_state_f32 = state0.float().clone() - ref_state_f32[:] = states_buffer_f32[:, k - 1] + for slot in range(batch): + if pnat_list[slot] > 0: + ref_state_f32[slot] = states_buffer_f32[slot, pnat_list[slot] - 1] + ref_state_after_replay = ref_state_f32.clone() ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) selective_state_update( - ref_state_f32, - x2, - dt2, - A, - B2, - C2, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=None, - out=ref_out, + ref_state_f32, x2, dt2, A, B2, C2, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=None, out=ref_out, ) + # Pre-call snapshots double as the "expected" baseline for untouched + # regions / untouched buffer. Kernel operates on the _test copies. test_state = state0.clone() - prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + state_pre = test_state.clone() + old_x_pre = old_x_init.clone() + old_B_pre = old_B_init.clone() + old_dt_pre = old_dt_init.clone() + old_dA_cumsum_pre = old_dA_cumsum_init.clone() + cache_buf_idx_pre = cache_buf_idx.clone() + + old_x_test = old_x_pre.clone() + old_B_test = old_B_pre.clone() + old_dt_test = old_dt_pre.clone() + old_dA_cumsum_test = old_dA_cumsum_pre.clone() + cache_buf_idx_test = cache_buf_idx_pre.clone() + + n_writes_t, slot_perm_t = _make_persistent_inputs( + prev_tokens, T, max_window, batch, None, device, + ) checkpointing_state_update( test_state, - old_x.clone(), - old_B.clone(), - old_dt.clone(), - old_dA_cumsum.clone(), - cache_buf_idx.clone(), + old_x_test, + old_B_test, + old_dt_test, + old_dA_cumsum_test, + cache_buf_idx_test, prev_tokens, x=x2, dt=dt2, @@ -1515,49 +1526,177 @@ def test_checkpointing_heads_per_block( B=B2, C=C2, out=test_out, + n_writes=n_writes_t, + slot_perm=slot_perm_t, D=D, dt_bias=dt_bias, dt_softplus=True, state_batch_indices=None, + mode=mode, + rectangle_for_nowrite=rectangle_nowrite, _heads_per_block=heads_per_block, ) + # ---------------- Output + state checks (existing coverage) ------------- torch.testing.assert_close( test_out, ref_out, rtol=2e-2, atol=1.0, msg=f"Output mismatch with HPB={heads_per_block}, T={T}, " - f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", + f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}, " + f"rect={rectangle_nowrite}, pnats={pnat_list}", ) - expected_state = states_buffer_f32[:, k - 1].to(state_dtype) - torch.testing.assert_close( - test_state, - expected_state, - rtol=2e-2, - atol=1.0, - msg=f"State mismatch with HPB={heads_per_block}, T={T}, " - f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", - ) + for slot in range(batch): + if pnat_means_write[slot]: + torch.testing.assert_close( + test_state[slot].float(), ref_state_after_replay[slot].float(), + rtol=2e-2, atol=1.0, + msg=( + f"Write slot {slot} (PNAT={pnat_list[slot]}): state mismatch " + f"(HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, state_dtype={state_dtype}, " + f"rect={rectangle_nowrite})" + ), + ) + else: + torch.testing.assert_close( + test_state[slot], state_pre[slot], + rtol=0, atol=0, + msg=( + f"Nowrite slot {slot} (PNAT={pnat_list[slot]}): state HBM " + f"modified (HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, state_dtype={state_dtype}, " + f"rect={rectangle_nowrite})" + ), + ) + + # ---------------- New checks: kernel write contract --------------------- + # For each slot: decide active/staging buffer, write offset, build the + # expected per-buffer tensor element-wise, compare against the actual + # kernel-modified tensor. The expected_* tensors are clones of the + # pre-call snapshot with only [target_buf, write_offset:write_end] + # overwritten — so the full-slot equality compares implicitly assert + # the "other buffer untouched" and "outside write range untouched" + # contracts. + for slot in range(batch): + pnat = pnat_list[slot] + is_write = pnat_means_write[slot] + active_buf = int(cache_buf_idx_pre[slot].item()) + staging_buf = 1 - active_buf + + if is_write: + target_buf = staging_buf + write_offset = 0 + else: + target_buf = active_buf + write_offset = pnat + write_end = write_offset + T + + # ----- old_x (single-buffer) ----- + expected_old_x_slot = old_x_pre[slot].clone() + expected_old_x_slot[write_offset:write_end] = x2[slot] + torch.testing.assert_close( + old_x_test[slot], expected_old_x_slot, + rtol=0, atol=0, + msg=( + f"old_x slot {slot} (PNAT={pnat}, is_write={is_write}, " + f"write_offset={write_offset}): mismatch " + f"(HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, rect={rectangle_nowrite})" + ), + ) + + # ----- old_B (double-buffer) ----- + expected_old_B_slot = old_B_pre[slot].clone() + expected_old_B_slot[target_buf, write_offset:write_end] = B2[slot] + torch.testing.assert_close( + old_B_test[slot], expected_old_B_slot, + rtol=0, atol=0, + msg=( + f"old_B slot {slot} (PNAT={pnat}, is_write={is_write}, " + f"target_buf={target_buf}, write_offset={write_offset}): " + f"mismatch (HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, rect={rectangle_nowrite})" + ), + ) + + # ----- old_dt (double-buffer (cache, 2, nheads, max_window)) ----- + # Kernel writes per-head processed dt at [target_buf, :, write_offset:write_end]. + # softplus on chip vs F.softplus host: small ULP diff possible; use + # tight but non-zero tolerance. + expected_old_dt_slot = old_dt_pre[slot].clone() + expected_old_dt_slot[target_buf, :, write_offset:write_end] = dt2_proc[slot].T + torch.testing.assert_close( + old_dt_test[slot], expected_old_dt_slot, + rtol=1e-5, atol=1e-5, + msg=( + f"old_dt slot {slot} (PNAT={pnat}, is_write={is_write}, " + f"target_buf={target_buf}, write_offset={write_offset}): " + f"mismatch (HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, rect={rectangle_nowrite})" + ), + ) + + # ----- old_dA_cumsum (double-buffer (cache, 2, nheads, max_window)) ----- + # WRITE: per-step cumsum starting from 0 (fresh staging buf). + # NOWRITE: continuous — cumsum offset by the prefix value at + # old_dA_cumsum_pre[slot, active_buf, head, PNAT-1] (or 0 if PNAT=0). + # This is the direct regression assertion for the bug just fixed. + expected_old_dAcs_slot = old_dA_cumsum_pre[slot].clone() + step_cumsum = dA_cumsum2_step[slot].T # (nheads, T) + if is_write: + expected_old_dAcs_slot[target_buf, :, write_offset:write_end] = step_cumsum + else: + if pnat > 0: + prefix = old_dA_cumsum_pre[slot, active_buf, :, pnat - 1] + else: + prefix = torch.zeros(nheads, device=device, dtype=torch.float32) + expected_old_dAcs_slot[target_buf, :, write_offset:write_end] = ( + step_cumsum + prefix[:, None] + ) + torch.testing.assert_close( + old_dA_cumsum_test[slot], expected_old_dAcs_slot, + rtol=1e-5, atol=1e-5, + msg=( + f"old_dA_cumsum slot {slot} (PNAT={pnat}, is_write={is_write}, " + f"target_buf={target_buf}, write_offset={write_offset}): " + f"mismatch (HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, rect={rectangle_nowrite})" + ), + ) # HPB > 1 multi-step test. Production chains decode steps; bugs in # buffer ordering or stale cache values accumulate across steps and can # be invisible in a single-step test. +# +# Divergent per-slot acceptance: slot 0 accepts all T tokens each step, +# slot 1 accepts a smaller fixed count. This forces the write/nowrite +# mask to differ between slots on multiple steps (n_writes ∈ {0, 1, 2} +# within an 8-step run and slot_perm hits both identity [0,1] and the +# swapped [1,0] order — exercising the kernel's per-slot dispatch). +# `rectangle_nowrite` forces the rectangle vs non-rectangle nowrite path +# (via `mode="persistent_main"` + `rectangle_for_nowrite=…` kwargs) +# rather than relying on the tuning table's mode pick. @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("T", [6, 16], ids=["T6", "T16"]) @pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) @pytest.mark.parametrize("paged_cache", [False, True], ids=["contig", "paged"]) +@pytest.mark.parametrize("rectangle_nowrite", [True, False], ids=["rect", "norect"]) +@pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_checkpointing_heads_per_block_multistep( - nheads, head_dim, d_state, ngroups, state_dtype, T, heads_per_block, paged_cache + nheads, head_dim, d_state, ngroups, state_dtype, T, heads_per_block, + paged_cache, rectangle_nowrite, mode, ): """ Chain N decode steps with HPB > 1 and verify each step's output matches a fresh reference. A bug that mixes up WRITE/READ buffers, writes wrong data to cache, or races in the two-loop structure would accumulate - across steps. + across steps. Per-slot acceptance diverges so write/nowrite masks + differ across slots — the n_writes=1 case (mixed batch) is exercised. """ batch = 2 device = "cuda" @@ -1581,11 +1720,9 @@ def test_checkpointing_heads_per_block_multistep( if paged_cache: cache_size = 4 state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) - slots = state_batch_indices else: cache_size = batch state_batch_indices = None - slots = slice(None) all_x = [] all_dt = [] @@ -1604,41 +1741,72 @@ def test_checkpointing_heads_per_block_multistep( cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype ) + # max_window = 2*np2(T) so the buffer has slack for several nowrite + # steps before overflow — required for the divergent-acceptance pattern + # to produce a chain of mixed write/nowrite steps within n_steps=8. + # For T=6 this is 16, for T=16 this is 32. + max_window = max(triton.next_power_of_2(2 * T), 16) + + # Per-slot acceptance counts. Slot 0 advances by 6 per step, slot 1 by + # 4 — divergence (PNAT trajectories desync, write_mask varies between + # slots, n_writes hits 0/1/2 within the 8-step run). Values constant + # across T because they exercise the kernel's per-slot dispatch + # independent of T; acc <= T is the only requirement. + accepted_per_slot = [6, 4] + accepted_tensor = torch.tensor(accepted_per_slot, device=device, dtype=torch.int32) + + # Per-slot reference: each slot's reference state advances by only its + # `accepted` tokens per step, not all T. selective_state_update doesn't + # natively support per-slot variable T, so run it once per slot per step + # with a single-slot batch view. ref_state = state_init.float().clone() ref_outs = [] - ref_slots = ( - state_batch_indices - if paged_cache - else torch.arange(batch, device=device, dtype=torch.int32) - ) for step in range(n_steps): out_step = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state, - all_x[step], - all_dt[step], - A, - all_B[step], - all_C[step], - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=ref_slots, - out=out_step, - ) + for s_local in range(batch): + acc = accepted_per_slot[s_local] + if acc == 0: + continue + c_idx = ( + state_batch_indices[s_local].item() + if state_batch_indices is not None else s_local + ) + s_state = ref_state[c_idx:c_idx + 1].clone() + s_x = all_x[step][s_local:s_local + 1, :acc].contiguous() + s_dt = all_dt[step][s_local:s_local + 1, :acc].contiguous() + s_B = all_B[step][s_local:s_local + 1, :acc].contiguous() + s_C = all_C[step][s_local:s_local + 1, :acc].contiguous() + s_out = torch.zeros(1, acc, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + s_state, s_x, s_dt, A, s_B, s_C, + D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=torch.tensor([0], device=device, dtype=torch.int32), + out=s_out, + ) + out_step[s_local, :acc] = s_out[0] + ref_state[c_idx] = s_state[0] ref_outs.append(out_step) test_state = state_init.clone() - old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.zeros(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - old_dA_cumsum = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_x = torch.zeros(cache_size, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.zeros(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.zeros(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.zeros(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) + # Per-active-slot PNAT tracker, advanced by `accepted` (not T) per step. + pnat_active = torch.zeros(batch, device=device, dtype=torch.int32) + for step in range(n_steps): - k = T if step > 0 else 0 - prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + prev_tokens = torch.zeros(cache_size, device=device, dtype=torch.int32) + if state_batch_indices is not None: + prev_tokens[state_batch_indices.long()] = pnat_active + else: + prev_tokens[:] = pnat_active test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + n_writes_t, slot_perm_t = _make_persistent_inputs( + prev_tokens, T, max_window, batch, state_batch_indices, device, + ) checkpointing_state_update( test_state, @@ -1654,27 +1822,49 @@ def test_checkpointing_heads_per_block_multistep( B=all_B[step], C=all_C[step], out=test_out, + n_writes=n_writes_t, + slot_perm=slot_perm_t, D=D, dt_bias=dt_bias, dt_softplus=True, state_batch_indices=state_batch_indices, _heads_per_block=heads_per_block, + mode=mode, + rectangle_for_nowrite=rectangle_nowrite, ) - if paged_cache: - cache_buf_idx[slots] = 1 - cache_buf_idx[slots] - else: - cache_buf_idx[:] = 1 - cache_buf_idx - - torch.testing.assert_close( - test_out, - ref_outs[step], - rtol=2e-2, - atol=2.0, - msg=f"Output mismatch at step {step} with HPB={heads_per_block}, " - f"T={T}, nheads={nheads}, ngroups={ngroups}, " - f"state_dtype={state_dtype}, paged_cache={paged_cache}", + # PNAT update uses `accepted` (per-slot), not T. Write step resets + # PNAT to `accepted` of this step (new buffer starts fresh); nowrite + # appends `accepted` to current PNAT. + write_mask = (pnat_active + T) > max_window + new_pnat_active = torch.where( + write_mask, accepted_tensor, pnat_active + accepted_tensor, ) + pnat_active = new_pnat_active + cache_active_idx = ( + state_batch_indices.long() if state_batch_indices is not None + else torch.arange(batch, device=device) + ) + write_slots = cache_active_idx[write_mask] + cache_buf_idx[write_slots] = 1 - cache_buf_idx[write_slots] + + # Per-slot output comparison: only the first `accepted` output tokens + # of each slot are meaningful (the rest are produced from "candidate" + # tokens that wouldn't be accepted in production). + for s_local in range(batch): + acc = accepted_per_slot[s_local] + if acc == 0: + continue + torch.testing.assert_close( + test_out[s_local, :acc], + ref_outs[step][s_local, :acc], + rtol=2e-2, + atol=2.0, + msg=f"Output mismatch at step {step}, slot {s_local} " + f"(acc={acc}) with HPB={heads_per_block}, T={T}, " + f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}, " + f"paged_cache={paged_cache}, rectangle={rectangle_nowrite}", + ) # ----- SR grid-bracket tests (fp8 and fp16) ----- From c578008a956ed42b22d354c7ea9c0c9fbf274dca Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 19 May 2026 09:30:04 -0700 Subject: [PATCH 55/89] mamba_checkpointing/slim: wire split-form loop_stages/cta_per_sm in tuning + bench-slim API alignment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split-form tuning resolution ============================ The slim wrapper's table-resolution code at line 2597-98 only resolved the unsplit `_cta_per_sm` and `_num_loop_stages` knobs from `_DEFAULT_TUNING`, but every pm entry in the table stores the per-half split forms (`_cta_per_sm_write`/`_cta_per_sm_nowrite`, `_num_loop_stages_write`/ `_num_loop_stages_nowrite`). Without the resolution lines, callers that relied on auto-resolution from the table got None on the split forms — pm's kernel constexpr `NUM_LOOP_STAGES` fell through to Triton's default in `tl.range`, and `NUM_PERSISTENT` fell through to the hardcoded `or 1` fallback (cta_per_sm=1 instead of the tuned per-batch value). Add the four missing `.get(...)` lines so the split-form values land in the wrapper kwargs. search_driver was unaffected (always passes explicit values via cell-list). Bench-slim wired to slim wrapper ================================ benchmark_replay_selective_state_update_slim.py's fast-import path was loading `checkpointing_state_update.py` (the non-slim wrapper) instead of `checkpointing_state_update_slim.py`. Switch the fast-import to load slim, and update the bench's variant_fn call to use slim's `n_writes`/`slot_perm` required kwargs (dropping the legacy `_n_writes_dev`/`_n_writes` plumbing). None-default plumbing for tuning knobs ====================================== Default `--modes` and `--rectangle-for-nowrite` to None so the wrapper resolves them from `_DEFAULT_TUNING` per (batch, dtype, sr) cell. Tag emission for ALL tuning knobs (M, W, S, pW, pS, H, R, CT, CPS, LS, FL, WS, TMARL, TMAWL, TMANL, TMAWS, RECT, MODE) now emits "auto" when the bench-level value is None — uniform key set across cells regardless of whether the caller pinned a knob or deferred to the table. Conditional mode-checks and n_writes_dev pre-alloc updated to include `None` so mode=None doesn't silently skip plumbing. `_kernels_per_iter_incremental` gets a `mode is None` branch (defaults to pm's K=3; pd-resolved cells will land in the skipped sidecar — acceptable for audit runs). The slim bench now exercises the slim wrapper end-to-end and the slim wrapper picks tuned values from `_DEFAULT_TUNING` for both pm and pd paths. search_driver workflows in scripts/ are unaffected. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...aseline_checkpointing_state_update_slim.py | 3182 +++++++++++++++++ ...aseline_checkpointing_state_update_slim.py | 3173 ++++++++++++++++ .../mamba/checkpointing_state_update_slim.py | 9 + ...mark_replay_selective_state_update_slim.py | 183 +- 4 files changed, 6467 insertions(+), 80 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/mamba/_v0_baseline_checkpointing_state_update_slim.py create mode 100644 tensorrt_llm/_torch/modules/mamba/_v3_baseline_checkpointing_state_update_slim.py diff --git a/tensorrt_llm/_torch/modules/mamba/_v0_baseline_checkpointing_state_update_slim.py b/tensorrt_llm/_torch/modules/mamba/_v0_baseline_checkpointing_state_update_slim.py new file mode 100644 index 000000000000..65c73cdd5aa7 --- /dev/null +++ b/tensorrt_llm/_torch/modules/mamba/_v0_baseline_checkpointing_state_update_slim.py @@ -0,0 +1,3182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +# +# Adapted from: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +# SPDX-FileCopyrightText: Copyright contributors to the sglang project +# +# Copyright (c) 2024, Tri Dao, Albert Gu. +# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID +from tensorrt_llm._utils import get_sm_version + +from .softplus import softplus + + +# Lazy global allocator for Triton TMA tensor descriptors. Required by any +# host- or device-built tensor_descriptor; without it Triton raises at first +# launch. See TMA backlog item #17 / scratch experiment notes. +_TMA_ALLOCATOR_SET = False + + +def _ensure_tma_allocator() -> None: + global _TMA_ALLOCATOR_SET + if _TMA_ALLOCATOR_SET: + return + + def _alloc_fn(size, alignment, stream): + # Triton expects an int8 buffer of `size` bytes; alignment is enforced + # by the allocator returning a buffer satisfying it (PyTorch's + # cudaMalloc-backed tensors are 256B-aligned, so we're fine). + return torch.empty(size, device="cuda", dtype=torch.int8) + + triton.set_allocator(_alloc_fn) + _TMA_ALLOCATOR_SET = True + + +@triton.jit +def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 pair → fp16x2 using Philox random bits. + + Uses PTX cvt.rs.f16x2.f32 which rounds each fp32 value to fp16 using + the random bits to break ties, avoiding systematic rounding bias that + accumulates over many decode steps with fp16 state. + + Adapted from flashinfer (Apache-2.0, vLLM/mamba lineage). + """ + return tl.inline_asm_elementwise( + asm="""{ + cvt.rs.f16x2.f32 $0, $2, $1, $3; + }""", + constraints=("=r,r,r,r,r"), + args=(x, rand), + dtype=tl.float16, + is_pure=True, + pack=2, + ) + + +@triton.jit +def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 quad → fp8 e4m3 using Philox random bits. + + Uses PTX cvt.rs.satfinite.e4m3x4.f32 which combines stochastic rounding + and saturating cast in a single op (output is final fp8, no separate + clamp needed). The reversed source-register order {$4,$3,$2,$1} is + load-bearing — PTX packs leftmost source into the high byte but Triton's + pack=4 is little-endian, so the natural {$1,$2,$3,$4} order would + silently shuffle every group of 4 contiguous outputs. + + Requires SM_100a+ (Blackwell B200). Caller must gate at the wrapper + level — this kernel does not check. + + Adapted from vLLM PR #40012 (Apache-2.0). + """ + return tl.inline_asm_elementwise( + asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", + constraints="=r,r,r,r,r,r,r,r,r", + args=(x, rand), + dtype=tl.float8e4nv, + is_pure=True, + pack=4, + ) + + +@triton.jit +def _bitrev32(x: tl.tensor) -> tl.tensor: + return tl.inline_asm_elementwise( + asm="brev.b32 $0, $1;", + constraints="=r,r", + args=(x,), + dtype=tl.uint32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _stochastic_round_int8_packed( + x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor +) -> tl.tensor: + """Stochastic rounding for int8 using one random uint32 per 4 values.""" + low = rand & 0x0000FFFF + high = (rand >> 16) & 0x0000FFFF + low_rev = _bitrev32(low) >> 16 + high_rev = _bitrev32(high) >> 16 + rand_pos = offs_n & 3 + rand16 = tl.where( + rand_pos == 0, + low, + tl.where(rand_pos == 1, low_rev, tl.where(rand_pos == 2, high, high_rev)), + ) + rand01 = rand16.to(tl.float32) * (1.0 / float(1 << 16)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +@triton.jit +def _stochastic_round_int16_packed( + x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor +) -> tl.tensor: + """Stochastic rounding for int16 using one random uint32 per 2 values.""" + rand_bits = tl.where((offs_n & 1) == 0, rand, _bitrev32(rand)) + rand01 = (rand_bits & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +# Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, +# old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. +# Grid: (batch, nheads // HEADS_PER_BLOCK). + + +@triton.jit() +def _replay_precompute_impl( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers (both buffers reachable via stride_*_dbuf). This + # kernel writes to either the active (= cache_buf_idx) or inactive + # (= 1 - cache_buf_idx) buffer depending on WRITE_CHECKPOINT — see + # comment block at top of kernel body. + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + # Double-buffer index (per cache slot) — selects this step's "active" + # buffer (= where the historical inputs for this step live). + cache_buf_idx_ptr, + # Per-request accepted-tokens count (already-cached old tokens at + # [0, PNAT) of the active buffer; new tokens this step go after them + # on no-checkpoint steps). + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides: (cache, 2, T, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Checkpointing flag — selects target buffer + offset for new-token + # cache writes. See "Cache write semantics" block below. + # Runtime (not constexpr): the only WRITE_CHECKPOINT-dependent code in + # this body is the write_buf/write_offset selection, which is plain + # arithmetic — no constexpr-shaped tile or whole-block gate. Letting + # it be runtime lets the dynamic dispatch kernel call us once with the + # per-slot needs_write flag instead of inlining two specializations. + write_checkpoint, +): + pid_b = tl.program_id(axis=0) + pid_hg = tl.program_id(axis=1) # head-group index + first_head = pid_hg * HEADS_PER_BLOCK + + # Resolve cache index for writes + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + + # --- Cache write semantics --- + # cache_buf_idx names this step's "active" buffer — the one with the + # historical inputs at [0, PNAT). The other buffer is "staging". + # + # Where do we write new tokens this step? + # WRITE_CHECKPOINT=False (no overflow): append to ACTIVE buffer at + # offset [PNAT : PNAT+T). Caller does NOT flip cache_buf_idx + # afterward; PNAT_next = PNAT + accepted. [0, PNAT) preserved. + # WRITE_CHECKPOINT=True (would overflow): write to STAGING buffer at + # [0, T). Caller flips cache_buf_idx afterward; next step's + # active = the one we just wrote. PNAT_next = accepted. Old + # data in the previous active buffer is folded into state via + # the replay update and discarded. This matches today's replay + # kernel behavior exactly. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if write_checkpoint: + write_buf = 1 - buf_active + write_offset = 0 + else: + write_buf = buf_active + write_offset = prev_num_accepted_tokens + + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + t_mask = offs_t < T + n_mask = offs_n < dstate + + # Causal mask is shared across all heads (depends only on offs_t) + causal_mask = offs_t[:, None] >= offs_t[None, :] + valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] + + # --- Vectorized pre-wait phase across HEADS_PER_BLOCK heads --- + # Compute dt, dA_cumsum, decay_vec as (H, T) tiles. Pre-compute + # scale_combo = decay_matrix * dt[:, None, :] as an (H, T, T) tile that + # stays in registers across gdc_wait — eliminates the post-wait reload + # of dt + dA_cumsum and the per-head loop. + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) + + # Load dt (H, T) + dt_addrs = ( + dt_ptr + pid_b * stride_dt_batch + + heads_block[:, None] * stride_dt_head + + offs_t[None, :] * stride_dt_T + ) + dt = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias[:, None] + if DT_SOFTPLUS: + dt = softplus(dt) + + A = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) # (H,) + dA_cumsum = tl.cumsum(A[:, None] * dt, axis=1) # (H, T) + decay_vec = tl.exp(dA_cumsum) # (H, T) + + # Cross-step continuity for old_dA_cumsum: when appending to active_buf at + # offset PNAT > 0, the previous step left a running cumsum at [0, PNAT) + # whose tail value lives at active_buf[head, PNAT-1]. Add that tail to + # this step's per-step-restarted cumsum before storing so the buffer + # holds one continuous cumsum across N back-to-back nowrites. Write path + # (write_buf = 1 - buf_active, write_offset = 0) starts fresh, no prefix. + # Both branches are on scalar runtime values (write_checkpoint and PNAT), + # uniform across the block — use scalar if to short-circuit the load. + if write_checkpoint or prev_num_accepted_tokens == 0: + prev_total = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) + else: + last_cumsum_ptrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T + ) + prev_total = tl.load(last_cumsum_ptrs).to(tl.float32) + + # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of write_buf. + old_dt_addrs = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block[:, None] * stride_old_dt_head + + (write_offset + offs_t)[None, :] * stride_old_dt_T + ) + tl.store(old_dt_addrs, dt, mask=t_mask[None, :]) + + old_dA_cumsum_addrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block[:, None] * stride_old_dA_cumsum_head + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T + ) + tl.store(old_dA_cumsum_addrs, dA_cumsum + prev_total[:, None], mask=t_mask[None, :]) + + # decay_vec scratch — always at offs_t. + decay_vec_addrs = ( + decay_vec_ptr + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) + tl.store(decay_vec_addrs, decay_vec, mask=t_mask[None, :]) + + # scale_combo (H, T, T) = exp(dA_cumsum[h, t1] - dA_cumsum[h, t2]) * dt[h, t2] + # Stays live across gdc_wait — used post-wait to compute CB_scaled. + decay_matrix = tl.exp(dA_cumsum[:, :, None] - dA_cumsum[:, None, :]) # (H, T, T) + scale_combo = decay_matrix * dt[:, None, :] # (H, T, T) + + # --- Wait for upstream kernel (external PDL) before loading B and C --- + # All dt processing above is independent of conv1d outputs. + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- + group_idx = first_head // nheads_ngroups_ratio + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_all = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_all = tl.load( + B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + # Compute raw CB once — shared across all heads in this block + raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) + + # Store B to cache at [write_offset : write_offset+T) of write_buf. + if first_head % nheads_ngroups_ratio == 0: + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + write_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_all, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # --- Vectorized post-wait phase: scale_combo (H, T, T) is still live in + # registers from pre-wait; multiply by raw_CB (T, T), apply causal mask, + # store as one (H, T, T) tile. --- + CB_scaled_block = tl.where( + valid_mask[None, :, :], + raw_CB[None, :, :] * scale_combo, + 0.0, + ) # (H, T, T) + cb_scaled_addrs = ( + cb_scaled_ptr + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_t[None, None, :] * stride_cb_j + ) # (H, T, T) + cb_store_mask = ( + (offs_t[None, :, None] < BLOCK_SIZE_T) + & (offs_t[None, None, :] < BLOCK_SIZE_T) + ) + tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) + + +# Replay-style precompute kernel. Thin wrapper around _replay_precompute_impl +# that carries the @triton.heuristics for constexpr derivation; called from +# the Python wrapper on the replay-style path (write or replay-nowrite). +@triton.jit() +def _rectangle_precompute_impl( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle + decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) + # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite + # path: read from buf_active at [0, PNAT), write new tokens at + # [PNAT, PNAT+T) of buf_active (same buffer). + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # rectangle K-axis bound + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides (rectangle: (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides: (cache, 2, T_max, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T_max) + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T_max) + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, +): + pid_b = tl.program_id(axis=0) + pid_hg = tl.program_id(axis=1) + first_head = pid_hg * HEADS_PER_BLOCK + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_buf = buf_active + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout: old at [0, PNAT) (mask is_old_k); + # new at [MAX-T, MAX) at compile-time shift K_NEW_SHIFT = MAX - T. + # PNAT + T <= MAX is guaranteed on the nowrite path → no overlap. + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis (output rows) + offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis (rectangle input cols) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + t_mask = offs_t < T + n_mask = offs_n < dstate + + # K-axis masks (approach C: runtime PNAT-offset instead of K_NEW_SHIFT) + # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Loop 1: per-head dt processing. dt → dt_processed → dA_cumsum → + # decay_vec_new (= exp(cumAdt_new)). Stored to write_buf for next step. + # decay_vec_full (= total_decay * decay_vec_new) is finalized in loop 2 + # once total_decay is loaded; loop 1 stores raw decay_vec_new to scratch. + for h_local in range(HEADS_PER_BLOCK): + head_idx = first_head + h_local + + dt_base = dt_ptr + pid_b * stride_dt_batch + head_idx * stride_dt_head + dt = tl.load(dt_base + offs_t * stride_dt_T, mask=t_mask, other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + head_idx * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias + if DT_SOFTPLUS: + dt = softplus(dt) + + A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) + dA_cumsum = tl.cumsum(A * dt, axis=0) + + # Cross-step continuity for old_dA_cumsum: rectangle precompute runs + # only on the nowrite path (write_buf == buf_active, write_offset == PNAT). + # Add the running tail from buf_active[head_idx, PNAT-1] so the buffer + # holds one continuous cumsum across back-to-back nowrites. PNAT is + # scalar/uniform, use scalar if to short-circuit the load at PNAT=0. + if prev_num_accepted_tokens == 0: + prev_total = 0.0 + else: + last_cumsum_ptr = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T + ) + prev_total = tl.load(last_cumsum_ptr).to(tl.float32) + + # Store dt and dA_cumsum to write_buf at [write_offset, write_offset+T) + # for next step's replay/rectangle use. + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + head_idx * stride_old_dt_head + ) + tl.store( + old_dt_base + (write_offset + offs_t) * stride_old_dt_T, + dt, + mask=t_mask, + ) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + head_idx * stride_old_dA_cumsum_head + ) + tl.store( + old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, + dA_cumsum + prev_total, + mask=t_mask, + ) + + # ---- Hoisted: cache-only loads independent of conv1d ---- + # old_B (group-level, BLOCK_K × BLOCK_DSTATE = ~8KB tile) and the + # decay_vec_full per-head pre-compute (which writes to DRAM and doesn't + # need cross-gdc_wait variables) are issued BEFORE gdc_wait so their + # HBM latency overlaps with conv1d. Per-head factor_dt/exp_diff stay + # below gdc_wait — they need cross-iteration spans, which Triton can't + # express without a DRAM round-trip; the per-head LOADS in the post- + # wait loop are small and cheap, so leave them. + group_idx = first_head // nheads_ngroups_ratio + + # Group-level: old B from active buffer at [0, PNAT) of the K-axis. + old_B_read_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + buf_active * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_load = tl.load( + old_B_read_base + + safe_old_k[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + mask=is_old_k[:, None] & n_mask[None, :], + other=0.0, + ) + + # Pre-wait: vectorized across HEADS_PER_BLOCK heads. Compute decay_vec_full + # (H, T) and combo = factor_dt * exp_diff (H, T, K). Store decay_vec_full; + # combo_block stays in registers across gdc_wait — used directly post-wait + # to compute rect_CB_scaled without a global memory roundtrip. + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) + + # Per-head bases (H,) — broadcast with offs_k or offs_t for 2D loads. + old_dt_read_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_active * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_read_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + ) + old_dt_write_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_write_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + ) + + # (H, K) loads at [0, PNAT) — old data from previous step. + hk_mask = is_old_k[None, :] # (1, K) + old_dt_all = tl.load( + old_dt_read_h[:, None] + safe_old_k[None, :] * stride_old_dt_T, + mask=hk_mask, other=0.0, + ).to(tl.float32) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, + mask=hk_mask, other=0.0, + ).to(tl.float32) + # (H, T) loads at [PNAT, PNAT+T) — this step's dA_cumsum_new from loop 1. + # With the cross-step continuity fix in loop 1, the values stored at + # [PNAT, PNAT+T) are the continuous cumsum (prefix + per-step new + # cumsum) — i.e., continuous_cumsum[PNAT..PNAT+T-1] in global indexing. + ht_mask = t_mask[None, :] # (1, T) + dA_cumsum_new = tl.load( + old_dA_cumsum_write_h[:, None] + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T, + mask=ht_mask, other=0.0, + ).to(tl.float32) + # (H, K) loads at K_NEW_SHIFT-shifted positions for new tokens. + hkn_mask = is_new_k[None, :] + dt_at_kn = tl.load( + old_dt_write_h[:, None] + + (write_offset + safe_k_new)[None, :] * stride_old_dt_T, + mask=hkn_mask, other=0.0, + ).to(tl.float32) + dA_cumsum_at_kn = tl.load( + old_dA_cumsum_write_h[:, None] + + (write_offset + safe_k_new)[None, :] * stride_old_dA_cumsum_T, + mask=hkn_mask, other=0.0, + ).to(tl.float32) + + # decay_vec_full[t] = exp(continuous_cumsum[PNAT+t]) — directly the + # continuous value now stored at buffer position write_offset+t. Was + # decomposed as total_decay * exp(per_step_new[t]) when the buffer held + # per-step (non-continuous) cumsum; with the continuity fix the value + # IS continuous_cumsum[PNAT+t] so no decomposition is needed. + decay_vec_full_block = tl.exp(dA_cumsum_new) # (H, T) + decay_vec_addrs = ( + decay_vec_ptr + + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) # (H, T) + tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) + + # combo_block = factor_dt * exp_diff — (H, T, K). Stays in registers + # across gdc_wait. With continuous cumsum in the buffer, s_k for any k + # (old or new) is simply -continuous_cumsum[k]; exp_diff[t, k] then + # equals exp(continuous_cumsum[PNAT+t] - continuous_cumsum[k]) — the + # decay weight for token k's contribution to the output at position + # PNAT+t. No need to subtract any "total" — the dA_cumsum_new[t] term + # already carries the full prefix. + # + # Numerical note: pre-fix this kernel computed `total - old_dA[k]` + # (small-minus-small) then summed `+ dA_cumsum_new[t]` (also small, + # per-step). Post-fix `s_k = -old_dA[k]` is large-magnitude positive + # and `dA_cumsum_new[t]` is large-magnitude negative; their sum + # cancels back to the same small value. Cancellation error is bounded + # by ulp(max_magnitude) ≈ 2^-23 · |continuous_cumsum| — negligible + # for max_window ≤ ~1024. Still one exp on the sum (not two muls of + # exps), so no overflow regression vs the original formulation. + factor_dt = tl.where(is_old_k[None, :], old_dt_all, dt_at_kn) # (H, K) + s_k = tl.where( + is_old_k[None, :], + -old_dA_cumsum_all, + -dA_cumsum_at_kn, + ) # (H, K) + # exp_diff (H, T, K) = exp(s_k (H, 1, K) + dA_cumsum_new (H, T, 1)). + exp_diff = tl.exp(s_k[:, None, :] + dA_cumsum_new[:, :, None]) + combo_block = factor_dt[:, None, :] * exp_diff # (H, T, K) + + # ---- gdc_wait: from here on we depend on conv1d's outputs ---- + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # Conv1d outputs: B and C + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + B_new_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_all = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_new_orig = tl.load( + B_new_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_new_shifted = tl.load( + B_new_base + safe_k_new[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=is_new_k[:, None] & n_mask[None, :], + other=0.0, + ) + # Disjoint masks: old at [0, PNAT), new at [K_NEW_SHIFT, K_NEW_SHIFT+T). + B_combined = old_B_load + B_new_shifted + raw_rect_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_combined).to(tl.bfloat16)) + + # Append new B to cache at [PNAT, PNAT+T) of write_buf (once per group). + if first_head % nheads_ngroups_ratio == 0: + old_B_write_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + write_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_write_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_new_orig, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # Causal mask (BLOCK_SIZE_T × BLOCK_SIZE_K, shared across heads). + # Approach C: new tokens at runtime [PNAT, PNAT+T) instead of K_NEW_SHIFT. + t_idx_2d = offs_t[:, None] + k_idx_2d = offs_k[None, :] + is_old_k_2d = k_idx_2d < prev_num_accepted_tokens + k_new_idx_2d = k_idx_2d - prev_num_accepted_tokens + is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) + causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] + + # Post-wait vectorized: combo_block (H, T, K) is still live in registers. + # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store as + # one (H, T, K) tile. + rect_CB_scaled_block = tl.where( + causal_combined[None, :, :], + raw_rect_CB[None, :, :] * combo_block, + 0.0, + ) # (H, T, K) + cb_scaled_addrs = ( + cb_scaled_ptr + + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_k[None, None, :] * stride_cb_j + ) # (H, T, K) + cb_store_mask_3d = ( + (offs_t[None, :, None] < BLOCK_SIZE_T) + & (offs_k[None, None, :] < BLOCK_SIZE_K) + ) # (1, T, K) → broadcasts to (H, T, K) + tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) + + +# Rectangle precompute kernel. Thin wrapper around _rectangle_precompute_impl +# that carries the @triton.heuristics for constexpr derivation; called from +# the Python wrapper on the rectangle nowrite path. +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _dynamic_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. + RECTANGLE: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + pid_b = tl.program_id(axis=0) + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH + # write_checkpoint is now runtime in replay precompute, so a single + # call site handles both write and nowrite for the replay branch. + # Take rectangle only when RECTANGLE is True AND this slot doesn't + # need write; everything else funnels into replay. + if needs_write_runtime or not RECTANGLE: + _replay_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + needs_write_runtime, + ) + else: + _rectangle_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + ) + + +# Main kernel: tl.dot replay + precomputed CB output. +# Grid: (cdiv(dim, M), batch, nheads). + + +@triton.jit() +def _persistent_main_impl( + # Per-work-unit indices (computed by the persistent wrapper). + # `pid_b` is the post-perm slot index (caller has already applied any + # slot permutation and slot_offset). + pid_m, + pid_b, + pid_h, + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or + # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor + # USE_TMA_STORE is enabled (kernel ignores it via constexpr). + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + # IS_DYNAMIC: kept in the signature for caller-side bookkeeping (the + # outer _persistent_main_kernel still inspects it to decide the slot- + # IS_DYNAMIC: when True (persistent_dynamic), is_write is per-slot from + # PNAT. When False (persistent_main), is_write is constexpr from + # WRITE_CHECKPOINT. See also WC_IS_CONSTEXPR below. + IS_DYNAMIC: tl.constexpr, + # WC_IS_CONSTEXPR: when True, force is_write = WRITE_CHECKPOINT (constexpr) + # regardless of IS_DYNAMIC. Callers in RECT=1 use this in the is_w=True + # arm of _persistent_main_kernel (we know all slots that reach this call + # need is_write=True because is_w was the PNAT-derived runtime check, and + # this arm only fires when is_w is True). Passing WRITE_CHECKPOINT=True + # as a literal at the call site + WC_IS_CONSTEXPR=True here lets the inner + # body DCE the nowrite path under IS_DYNAMIC=True too — same codegen + # quality as persistent_main mode (-3.7% measured at b=1024 dyn-shape). + # When False (RECT=0 callers, where both write and nowrite slots are + # dispatched to ONE call), use the original runtime is_write under + # IS_DYNAMIC=True; avoids the binary-doubling regression that two + # specialized calls would cause. + WC_IS_CONSTEXPR: tl.constexpr = False, + # TMA flags — picked inside body based on is_write. When is_write is + # constexpr (either IS_DYNAMIC=False or WC_IS_CONSTEXPR=True), the + # use_tma_load = USE_TMA_LOAD_WRITE if is_write else USE_TMA_LOAD_NOWRITE + # ternary constexpr-folds and only one TMA load form survives. + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # IS_DYNAMIC: kernel-mode label, used by the OUTER _persistent_main_kernel + # to decide slot-range derivation and outer is_w dispatch strategy + # (constexpr WC for persistent_main; runtime is_w split -> 2 specialized + # impl calls for persistent_dynamic). Inside this impl, IS_DYNAMIC is + # NOT consulted at runtime -- WRITE_CHECKPOINT is the only constexpr that + # gates the write/nowrite codegen, in BOTH modes. + + # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized + # state dtype (int8 / int16 / float8e4nv) and only those. + tl.static_assert( + (QUANT_MAX > 0.0) + == ( + (state_ptr.dtype.element_ty == tl.int8) + or (state_ptr.dtype.element_ty == tl.int16) + or (state_ptr.dtype.element_ty == tl.float8e4nv) + ), + "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", + ) + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + # Resolve is_write: see WC_IS_CONSTEXPR / IS_DYNAMIC docs in the param + # list above. Three cases: + # - WC_IS_CONSTEXPR=True (RECT=1 is_w=True arm callers): use WC + # constexpr. Caller knows the slot needs write; inner DCEs nowrite + # paths. Avoids the binary-doubling overhead that calling the impl + # twice would cause, while still constexpr-DCEing the nowrite half. + # - IS_DYNAMIC=True (RECT=0 caller, persistent_dynamic): runtime + # branch on PNAT. Both write and nowrite codegen live in one body + # (no bloat) — same as the pre-refactor behavior. + # - IS_DYNAMIC=False (persistent_main): WC constexpr from caller. + if WC_IS_CONSTEXPR: + is_write: tl.constexpr = WRITE_CHECKPOINT + elif IS_DYNAMIC: + is_write = (prev_num_accepted_tokens + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_write = WRITE_CHECKPOINT + if is_write: + write_buf = 1 - active_buf # noqa: F841 + write_offset = 0 + else: + write_buf = active_buf # noqa: F841 + write_offset = prev_num_accepted_tokens + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # Load state. state_tma_descriptor is a host-built tensor_descriptor + # over a flat (cache*nheads*dim, dstate) view of state when any TMA + # path is enabled; raw `state_ptr` is the underlying tensor and is + # always passed. state_ptrs / state_ptr_raw are the raw-pointer view + # used for !TMA load and store paths. offs_y is the flat row index + # for TMA load/store; computed unconditionally (cheap int math; DCE'd + # when no TMA path is reachable). + state_mask = m_mask[:, None] & n_mask[None, :] + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + # Load state. Branch on is_write (constexpr = WRITE_CHECKPOINT in BOTH + # modes after the outer-dispatch refactor), then constexpr-pick TMA-vs- + # tl.load per side. Outer `if` DCE's, only the matching side's + # constexpr-gated load survives -- same compile-time picking for both + # persistent_main and persistent_dynamic (the latter dispatches at the + # outer kernel level so each impl instance sees a constexpr WC). + if is_write: + if USE_TMA_LOAD_WRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + else: + if USE_TMA_LOAD_NOWRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, + other=1.0, + ).to(tl.float32) + state = state * decode_scale[:, None] + + # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) + group_idx = pid_h // nheads_ngroups_ratio + + old_window_mask = offs_window < prev_num_accepted_tokens + + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + active_buf * stride_old_dt_dbuf + + pid_h * stride_old_dt_head + ) + old_dt_all = tl.load( + old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 + ).to(tl.float32) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + active_buf * stride_old_dA_cumsum_dbuf + + pid_h * stride_old_dA_cumsum_head + ) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, + mask=old_window_mask, other=0.0, + ).to(tl.float32) + + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( + tl.float32 + ) + + coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all + + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + old_x_all = tl.load( + old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + mask=old_window_mask[:, None] & m_mask[None, :], + other=0.0, + ) + + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + active_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_all = tl.load( + old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + mask=old_window_mask[:, None] & n_mask[None, :], + other=0.0, + ).to(tl.float32) + + dB_scaled = coeff[:, None] * old_B_all + + total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) + state *= total_decay + + state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) + + if is_write: + if USE_RS_ROUNDING: + # Generate random tensor for stochastic rounding. The amount of + # randomness needed depends on the SR codegen path: + # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) + # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) + # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs + # int16 SR (24b + bitrev32): 1 b32 per 2 outputs + # The PTX cvt.rs.* instructions consume a single 32-bit random + # and split the bits internally for 2 or 4 conversions. Generate + # only what's actually consumed and broadcast to fill the unused + # slots — saves Philox rounds proportionally. + if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: + RAND_DIVISOR: tl.constexpr = 4 # fp8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: + RAND_DIVISOR: tl.constexpr = 4 # int8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: + RAND_DIVISOR: tl.constexpr = 2 # int16 SR + elif QUANT_MAX == 0.0: + RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) + else: + RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized + + rand_seed = tl.load(rand_seed_ptr) + base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head + # Number of unique randoms per row = dstate / RAND_DIVISOR. + # randint4x emits 4 randoms per offset, so use that / 4 offsets. + offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) + rand_offsets_q = ( + base_rand + + offs_m[:, None] * stride_state_dim + + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) + ) # (M, dstate / (4*RAND_DIVISOR)) + if PHILOX_ROUNDS > 0: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + else: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) + r01 = tl.join(r0, r1) + r23 = tl.join(r2, r3) + r0123 = tl.join(r01, r23) + rand_compact = tl.reshape( + r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) + ) + # Broadcast each unique rand to RAND_DIVISOR adjacent positions. + # Pack-group (pack=2 fp16 / pack=4 fp8) consumes adjacent positions; + # the unique rand lands at the asm's read slot; duplicates feed + # the dead slots. Triton's broadcast_to is stride-0 in IR. + if RAND_DIVISOR > 1: + rand_3d = rand_compact[:, :, None] + rand_3d = tl.broadcast_to( + rand_3d, + (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), + ) + rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + else: + rand = rand_compact + + if QUANT_MAX > 0.0: + amax = tl.max(tl.abs(state), axis=1) + encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) + decode_scale = 1.0 / encode_scale + state_scales_ptrs = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + + offs_m * stride_state_scales_dim + ) + tl.store(state_scales_ptrs, decode_scale, mask=m_mask) + state_q = state * encode_scale[:, None] + if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): + _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_fp8sr) + else: + tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) + else: + if USE_RS_ROUNDING: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized SR fall-through expects int8 or int16; " + "fp8 SR is handled by the prior branch.", + ) + if state_ptrs.dtype.element_ty == tl.int8: + state_q = _stochastic_round_int8_packed( + state_q, rand, offs_n[None, :] + ) + else: + state_q = _stochastic_round_int16_packed( + state_q, rand, offs_n[None, :] + ) + elif state_ptrs.dtype.element_ty != tl.float8e4nv: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized RN with explicit round() expects int8 or int16.", + ) + state_q = tl.extra.cuda.libdevice.round(state_q) + state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) + _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_cast) + else: + tl.store(state_ptrs, _state_q_cast, mask=state_mask) + elif USE_RS_ROUNDING: + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "Non-quantized SR only supports fp16 state.", + ) + _state_sr = _stochastic_round_fp16x2(state, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_sr) + else: + tl.store(state_ptrs, _state_sr, mask=state_mask) + else: + _state_cast = state.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_cast) + else: + tl.store(state_ptrs, _state_cast, mask=state_mask) + + # Phase 2: Output using precomputed CB_scaled and decay_vec + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + x_all = tl.load( + x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_base + + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_all, + mask=t_mask[:, None] & m_mask[None, :], + ) + x_all = x_all.to(tl.float32) + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + tl.float32 + ) + + init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) + out_all = init_out + cb_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# `_persistent_rectangle_impl`: rectangle nowrite path for the persistent +# kernel. Body is a copy of `_rectangle_main_impl` with `pid_m`/`pid_b`/`pid_h` +# lifted to args (same pattern as `_persistent_main_impl` vs `_replay_main_impl`). +# Called only for nowrite slots when the kernel runs with RECTANGLE=True. +# Dropped from the rect impl: LAUNCH_DEPENDENT_KERNELS / USE_PERM / REVERSE_PERM +# (kernel-level, signalled once at top); slot_perm_ptr (kernel resolves perm). +@triton.jit() +def _persistent_rectangle_impl( + # Per-work-unit indices (computed by the persistent wrapper). + pid_m, + pid_b, + pid_h, + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as + # replay path). Used when USE_TMA_LOAD; ignored otherwise. + state_tma_descriptor, + state_scales_ptr, # only consulted when QUANT_MAX > 0 + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides (rectangle (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + QUANT_MAX: tl.constexpr, + USE_TMA_LOAD: tl.constexpr = False, +): + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout (matches precompute). + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_k = tl.arange(0, BLOCK_SIZE_K) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # K-axis masks (approach C: PNAT-runtime offset, matches precompute). + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Load state. Quant scale hoist: defer `* decode_scale` post-matmul. + if USE_TMA_LOAD: + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state = state_tma_descriptor.load([offs_y, 0]) + else: + state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, other=1.0, + ).to(tl.float32) + else: + state = state.to(tl.float32) + + # Group / pointer offset setup + group_idx = pid_h // nheads_ngroups_ratio + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + # Hoist: old_x doesn't depend on conv1d/precompute; load before gdc_wait. + old_x_load = tl.load( + old_x_base + + safe_old_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + mask=is_old_k[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + x_K = tl.load( + x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=is_new_k[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_base + + offs_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_K, + mask=is_new_k[:, None] & m_mask[None, :], + ) + + x_K_f32 = x_K.to(tl.float32) + x_combined = old_x_load + x_K_f32 + + if HAS_D or HAS_Z: + sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) + x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) + else: + x_all = x_K_f32 # placeholder; unused + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec_full = tl.load( + decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 + ).to(tl.float32) + + state_out = ( + tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) + * decay_vec_full[:, None] + ) + if QUANT_MAX > 0.0: + state_out = state_out * decode_scale[None, :] + + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) + + out_all = state_out + token_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# Persistent main kernel: 1D grid, persistent CTA loop. +# Heuristics mirror those of `_checkpointing_main_kernel`. +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"NUM_PID_M_BLOCKS": lambda args: triton.cdiv(args["dim"], args["BLOCK_SIZE_M"])} +) +@triton.jit() +def _persistent_main_kernel( + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view. + # Shared across BOTH the replay path (consumed by _persistent_main_impl + # when USE_TMA_LOAD_*/STORE) AND the rectangle path (consumed by + # _persistent_rectangle_impl when USE_TMA_LOAD) — same descriptor, same + # block_shape, just gated by separate constexprs per impl. Wrapper sets + # this to a TensorDescriptor when ANY of the three TMA flags is on, else + # to `state_ptr` (raw); each impl ignores it via its own constexpr when + # not consuming it. + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + rand_seed_ptr, + pad_slot_id, + # Persistent-loop work-distribution scalars. Caller pre-sorts the batch + # write-first; the kernel uses (n_writes, batch_total, WRITE_CHECKPOINT) + # to derive its own slot range. Write half processes [0, n_writes), + # nowrite half processes [n_writes, batch_total). + # + # n_writes_ptr is a device pointer to a (1,) int32 tensor. Reading + # from device memory (rather than taking a Python int kernel arg) is + # required so mix-mode benchmarking can vary n_writes per iter inside + # a captured CUDA graph — the source tensor's contents change, the + # pointer doesn't. Cost: one int load per kernel launch (~negligible). + # When IS_DYNAMIC=True the value is unused (Triton DCEs the load). + n_writes_ptr, # int32 *: device-side count of write-mode slots + batch_total, # int32: total slot count + nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + USE_PERM: tl.constexpr, + # NUM_PERSISTENT: runtime int (not constexpr). Used ONLY as the loop + # stride in `tl.range(pid, total_work, NUM_PERSISTENT, ...)`. Making it + # runtime collapses the cta_per_sm tuning dim from the kernel's compile + # signature: 8 CPS values used to mean 8x recompiles; now they share one + # compiled kernel. Work decomposition (pid_m, pid_b_local, pid_h) does + # NOT depend on NUM_PERSISTENT — it uses constexpr NUM_PID_M_BLOCKS and + # runtime n_slots_local — so loop unrolling and flatten=/num_stages=/ + # warp_specialize= optimizations on `tl.range` operate independently of + # the stride value. + NUM_PERSISTENT, + NUM_LOOP_STAGES: tl.constexpr, + NUM_PID_M_BLOCKS: tl.constexpr, + FLATTEN: tl.constexpr, + WARP_SPECIALIZE: tl.constexpr, + IS_DYNAMIC: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr = 16, # rectangle K-axis (heuristic-derived) + RECTANGLE: tl.constexpr = False, # when True, dispatch nowrite slots to _persistent_rectangle_impl + # 3 TMA toggles per the 3 live paths per-compilation: + # USE_TMA_LOAD_WRITE — replay-style state load when is_write + # USE_TMA_LOAD_NOWRITE — nowrite-path state load (rect when RECTANGLE, + # else replay-nowrite) + # USE_TMA_STORE — replay-style state store (only fires on write + # path; no-op when not is_write) + # Wrapper picks USE_TMA_LOAD_NOWRITE = _use_tma_rect_load (if rectangle) + # or _use_tma_replay_nowrite_load (if not). + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # PDL signal: fire once at kernel entry (not per work unit). + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + # Load runtime n_writes from device memory. Read once at kernel entry; + # used only by the !IS_DYNAMIC slot-range derivation below. Triton + # DCEs the load when IS_DYNAMIC=True (n_writes is dead there). + n_writes = tl.load(n_writes_ptr) + + # Derive this kernel's slot range. Two modes: + # IS_DYNAMIC=False (persistent_main): caller pre-sorts and splits halves; + # slot range is [0, n_writes) when WRITE_CHECKPOINT else [n_writes, batch_total) + # IS_DYNAMIC=True (persistent_dynamic): single launch covers full batch; + # each work-item dispatches via runtime PNAT check inside the impl. + if IS_DYNAMIC: + slot_lo = 0 + slot_hi = batch_total + else: + if WRITE_CHECKPOINT: + slot_lo = 0 + slot_hi = n_writes + else: + slot_lo = n_writes + slot_hi = batch_total + n_slots_local = slot_hi - slot_lo + + pid = tl.program_id(axis=0) + total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads + + # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) + # with pid_m varying fastest (M-tile cache locality on state load), then + # slot, then head — mirrors the existing 3D grid's axis ordering + # (axis=0 fastest = pid_m). + for tile_id in tl.range( + pid, total_work, NUM_PERSISTENT, + flatten=FLATTEN, num_stages=NUM_LOOP_STAGES, warp_specialize=WARP_SPECIALIZE, + ): + pid_m = tile_id % NUM_PID_M_BLOCKS + pid_b_local = (tile_id // NUM_PID_M_BLOCKS) % n_slots_local + pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) + # Translate local slot index → global slot index. When USE_PERM is + # set, the caller-provided slot_perm gives the original slot index + # for the post-sort position. + pid_b_grid = pid_b_local + slot_lo + if USE_PERM: + pid_b = tl.load(slot_perm_ptr + pid_b_grid) + else: + pid_b = pid_b_grid + + # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle + # impl. Pre-load cache_batch_idx + PNAT here only for the RECTANGLE + # path's branch decision. Both impls re-load and handle pad_slot_id + # internally (Triton's L1 cache makes the duplicate loads ~free). + if RECTANGLE: + if HAS_CACHE_BATCH_INDICES: + cbi_pre = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + is_pad = cbi_pre == pad_slot_id + else: + cbi_pre = pid_b.to(tl.int64) + is_pad = False + if not is_pad: + pnat_pre = tl.load(prev_num_accepted_tokens_ptr + cbi_pre) + if IS_DYNAMIC: + is_w = (pnat_pre + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_w = WRITE_CHECKPOINT + if is_w: + # Pass WRITE_CHECKPOINT=True constexpr to specialize this + # impl call for the write path. Under IS_DYNAMIC=True, the + # kernel-level WRITE_CHECKPOINT is False (launcher default), + # but the OUTER is_w branch we are inside narrows the + # runtime path to writes-only, so we override to True here + # so the impl's constexpr-gated `if is_write:` blocks DCE + # to the write-only codegen. Under IS_DYNAMIC=False + # (persistent_main), the kernel-level WRITE_CHECKPOINT is + # itself True for this half (write half launches with + # WC=True), and the outer is_w = WRITE_CHECKPOINT = True + # constexpr-folds; passing literal True here is consistent + # and constexpr-equivalent. + _persistent_main_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, + True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) + True, # WC_IS_CONSTEXPR — force inner to use WC constexpr + # 3 TMA flags: write-load fires here (we're in the + # is_write branch), nowrite-load is dead (no slot + # reaches it), store fires (write path). + USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + ) + else: + # Rectangle nowrite: pass state_ptr (raw, always) + + # state_tma_descriptor (the single unified descriptor — + # same memory replay paths use). Rect impl gates use + # of the descriptor via its USE_TMA_LOAD constexpr. + _persistent_rectangle_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, + LAUNCH_WITH_PDL, QUANT_MAX, + USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle + ) + # else: pad slot — skip both impls (both would early-return anyway) + else: + # No rectangle path — single _persistent_main_impl call covers + # both write and nowrite slots via WC constexpr (non-dynamic) or + # runtime is_write (IS_DYNAMIC=True). Pass all 3 TMA flags; + # impl picks USE_TMA_LOAD_WRITE vs USE_TMA_LOAD_NOWRITE based on + # its computed is_write — constexpr-folds when is_write is + # constexpr (non-dyn), runtime branch when IS_DYNAMIC=True. + # (Reverted from outer two-call dispatch: that doubled the + # compiled body size under IS_DYNAMIC=True and regressed RECT=0 + # perf by ~+24%.) + _persistent_main_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, + WRITE_CHECKPOINT, IS_DYNAMIC, + False, # WC_IS_CONSTEXPR=False — RECT=0 has both write/nowrite slots in one call + USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + ) + + +# ============================================================================ +# Python wrapper +# ============================================================================ + + +_QUANT_MAX_BY_DTYPE = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, +} + + +# --------------------------------------------------------------------------- +# Default tunings — looked up by (effective_batch, dtype, sr) when the caller +# leaves mode/knobs as None. +# +# Effective batch = raw_batch × nheads_per_rank. Our sweep was at TP=8 with +# the standard Mamba2 nheads; at call time we compute it from the input +# tensor shape so callers at other TP / nheads pick up the right cell. +# +# Schema: dict[(dtype_str, sr_str)] → list[(eff_batch_threshold, mode, knobs)] +# sorted by threshold ascending. Lookup finds the first threshold ≥ eff_b +# (so missing intermediate batches fall up to the next tuned cell). If +# eff_b exceeds the largest threshold, use the largest entry. +# +# Each `knobs` dict only contains keys for the chosen mode; the wrapper +# unpacks them with the same name as the matching kwargs. Caller-provided +# kwargs always win over table values. +# +# This table is intentionally NOT parameterized by T or max_window. Our +# sweep was T=6, max_window=16. Callers outside that regime silently get +# the same numbers — they may be suboptimal but they're correct. +# +# Source: audit_v2.py --emit-tuning. Auto-generated from per-cell search +# winners (best of pd / pm by bucket_expected_renorm). Sweep was TP=8 with +# NHEADS=128 → nheads_per_rank=16; thresholds are in effective_batch units. +# Missing dtype/SR combos (fp16/RN, int8/RN, fp8/*) fall back via the +# _resolve_tuning chain — RN→SR for same dtype, then fp8→int8/SR. +_DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { + ("fp32", "RN"): [ + ( 16, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 8, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.22us + ( 32, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.17us + ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.08us + ( 128, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=8, score=9.00us + ( 256, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=16, score=10.92us + ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=13.53us + ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=19.50us + ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=30.28us + ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=50.32us + ( 8192, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=90.99us + (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=171.69us + ], + ("fp16", "SR"): [ + ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.16us + ( 32, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.01us + ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=7.95us + ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 4, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=8.87us + ( 256, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 2, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.28us + ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=12.90us + ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=16.71us + ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=25.71us + ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=39.80us + ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=71.34us + (16384, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=133.51us + ], + ("int8", "SR"): [ + ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.34us + ( 32, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.36us + ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 2, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.40us + ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 10, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=9.37us + ( 256, "persistent_dynamic", {'_block_size_m': 16, '_cta_per_sm': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages': 1, '_num_stages': 4, '_num_warps': 1, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.02us + ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=32, score=13.15us + ( 1024, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=17.82us + ( 2048, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=27.01us + ( 4096, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=43.23us + ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=77.01us + (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=140.43us + ], +} + + +# Knob names that map between the modes' single-value (pd) and split-value +# (pm) namespaces. Used by `_bridge_tuning_knobs` when caller forces a mode +# different from the table's recommendation. +_PD_TO_PM_SPLIT_MAP = { # pd unsplit knob → (pm_write_knob, pm_nowrite_knob) + "_block_size_m": ("_block_size_m_write", "_block_size_m_nowrite"), + "_num_warps": ("_num_warps_write", "_num_warps_nowrite"), + "_num_stages": ("_num_stages_write", "_num_stages_nowrite"), + # CPS / LS are persistent-loop knobs; pd uses _cta_per_sm + _num_loop_stages + # as unsplit, pm uses _cta_per_sm_write/_nowrite + _num_loop_stages_write/_nowrite. + "_cta_per_sm": ("_cta_per_sm_write", "_cta_per_sm_nowrite"), + "_num_loop_stages": ("_num_loop_stages_write", "_num_loop_stages_nowrite"), +} + + +def _bridge_tuning_knobs(knobs: dict, from_mode: str, to_mode: str) -> dict: + """Convert a tuning dict between pd ↔ pm knob namespaces. + + pd → pm: copy each unsplit value to both write/nowrite split knobs; drop + the unsplit form (pm doesn't read it). + pm → pd: take the nowrite split value as the unsplit knob; drop the + write/nowrite split forms (pd doesn't read them). + Shape knobs that exist in both modes (_heads_per_block, _flatten, + _warp_specialize, TMA flags, rectangle_for_nowrite) carry over unchanged. + """ + out = dict(knobs) + if from_mode == "persistent_dynamic" and to_mode == "persistent_main": + for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): + if unsplit in out: + out.setdefault(pm_w, out[unsplit]) + out.setdefault(pm_nw, out[unsplit]) + del out[unsplit] + elif from_mode == "persistent_main" and to_mode == "persistent_dynamic": + for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): + if pm_nw in out: + out.setdefault(unsplit, out[pm_nw]) + out.pop(pm_w, None) + out.pop(pm_nw, None) + return out + + +def _resolve_tuning( + batch: int, nheads_per_rank: int, dt_str: str, sr_str: str, +) -> tuple[str, dict] | None: + """Look up the default mode + knobs for this (eff_batch, dt, sr) cell. + + Returns (mode, knobs_dict) or None if the table has no entry covering + this dtype/sr (including the fp8→int8/SR and dtype/RN→dtype/SR fallbacks). + Returning None lets the wrapper fall back to caller-provided kwargs or + kernel-side defaults. + """ + eff_b = batch * max(1, nheads_per_rank) + # Lookup chain. Order: + # 1. Exact (dt, sr). + # 2. (dt, SR) if RN missing for that dtype. + # 3. Cross-dtype fallback for dtypes we haven't tuned: + # bf16 / int16 → fp16/SR + # fp8 → int8/SR + # Unknown dtype → raise. + valid_dtypes = {"fp32", "fp16", "bf16", "int8", "int16", "fp8"} + if dt_str not in valid_dtypes: + raise ValueError( + f"checkpointing_state_update: unsupported state dtype {dt_str!r}; " + f"expected one of {sorted(valid_dtypes)}" + ) + keys_to_try = [(dt_str, sr_str)] + if sr_str == "RN": + keys_to_try.append((dt_str, "SR")) + if dt_str in ("bf16", "int16"): + keys_to_try.append(("fp16", "SR")) + elif dt_str == "fp8": + keys_to_try.append(("int8", "SR")) + entries = None + for k in keys_to_try: + if k in _DEFAULT_TUNING: + entries = _DEFAULT_TUNING[k] + break + if entries is None: + return None + # Find first threshold ≥ eff_b; if none, use largest entry. + for thresh, mode, knobs in entries: + if eff_b <= thresh: + return mode, dict(knobs) + thresh, mode, knobs = entries[-1] + return mode, dict(knobs) + + +def checkpointing_state_update( + state: torch.Tensor, + old_x: torch.Tensor, + old_B: torch.Tensor, + old_dt: torch.Tensor, + old_dA_cumsum: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + out: torch.Tensor, + # Required persistent-mode plumbing (REQUIRED for both pd and pm; pd + # ignores both internally but the wrapper still demands them): + # n_writes : (1,) int32 device tensor with the count of write-mode + # slots in the batch. pm uses it to size the two halves; + # pd ignores it (per-slot runtime PNAT check). + # slot_perm : (batch,) int32 device tensor remapping grid pid → slot. + # pm uses it to cluster writes first (kernel grid step is + # write_half then nowrite_half); pd ignores it. Callers + # that don't care about ordering should pass arange(batch). + n_writes: torch.Tensor, + slot_perm: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + pad_slot_id: int = PAD_SLOT_ID, + rand_seed: torch.Tensor | None = None, + philox_rounds: int = 10, + state_scales: torch.Tensor | None = None, + launch_with_pdl=False, + use_internal_pdl=True, + write_checkpoint: bool = True, + rectangle_for_nowrite: bool | None = None, + mode: str | None = None, + _block_size_m: int | None = None, + _num_warps: int | None = None, + _num_stages: int | None = None, + _precompute_num_warps: int | None = None, + _precompute_num_stages: int | None = None, + _heads_per_block: int | None = None, + _maxnreg: int | None = None, + _num_ctas: int | None = None, + # Per-main knobs (override shared values for one half of the dl-family / + # persistent_main launches). Default None = tied to the shared value + # (backward compat). The two main kernels (write vs nowrite) have + # different per-slot work — write does a state shift + store, nowrite + # just appends — so the optimum (M, W, S, H) can differ. Precompute + # knobs are intentionally NOT split: shared precompute wins (cheaper + # launch, hotter precompute outputs in L2). Persistent CPS / LS knobs + # are also split per-main since the two persistent_main launches have + # different grid sizes. + _block_size_m_write: int | None = None, + _block_size_m_nowrite: int | None = None, + _num_warps_write: int | None = None, + _num_warps_nowrite: int | None = None, + _num_stages_write: int | None = None, + _num_stages_nowrite: int | None = None, + # Note: heads_per_block / precompute_num_warps are NOT split — they only + # affect the precompute kernel, which is shared across write/nowrite. + # TMA state-tensor toggles — 4 independent paths (see CHECKPOINTING_DESIGN.md + # item #17 for measured perf profiles). Each is False=raw load/store, True= + # use a host-built TMA tensor_descriptor for that path. + _use_tma_rect_load: bool | None = None, # rect kernel's state load (nowrite-only) + _use_tma_replay_write_load: bool | None = None, # replay-style state load when WC=True + _use_tma_replay_write_store: bool | None = None, # replay-style state store when WC=True + _use_tma_replay_nowrite_load: bool | None = None, # replay-style state load when WC=False + # Persistent-mode tuning kwargs (consulted for both pd and pm; pd uses + # _cta_per_sm / _num_loop_stages, pm uses the _write/_nowrite splits): + # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally + # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. + # _num_loop_stages : int — `num_stages` arg on the inner `tl.range(...)` + # persistent loop. Note: this is loop-level, NOT the kernel-arg + # `num_stages` (which only pipelines dot-feeding loads). + # _flatten : bool — `flatten` arg on `tl.range(...)`. + # _warp_specialize : bool — `warp_specialize` arg on `tl.range(...)`. + _cta_per_sm: int | None = None, + _num_loop_stages: int | None = None, + _flatten: bool | None = None, + _warp_specialize: bool | None = None, + # Per-main persistent-specific knobs. Same rationale as the BLOCK_SIZE_M + # split above: the two persistent_main launches (write half vs nowrite + # half) have different grid sizes and per-work-item costs, so they may + # want different cta_per_sm / num_loop_stages. + _cta_per_sm_write: int | None = None, + _cta_per_sm_nowrite: int | None = None, + _num_loop_stages_write: int | None = None, + _num_loop_stages_nowrite: int | None = None, +): + """ + Replay SSM state update with precomputed CB and tl.dot fast-forward. + + Two-kernel architecture: + 1. Precompute kernel: computes CB_scaled and decay_vec from B, C, dt, A. + Writes processed dt/dA_cumsum/B to double-buffered cache for next step. + 2. Main kernel: replays old tokens via tl.dot fast-forward on cached data, + then computes output using precomputed CB_scaled and new x/C inputs. + + PDL (Programmatic Dependent Launch) chain: + conv1d → (external PDL) → precompute → (internal PDL) → main + External PDL: precompute starts while conv1d is running; gdc_wait() + in precompute blocks until conv1d completes before loading B/C. + Internal PDL: main starts while precompute is running; main's replay + phase uses only cached data from the previous step. gdc_wait() in + main blocks until precompute completes before loading conv1d outputs + (x, C) and precompute outputs (CB_scaled, decay_vec). + + Uses double-buffered cache tensors. cache_buf_idx[slot] indicates which + buffer (0 or 1) to READ from for replay. The WRITE buffer is 1 - read. + Caller must flip cache_buf_idx[slot] after each call. + + Arguments: + state: (cache, nheads, dim, dstate) in-place. After the call, contains + the state after replaying prev_num_accepted_tokens old tokens. + old_x: (cache, T, nheads, dim) bf16 — old x cache (single-buffered). + old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. + old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. + old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. + cache_buf_idx: (cache,) int32 — which buffer to read (0 or 1). + prev_num_accepted_tokens: (cache,) int32. + x: (batch, T, nheads, dim) new token inputs. + dt: (batch, T, nheads, dim) with stride(-1)==0 (tie_hdim). + A: (nheads, dim, dstate) with stride(-1)==0, stride(-2)==0 (tie_hdim). + B: (batch, T, ngroups, dstate). + C: (batch, T, ngroups, dstate). + out: (batch, T, nheads, dim) preallocated output. + D: (nheads, dim) optional feed-through parameter. + z: (batch, T, nheads, dim) optional silu gate. + dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). + state_batch_indices: (batch,) optional cache slot mapping. + rand_seed: optional single-element int64 CUDA tensor for Philox PRNG seed. + When provided, state is stochastically rounded on store. Supported + for state.dtype in (fp16, int8, int16, fp8_e4m3fn); other dtypes + silently use deterministic rounding. fp16+SR and fp8+SR both + require sm_100a (Blackwell B200+) — wrapper asserts this loudly. + philox_rounds: number of Philox PRNG rounds (default 10). + state_scales: required when state.dtype in (int8, int16, fp8_e4m3fn). + Shape (cache_size, nheads, dim), fp32. Per-(head, dim) channel + decode scale (= 1 / encode_scale). The kernel writes scales on + checkpoint steps and reads them on load (broadcast over dstate). + Ignored for non-quantized state dtypes. + launch_with_pdl: enable external PDL (conv1d → precompute chain). + Defaults False; caller opts in when the upstream chain is PDL-safe. + Ignored on hardware that doesn't support PDL (sm < 90). + use_internal_pdl: enable internal PDL (precompute → main overlap). + Defaults True; override for testing only. + Ignored on hardware that doesn't support PDL (sm < 90). + + _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, + _precompute_num_warps, _precompute_num_stages, _heads_per_block, + _maxnreg, _num_ctas) are benchmark-only overrides; production callers + should leave them None to use the heuristic-tuned defaults. + """ + # PDL needs sm >= 90. + if get_sm_version() < 90: + launch_with_pdl = False + use_internal_pdl = False + + # Mode selection: + # mode=None (default): look up the table-tuned mode + knobs for this + # (effective_batch, dtype, sr) cell. See `_resolve_tuning` above. + # mode="persistent_dynamic": single persistent-CTA kernel covering the + # full batch. Each work-item dispatches via runtime PNAT check + # (is_write = (pnat + T) > MAX). No write/nowrite split. + # slot_perm is honored but optional. write_checkpoint is ignored. + # mode="persistent_main": persistent-CTA kernel with two launches + # (write half + nowrite half). Caller MUST pre-sort slot_perm + # write-first; the n_writes tensor partitions the persistent loop + # into the two halves with the right WRITE_CHECKPOINT constexpr + # each time. RECTANGLE constexpr (= rectangle_for_nowrite) picks + # rect vs replay for the nowrite half. write_checkpoint is ignored. + # Note: mode-and-knob resolution from the default-tuning table happens + # below, after we have `batch` and `nheads`. + + # --- Hardware support gates --- + # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX + # instructions (Ada Lovelace introduced them; Hopper/Blackwell carry them). + if state.dtype == torch.float8_e4m3fn: + assert get_sm_version() >= 89, ( + "fp8_e4m3fn state requires SM 89+ (Ada Lovelace / Hopper / Blackwell) " + f"for fp32↔fp8 cvt PTX instructions; current SM is {get_sm_version()}." + ) + + # PTX cvt.rs.* (stochastic rounding) family lands on Blackwell only. + # Wrapper fails loud; framework decides fall-back (e.g. drop SR, use RN). + # int8 / int16 SR uses pure-Triton libdevice.floor + uniform noise — no + # PTX SR instruction needed, runs anywhere. + if rand_seed is not None: + if state.dtype == torch.float16: + assert get_sm_version() >= 100, ( + "fp16 stochastic rounding (PTX cvt.rs.f16x2.f32) requires " + f"sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + ) + elif state.dtype == torch.float8_e4m3fn: + assert get_sm_version() >= 100, ( + "fp8 stochastic rounding (PTX cvt.rs.satfinite.e4m3x4.f32) " + f"requires sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + ) + + # --- Unsqueeze inputs to canonical shapes --- + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if x.dim() == 3: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if dt.dim() == 3: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if B.dim() == 3: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if C.dim() == 3: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None: + if z.dim() == 2: + z = z.unsqueeze(1) + if z.dim() == 3: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + if out.dim() == 2: + out = out.unsqueeze(1) + if out.dim() == 3: + out = out.unsqueeze(1) + + cache_size, nheads, dim, dstate = state.shape + batch, T, _, _ = x.shape + ngroups = B.shape[2] + assert nheads % ngroups == 0 + + # --- Quantization plumbing (needed for SR/RN classification below) --- + # QUANT_MAX > 0 ⇔ state is int8 / int16 / fp8_e4m3fn. Kernel-entry + # static_assert on the Triton side mirrors this invariant. + quant_max = _QUANT_MAX_BY_DTYPE.get(state.dtype, 0.0) + is_quantized = quant_max > 0.0 + + # --- Default-tuning lookup --- + # Resolve (mode, knobs) from the table when caller leaves them None. + # Caller-provided kwargs always win. If the caller forces a mode that + # differs from the table's recommendation for this cell, we BRIDGE the + # table's knobs into the forced mode's knob namespace rather than fall + # back to (likely-terrible) kernel defaults: + # table pd → forced pm: copy each unsplit pd knob (M, W, S, CPS, LS) + # to both write and nowrite split knobs. + # table pm → forced pd: take the nowrite split values (Mnw, Wnw, Snw, + # CPSnw, LSnw) as the unsplit knobs. + # Empty table → no-op (caller passes whatever, mode falls back to pd). + _dt_str = { + torch.float32: "fp32", + torch.float16: "fp16", + torch.bfloat16: "bf16", + torch.int8: "int8", + torch.int16: "int16", + torch.float8_e4m3fn: "fp8", + }.get(state.dtype, str(state.dtype)) + _sr_str = "SR" if (rand_seed is not None and is_quantized) else "RN" + _table_entry = _resolve_tuning(batch, nheads, _dt_str, _sr_str) + if _table_entry is not None: + _table_mode, _table_knobs = _table_entry + if mode is None: + mode = _table_mode + if mode != _table_mode: + # Bridge across modes — see header comment above. + _table_knobs = _bridge_tuning_knobs(_table_knobs, _table_mode, mode) + # Fill None-valued kwargs from table. We can't reliably mutate + # locals() for re-read, so re-bind each kwarg explicitly. + if rectangle_for_nowrite is None and "rectangle_for_nowrite" in _table_knobs: + rectangle_for_nowrite = bool(_table_knobs["rectangle_for_nowrite"]) + _block_size_m = _block_size_m if _block_size_m is not None else _table_knobs.get("_block_size_m") + _num_warps = _num_warps if _num_warps is not None else _table_knobs.get("_num_warps") + _num_stages = _num_stages if _num_stages is not None else _table_knobs.get("_num_stages") + _heads_per_block = _heads_per_block if _heads_per_block is not None else _table_knobs.get("_heads_per_block") + _precompute_num_warps = _precompute_num_warps if _precompute_num_warps is not None else _table_knobs.get("_precompute_num_warps") + _precompute_num_stages = _precompute_num_stages if _precompute_num_stages is not None else _table_knobs.get("_precompute_num_stages") + _block_size_m_write = _block_size_m_write if _block_size_m_write is not None else _table_knobs.get("_block_size_m_write") + _block_size_m_nowrite = _block_size_m_nowrite if _block_size_m_nowrite is not None else _table_knobs.get("_block_size_m_nowrite") + _num_warps_write = _num_warps_write if _num_warps_write is not None else _table_knobs.get("_num_warps_write") + _num_warps_nowrite = _num_warps_nowrite if _num_warps_nowrite is not None else _table_knobs.get("_num_warps_nowrite") + _num_stages_write = _num_stages_write if _num_stages_write is not None else _table_knobs.get("_num_stages_write") + _num_stages_nowrite = _num_stages_nowrite if _num_stages_nowrite is not None else _table_knobs.get("_num_stages_nowrite") + _cta_per_sm = _cta_per_sm if _cta_per_sm is not None else _table_knobs.get("_cta_per_sm") + _num_loop_stages = _num_loop_stages if _num_loop_stages is not None else _table_knobs.get("_num_loop_stages") + # Split-form resolution for pm's per-half knobs. Without these the + # table's _num_loop_stages_{write,nowrite} and _cta_per_sm_{write, + # nowrite} values are dead — pm reads the split forms but the + # wrapper would leave them None, falling through to Triton defaults + # (or our hardcoded `or 1` / `or 2` per-mode fallbacks). + _num_loop_stages_write = _num_loop_stages_write if _num_loop_stages_write is not None else _table_knobs.get("_num_loop_stages_write") + _num_loop_stages_nowrite = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _table_knobs.get("_num_loop_stages_nowrite") + _cta_per_sm_write = _cta_per_sm_write if _cta_per_sm_write is not None else _table_knobs.get("_cta_per_sm_write") + _cta_per_sm_nowrite = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _table_knobs.get("_cta_per_sm_nowrite") + _flatten = _flatten if _flatten is not None else _table_knobs.get("_flatten") + _warp_specialize = _warp_specialize if _warp_specialize is not None else _table_knobs.get("_warp_specialize") + _use_tma_rect_load = _use_tma_rect_load or bool(_table_knobs.get("_use_tma_rect_load", False)) + _use_tma_replay_write_load = _use_tma_replay_write_load or bool(_table_knobs.get("_use_tma_replay_write_load", False)) + _use_tma_replay_write_store = _use_tma_replay_write_store or bool(_table_knobs.get("_use_tma_replay_write_store", False)) + _use_tma_replay_nowrite_load = _use_tma_replay_nowrite_load or bool(_table_knobs.get("_use_tma_replay_nowrite_load", False)) + # Final defaults if neither caller nor table set them (empty table case). + if mode is None: + mode = "persistent_dynamic" + if rectangle_for_nowrite is None: + rectangle_for_nowrite = False + assert mode in ("persistent_dynamic", "persistent_main"), ( + f"unknown mode {mode!r}; expected 'persistent_dynamic' or 'persistent_main'" + ) + if is_quantized: + assert state_scales is not None, ( + f"state.dtype={state.dtype} requires state_scales tensor " + "(shape (cache_size, nheads, dim), fp32)." + ) + assert state_scales.shape == (cache_size, nheads, dim), ( + f"state_scales shape mismatch: expected {(cache_size, nheads, dim)}, " + f"got {state_scales.shape}." + ) + assert state_scales.dtype == torch.float32, ( + f"state_scales must be fp32, got {state_scales.dtype}." + ) + assert state_scales.device == state.device + + # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the + # placeholder degenerate case max_window = T (every step is a checkpoint + # step). For real replay-style checkpointing, max_window > T and + # `prev_num_accepted_tokens` can be 0..max_window. Window-axis kernel + # tiles (BLOCK_SIZE_WINDOW, BLOCK_SIZE_K) are derived independently from + # MAX_REPLAY_BUFFER_LENGTH so max_window can exceed BLOCK_SIZE_T freely. + max_window = old_x.shape[1] + assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" + + assert x.shape == (batch, T, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + assert B.shape == (batch, T, ngroups, dstate) + assert C.shape == B.shape + assert old_x.shape == (cache_size, max_window, nheads, dim) + assert old_B.shape == (cache_size, 2, max_window, ngroups, dstate) + assert old_dt.shape == (cache_size, 2, nheads, max_window) + assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_window) + assert cache_buf_idx.shape == (cache_size,) + assert prev_num_accepted_tokens.shape == (cache_size,) + + tie_hdim = ( + A.stride(-1) == 0 + and A.stride(-2) == 0 + and dt.stride(-1) == 0 + and (dt_bias is None or dt_bias.stride(-1) == 0) + ) + assert tie_hdim + + device = x.device + BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) + # Rectangle K-axis bound = window (max_window). Computed unconditionally + # so the launch sites can refer to it; only used on the rectangle path. + BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) + + # Allocate precomputed intermediates (per-call, not cached). Always + # allocate (T, K) — the largest layout that any path uses. Replay-style + # paths only touch the first T columns; rectangle/dynamic use the full K. + # The few extra unused columns per row are negligible (~6KB per layer at + # production sizes) and let the dispatch helpers share one buffer. + cb_scaled = torch.empty( + batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 + ) + decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) + + z_strides = ( + (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) if z is not None else (0, 0, 0, 0) + ) + + # Kernel tuning: BLOCK_SIZE_M, num_warps, HEADS_PER_BLOCK, precompute_num_warps. + # Dtype-aware heuristic from B200 sweeps (batch 1-512, T=6/32, TP=8, conv1d + + # chained PDL). Keyed on total_heads, BLOCK_SIZE_T, and state dtype; 16-bit + # states prefer different tiles from fp32 due to lower bandwidth. Philox + # gets its own branch — stochastic rounding shifts compute toward CUDA cores, + # so small-batch configs want more warps to hide the extra work. + total_heads = batch * nheads + heads_per_group = nheads // ngroups + state_is_16bit = state.dtype in (torch.float16, torch.bfloat16) + use_philox = rand_seed is not None + if BLOCK_SIZE_T <= 16: + if use_philox and state_is_16bit: + # Philox: more warps at small batch to hide CUDA core work. + # At large batch, converges to non-Philox fp16 config. + if total_heads <= 16: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 4, 4, 4, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + elif state_is_16bit: + if total_heads <= 16: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 4, 1 + elif total_heads <= 64: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 2, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 1, + min(2, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + else: # fp32 state (no Philox — fp32 doesn't need stochastic rounding) + if total_heads <= 32: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 4, 1 + elif total_heads <= 64: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + elif total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 2, 2, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 1, 2, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 2, + min(2, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + else: # T > 16 + if state_is_16bit: + if total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 16, + 1, + 4, + min(2, heads_per_group), + ) + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 1, + min(4, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 4, + min(2, heads_per_group), + ) + else: # fp32 state + if total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 2, + 4, + min(2, heads_per_group), + ) + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 2, + min(4, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 4, + min(2, heads_per_group), + ) + if _block_size_m is not None: + BLOCK_SIZE_M = _block_size_m + if _num_warps is not None: + num_warps = _num_warps + if _heads_per_block is not None: + # Cap at heads_per_group: HEADS_PER_BLOCK divides the kernel's head + # axis, so a table value larger than the model's heads-per-group + # would overshoot. Protects callers running smaller models than + # the one we tuned against. + heads_per_block = min(_heads_per_block, heads_per_group) + if _precompute_num_warps is not None: + precompute_num_warps = _precompute_num_warps + + # Per-main knob resolution: each _*_{write,nowrite} arg, if not None, + # overrides the corresponding shared value for ONE main launch only. + # Default (None) = tied to shared value (current behavior). + BLOCK_SIZE_M_WRITE = _block_size_m_write if _block_size_m_write is not None else BLOCK_SIZE_M + BLOCK_SIZE_M_NOWRITE = _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M + NUM_WARPS_WRITE = _num_warps_write if _num_warps_write is not None else num_warps + NUM_WARPS_NOWRITE = _num_warps_nowrite if _num_warps_nowrite is not None else num_warps + NUM_STAGES_WRITE = _num_stages_write if _num_stages_write is not None else _num_stages + NUM_STAGES_NOWRITE = _num_stages_nowrite if _num_stages_nowrite is not None else _num_stages + # Persistent-only per-main: + CTA_PER_SM_WRITE = _cta_per_sm_write if _cta_per_sm_write is not None else _cta_per_sm + CTA_PER_SM_NOWRITE = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _cta_per_sm + NUM_LOOP_STAGES_WRITE = _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages + NUM_LOOP_STAGES_NOWRITE = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages + + HAS_CACHE_BATCH_INDICES = state_batch_indices is not None + + assert nheads % heads_per_block == 0, ( + f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" + ) + assert heads_per_block <= heads_per_group, ( + f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" + ) + + # state_scales pointer + strides: real tensor when quantized, otherwise + # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). + if is_quantized: + state_scales_arg = state_scales + state_scales_strides = ( + state_scales.stride(0), + state_scales.stride(1), + state_scales.stride(2), + ) + else: + state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 + state_scales_strides = (0, 0, 0) + + # Per-path TMA descriptors for state — write-side and nowrite-side. Each + # kernel launch consumes the descriptor whose block_shape[0] matches its + # BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need + # distinct descriptors; otherwise the descriptor's block_shape[0] would + # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot / arithmetic + # on the loaded tile fails shape inference at compile time + # ("Cannot make_shape_compatible: incompatible dimensions"). When Mw == + # Mnw (tied, the common case) the two descriptors are the same object. + # Same memory (state's flat 2D view, shape (cache*nheads*dim, dstate)) + # and same dstate block_shape — only block_shape[0] differs. + # When no TMA flag is on, both variables hold the raw `state` tensor as a + # dummy; kernels never reference it because their constexprs are all + # False (Triton DCEs the dead branches). + # `triton.set_allocator()` must run before any descriptor-using launch. + if (_use_tma_rect_load or _use_tma_replay_write_load + or _use_tma_replay_write_store or _use_tma_replay_nowrite_load): + from triton.tools.tensor_descriptor import TensorDescriptor + _ensure_tma_allocator() + assert state.is_contiguous(), "TMA state requires contiguous state" + assert state.stride(-1) == 1, "TMA state requires inner stride 1" + _state_flat = state.view(-1, state.shape[-1]) + _dstate_pow2 = triton.next_power_of_2(dstate) + state_tma_descriptor_write = TensorDescriptor.from_tensor( + _state_flat, block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], + ) + if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: + state_tma_descriptor_nowrite = state_tma_descriptor_write + else: + state_tma_descriptor_nowrite = TensorDescriptor.from_tensor( + _state_flat, block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], + ) + else: + state_tma_descriptor_write = state # dummy; all consuming constexprs False + state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False + + # Slot permutation — pointer + USE_PERM gate. Always required: the + # persistent_main kernel reads pid_b through slot_perm to walk the + # write-first sorted batch. persistent_dynamic forces USE_PERM=False + # at the call site (see launch_persistent_dynamic_main below) so the + # perm value doesn't matter for pd, but the tensor must still be valid. + assert isinstance(slot_perm, torch.Tensor), ( + f"slot_perm must be a torch.Tensor, got {type(slot_perm).__name__}" + ) + assert slot_perm.device == device, ( + f"slot_perm must be on device {device}, got {slot_perm.device}" + ) + assert slot_perm.dtype in (torch.int32, torch.int64), ( + f"slot_perm must be int32/int64, got {slot_perm.dtype}" + ) + assert slot_perm.shape == (batch,), ( + f"slot_perm must have shape (batch={batch},), got {tuple(slot_perm.shape)}" + ) + assert isinstance(n_writes, torch.Tensor), ( + f"n_writes must be a torch.Tensor, got {type(n_writes).__name__}" + ) + assert n_writes.device == device, ( + f"n_writes must be on device {device}, got {n_writes.device}" + ) + assert n_writes.dtype == torch.int32, ( + f"n_writes must be int32, got {n_writes.dtype}" + ) + assert n_writes.shape == (1,), ( + f"n_writes must have shape (1,), got {tuple(n_writes.shape)}" + ) + slot_perm_arg = slot_perm + use_perm = True + + precomp_grid = (batch, nheads // heads_per_block) + d_strides = (D.stride(0), D.stride(1)) if D is not None else (0, 0) + + # ---- Launch helpers (close over locals) ------------------------------- + # Each helper is a thin closure that calls one Triton kernel with the + # full positional + kwarg argument list. Mode-dependent constexprs + # (write_checkpoint, early_out, rectangle) are passed in. + + def launch_dynamic_precompute(rectangle: bool): + _dynamic_precompute_kernel[precomp_grid]( + dt, dt_bias, A, B, C, + cb_scaled, decay_vec, + old_B, old_dt, old_dA_cumsum, + cache_buf_idx, prev_num_accepted_tokens, + state_batch_indices, pad_slot_id, + T, max_window, dstate, nheads // ngroups, + dt.stride(0), dt.stride(1), dt.stride(2), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), B.stride(1), B.stride(2), B.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + dt_softplus, + HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, + LAUNCH_WITH_PDL=launch_with_pdl, + LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, + HEADS_PER_BLOCK=heads_per_block, + RECTANGLE=rectangle, + num_warps=precompute_num_warps, + **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), + launch_pdl=launch_with_pdl, + ) + + # ---- launch_persistent_main ------------------------------------------ + # Persistent-CTA main kernel. Single launch covers `n_slots` slots + # starting at `slot_offset`. Caller invokes twice: once for the write + # half (slot_offset=0, n_slots=n_writes, write_checkpoint=True) and + # once for the nowrite half (slot_offset=n_writes, + # n_slots=batch-n_writes, write_checkpoint=False). Hard-sort + # contract: caller has pre-sorted slots so [0, n_writes) are writes + # and [n_writes, batch) are nowrites. + + # Resolve persistent-mode bench knobs. Defaults: cta_per_sm = 1 + # (one CTA per SM, matches upstream `_p_matmul_ogs.py`); num_loop_stages + # = 2 (matches in-tree `swiglu` precedent for non-dot persistent loops); + # flatten = True (canonical Triton 3.6 idiom); warp_specialize = False. + _num_sms = torch.cuda.get_device_properties(device).multi_processor_count + cta_per_sm_arg = _cta_per_sm if _cta_per_sm else 1 + num_persistent_arg = cta_per_sm_arg * _num_sms + num_loop_stages_arg = _num_loop_stages if _num_loop_stages else 2 + flatten_arg = True if _flatten is None else bool(_flatten) + warp_specialize_arg = False if _warp_specialize is None else bool(_warp_specialize) + # Per-launch work-item count. At small batch, total_work may be < the + # full persistent grid; capping `grid` at `min(NUM_PERSISTENT, total_work)` + # avoids launching empty CTAs that pay setup cost for no work. Correctness: + # the kernel's `tl.range(pid, total_work, NUM_PERSISTENT)` ensures each + # tile_id is covered exactly once across all live pids in [0, grid) when + # grid <= NUM_PERSISTENT (each CTA does 1 tile; loop step >= total_work + # exits immediately) AND when grid == NUM_PERSISTENT (each CTA loops over + # multiple tiles). NUM_PERSISTENT is now a runtime int (see kernel def + # docstring at _persistent_main_kernel) so changing cta_per_sm does NOT + # trigger a new Triton compile — same kernel binary, different loop step. + # (Named UPPERCASE for historical Triton-style consistency only; not + # constexpr.) + _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + + def launch_persistent_main(write_checkpoint: bool, + *, + launch_dependent_kernels: bool = False, + rectangle: bool = False): + # `n_writes` (wrapper-level) is the (1,) int32 device tensor with the + # write count. Both halves always launch; the kernel's runtime PNAT + # check iterates only the slots that belong to its half. + _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE + _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE + _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE + _cps = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE + _cps = _cps if _cps else 1 + _nls = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE + _nls = _nls if _nls else 2 + _num_persistent = _cps * _num_sms + _num_pid_m_local = (dim + _bsm - 1) // _bsm + # Grid sizing: cap at min(full persistent grid, upper-bound total work). + # We use `batch` as the upper bound on slots-per-half — overcounting + # by a few CTAs is fine since the kernel's runtime check only + # iterates the slots that actually belong to its half. + _total_work_launch = max(1, batch * _num_pid_m_local * nheads) + grid = (min(_num_persistent, _total_work_launch),) + # Per-path TMA descriptor — block_shape[0] must match _bsm. + _desc = (state_tma_descriptor_write if write_checkpoint + else state_tma_descriptor_nowrite) + _persistent_main_kernel[grid]( + state, _desc, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + n_writes, batch, nheads, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + _bsm, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=write_checkpoint, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + NUM_PERSISTENT=_num_persistent, + NUM_LOOP_STAGES=_nls, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=False, + RECTANGLE=rectangle, + # 3 TMA flags. IS_DYNAMIC=False: WC fixed per launch; impl + # constexpr-folds the LOAD pick. When WC=True (write half), + # NOWRITE_LOAD is dummy False; when WC=False, WRITE_LOAD/STORE + # dummy False. NOWRITE_LOAD picks rect-load (RECTANGLE) or + # replay-nowrite-load. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), + USE_TMA_LOAD_NOWRITE=bool( + (_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load) + and not write_checkpoint + ), + USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), + num_warps=_nw, + **({"num_stages": _ns} if _ns else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, + launch_dependent_kernels: bool = False, + rectangle: bool = False): + # Single-launch persistent kernel covering the whole batch with + # runtime per-slot WRITE_CHECKPOINT branch. No half-split, no + # n_writes needed (the kernel ignores n_writes_dev when + # IS_DYNAMIC=True; Triton DCEs the load). is_write is computed + # at runtime per work-item from the loaded PNAT. + # We still pass `n_writes_dev` (the same tensor the persistent_main + # path uses) so the kernel signature is uniform; the value is + # immaterial. + # Grid sizing: cap at total_work (= batch * num_pid_m * nheads) for + # the dynamic case (full-batch coverage); see launch_persistent_main + # comment for correctness rationale. + _total_work_launch = max(1, batch * _num_pid_m * nheads) + grid = (min(num_persistent_arg, _total_work_launch),) + # Persistent-dynamic kernel uses a single BLOCK_SIZE_M (same as the + # wrapper's BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE tied convention), so + # the write-side descriptor matches. Both write and nowrite slots + # in this kernel share that BSM. + _persistent_main_kernel[grid]( + state, state_tma_descriptor_write, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + n_writes, batch, nheads, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=False, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + # persistent_dynamic forces USE_PERM=False regardless of caller- + # provided slot_perm — our pd tuning runs all happened with + # SORT=0 (no slot_perm passed), so honoring slot_perm here would + # silently shift pd to an untimed code path. Revisit if/when + # we benchmark pd with slot_perm. + USE_PERM=False, + NUM_PERSISTENT=num_persistent_arg, + NUM_LOOP_STAGES=num_loop_stages_arg, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=True, + RECTANGLE=rectangle, + # 3 TMA flags. IS_DYNAMIC=True: is_write is runtime per slot; + # impl's load TMA picks per-slot (constexpr ternary becomes a + # runtime branch — both load forms emitted, ~negligible cost). + # NOWRITE_LOAD picks rect-load when RECTANGLE, else + # replay-nowrite-load. STORE only fires on runtime is_write. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load), + USE_TMA_LOAD_NOWRITE=bool( + _use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load + ), + USE_TMA_STORE=bool(_use_tma_replay_write_store), + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + # ---- Mode dispatch ---------------------------------------------------- + with torch.cuda.device(device.index): + if mode == "persistent_dynamic": + # Single-launch persistent kernel covering the full batch. Each + # work-item dispatches via runtime PNAT check. Kernel ignores + # n_writes (Triton DCEs the load) when IS_DYNAMIC=True; we still + # pass the wrapper-provided tensor as required by the signature. + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_persistent_dynamic_main( + n_writes, + launch_dependent_kernels=False, + rectangle=rectangle_for_nowrite, + ) + elif mode == "persistent_main": + # Persistent-CTA main kernel. One shared dynamic_precompute + # (per-slot dispatch via PNAT) feeds two persistent_main + # launches (write half + nowrite half). Both halves ALWAYS + # launch; the kernel's runtime check iterates only the slots + # belonging to its half (write: [0, n_writes), nowrite: + # [n_writes, batch)). + # + # Caller-provided contract: `n_writes` is a (1,) int32 device + # tensor (the kernel reads it at runtime, after the precompute); + # `slot_perm` is a (batch,) int32 device tensor pre-sorted + # write-first. + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_persistent_main( + write_checkpoint=True, + launch_dependent_kernels=True, + rectangle=False, # write always replay-style + ) + launch_persistent_main( + write_checkpoint=False, + launch_dependent_kernels=False, + rectangle=rectangle_for_nowrite, + ) + else: + raise ValueError( + f"mode={mode!r} is not supported. Supported modes: " + f"'persistent_dynamic', 'persistent_main'." + ) diff --git a/tensorrt_llm/_torch/modules/mamba/_v3_baseline_checkpointing_state_update_slim.py b/tensorrt_llm/_torch/modules/mamba/_v3_baseline_checkpointing_state_update_slim.py new file mode 100644 index 000000000000..961dbf69bd7e --- /dev/null +++ b/tensorrt_llm/_torch/modules/mamba/_v3_baseline_checkpointing_state_update_slim.py @@ -0,0 +1,3173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +# +# Adapted from: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +# SPDX-FileCopyrightText: Copyright contributors to the sglang project +# +# Copyright (c) 2024, Tri Dao, Albert Gu. +# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID +from tensorrt_llm._utils import get_sm_version + +from .softplus import softplus + + +# Lazy global allocator for Triton TMA tensor descriptors. Required by any +# host- or device-built tensor_descriptor; without it Triton raises at first +# launch. See TMA backlog item #17 / scratch experiment notes. +_TMA_ALLOCATOR_SET = False + + +def _ensure_tma_allocator() -> None: + global _TMA_ALLOCATOR_SET + if _TMA_ALLOCATOR_SET: + return + + def _alloc_fn(size, alignment, stream): + # Triton expects an int8 buffer of `size` bytes; alignment is enforced + # by the allocator returning a buffer satisfying it (PyTorch's + # cudaMalloc-backed tensors are 256B-aligned, so we're fine). + return torch.empty(size, device="cuda", dtype=torch.int8) + + triton.set_allocator(_alloc_fn) + _TMA_ALLOCATOR_SET = True + + +@triton.jit +def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 pair → fp16x2 using Philox random bits. + + Uses PTX cvt.rs.f16x2.f32 which rounds each fp32 value to fp16 using + the random bits to break ties, avoiding systematic rounding bias that + accumulates over many decode steps with fp16 state. + + Adapted from flashinfer (Apache-2.0, vLLM/mamba lineage). + """ + return tl.inline_asm_elementwise( + asm="""{ + cvt.rs.f16x2.f32 $0, $2, $1, $3; + }""", + constraints=("=r,r,r,r,r"), + args=(x, rand), + dtype=tl.float16, + is_pure=True, + pack=2, + ) + + +@triton.jit +def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 quad → fp8 e4m3 using Philox random bits. + + Uses PTX cvt.rs.satfinite.e4m3x4.f32 which combines stochastic rounding + and saturating cast in a single op (output is final fp8, no separate + clamp needed). The reversed source-register order {$4,$3,$2,$1} is + load-bearing — PTX packs leftmost source into the high byte but Triton's + pack=4 is little-endian, so the natural {$1,$2,$3,$4} order would + silently shuffle every group of 4 contiguous outputs. + + Requires SM_100a+ (Blackwell B200). Caller must gate at the wrapper + level — this kernel does not check. + + Adapted from vLLM PR #40012 (Apache-2.0). + """ + return tl.inline_asm_elementwise( + asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", + constraints="=r,r,r,r,r,r,r,r,r", + args=(x, rand), + dtype=tl.float8e4nv, + is_pure=True, + pack=4, + ) + + +@triton.jit +def _bitrev32(x: tl.tensor) -> tl.tensor: + return tl.inline_asm_elementwise( + asm="brev.b32 $0, $1;", + constraints="=r,r", + args=(x,), + dtype=tl.uint32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _stochastic_round_int8_packed( + x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor +) -> tl.tensor: + """Stochastic rounding for int8 using one random uint32 per 4 values.""" + low = rand & 0x0000FFFF + high = (rand >> 16) & 0x0000FFFF + low_rev = _bitrev32(low) >> 16 + high_rev = _bitrev32(high) >> 16 + rand_pos = offs_n & 3 + rand16 = tl.where( + rand_pos == 0, + low, + tl.where(rand_pos == 1, low_rev, tl.where(rand_pos == 2, high, high_rev)), + ) + rand01 = rand16.to(tl.float32) * (1.0 / float(1 << 16)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +@triton.jit +def _stochastic_round_int16_packed( + x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor +) -> tl.tensor: + """Stochastic rounding for int16 using one random uint32 per 2 values.""" + rand_bits = tl.where((offs_n & 1) == 0, rand, _bitrev32(rand)) + rand01 = (rand_bits & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +# Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, +# old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. +# Grid: (batch, nheads // HEADS_PER_BLOCK). + + +@triton.jit() +def _replay_precompute_impl( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers (both buffers reachable via stride_*_dbuf). This + # kernel writes to either the active (= cache_buf_idx) or inactive + # (= 1 - cache_buf_idx) buffer depending on WRITE_CHECKPOINT — see + # comment block at top of kernel body. + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + # Double-buffer index (per cache slot) — selects this step's "active" + # buffer (= where the historical inputs for this step live). + cache_buf_idx_ptr, + # Per-request accepted-tokens count (already-cached old tokens at + # [0, PNAT) of the active buffer; new tokens this step go after them + # on no-checkpoint steps). + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides: (cache, 2, T, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Checkpointing flag — selects target buffer + offset for new-token + # cache writes. See "Cache write semantics" block below. + # Runtime (not constexpr): the only WRITE_CHECKPOINT-dependent code in + # this body is the write_buf/write_offset selection, which is plain + # arithmetic — no constexpr-shaped tile or whole-block gate. Letting + # it be runtime lets the dynamic dispatch kernel call us once with the + # per-slot needs_write flag instead of inlining two specializations. + write_checkpoint, +): + pid_b = tl.program_id(axis=0) + pid_hg = tl.program_id(axis=1) # head-group index + first_head = pid_hg * HEADS_PER_BLOCK + + # Resolve cache index for writes + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + + # --- Cache write semantics --- + # cache_buf_idx names this step's "active" buffer — the one with the + # historical inputs at [0, PNAT). The other buffer is "staging". + # + # Where do we write new tokens this step? + # WRITE_CHECKPOINT=False (no overflow): append to ACTIVE buffer at + # offset [PNAT : PNAT+T). Caller does NOT flip cache_buf_idx + # afterward; PNAT_next = PNAT + accepted. [0, PNAT) preserved. + # WRITE_CHECKPOINT=True (would overflow): write to STAGING buffer at + # [0, T). Caller flips cache_buf_idx afterward; next step's + # active = the one we just wrote. PNAT_next = accepted. Old + # data in the previous active buffer is folded into state via + # the replay update and discarded. This matches today's replay + # kernel behavior exactly. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if write_checkpoint: + write_buf = 1 - buf_active + write_offset = 0 + else: + write_buf = buf_active + write_offset = prev_num_accepted_tokens + + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + t_mask = offs_t < T + n_mask = offs_n < dstate + + # Causal mask is shared across all heads (depends only on offs_t) + causal_mask = offs_t[:, None] >= offs_t[None, :] + valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] + + # --- Vectorized pre-wait phase across HEADS_PER_BLOCK heads --- + # Compute dt, dA_cumsum, decay_vec as (H, T) tiles. Pre-compute + # scale_combo = decay_matrix * dt[:, None, :] as an (H, T, T) tile that + # stays in registers across gdc_wait — eliminates the post-wait reload + # of dt + dA_cumsum and the per-head loop. + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) + + # Load dt (H, T) + dt_addrs = ( + dt_ptr + pid_b * stride_dt_batch + + heads_block[:, None] * stride_dt_head + + offs_t[None, :] * stride_dt_T + ) + dt = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias[:, None] + if DT_SOFTPLUS: + dt = softplus(dt) + + A = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) # (H,) + dA_cumsum = tl.cumsum(A[:, None] * dt, axis=1) # (H, T) + decay_vec = tl.exp(dA_cumsum) # (H, T) + + # Cross-step continuity for old_dA_cumsum: when appending to active_buf at + # offset PNAT > 0, the previous step left a running cumsum at [0, PNAT) + # whose tail value lives at active_buf[head, PNAT-1]. Add that tail to + # this step's per-step-restarted cumsum before storing so the buffer + # holds one continuous cumsum across N back-to-back nowrites. Write path + # (write_buf = 1 - buf_active, write_offset = 0) starts fresh, no prefix. + # Both branches are on scalar runtime values (write_checkpoint and PNAT), + # uniform across the block — use scalar if to short-circuit the load. + if write_checkpoint or prev_num_accepted_tokens == 0: + prev_total = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) + else: + last_cumsum_ptrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T + ) + prev_total = tl.load(last_cumsum_ptrs).to(tl.float32) + + # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of write_buf. + old_dt_addrs = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block[:, None] * stride_old_dt_head + + (write_offset + offs_t)[None, :] * stride_old_dt_T + ) + tl.store(old_dt_addrs, dt, mask=t_mask[None, :]) + + old_dA_cumsum_addrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block[:, None] * stride_old_dA_cumsum_head + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T + ) + tl.store(old_dA_cumsum_addrs, dA_cumsum + prev_total[:, None], mask=t_mask[None, :]) + + # decay_vec scratch — always at offs_t. + decay_vec_addrs = ( + decay_vec_ptr + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) + tl.store(decay_vec_addrs, decay_vec, mask=t_mask[None, :]) + + # scale_combo (H, T, T) = exp(dA_cumsum[h, t1] - dA_cumsum[h, t2]) * dt[h, t2] + # Stays live across gdc_wait — used post-wait to compute CB_scaled. + decay_matrix = tl.exp(dA_cumsum[:, :, None] - dA_cumsum[:, None, :]) # (H, T, T) + scale_combo = decay_matrix * dt[:, None, :] # (H, T, T) + + # --- Wait for upstream kernel (external PDL) before loading B and C --- + # All dt processing above is independent of conv1d outputs. + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- + group_idx = first_head // nheads_ngroups_ratio + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_all = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_all = tl.load( + B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + # Compute raw CB once — shared across all heads in this block + raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) + + # Store B to cache at [write_offset : write_offset+T) of write_buf. + if first_head % nheads_ngroups_ratio == 0: + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + write_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_all, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # --- Vectorized post-wait phase: scale_combo (H, T, T) is still live in + # registers from pre-wait; multiply by raw_CB (T, T), apply causal mask, + # store as one (H, T, T) tile. --- + CB_scaled_block = tl.where( + valid_mask[None, :, :], + raw_CB[None, :, :] * scale_combo, + 0.0, + ) # (H, T, T) + cb_scaled_addrs = ( + cb_scaled_ptr + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_t[None, None, :] * stride_cb_j + ) # (H, T, T) + cb_store_mask = ( + (offs_t[None, :, None] < BLOCK_SIZE_T) + & (offs_t[None, None, :] < BLOCK_SIZE_T) + ) + tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) + + +# Replay-style precompute kernel. Thin wrapper around _replay_precompute_impl +# that carries the @triton.heuristics for constexpr derivation; called from +# the Python wrapper on the replay-style path (write or replay-nowrite). +@triton.jit() +def _rectangle_precompute_impl( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle + decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) + # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite + # path: read from buf_active at [0, PNAT), write new tokens at + # [PNAT, PNAT+T) of buf_active (same buffer). + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # rectangle K-axis bound + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides (rectangle: (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides: (cache, 2, T_max, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T_max) + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T_max) + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, +): + pid_b = tl.program_id(axis=0) + pid_hg = tl.program_id(axis=1) + first_head = pid_hg * HEADS_PER_BLOCK + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_buf = buf_active + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout: old at [0, PNAT) (mask is_old_k); + # new at [MAX-T, MAX) at compile-time shift K_NEW_SHIFT = MAX - T. + # PNAT + T <= MAX is guaranteed on the nowrite path → no overlap. + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis (output rows) + offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis (rectangle input cols) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + t_mask = offs_t < T + n_mask = offs_n < dstate + + # K-axis masks (approach C: runtime PNAT-offset instead of K_NEW_SHIFT) + # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # V2: Loop 1 vectorized across HEADS_PER_BLOCK — single (H, T) tile, + # no per-head loop. Mirrors the replay precompute's pre-wait layout. + offs_h_lp1 = tl.arange(0, HEADS_PER_BLOCK) + heads_block_lp1 = first_head + offs_h_lp1 # (H,) + + dt_addrs_v = ( + dt_ptr + pid_b * stride_dt_batch + + heads_block_lp1[:, None] * stride_dt_head + + offs_t[None, :] * stride_dt_T + ) + dt_v = tl.load(dt_addrs_v, mask=t_mask[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias_v = tl.load(dt_bias_ptr + heads_block_lp1 * stride_dt_bias_head).to(tl.float32) + dt_v = dt_v + dt_bias_v[:, None] + if DT_SOFTPLUS: + dt_v = softplus(dt_v) + + A_v = tl.load(A_ptr + heads_block_lp1 * stride_A_head).to(tl.float32) # (H,) + dA_cumsum_v = tl.cumsum(A_v[:, None] * dt_v, axis=1) # (H, T) + + # Cross-step continuity: hoisted (H,) prefix load. + if prev_num_accepted_tokens == 0: + prev_total_v = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) + else: + prev_total_ptrs_v = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block_lp1 * stride_old_dA_cumsum_head + + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T + ) + prev_total_v = tl.load(prev_total_ptrs_v).to(tl.float32) + + # Coalesced (H, T) stores — replaces HPB per-head scalar stores. + old_dt_addrs_v = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block_lp1[:, None] * stride_old_dt_head + + (write_offset + offs_t)[None, :] * stride_old_dt_T + ) + tl.store(old_dt_addrs_v, dt_v, mask=t_mask[None, :]) + + old_dA_cumsum_addrs_v = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block_lp1[:, None] * stride_old_dA_cumsum_head + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T + ) + tl.store( + old_dA_cumsum_addrs_v, + dA_cumsum_v + prev_total_v[:, None], + mask=t_mask[None, :], + ) + + # ---- Hoisted: cache-only loads independent of conv1d ---- + # old_B (group-level, BLOCK_K × BLOCK_DSTATE = ~8KB tile) and the + # decay_vec_full per-head pre-compute (which writes to DRAM and doesn't + # need cross-gdc_wait variables) are issued BEFORE gdc_wait so their + # HBM latency overlaps with conv1d. Per-head factor_dt/exp_diff stay + # below gdc_wait — they need cross-iteration spans, which Triton can't + # express without a DRAM round-trip; the per-head LOADS in the post- + # wait loop are small and cheap, so leave them. + group_idx = first_head // nheads_ngroups_ratio + + # Group-level: old B from active buffer at [0, PNAT) of the K-axis. + old_B_read_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + buf_active * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_load = tl.load( + old_B_read_base + + safe_old_k[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + mask=is_old_k[:, None] & n_mask[None, :], + other=0.0, + ) + + # Pre-wait: vectorized across HEADS_PER_BLOCK heads. Compute decay_vec_full + # (H, T) and combo = factor_dt * exp_diff (H, T, K). Store decay_vec_full; + # combo_block stays in registers across gdc_wait — used directly post-wait + # to compute rect_CB_scaled without a global memory roundtrip. + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) + + # Per-head bases (H,) — broadcast with offs_k or offs_t for 2D loads. + old_dt_read_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_active * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_read_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + ) + old_dt_write_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_write_h = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + ) + + # (H, K) loads at [0, PNAT) — old data from previous step. + hk_mask = is_old_k[None, :] # (1, K) + old_dt_all = tl.load( + old_dt_read_h[:, None] + safe_old_k[None, :] * stride_old_dt_T, + mask=hk_mask, other=0.0, + ).to(tl.float32) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, + mask=hk_mask, other=0.0, + ).to(tl.float32) + # V3: use loop-1 registers directly instead of reloading dA_cumsum_new + # from buffer. dA_cumsum_v + prev_total_v[:, None] IS what the buffer + # holds at positions [PNAT, PNAT+T). Saves an (H, T) DRAM round-trip + # per kernel call. + ht_mask = t_mask[None, :] # (1, T) + dA_cumsum_new = dA_cumsum_v + prev_total_v[:, None] # (H, T) + # (H, K) loads at K_NEW_SHIFT-shifted positions for new tokens. + hkn_mask = is_new_k[None, :] + dt_at_kn = tl.load( + old_dt_write_h[:, None] + + (write_offset + safe_k_new)[None, :] * stride_old_dt_T, + mask=hkn_mask, other=0.0, + ).to(tl.float32) + dA_cumsum_at_kn = tl.load( + old_dA_cumsum_write_h[:, None] + + (write_offset + safe_k_new)[None, :] * stride_old_dA_cumsum_T, + mask=hkn_mask, other=0.0, + ).to(tl.float32) + + # decay_vec_full[t] = exp(continuous_cumsum[PNAT+t]) — directly the + # continuous value now stored at buffer position write_offset+t. Was + # decomposed as total_decay * exp(per_step_new[t]) when the buffer held + # per-step (non-continuous) cumsum; with the continuity fix the value + # IS continuous_cumsum[PNAT+t] so no decomposition is needed. + decay_vec_full_block = tl.exp(dA_cumsum_new) # (H, T) + decay_vec_addrs = ( + decay_vec_ptr + + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) # (H, T) + tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) + + # combo_block = factor_dt * exp_diff — (H, T, K). Stays in registers + # across gdc_wait. With continuous cumsum in the buffer, s_k for any k + # (old or new) is simply -continuous_cumsum[k]; exp_diff[t, k] then + # equals exp(continuous_cumsum[PNAT+t] - continuous_cumsum[k]) — the + # decay weight for token k's contribution to the output at position + # PNAT+t. No need to subtract any "total" — the dA_cumsum_new[t] term + # already carries the full prefix. + # + # Numerical note: pre-fix this kernel computed `total - old_dA[k]` + # (small-minus-small) then summed `+ dA_cumsum_new[t]` (also small, + # per-step). Post-fix `s_k = -old_dA[k]` is large-magnitude positive + # and `dA_cumsum_new[t]` is large-magnitude negative; their sum + # cancels back to the same small value. Cancellation error is bounded + # by ulp(max_magnitude) ≈ 2^-23 · |continuous_cumsum| — negligible + # for max_window ≤ ~1024. Still one exp on the sum (not two muls of + # exps), so no overflow regression vs the original formulation. + factor_dt = tl.where(is_old_k[None, :], old_dt_all, dt_at_kn) # (H, K) + s_k = tl.where( + is_old_k[None, :], + -old_dA_cumsum_all, + -dA_cumsum_at_kn, + ) # (H, K) + # exp_diff (H, T, K) = exp(s_k (H, 1, K) + dA_cumsum_new (H, T, 1)). + exp_diff = tl.exp(s_k[:, None, :] + dA_cumsum_new[:, :, None]) + combo_block = factor_dt[:, None, :] * exp_diff # (H, T, K) + + # ---- gdc_wait: from here on we depend on conv1d's outputs ---- + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + # Conv1d outputs: B and C + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + B_new_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_all = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_new_orig = tl.load( + B_new_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + B_new_shifted = tl.load( + B_new_base + safe_k_new[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=is_new_k[:, None] & n_mask[None, :], + other=0.0, + ) + # Disjoint masks: old at [0, PNAT), new at [K_NEW_SHIFT, K_NEW_SHIFT+T). + B_combined = old_B_load + B_new_shifted + raw_rect_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_combined).to(tl.bfloat16)) + + # Append new B to cache at [PNAT, PNAT+T) of write_buf (once per group). + if first_head % nheads_ngroups_ratio == 0: + old_B_write_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + write_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_write_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_new_orig, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # Causal mask (BLOCK_SIZE_T × BLOCK_SIZE_K, shared across heads). + # Approach C: new tokens at runtime [PNAT, PNAT+T) instead of K_NEW_SHIFT. + t_idx_2d = offs_t[:, None] + k_idx_2d = offs_k[None, :] + is_old_k_2d = k_idx_2d < prev_num_accepted_tokens + k_new_idx_2d = k_idx_2d - prev_num_accepted_tokens + is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) + causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] + + # Post-wait vectorized: combo_block (H, T, K) is still live in registers. + # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store as + # one (H, T, K) tile. + rect_CB_scaled_block = tl.where( + causal_combined[None, :, :], + raw_rect_CB[None, :, :] * combo_block, + 0.0, + ) # (H, T, K) + cb_scaled_addrs = ( + cb_scaled_ptr + + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_k[None, None, :] * stride_cb_j + ) # (H, T, K) + cb_store_mask_3d = ( + (offs_t[None, :, None] < BLOCK_SIZE_T) + & (offs_k[None, None, :] < BLOCK_SIZE_K) + ) # (1, T, K) → broadcasts to (H, T, K) + tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) + + +# Rectangle precompute kernel. Thin wrapper around _rectangle_precompute_impl +# that carries the @triton.heuristics for constexpr derivation; called from +# the Python wrapper on the rectangle nowrite path. +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.jit() +def _dynamic_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. + RECTANGLE: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + pid_b = tl.program_id(axis=0) + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH + # write_checkpoint is now runtime in replay precompute, so a single + # call site handles both write and nowrite for the replay branch. + # Take rectangle only when RECTANGLE is True AND this slot doesn't + # need write; everything else funnels into replay. + if needs_write_runtime or not RECTANGLE: + _replay_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + needs_write_runtime, + ) + else: + _rectangle_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + pad_slot_id, + T, + MAX_REPLAY_BUFFER_LENGTH, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + ) + + +# Main kernel: tl.dot replay + precomputed CB output. +# Grid: (cdiv(dim, M), batch, nheads). + + +@triton.jit() +def _persistent_main_impl( + # Per-work-unit indices (computed by the persistent wrapper). + # `pid_b` is the post-perm slot index (caller has already applied any + # slot permutation and slot_offset). + pid_m, + pid_b, + pid_h, + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or + # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor + # USE_TMA_STORE is enabled (kernel ignores it via constexpr). + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + rand_seed_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + # IS_DYNAMIC: kept in the signature for caller-side bookkeeping (the + # outer _persistent_main_kernel still inspects it to decide the slot- + # IS_DYNAMIC: when True (persistent_dynamic), is_write is per-slot from + # PNAT. When False (persistent_main), is_write is constexpr from + # WRITE_CHECKPOINT. See also WC_IS_CONSTEXPR below. + IS_DYNAMIC: tl.constexpr, + # WC_IS_CONSTEXPR: when True, force is_write = WRITE_CHECKPOINT (constexpr) + # regardless of IS_DYNAMIC. Callers in RECT=1 use this in the is_w=True + # arm of _persistent_main_kernel (we know all slots that reach this call + # need is_write=True because is_w was the PNAT-derived runtime check, and + # this arm only fires when is_w is True). Passing WRITE_CHECKPOINT=True + # as a literal at the call site + WC_IS_CONSTEXPR=True here lets the inner + # body DCE the nowrite path under IS_DYNAMIC=True too — same codegen + # quality as persistent_main mode (-3.7% measured at b=1024 dyn-shape). + # When False (RECT=0 callers, where both write and nowrite slots are + # dispatched to ONE call), use the original runtime is_write under + # IS_DYNAMIC=True; avoids the binary-doubling regression that two + # specialized calls would cause. + WC_IS_CONSTEXPR: tl.constexpr = False, + # TMA flags — picked inside body based on is_write. When is_write is + # constexpr (either IS_DYNAMIC=False or WC_IS_CONSTEXPR=True), the + # use_tma_load = USE_TMA_LOAD_WRITE if is_write else USE_TMA_LOAD_NOWRITE + # ternary constexpr-folds and only one TMA load form survives. + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # IS_DYNAMIC: kernel-mode label, used by the OUTER _persistent_main_kernel + # to decide slot-range derivation and outer is_w dispatch strategy + # (constexpr WC for persistent_main; runtime is_w split -> 2 specialized + # impl calls for persistent_dynamic). Inside this impl, IS_DYNAMIC is + # NOT consulted at runtime -- WRITE_CHECKPOINT is the only constexpr that + # gates the write/nowrite codegen, in BOTH modes. + + # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized + # state dtype (int8 / int16 / float8e4nv) and only those. + tl.static_assert( + (QUANT_MAX > 0.0) + == ( + (state_ptr.dtype.element_ty == tl.int8) + or (state_ptr.dtype.element_ty == tl.int16) + or (state_ptr.dtype.element_ty == tl.float8e4nv) + ), + "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", + ) + + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + # Resolve is_write: see WC_IS_CONSTEXPR / IS_DYNAMIC docs in the param + # list above. Three cases: + # - WC_IS_CONSTEXPR=True (RECT=1 is_w=True arm callers): use WC + # constexpr. Caller knows the slot needs write; inner DCEs nowrite + # paths. Avoids the binary-doubling overhead that calling the impl + # twice would cause, while still constexpr-DCEing the nowrite half. + # - IS_DYNAMIC=True (RECT=0 caller, persistent_dynamic): runtime + # branch on PNAT. Both write and nowrite codegen live in one body + # (no bloat) — same as the pre-refactor behavior. + # - IS_DYNAMIC=False (persistent_main): WC constexpr from caller. + if WC_IS_CONSTEXPR: + is_write: tl.constexpr = WRITE_CHECKPOINT + elif IS_DYNAMIC: + is_write = (prev_num_accepted_tokens + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_write = WRITE_CHECKPOINT + if is_write: + write_buf = 1 - active_buf # noqa: F841 + write_offset = 0 + else: + write_buf = active_buf # noqa: F841 + write_offset = prev_num_accepted_tokens + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # Load state. state_tma_descriptor is a host-built tensor_descriptor + # over a flat (cache*nheads*dim, dstate) view of state when any TMA + # path is enabled; raw `state_ptr` is the underlying tensor and is + # always passed. state_ptrs / state_ptr_raw are the raw-pointer view + # used for !TMA load and store paths. offs_y is the flat row index + # for TMA load/store; computed unconditionally (cheap int math; DCE'd + # when no TMA path is reachable). + state_mask = m_mask[:, None] & n_mask[None, :] + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + # Load state. Branch on is_write (constexpr = WRITE_CHECKPOINT in BOTH + # modes after the outer-dispatch refactor), then constexpr-pick TMA-vs- + # tl.load per side. Outer `if` DCE's, only the matching side's + # constexpr-gated load survives -- same compile-time picking for both + # persistent_main and persistent_dynamic (the latter dispatches at the + # outer kernel level so each impl instance sees a constexpr WC). + if is_write: + if USE_TMA_LOAD_WRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + else: + if USE_TMA_LOAD_NOWRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, + other=1.0, + ).to(tl.float32) + state = state * decode_scale[:, None] + + # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) + group_idx = pid_h // nheads_ngroups_ratio + + old_window_mask = offs_window < prev_num_accepted_tokens + + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + active_buf * stride_old_dt_dbuf + + pid_h * stride_old_dt_head + ) + old_dt_all = tl.load( + old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 + ).to(tl.float32) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + active_buf * stride_old_dA_cumsum_dbuf + + pid_h * stride_old_dA_cumsum_head + ) + old_dA_cumsum_all = tl.load( + old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, + mask=old_window_mask, other=0.0, + ).to(tl.float32) + + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( + tl.float32 + ) + + coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all + + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + old_x_all = tl.load( + old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + mask=old_window_mask[:, None] & m_mask[None, :], + other=0.0, + ) + + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + active_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + old_B_all = tl.load( + old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + mask=old_window_mask[:, None] & n_mask[None, :], + other=0.0, + ).to(tl.float32) + + dB_scaled = coeff[:, None] * old_B_all + + total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) + state *= total_decay + + state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) + + if is_write: + if USE_RS_ROUNDING: + # Generate random tensor for stochastic rounding. The amount of + # randomness needed depends on the SR codegen path: + # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) + # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) + # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs + # int16 SR (24b + bitrev32): 1 b32 per 2 outputs + # The PTX cvt.rs.* instructions consume a single 32-bit random + # and split the bits internally for 2 or 4 conversions. Generate + # only what's actually consumed and broadcast to fill the unused + # slots — saves Philox rounds proportionally. + if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: + RAND_DIVISOR: tl.constexpr = 4 # fp8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: + RAND_DIVISOR: tl.constexpr = 4 # int8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: + RAND_DIVISOR: tl.constexpr = 2 # int16 SR + elif QUANT_MAX == 0.0: + RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) + else: + RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized + + rand_seed = tl.load(rand_seed_ptr) + base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head + # Number of unique randoms per row = dstate / RAND_DIVISOR. + # randint4x emits 4 randoms per offset, so use that / 4 offsets. + offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) + rand_offsets_q = ( + base_rand + + offs_m[:, None] * stride_state_dim + + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) + ) # (M, dstate / (4*RAND_DIVISOR)) + if PHILOX_ROUNDS > 0: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + else: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) + r01 = tl.join(r0, r1) + r23 = tl.join(r2, r3) + r0123 = tl.join(r01, r23) + rand_compact = tl.reshape( + r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) + ) + # Broadcast each unique rand to RAND_DIVISOR adjacent positions. + # Pack-group (pack=2 fp16 / pack=4 fp8) consumes adjacent positions; + # the unique rand lands at the asm's read slot; duplicates feed + # the dead slots. Triton's broadcast_to is stride-0 in IR. + if RAND_DIVISOR > 1: + rand_3d = rand_compact[:, :, None] + rand_3d = tl.broadcast_to( + rand_3d, + (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), + ) + rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + else: + rand = rand_compact + + if QUANT_MAX > 0.0: + amax = tl.max(tl.abs(state), axis=1) + encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) + decode_scale = 1.0 / encode_scale + state_scales_ptrs = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + + offs_m * stride_state_scales_dim + ) + tl.store(state_scales_ptrs, decode_scale, mask=m_mask) + state_q = state * encode_scale[:, None] + if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): + _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_fp8sr) + else: + tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) + else: + if USE_RS_ROUNDING: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized SR fall-through expects int8 or int16; " + "fp8 SR is handled by the prior branch.", + ) + if state_ptrs.dtype.element_ty == tl.int8: + state_q = _stochastic_round_int8_packed( + state_q, rand, offs_n[None, :] + ) + else: + state_q = _stochastic_round_int16_packed( + state_q, rand, offs_n[None, :] + ) + elif state_ptrs.dtype.element_ty != tl.float8e4nv: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized RN with explicit round() expects int8 or int16.", + ) + state_q = tl.extra.cuda.libdevice.round(state_q) + state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) + _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_cast) + else: + tl.store(state_ptrs, _state_q_cast, mask=state_mask) + elif USE_RS_ROUNDING: + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "Non-quantized SR only supports fp16 state.", + ) + _state_sr = _stochastic_round_fp16x2(state, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_sr) + else: + tl.store(state_ptrs, _state_sr, mask=state_mask) + else: + _state_cast = state.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_cast) + else: + tl.store(state_ptrs, _state_cast, mask=state_mask) + + # Phase 2: Output using precomputed CB_scaled and decay_vec + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + x_all = tl.load( + x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_base + + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_all, + mask=t_mask[:, None] & m_mask[None, :], + ) + x_all = x_all.to(tl.float32) + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + tl.float32 + ) + + init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) + out_all = init_out + cb_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# `_persistent_rectangle_impl`: rectangle nowrite path for the persistent +# kernel. Body is a copy of `_rectangle_main_impl` with `pid_m`/`pid_b`/`pid_h` +# lifted to args (same pattern as `_persistent_main_impl` vs `_replay_main_impl`). +# Called only for nowrite slots when the kernel runs with RECTANGLE=True. +# Dropped from the rect impl: LAUNCH_DEPENDENT_KERNELS / USE_PERM / REVERSE_PERM +# (kernel-level, signalled once at top); slot_perm_ptr (kernel resolves perm). +@triton.jit() +def _persistent_rectangle_impl( + # Per-work-unit indices (computed by the persistent wrapper). + pid_m, + pid_b, + pid_h, + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as + # replay path). Used when USE_TMA_LOAD; ignored otherwise. + state_tma_descriptor, + state_scales_ptr, # only consulted when QUANT_MAX > 0 + old_x_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + pad_slot_id, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides (rectangle (batch, nheads, T, K)) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + QUANT_MAX: tl.constexpr, + USE_TMA_LOAD: tl.constexpr = False, +): + if HAS_CACHE_BATCH_INDICES: + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + if cache_batch_idx == pad_slot_id: + return + else: + cache_batch_idx = pid_b.to(tl.int64) + + # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_offset = prev_num_accepted_tokens + + # Static rectangle K-axis layout (matches precompute). + K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_k = tl.arange(0, BLOCK_SIZE_K) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # K-axis masks (approach C: PNAT-runtime offset, matches precompute). + is_old_k = offs_k < prev_num_accepted_tokens + safe_old_k = tl.where(is_old_k, offs_k, 0) + k_new_idx = offs_k - prev_num_accepted_tokens + is_new_k = (k_new_idx >= 0) & (k_new_idx < T) + safe_k_new = tl.where(is_new_k, k_new_idx, 0) + + # Load state. Quant scale hoist: defer `* decode_scale` post-matmul. + if USE_TMA_LOAD: + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state = state_tma_descriptor.load([offs_y, 0]) + else: + state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, other=1.0, + ).to(tl.float32) + else: + state = state.to(tl.float32) + + # Group / pointer offset setup + group_idx = pid_h // nheads_ngroups_ratio + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + # Hoist: old_x doesn't depend on conv1d/precompute; load before gdc_wait. + old_x_load = tl.load( + old_x_base + + safe_old_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + mask=is_old_k[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + + if LAUNCH_WITH_PDL: + tl.extra.cuda.gdc_wait() + + C_all = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + x_K = tl.load( + x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=is_new_k[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_base + + offs_k[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + x_K, + mask=is_new_k[:, None] & m_mask[None, :], + ) + + x_K_f32 = x_K.to(tl.float32) + x_combined = old_x_load + x_K_f32 + + if HAS_D or HAS_Z: + sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) + x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) + else: + x_all = x_K_f32 # placeholder; unused + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec_full = tl.load( + decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 + ).to(tl.float32) + + state_out = ( + tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) + * decay_vec_full[:, None] + ) + if QUANT_MAX > 0.0: + state_out = state_out * decode_scale[None, :] + + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) + + out_all = state_out + token_out + + if HAS_D: + out_all = out_all + x_all * D[None, :] + + if HAS_Z: + z_all = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], other=0.0, + ).to(tl.float32) + out_all_z = out_all * z_all * tl.sigmoid(z_all) + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) + else: + out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) + + +# Persistent main kernel: 1D grid, persistent CTA loop. +# Heuristics mirror those of `_checkpointing_main_kernel`. +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics( + {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} +) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +@triton.heuristics( + {"BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} +) +@triton.heuristics( + {"NUM_PID_M_BLOCKS": lambda args: triton.cdiv(args["dim"], args["BLOCK_SIZE_M"])} +) +@triton.jit() +def _persistent_main_kernel( + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view. + # Shared across BOTH the replay path (consumed by _persistent_main_impl + # when USE_TMA_LOAD_*/STORE) AND the rectangle path (consumed by + # _persistent_rectangle_impl when USE_TMA_LOAD) — same descriptor, same + # block_shape, just gated by separate constexprs per impl. Wrapper sets + # this to a TensorDescriptor when ANY of the three TMA flags is on, else + # to `state_ptr` (raw); each impl ignores it via its own constexpr when + # not consuming it. + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + slot_perm_ptr, + rand_seed_ptr, + pad_slot_id, + # Persistent-loop work-distribution scalars. Caller pre-sorts the batch + # write-first; the kernel uses (n_writes, batch_total, WRITE_CHECKPOINT) + # to derive its own slot range. Write half processes [0, n_writes), + # nowrite half processes [n_writes, batch_total). + # + # n_writes_ptr is a device pointer to a (1,) int32 tensor. Reading + # from device memory (rather than taking a Python int kernel arg) is + # required so mix-mode benchmarking can vary n_writes per iter inside + # a captured CUDA graph — the source tensor's contents change, the + # pointer doesn't. Cost: one int load per kernel launch (~negligible). + # When IS_DYNAMIC=True the value is unused (Triton DCEs the load). + n_writes_ptr, # int32 *: device-side count of write-mode slots + batch_total, # int32: total slot count + nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides + stride_old_x_cache, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + HAS_CACHE_BATCH_INDICES: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + USE_PERM: tl.constexpr, + # NUM_PERSISTENT: runtime int (not constexpr). Used ONLY as the loop + # stride in `tl.range(pid, total_work, NUM_PERSISTENT, ...)`. Making it + # runtime collapses the cta_per_sm tuning dim from the kernel's compile + # signature: 8 CPS values used to mean 8x recompiles; now they share one + # compiled kernel. Work decomposition (pid_m, pid_b_local, pid_h) does + # NOT depend on NUM_PERSISTENT — it uses constexpr NUM_PID_M_BLOCKS and + # runtime n_slots_local — so loop unrolling and flatten=/num_stages=/ + # warp_specialize= optimizations on `tl.range` operate independently of + # the stride value. + NUM_PERSISTENT, + NUM_LOOP_STAGES: tl.constexpr, + NUM_PID_M_BLOCKS: tl.constexpr, + FLATTEN: tl.constexpr, + WARP_SPECIALIZE: tl.constexpr, + IS_DYNAMIC: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr = 16, # rectangle K-axis (heuristic-derived) + RECTANGLE: tl.constexpr = False, # when True, dispatch nowrite slots to _persistent_rectangle_impl + # 3 TMA toggles per the 3 live paths per-compilation: + # USE_TMA_LOAD_WRITE — replay-style state load when is_write + # USE_TMA_LOAD_NOWRITE — nowrite-path state load (rect when RECTANGLE, + # else replay-nowrite) + # USE_TMA_STORE — replay-style state store (only fires on write + # path; no-op when not is_write) + # Wrapper picks USE_TMA_LOAD_NOWRITE = _use_tma_rect_load (if rectangle) + # or _use_tma_replay_nowrite_load (if not). + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # PDL signal: fire once at kernel entry (not per work unit). + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + # Load runtime n_writes from device memory. Read once at kernel entry; + # used only by the !IS_DYNAMIC slot-range derivation below. Triton + # DCEs the load when IS_DYNAMIC=True (n_writes is dead there). + n_writes = tl.load(n_writes_ptr) + + # Derive this kernel's slot range. Two modes: + # IS_DYNAMIC=False (persistent_main): caller pre-sorts and splits halves; + # slot range is [0, n_writes) when WRITE_CHECKPOINT else [n_writes, batch_total) + # IS_DYNAMIC=True (persistent_dynamic): single launch covers full batch; + # each work-item dispatches via runtime PNAT check inside the impl. + if IS_DYNAMIC: + slot_lo = 0 + slot_hi = batch_total + else: + if WRITE_CHECKPOINT: + slot_lo = 0 + slot_hi = n_writes + else: + slot_lo = n_writes + slot_hi = batch_total + n_slots_local = slot_hi - slot_lo + + pid = tl.program_id(axis=0) + total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads + + # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) + # with pid_m varying fastest (M-tile cache locality on state load), then + # slot, then head — mirrors the existing 3D grid's axis ordering + # (axis=0 fastest = pid_m). + for tile_id in tl.range( + pid, total_work, NUM_PERSISTENT, + flatten=FLATTEN, num_stages=NUM_LOOP_STAGES, warp_specialize=WARP_SPECIALIZE, + ): + pid_m = tile_id % NUM_PID_M_BLOCKS + pid_b_local = (tile_id // NUM_PID_M_BLOCKS) % n_slots_local + pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) + # Translate local slot index → global slot index. When USE_PERM is + # set, the caller-provided slot_perm gives the original slot index + # for the post-sort position. + pid_b_grid = pid_b_local + slot_lo + if USE_PERM: + pid_b = tl.load(slot_perm_ptr + pid_b_grid) + else: + pid_b = pid_b_grid + + # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle + # impl. Pre-load cache_batch_idx + PNAT here only for the RECTANGLE + # path's branch decision. Both impls re-load and handle pad_slot_id + # internally (Triton's L1 cache makes the duplicate loads ~free). + if RECTANGLE: + if HAS_CACHE_BATCH_INDICES: + cbi_pre = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + is_pad = cbi_pre == pad_slot_id + else: + cbi_pre = pid_b.to(tl.int64) + is_pad = False + if not is_pad: + pnat_pre = tl.load(prev_num_accepted_tokens_ptr + cbi_pre) + if IS_DYNAMIC: + is_w = (pnat_pre + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_w = WRITE_CHECKPOINT + if is_w: + # Pass WRITE_CHECKPOINT=True constexpr to specialize this + # impl call for the write path. Under IS_DYNAMIC=True, the + # kernel-level WRITE_CHECKPOINT is False (launcher default), + # but the OUTER is_w branch we are inside narrows the + # runtime path to writes-only, so we override to True here + # so the impl's constexpr-gated `if is_write:` blocks DCE + # to the write-only codegen. Under IS_DYNAMIC=False + # (persistent_main), the kernel-level WRITE_CHECKPOINT is + # itself True for this half (write half launches with + # WC=True), and the outer is_w = WRITE_CHECKPOINT = True + # constexpr-folds; passing literal True here is consistent + # and constexpr-equivalent. + _persistent_main_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, + True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) + True, # WC_IS_CONSTEXPR — force inner to use WC constexpr + # 3 TMA flags: write-load fires here (we're in the + # is_write branch), nowrite-load is dead (no slot + # reaches it), store fires (write path). + USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + ) + else: + # Rectangle nowrite: pass state_ptr (raw, always) + + # state_tma_descriptor (the single unified descriptor — + # same memory replay paths use). Rect impl gates use + # of the descriptor via its USE_TMA_LOAD constexpr. + _persistent_rectangle_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, + LAUNCH_WITH_PDL, QUANT_MAX, + USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle + ) + # else: pad slot — skip both impls (both would early-return anyway) + else: + # No rectangle path — single _persistent_main_impl call covers + # both write and nowrite slots via WC constexpr (non-dynamic) or + # runtime is_write (IS_DYNAMIC=True). Pass all 3 TMA flags; + # impl picks USE_TMA_LOAD_WRITE vs USE_TMA_LOAD_NOWRITE based on + # its computed is_write — constexpr-folds when is_write is + # constexpr (non-dyn), runtime branch when IS_DYNAMIC=True. + # (Reverted from outer two-call dispatch: that doubled the + # compiled body size under IS_DYNAMIC=True and regressed RECT=0 + # perf by ~+24%.) + _persistent_main_impl( + pid_m, pid_b, pid_h, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, + WRITE_CHECKPOINT, IS_DYNAMIC, + False, # WC_IS_CONSTEXPR=False — RECT=0 has both write/nowrite slots in one call + USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + ) + + +# ============================================================================ +# Python wrapper +# ============================================================================ + + +_QUANT_MAX_BY_DTYPE = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, +} + + +# --------------------------------------------------------------------------- +# Default tunings — looked up by (effective_batch, dtype, sr) when the caller +# leaves mode/knobs as None. +# +# Effective batch = raw_batch × nheads_per_rank. Our sweep was at TP=8 with +# the standard Mamba2 nheads; at call time we compute it from the input +# tensor shape so callers at other TP / nheads pick up the right cell. +# +# Schema: dict[(dtype_str, sr_str)] → list[(eff_batch_threshold, mode, knobs)] +# sorted by threshold ascending. Lookup finds the first threshold ≥ eff_b +# (so missing intermediate batches fall up to the next tuned cell). If +# eff_b exceeds the largest threshold, use the largest entry. +# +# Each `knobs` dict only contains keys for the chosen mode; the wrapper +# unpacks them with the same name as the matching kwargs. Caller-provided +# kwargs always win over table values. +# +# This table is intentionally NOT parameterized by T or max_window. Our +# sweep was T=6, max_window=16. Callers outside that regime silently get +# the same numbers — they may be suboptimal but they're correct. +# +# Source: audit_v2.py --emit-tuning. Auto-generated from per-cell search +# winners (best of pd / pm by bucket_expected_renorm). Sweep was TP=8 with +# NHEADS=128 → nheads_per_rank=16; thresholds are in effective_batch units. +# Missing dtype/SR combos (fp16/RN, int8/RN, fp8/*) fall back via the +# _resolve_tuning chain — RN→SR for same dtype, then fp8→int8/SR. +_DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { + ("fp32", "RN"): [ + ( 16, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 8, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.22us + ( 32, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.17us + ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.08us + ( 128, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=8, score=9.00us + ( 256, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=16, score=10.92us + ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=13.53us + ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=19.50us + ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=30.28us + ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=50.32us + ( 8192, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=90.99us + (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=171.69us + ], + ("fp16", "SR"): [ + ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.16us + ( 32, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.01us + ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=7.95us + ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 4, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=8.87us + ( 256, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 2, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.28us + ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=12.90us + ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=16.71us + ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=25.71us + ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=39.80us + ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=71.34us + (16384, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=133.51us + ], + ("int8", "SR"): [ + ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.34us + ( 32, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.36us + ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 2, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.40us + ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 10, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=9.37us + ( 256, "persistent_dynamic", {'_block_size_m': 16, '_cta_per_sm': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages': 1, '_num_stages': 4, '_num_warps': 1, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.02us + ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=32, score=13.15us + ( 1024, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=17.82us + ( 2048, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=27.01us + ( 4096, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=43.23us + ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=77.01us + (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=140.43us + ], +} + + +# Knob names that map between the modes' single-value (pd) and split-value +# (pm) namespaces. Used by `_bridge_tuning_knobs` when caller forces a mode +# different from the table's recommendation. +_PD_TO_PM_SPLIT_MAP = { # pd unsplit knob → (pm_write_knob, pm_nowrite_knob) + "_block_size_m": ("_block_size_m_write", "_block_size_m_nowrite"), + "_num_warps": ("_num_warps_write", "_num_warps_nowrite"), + "_num_stages": ("_num_stages_write", "_num_stages_nowrite"), + # CPS / LS are persistent-loop knobs; pd uses _cta_per_sm + _num_loop_stages + # as unsplit, pm uses _cta_per_sm_write/_nowrite + _num_loop_stages_write/_nowrite. + "_cta_per_sm": ("_cta_per_sm_write", "_cta_per_sm_nowrite"), + "_num_loop_stages": ("_num_loop_stages_write", "_num_loop_stages_nowrite"), +} + + +def _bridge_tuning_knobs(knobs: dict, from_mode: str, to_mode: str) -> dict: + """Convert a tuning dict between pd ↔ pm knob namespaces. + + pd → pm: copy each unsplit value to both write/nowrite split knobs; drop + the unsplit form (pm doesn't read it). + pm → pd: take the nowrite split value as the unsplit knob; drop the + write/nowrite split forms (pd doesn't read them). + Shape knobs that exist in both modes (_heads_per_block, _flatten, + _warp_specialize, TMA flags, rectangle_for_nowrite) carry over unchanged. + """ + out = dict(knobs) + if from_mode == "persistent_dynamic" and to_mode == "persistent_main": + for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): + if unsplit in out: + out.setdefault(pm_w, out[unsplit]) + out.setdefault(pm_nw, out[unsplit]) + del out[unsplit] + elif from_mode == "persistent_main" and to_mode == "persistent_dynamic": + for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): + if pm_nw in out: + out.setdefault(unsplit, out[pm_nw]) + out.pop(pm_w, None) + out.pop(pm_nw, None) + return out + + +def _resolve_tuning( + batch: int, nheads_per_rank: int, dt_str: str, sr_str: str, +) -> tuple[str, dict] | None: + """Look up the default mode + knobs for this (eff_batch, dt, sr) cell. + + Returns (mode, knobs_dict) or None if the table has no entry covering + this dtype/sr (including the fp8→int8/SR and dtype/RN→dtype/SR fallbacks). + Returning None lets the wrapper fall back to caller-provided kwargs or + kernel-side defaults. + """ + eff_b = batch * max(1, nheads_per_rank) + # Lookup chain. Order: + # 1. Exact (dt, sr). + # 2. (dt, SR) if RN missing for that dtype. + # 3. Cross-dtype fallback for dtypes we haven't tuned: + # bf16 / int16 → fp16/SR + # fp8 → int8/SR + # Unknown dtype → raise. + valid_dtypes = {"fp32", "fp16", "bf16", "int8", "int16", "fp8"} + if dt_str not in valid_dtypes: + raise ValueError( + f"checkpointing_state_update: unsupported state dtype {dt_str!r}; " + f"expected one of {sorted(valid_dtypes)}" + ) + keys_to_try = [(dt_str, sr_str)] + if sr_str == "RN": + keys_to_try.append((dt_str, "SR")) + if dt_str in ("bf16", "int16"): + keys_to_try.append(("fp16", "SR")) + elif dt_str == "fp8": + keys_to_try.append(("int8", "SR")) + entries = None + for k in keys_to_try: + if k in _DEFAULT_TUNING: + entries = _DEFAULT_TUNING[k] + break + if entries is None: + return None + # Find first threshold ≥ eff_b; if none, use largest entry. + for thresh, mode, knobs in entries: + if eff_b <= thresh: + return mode, dict(knobs) + thresh, mode, knobs = entries[-1] + return mode, dict(knobs) + + +def checkpointing_state_update( + state: torch.Tensor, + old_x: torch.Tensor, + old_B: torch.Tensor, + old_dt: torch.Tensor, + old_dA_cumsum: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + out: torch.Tensor, + # Required persistent-mode plumbing (REQUIRED for both pd and pm; pd + # ignores both internally but the wrapper still demands them): + # n_writes : (1,) int32 device tensor with the count of write-mode + # slots in the batch. pm uses it to size the two halves; + # pd ignores it (per-slot runtime PNAT check). + # slot_perm : (batch,) int32 device tensor remapping grid pid → slot. + # pm uses it to cluster writes first (kernel grid step is + # write_half then nowrite_half); pd ignores it. Callers + # that don't care about ordering should pass arange(batch). + n_writes: torch.Tensor, + slot_perm: torch.Tensor, + D: torch.Tensor | None = None, + z: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + pad_slot_id: int = PAD_SLOT_ID, + rand_seed: torch.Tensor | None = None, + philox_rounds: int = 10, + state_scales: torch.Tensor | None = None, + launch_with_pdl=False, + use_internal_pdl=True, + write_checkpoint: bool = True, + rectangle_for_nowrite: bool | None = None, + mode: str | None = None, + _block_size_m: int | None = None, + _num_warps: int | None = None, + _num_stages: int | None = None, + _precompute_num_warps: int | None = None, + _precompute_num_stages: int | None = None, + _heads_per_block: int | None = None, + _maxnreg: int | None = None, + _num_ctas: int | None = None, + # Per-main knobs (override shared values for one half of the dl-family / + # persistent_main launches). Default None = tied to the shared value + # (backward compat). The two main kernels (write vs nowrite) have + # different per-slot work — write does a state shift + store, nowrite + # just appends — so the optimum (M, W, S, H) can differ. Precompute + # knobs are intentionally NOT split: shared precompute wins (cheaper + # launch, hotter precompute outputs in L2). Persistent CPS / LS knobs + # are also split per-main since the two persistent_main launches have + # different grid sizes. + _block_size_m_write: int | None = None, + _block_size_m_nowrite: int | None = None, + _num_warps_write: int | None = None, + _num_warps_nowrite: int | None = None, + _num_stages_write: int | None = None, + _num_stages_nowrite: int | None = None, + # Note: heads_per_block / precompute_num_warps are NOT split — they only + # affect the precompute kernel, which is shared across write/nowrite. + # TMA state-tensor toggles — 4 independent paths (see CHECKPOINTING_DESIGN.md + # item #17 for measured perf profiles). Each is False=raw load/store, True= + # use a host-built TMA tensor_descriptor for that path. + _use_tma_rect_load: bool | None = None, # rect kernel's state load (nowrite-only) + _use_tma_replay_write_load: bool | None = None, # replay-style state load when WC=True + _use_tma_replay_write_store: bool | None = None, # replay-style state store when WC=True + _use_tma_replay_nowrite_load: bool | None = None, # replay-style state load when WC=False + # Persistent-mode tuning kwargs (consulted for both pd and pm; pd uses + # _cta_per_sm / _num_loop_stages, pm uses the _write/_nowrite splits): + # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally + # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. + # _num_loop_stages : int — `num_stages` arg on the inner `tl.range(...)` + # persistent loop. Note: this is loop-level, NOT the kernel-arg + # `num_stages` (which only pipelines dot-feeding loads). + # _flatten : bool — `flatten` arg on `tl.range(...)`. + # _warp_specialize : bool — `warp_specialize` arg on `tl.range(...)`. + _cta_per_sm: int | None = None, + _num_loop_stages: int | None = None, + _flatten: bool | None = None, + _warp_specialize: bool | None = None, + # Per-main persistent-specific knobs. Same rationale as the BLOCK_SIZE_M + # split above: the two persistent_main launches (write half vs nowrite + # half) have different grid sizes and per-work-item costs, so they may + # want different cta_per_sm / num_loop_stages. + _cta_per_sm_write: int | None = None, + _cta_per_sm_nowrite: int | None = None, + _num_loop_stages_write: int | None = None, + _num_loop_stages_nowrite: int | None = None, +): + """ + Replay SSM state update with precomputed CB and tl.dot fast-forward. + + Two-kernel architecture: + 1. Precompute kernel: computes CB_scaled and decay_vec from B, C, dt, A. + Writes processed dt/dA_cumsum/B to double-buffered cache for next step. + 2. Main kernel: replays old tokens via tl.dot fast-forward on cached data, + then computes output using precomputed CB_scaled and new x/C inputs. + + PDL (Programmatic Dependent Launch) chain: + conv1d → (external PDL) → precompute → (internal PDL) → main + External PDL: precompute starts while conv1d is running; gdc_wait() + in precompute blocks until conv1d completes before loading B/C. + Internal PDL: main starts while precompute is running; main's replay + phase uses only cached data from the previous step. gdc_wait() in + main blocks until precompute completes before loading conv1d outputs + (x, C) and precompute outputs (CB_scaled, decay_vec). + + Uses double-buffered cache tensors. cache_buf_idx[slot] indicates which + buffer (0 or 1) to READ from for replay. The WRITE buffer is 1 - read. + Caller must flip cache_buf_idx[slot] after each call. + + Arguments: + state: (cache, nheads, dim, dstate) in-place. After the call, contains + the state after replaying prev_num_accepted_tokens old tokens. + old_x: (cache, T, nheads, dim) bf16 — old x cache (single-buffered). + old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. + old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. + old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. + cache_buf_idx: (cache,) int32 — which buffer to read (0 or 1). + prev_num_accepted_tokens: (cache,) int32. + x: (batch, T, nheads, dim) new token inputs. + dt: (batch, T, nheads, dim) with stride(-1)==0 (tie_hdim). + A: (nheads, dim, dstate) with stride(-1)==0, stride(-2)==0 (tie_hdim). + B: (batch, T, ngroups, dstate). + C: (batch, T, ngroups, dstate). + out: (batch, T, nheads, dim) preallocated output. + D: (nheads, dim) optional feed-through parameter. + z: (batch, T, nheads, dim) optional silu gate. + dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). + state_batch_indices: (batch,) optional cache slot mapping. + rand_seed: optional single-element int64 CUDA tensor for Philox PRNG seed. + When provided, state is stochastically rounded on store. Supported + for state.dtype in (fp16, int8, int16, fp8_e4m3fn); other dtypes + silently use deterministic rounding. fp16+SR and fp8+SR both + require sm_100a (Blackwell B200+) — wrapper asserts this loudly. + philox_rounds: number of Philox PRNG rounds (default 10). + state_scales: required when state.dtype in (int8, int16, fp8_e4m3fn). + Shape (cache_size, nheads, dim), fp32. Per-(head, dim) channel + decode scale (= 1 / encode_scale). The kernel writes scales on + checkpoint steps and reads them on load (broadcast over dstate). + Ignored for non-quantized state dtypes. + launch_with_pdl: enable external PDL (conv1d → precompute chain). + Defaults False; caller opts in when the upstream chain is PDL-safe. + Ignored on hardware that doesn't support PDL (sm < 90). + use_internal_pdl: enable internal PDL (precompute → main overlap). + Defaults True; override for testing only. + Ignored on hardware that doesn't support PDL (sm < 90). + + _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, + _precompute_num_warps, _precompute_num_stages, _heads_per_block, + _maxnreg, _num_ctas) are benchmark-only overrides; production callers + should leave them None to use the heuristic-tuned defaults. + """ + # PDL needs sm >= 90. + if get_sm_version() < 90: + launch_with_pdl = False + use_internal_pdl = False + + # Mode selection: + # mode=None (default): look up the table-tuned mode + knobs for this + # (effective_batch, dtype, sr) cell. See `_resolve_tuning` above. + # mode="persistent_dynamic": single persistent-CTA kernel covering the + # full batch. Each work-item dispatches via runtime PNAT check + # (is_write = (pnat + T) > MAX). No write/nowrite split. + # slot_perm is honored but optional. write_checkpoint is ignored. + # mode="persistent_main": persistent-CTA kernel with two launches + # (write half + nowrite half). Caller MUST pre-sort slot_perm + # write-first; the n_writes tensor partitions the persistent loop + # into the two halves with the right WRITE_CHECKPOINT constexpr + # each time. RECTANGLE constexpr (= rectangle_for_nowrite) picks + # rect vs replay for the nowrite half. write_checkpoint is ignored. + # Note: mode-and-knob resolution from the default-tuning table happens + # below, after we have `batch` and `nheads`. + + # --- Hardware support gates --- + # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX + # instructions (Ada Lovelace introduced them; Hopper/Blackwell carry them). + if state.dtype == torch.float8_e4m3fn: + assert get_sm_version() >= 89, ( + "fp8_e4m3fn state requires SM 89+ (Ada Lovelace / Hopper / Blackwell) " + f"for fp32↔fp8 cvt PTX instructions; current SM is {get_sm_version()}." + ) + + # PTX cvt.rs.* (stochastic rounding) family lands on Blackwell only. + # Wrapper fails loud; framework decides fall-back (e.g. drop SR, use RN). + # int8 / int16 SR uses pure-Triton libdevice.floor + uniform noise — no + # PTX SR instruction needed, runs anywhere. + if rand_seed is not None: + if state.dtype == torch.float16: + assert get_sm_version() >= 100, ( + "fp16 stochastic rounding (PTX cvt.rs.f16x2.f32) requires " + f"sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + ) + elif state.dtype == torch.float8_e4m3fn: + assert get_sm_version() >= 100, ( + "fp8 stochastic rounding (PTX cvt.rs.satfinite.e4m3x4.f32) " + f"requires sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + ) + + # --- Unsqueeze inputs to canonical shapes --- + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if x.dim() == 3: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if dt.dim() == 3: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if B.dim() == 3: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if C.dim() == 3: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None: + if z.dim() == 2: + z = z.unsqueeze(1) + if z.dim() == 3: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + if out.dim() == 2: + out = out.unsqueeze(1) + if out.dim() == 3: + out = out.unsqueeze(1) + + cache_size, nheads, dim, dstate = state.shape + batch, T, _, _ = x.shape + ngroups = B.shape[2] + assert nheads % ngroups == 0 + + # --- Quantization plumbing (needed for SR/RN classification below) --- + # QUANT_MAX > 0 ⇔ state is int8 / int16 / fp8_e4m3fn. Kernel-entry + # static_assert on the Triton side mirrors this invariant. + quant_max = _QUANT_MAX_BY_DTYPE.get(state.dtype, 0.0) + is_quantized = quant_max > 0.0 + + # --- Default-tuning lookup --- + # Resolve (mode, knobs) from the table when caller leaves them None. + # Caller-provided kwargs always win. If the caller forces a mode that + # differs from the table's recommendation for this cell, we BRIDGE the + # table's knobs into the forced mode's knob namespace rather than fall + # back to (likely-terrible) kernel defaults: + # table pd → forced pm: copy each unsplit pd knob (M, W, S, CPS, LS) + # to both write and nowrite split knobs. + # table pm → forced pd: take the nowrite split values (Mnw, Wnw, Snw, + # CPSnw, LSnw) as the unsplit knobs. + # Empty table → no-op (caller passes whatever, mode falls back to pd). + _dt_str = { + torch.float32: "fp32", + torch.float16: "fp16", + torch.bfloat16: "bf16", + torch.int8: "int8", + torch.int16: "int16", + torch.float8_e4m3fn: "fp8", + }.get(state.dtype, str(state.dtype)) + _sr_str = "SR" if (rand_seed is not None and is_quantized) else "RN" + _table_entry = _resolve_tuning(batch, nheads, _dt_str, _sr_str) + if _table_entry is not None: + _table_mode, _table_knobs = _table_entry + if mode is None: + mode = _table_mode + if mode != _table_mode: + # Bridge across modes — see header comment above. + _table_knobs = _bridge_tuning_knobs(_table_knobs, _table_mode, mode) + # Fill None-valued kwargs from table. We can't reliably mutate + # locals() for re-read, so re-bind each kwarg explicitly. + if rectangle_for_nowrite is None and "rectangle_for_nowrite" in _table_knobs: + rectangle_for_nowrite = bool(_table_knobs["rectangle_for_nowrite"]) + _block_size_m = _block_size_m if _block_size_m is not None else _table_knobs.get("_block_size_m") + _num_warps = _num_warps if _num_warps is not None else _table_knobs.get("_num_warps") + _num_stages = _num_stages if _num_stages is not None else _table_knobs.get("_num_stages") + _heads_per_block = _heads_per_block if _heads_per_block is not None else _table_knobs.get("_heads_per_block") + _precompute_num_warps = _precompute_num_warps if _precompute_num_warps is not None else _table_knobs.get("_precompute_num_warps") + _precompute_num_stages = _precompute_num_stages if _precompute_num_stages is not None else _table_knobs.get("_precompute_num_stages") + _block_size_m_write = _block_size_m_write if _block_size_m_write is not None else _table_knobs.get("_block_size_m_write") + _block_size_m_nowrite = _block_size_m_nowrite if _block_size_m_nowrite is not None else _table_knobs.get("_block_size_m_nowrite") + _num_warps_write = _num_warps_write if _num_warps_write is not None else _table_knobs.get("_num_warps_write") + _num_warps_nowrite = _num_warps_nowrite if _num_warps_nowrite is not None else _table_knobs.get("_num_warps_nowrite") + _num_stages_write = _num_stages_write if _num_stages_write is not None else _table_knobs.get("_num_stages_write") + _num_stages_nowrite = _num_stages_nowrite if _num_stages_nowrite is not None else _table_knobs.get("_num_stages_nowrite") + _cta_per_sm = _cta_per_sm if _cta_per_sm is not None else _table_knobs.get("_cta_per_sm") + _num_loop_stages = _num_loop_stages if _num_loop_stages is not None else _table_knobs.get("_num_loop_stages") + # Split-form resolution for pm's per-half knobs. Without these the + # table's _num_loop_stages_{write,nowrite} and _cta_per_sm_{write, + # nowrite} values are dead — pm reads the split forms but the + # wrapper would leave them None, falling through to Triton defaults + # (or our hardcoded `or 1` / `or 2` per-mode fallbacks). + _num_loop_stages_write = _num_loop_stages_write if _num_loop_stages_write is not None else _table_knobs.get("_num_loop_stages_write") + _num_loop_stages_nowrite = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _table_knobs.get("_num_loop_stages_nowrite") + _cta_per_sm_write = _cta_per_sm_write if _cta_per_sm_write is not None else _table_knobs.get("_cta_per_sm_write") + _cta_per_sm_nowrite = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _table_knobs.get("_cta_per_sm_nowrite") + _flatten = _flatten if _flatten is not None else _table_knobs.get("_flatten") + _warp_specialize = _warp_specialize if _warp_specialize is not None else _table_knobs.get("_warp_specialize") + _use_tma_rect_load = _use_tma_rect_load or bool(_table_knobs.get("_use_tma_rect_load", False)) + _use_tma_replay_write_load = _use_tma_replay_write_load or bool(_table_knobs.get("_use_tma_replay_write_load", False)) + _use_tma_replay_write_store = _use_tma_replay_write_store or bool(_table_knobs.get("_use_tma_replay_write_store", False)) + _use_tma_replay_nowrite_load = _use_tma_replay_nowrite_load or bool(_table_knobs.get("_use_tma_replay_nowrite_load", False)) + # Final defaults if neither caller nor table set them (empty table case). + if mode is None: + mode = "persistent_dynamic" + if rectangle_for_nowrite is None: + rectangle_for_nowrite = False + assert mode in ("persistent_dynamic", "persistent_main"), ( + f"unknown mode {mode!r}; expected 'persistent_dynamic' or 'persistent_main'" + ) + if is_quantized: + assert state_scales is not None, ( + f"state.dtype={state.dtype} requires state_scales tensor " + "(shape (cache_size, nheads, dim), fp32)." + ) + assert state_scales.shape == (cache_size, nheads, dim), ( + f"state_scales shape mismatch: expected {(cache_size, nheads, dim)}, " + f"got {state_scales.shape}." + ) + assert state_scales.dtype == torch.float32, ( + f"state_scales must be fp32, got {state_scales.dtype}." + ) + assert state_scales.device == state.device + + # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the + # placeholder degenerate case max_window = T (every step is a checkpoint + # step). For real replay-style checkpointing, max_window > T and + # `prev_num_accepted_tokens` can be 0..max_window. Window-axis kernel + # tiles (BLOCK_SIZE_WINDOW, BLOCK_SIZE_K) are derived independently from + # MAX_REPLAY_BUFFER_LENGTH so max_window can exceed BLOCK_SIZE_T freely. + max_window = old_x.shape[1] + assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" + + assert x.shape == (batch, T, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + assert B.shape == (batch, T, ngroups, dstate) + assert C.shape == B.shape + assert old_x.shape == (cache_size, max_window, nheads, dim) + assert old_B.shape == (cache_size, 2, max_window, ngroups, dstate) + assert old_dt.shape == (cache_size, 2, nheads, max_window) + assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_window) + assert cache_buf_idx.shape == (cache_size,) + assert prev_num_accepted_tokens.shape == (cache_size,) + + tie_hdim = ( + A.stride(-1) == 0 + and A.stride(-2) == 0 + and dt.stride(-1) == 0 + and (dt_bias is None or dt_bias.stride(-1) == 0) + ) + assert tie_hdim + + device = x.device + BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) + # Rectangle K-axis bound = window (max_window). Computed unconditionally + # so the launch sites can refer to it; only used on the rectangle path. + BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) + + # Allocate precomputed intermediates (per-call, not cached). Always + # allocate (T, K) — the largest layout that any path uses. Replay-style + # paths only touch the first T columns; rectangle/dynamic use the full K. + # The few extra unused columns per row are negligible (~6KB per layer at + # production sizes) and let the dispatch helpers share one buffer. + cb_scaled = torch.empty( + batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 + ) + decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) + + z_strides = ( + (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) if z is not None else (0, 0, 0, 0) + ) + + # Kernel tuning: BLOCK_SIZE_M, num_warps, HEADS_PER_BLOCK, precompute_num_warps. + # Dtype-aware heuristic from B200 sweeps (batch 1-512, T=6/32, TP=8, conv1d + + # chained PDL). Keyed on total_heads, BLOCK_SIZE_T, and state dtype; 16-bit + # states prefer different tiles from fp32 due to lower bandwidth. Philox + # gets its own branch — stochastic rounding shifts compute toward CUDA cores, + # so small-batch configs want more warps to hide the extra work. + total_heads = batch * nheads + heads_per_group = nheads // ngroups + state_is_16bit = state.dtype in (torch.float16, torch.bfloat16) + use_philox = rand_seed is not None + if BLOCK_SIZE_T <= 16: + if use_philox and state_is_16bit: + # Philox: more warps at small batch to hide CUDA core work. + # At large batch, converges to non-Philox fp16 config. + if total_heads <= 16: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 4, 4, 4, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + elif state_is_16bit: + if total_heads <= 16: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 4, 1 + elif total_heads <= 64: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 2, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 1, + min(2, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + else: # fp32 state (no Philox — fp32 doesn't need stochastic rounding) + if total_heads <= 32: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 4, 1 + elif total_heads <= 64: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 + elif total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 2, 2, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 1, 2, 1 + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 2, + min(2, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 + else: # T > 16 + if state_is_16bit: + if total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 16, + 1, + 4, + min(2, heads_per_group), + ) + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 1, + min(4, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 1, + 4, + min(2, heads_per_group), + ) + else: # fp32 state + if total_heads <= 128: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 + elif total_heads <= 256: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 32, + 2, + 4, + min(2, heads_per_group), + ) + elif total_heads <= 512: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 2, + min(4, heads_per_group), + ) + else: + BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( + 64, + 2, + 4, + min(2, heads_per_group), + ) + if _block_size_m is not None: + BLOCK_SIZE_M = _block_size_m + if _num_warps is not None: + num_warps = _num_warps + if _heads_per_block is not None: + # Cap at heads_per_group: HEADS_PER_BLOCK divides the kernel's head + # axis, so a table value larger than the model's heads-per-group + # would overshoot. Protects callers running smaller models than + # the one we tuned against. + heads_per_block = min(_heads_per_block, heads_per_group) + if _precompute_num_warps is not None: + precompute_num_warps = _precompute_num_warps + + # Per-main knob resolution: each _*_{write,nowrite} arg, if not None, + # overrides the corresponding shared value for ONE main launch only. + # Default (None) = tied to shared value (current behavior). + BLOCK_SIZE_M_WRITE = _block_size_m_write if _block_size_m_write is not None else BLOCK_SIZE_M + BLOCK_SIZE_M_NOWRITE = _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M + NUM_WARPS_WRITE = _num_warps_write if _num_warps_write is not None else num_warps + NUM_WARPS_NOWRITE = _num_warps_nowrite if _num_warps_nowrite is not None else num_warps + NUM_STAGES_WRITE = _num_stages_write if _num_stages_write is not None else _num_stages + NUM_STAGES_NOWRITE = _num_stages_nowrite if _num_stages_nowrite is not None else _num_stages + # Persistent-only per-main: + CTA_PER_SM_WRITE = _cta_per_sm_write if _cta_per_sm_write is not None else _cta_per_sm + CTA_PER_SM_NOWRITE = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _cta_per_sm + NUM_LOOP_STAGES_WRITE = _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages + NUM_LOOP_STAGES_NOWRITE = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages + + HAS_CACHE_BATCH_INDICES = state_batch_indices is not None + + assert nheads % heads_per_block == 0, ( + f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" + ) + assert heads_per_block <= heads_per_group, ( + f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" + ) + + # state_scales pointer + strides: real tensor when quantized, otherwise + # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). + if is_quantized: + state_scales_arg = state_scales + state_scales_strides = ( + state_scales.stride(0), + state_scales.stride(1), + state_scales.stride(2), + ) + else: + state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 + state_scales_strides = (0, 0, 0) + + # Per-path TMA descriptors for state — write-side and nowrite-side. Each + # kernel launch consumes the descriptor whose block_shape[0] matches its + # BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need + # distinct descriptors; otherwise the descriptor's block_shape[0] would + # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot / arithmetic + # on the loaded tile fails shape inference at compile time + # ("Cannot make_shape_compatible: incompatible dimensions"). When Mw == + # Mnw (tied, the common case) the two descriptors are the same object. + # Same memory (state's flat 2D view, shape (cache*nheads*dim, dstate)) + # and same dstate block_shape — only block_shape[0] differs. + # When no TMA flag is on, both variables hold the raw `state` tensor as a + # dummy; kernels never reference it because their constexprs are all + # False (Triton DCEs the dead branches). + # `triton.set_allocator()` must run before any descriptor-using launch. + if (_use_tma_rect_load or _use_tma_replay_write_load + or _use_tma_replay_write_store or _use_tma_replay_nowrite_load): + from triton.tools.tensor_descriptor import TensorDescriptor + _ensure_tma_allocator() + assert state.is_contiguous(), "TMA state requires contiguous state" + assert state.stride(-1) == 1, "TMA state requires inner stride 1" + _state_flat = state.view(-1, state.shape[-1]) + _dstate_pow2 = triton.next_power_of_2(dstate) + state_tma_descriptor_write = TensorDescriptor.from_tensor( + _state_flat, block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], + ) + if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: + state_tma_descriptor_nowrite = state_tma_descriptor_write + else: + state_tma_descriptor_nowrite = TensorDescriptor.from_tensor( + _state_flat, block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], + ) + else: + state_tma_descriptor_write = state # dummy; all consuming constexprs False + state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False + + # Slot permutation — pointer + USE_PERM gate. Always required: the + # persistent_main kernel reads pid_b through slot_perm to walk the + # write-first sorted batch. persistent_dynamic forces USE_PERM=False + # at the call site (see launch_persistent_dynamic_main below) so the + # perm value doesn't matter for pd, but the tensor must still be valid. + assert isinstance(slot_perm, torch.Tensor), ( + f"slot_perm must be a torch.Tensor, got {type(slot_perm).__name__}" + ) + assert slot_perm.device == device, ( + f"slot_perm must be on device {device}, got {slot_perm.device}" + ) + assert slot_perm.dtype in (torch.int32, torch.int64), ( + f"slot_perm must be int32/int64, got {slot_perm.dtype}" + ) + assert slot_perm.shape == (batch,), ( + f"slot_perm must have shape (batch={batch},), got {tuple(slot_perm.shape)}" + ) + assert isinstance(n_writes, torch.Tensor), ( + f"n_writes must be a torch.Tensor, got {type(n_writes).__name__}" + ) + assert n_writes.device == device, ( + f"n_writes must be on device {device}, got {n_writes.device}" + ) + assert n_writes.dtype == torch.int32, ( + f"n_writes must be int32, got {n_writes.dtype}" + ) + assert n_writes.shape == (1,), ( + f"n_writes must have shape (1,), got {tuple(n_writes.shape)}" + ) + slot_perm_arg = slot_perm + use_perm = True + + precomp_grid = (batch, nheads // heads_per_block) + d_strides = (D.stride(0), D.stride(1)) if D is not None else (0, 0) + + # ---- Launch helpers (close over locals) ------------------------------- + # Each helper is a thin closure that calls one Triton kernel with the + # full positional + kwarg argument list. Mode-dependent constexprs + # (write_checkpoint, early_out, rectangle) are passed in. + + def launch_dynamic_precompute(rectangle: bool): + _dynamic_precompute_kernel[precomp_grid]( + dt, dt_bias, A, B, C, + cb_scaled, decay_vec, + old_B, old_dt, old_dA_cumsum, + cache_buf_idx, prev_num_accepted_tokens, + state_batch_indices, pad_slot_id, + T, max_window, dstate, nheads // ngroups, + dt.stride(0), dt.stride(1), dt.stride(2), + dt_bias.stride(0) if dt_bias is not None else 0, + A.stride(0), + B.stride(0), B.stride(1), B.stride(2), B.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + dt_softplus, + HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, + LAUNCH_WITH_PDL=launch_with_pdl, + LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, + HEADS_PER_BLOCK=heads_per_block, + RECTANGLE=rectangle, + num_warps=precompute_num_warps, + **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), + launch_pdl=launch_with_pdl, + ) + + # ---- launch_persistent_main ------------------------------------------ + # Persistent-CTA main kernel. Single launch covers `n_slots` slots + # starting at `slot_offset`. Caller invokes twice: once for the write + # half (slot_offset=0, n_slots=n_writes, write_checkpoint=True) and + # once for the nowrite half (slot_offset=n_writes, + # n_slots=batch-n_writes, write_checkpoint=False). Hard-sort + # contract: caller has pre-sorted slots so [0, n_writes) are writes + # and [n_writes, batch) are nowrites. + + # Resolve persistent-mode bench knobs. Defaults: cta_per_sm = 1 + # (one CTA per SM, matches upstream `_p_matmul_ogs.py`); num_loop_stages + # = 2 (matches in-tree `swiglu` precedent for non-dot persistent loops); + # flatten = True (canonical Triton 3.6 idiom); warp_specialize = False. + _num_sms = torch.cuda.get_device_properties(device).multi_processor_count + cta_per_sm_arg = _cta_per_sm if _cta_per_sm else 1 + num_persistent_arg = cta_per_sm_arg * _num_sms + num_loop_stages_arg = _num_loop_stages if _num_loop_stages else 2 + flatten_arg = True if _flatten is None else bool(_flatten) + warp_specialize_arg = False if _warp_specialize is None else bool(_warp_specialize) + # Per-launch work-item count. At small batch, total_work may be < the + # full persistent grid; capping `grid` at `min(NUM_PERSISTENT, total_work)` + # avoids launching empty CTAs that pay setup cost for no work. Correctness: + # the kernel's `tl.range(pid, total_work, NUM_PERSISTENT)` ensures each + # tile_id is covered exactly once across all live pids in [0, grid) when + # grid <= NUM_PERSISTENT (each CTA does 1 tile; loop step >= total_work + # exits immediately) AND when grid == NUM_PERSISTENT (each CTA loops over + # multiple tiles). NUM_PERSISTENT is now a runtime int (see kernel def + # docstring at _persistent_main_kernel) so changing cta_per_sm does NOT + # trigger a new Triton compile — same kernel binary, different loop step. + # (Named UPPERCASE for historical Triton-style consistency only; not + # constexpr.) + _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + + def launch_persistent_main(write_checkpoint: bool, + *, + launch_dependent_kernels: bool = False, + rectangle: bool = False): + # `n_writes` (wrapper-level) is the (1,) int32 device tensor with the + # write count. Both halves always launch; the kernel's runtime PNAT + # check iterates only the slots that belong to its half. + _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE + _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE + _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE + _cps = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE + _cps = _cps if _cps else 1 + _nls = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE + _nls = _nls if _nls else 2 + _num_persistent = _cps * _num_sms + _num_pid_m_local = (dim + _bsm - 1) // _bsm + # Grid sizing: cap at min(full persistent grid, upper-bound total work). + # We use `batch` as the upper bound on slots-per-half — overcounting + # by a few CTAs is fine since the kernel's runtime check only + # iterates the slots that actually belong to its half. + _total_work_launch = max(1, batch * _num_pid_m_local * nheads) + grid = (min(_num_persistent, _total_work_launch),) + # Per-path TMA descriptor — block_shape[0] must match _bsm. + _desc = (state_tma_descriptor_write if write_checkpoint + else state_tma_descriptor_nowrite) + _persistent_main_kernel[grid]( + state, _desc, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + n_writes, batch, nheads, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + _bsm, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=write_checkpoint, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + USE_PERM=use_perm, + NUM_PERSISTENT=_num_persistent, + NUM_LOOP_STAGES=_nls, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=False, + RECTANGLE=rectangle, + # 3 TMA flags. IS_DYNAMIC=False: WC fixed per launch; impl + # constexpr-folds the LOAD pick. When WC=True (write half), + # NOWRITE_LOAD is dummy False; when WC=False, WRITE_LOAD/STORE + # dummy False. NOWRITE_LOAD picks rect-load (RECTANGLE) or + # replay-nowrite-load. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), + USE_TMA_LOAD_NOWRITE=bool( + (_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load) + and not write_checkpoint + ), + USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), + num_warps=_nw, + **({"num_stages": _ns} if _ns else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, + launch_dependent_kernels: bool = False, + rectangle: bool = False): + # Single-launch persistent kernel covering the whole batch with + # runtime per-slot WRITE_CHECKPOINT branch. No half-split, no + # n_writes needed (the kernel ignores n_writes_dev when + # IS_DYNAMIC=True; Triton DCEs the load). is_write is computed + # at runtime per work-item from the loaded PNAT. + # We still pass `n_writes_dev` (the same tensor the persistent_main + # path uses) so the kernel signature is uniform; the value is + # immaterial. + # Grid sizing: cap at total_work (= batch * num_pid_m * nheads) for + # the dynamic case (full-batch coverage); see launch_persistent_main + # comment for correctness rationale. + _total_work_launch = max(1, batch * _num_pid_m * nheads) + grid = (min(num_persistent_arg, _total_work_launch),) + # Persistent-dynamic kernel uses a single BLOCK_SIZE_M (same as the + # wrapper's BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE tied convention), so + # the write-side descriptor matches. Both write and nowrite slots + # in this kernel share that BSM. + _persistent_main_kernel[grid]( + state, state_tma_descriptor_write, state_scales_arg, old_x, + old_B, old_dt, old_dA_cumsum, + prev_num_accepted_tokens, cache_buf_idx, + x, C, D, z, out, + cb_scaled, decay_vec, + state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, + n_writes, batch, nheads, + T, max_window, dim, dstate, nheads // ngroups, + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_B.stride(0), old_B.stride(1), old_B.stride(2), + old_B.stride(3), old_B.stride(4), + old_dt.stride(0), old_dt.stride(1), + old_dt.stride(2), old_dt.stride(3), + old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + x.stride(0), x.stride(1), x.stride(2), x.stride(3), + C.stride(0), C.stride(1), C.stride(2), C.stride(3), + d_strides[0], d_strides[1], + z_strides[0], z_strides[1], z_strides[2], z_strides[3], + out.stride(0), out.stride(1), out.stride(2), out.stride(3), + cb_scaled.stride(0), cb_scaled.stride(1), + cb_scaled.stride(2), cb_scaled.stride(3), + decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + BLOCK_SIZE_M, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=False, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + # persistent_dynamic forces USE_PERM=False regardless of caller- + # provided slot_perm — our pd tuning runs all happened with + # SORT=0 (no slot_perm passed), so honoring slot_perm here would + # silently shift pd to an untimed code path. Revisit if/when + # we benchmark pd with slot_perm. + USE_PERM=False, + NUM_PERSISTENT=num_persistent_arg, + NUM_LOOP_STAGES=num_loop_stages_arg, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=True, + RECTANGLE=rectangle, + # 3 TMA flags. IS_DYNAMIC=True: is_write is runtime per slot; + # impl's load TMA picks per-slot (constexpr ternary becomes a + # runtime branch — both load forms emitted, ~negligible cost). + # NOWRITE_LOAD picks rect-load when RECTANGLE, else + # replay-nowrite-load. STORE only fires on runtime is_write. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load), + USE_TMA_LOAD_NOWRITE=bool( + _use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load + ), + USE_TMA_STORE=bool(_use_tma_replay_write_store), + num_warps=num_warps, + **({"num_stages": _num_stages} if _num_stages else {}), + **({"num_ctas": _num_ctas} if _num_ctas else {}), + **({"maxnreg": _maxnreg} if _maxnreg else {}), + launch_pdl=use_internal_pdl, + ) + + # ---- Mode dispatch ---------------------------------------------------- + with torch.cuda.device(device.index): + if mode == "persistent_dynamic": + # Single-launch persistent kernel covering the full batch. Each + # work-item dispatches via runtime PNAT check. Kernel ignores + # n_writes (Triton DCEs the load) when IS_DYNAMIC=True; we still + # pass the wrapper-provided tensor as required by the signature. + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_persistent_dynamic_main( + n_writes, + launch_dependent_kernels=False, + rectangle=rectangle_for_nowrite, + ) + elif mode == "persistent_main": + # Persistent-CTA main kernel. One shared dynamic_precompute + # (per-slot dispatch via PNAT) feeds two persistent_main + # launches (write half + nowrite half). Both halves ALWAYS + # launch; the kernel's runtime check iterates only the slots + # belonging to its half (write: [0, n_writes), nowrite: + # [n_writes, batch)). + # + # Caller-provided contract: `n_writes` is a (1,) int32 device + # tensor (the kernel reads it at runtime, after the precompute); + # `slot_perm` is a (batch,) int32 device tensor pre-sorted + # write-first. + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_persistent_main( + write_checkpoint=True, + launch_dependent_kernels=True, + rectangle=False, # write always replay-style + ) + launch_persistent_main( + write_checkpoint=False, + launch_dependent_kernels=False, + rectangle=rectangle_for_nowrite, + ) + else: + raise ValueError( + f"mode={mode!r} is not supported. Supported modes: " + f"'persistent_dynamic', 'persistent_main'." + ) diff --git a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py index 9dad703a8cde..65c73cdd5aa7 100644 --- a/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py +++ b/tensorrt_llm/_torch/modules/mamba/checkpointing_state_update_slim.py @@ -2596,6 +2596,15 @@ def checkpointing_state_update( _num_stages_nowrite = _num_stages_nowrite if _num_stages_nowrite is not None else _table_knobs.get("_num_stages_nowrite") _cta_per_sm = _cta_per_sm if _cta_per_sm is not None else _table_knobs.get("_cta_per_sm") _num_loop_stages = _num_loop_stages if _num_loop_stages is not None else _table_knobs.get("_num_loop_stages") + # Split-form resolution for pm's per-half knobs. Without these the + # table's _num_loop_stages_{write,nowrite} and _cta_per_sm_{write, + # nowrite} values are dead — pm reads the split forms but the + # wrapper would leave them None, falling through to Triton defaults + # (or our hardcoded `or 1` / `or 2` per-mode fallbacks). + _num_loop_stages_write = _num_loop_stages_write if _num_loop_stages_write is not None else _table_knobs.get("_num_loop_stages_write") + _num_loop_stages_nowrite = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _table_knobs.get("_num_loop_stages_nowrite") + _cta_per_sm_write = _cta_per_sm_write if _cta_per_sm_write is not None else _table_knobs.get("_cta_per_sm_write") + _cta_per_sm_nowrite = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _table_knobs.get("_cta_per_sm_nowrite") _flatten = _flatten if _flatten is not None else _table_knobs.get("_flatten") _warp_specialize = _warp_specialize if _warp_specialize is not None else _table_knobs.get("_warp_specialize") _use_tma_rect_load = _use_tma_rect_load or bool(_table_knobs.get("_use_tma_rect_load", False)) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py index 84d249ddfbb3..a6dceb85a320 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update_slim.py @@ -196,7 +196,9 @@ def _load(mod_name: str, file_name: str): _load("softplus", "softplus.py") # 3. The actual kernels replay_mod = _load("replay_selective_state_update", "replay_selective_state_update.py") - checkpoint_mod = _load("checkpointing_state_update", "checkpointing_state_update.py") + # _slim variant: bench loads the slim kernel file to match the bench + # filename suffix. --full-import path also loads slim (line 214). + checkpoint_mod = _load("checkpointing_state_update_slim", "checkpointing_state_update_slim.py") base_mod = _load("selective_state_update", "selective_state_update.py") conv1d_mod = _load("causal_conv1d_triton", "causal_conv1d_triton.py") @@ -625,8 +627,14 @@ def _kernels_per_iter_incremental( k = 2 # 1 dynamic_precomp + 1 persistent_main elif mode == "persistent_main": k = 2 if persistent_skip_empty else 3 # see docstring - else: - raise ValueError(f"_kernels_per_iter_incremental: unknown mode {mode!r}") + elif mode is None: + # mode=None: wrapper resolves from _DEFAULT_TUNING per-cell. Most + # tuning entries pick persistent_main (3 kernels: precompute + + # write-main + nowrite-main; slim no longer host-skips empty halves), + # so default to pm's count. pd-picking cells will skip on K mismatch + # (CUPTI sees only 2); their results land in the skipped sidecar + # rather than as bench rows. Acceptable for an audit run. + k = 3 if with_conv1d: k += 1 return k @@ -2944,13 +2952,21 @@ def _set(v): # flag set to a non-zero sweep value # it once here. _n_writes_dev_pure: torch.Tensor | None = None _host_n_writes_pure: int | None = None - if mode in ("persistent_main", "persistent_dynamic") and scenario_n_writes_dev is None: + # mode=None means the wrapper resolves from the tuning table — + # we don't know whether it'll pick pm or pd at this point, so + # pre-allocate as if pm (which needs the value). pd will + # ignore it via IS_DYNAMIC DCE. + if mode in ("persistent_main", "persistent_dynamic", None) and scenario_n_writes_dev is None: _n_writes_dev_pure = torch.zeros(1, dtype=torch.int32, device=state_work.device) - if mode == "persistent_main": + if mode in ("persistent_main", None): scn_fill = scn["fill"] - is_write_scenario_local = (scn_fill + mtp_len) > max_window - _host_n_writes_pure = batch if is_write_scenario_local else 0 - _n_writes_dev_pure.fill_(_host_n_writes_pure) + if scn_fill is not None: + is_write_scenario_local = (scn_fill + mtp_len) > max_window + _host_n_writes_pure = batch if is_write_scenario_local else 0 + _n_writes_dev_pure.fill_(_host_n_writes_pure) + # else: mix scenario — n_writes varies per iter, set per-iter + # later by the mix replay path; leave _n_writes_dev_pure as + # zeros placeholder. def _run_incr( block_size_m=block_size_m, @@ -2985,8 +3001,15 @@ def _run_incr( extra_kwargs["write_checkpoint"] = write_checkpoint extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite extra_kwargs["mode"] = mode - if sort_slots: - extra_kwargs["slot_perm"] = slot_perm_buf + # slim wrapper REQUIRES n_writes (1,) int32 device tensor + # and slot_perm (batch,) int32 device tensor. Always pass + # them (pd will ignore the count internally, pm splits). + extra_kwargs["slot_perm"] = slot_perm_buf + extra_kwargs["n_writes"] = ( + scenario_n_writes_dev + if scenario_n_writes_dev is not None + else _n_writes_dev_pure + ) # reverse_nowrite kwarg dropped from the slim kernel # wrapper (was a maindl/dlgrouped feature). The sweep # axis is retained here only so cell-list cells emitted @@ -3008,11 +3031,14 @@ def _run_incr( # n_writes is either 0 (all nowrite) or batch (all # write) depending on whether PNAT+T overflows the # window. Mix scenarios are skipped earlier. - if mode in ("persistent_main", "persistent_dynamic"): + if mode in ("persistent_main", "persistent_dynamic", None): # Per-cell sweep values for persistent-only knobs. # Apply to both persistent variants. _parse_sweep # returns [None] when the user didn't pass the flag, # in which case we leave the wrapper's defaults. + # mode=None means "let the wrapper resolve from the + # tuning table" — wrapper picks pm or pd, both consume + # these knobs, so propagate caller-explicit overrides. if cta_per_sm is not None: extra_kwargs["_cta_per_sm"] = cta_per_sm if num_loop_stages is not None: @@ -3044,24 +3070,10 @@ def _run_incr( "modes. Re-run with --sort-slots 1 or " "--hardcode-sort 1." ) - # n_writes plumbing: pure scenarios pass an int (host - # knows the value, wrapper host-skips empty halves); - # mix scenarios pass only a (1,) device tensor updated - # per iter by the benchmark pre-iter path (wrapper - # cannot host-skip since host_n_writes is unknown). - if scenario_n_writes_dev is not None: - # Mix: caller-allocated tensor, updated per iter. - extra_kwargs["_n_writes_dev"] = scenario_n_writes_dev - elif mode == "persistent_main": - # Pure pm: pre-allocated tensor + host int. - extra_kwargs["_n_writes"] = _host_n_writes_pure - extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure - elif mode == "persistent_dynamic": - # persistent_dynamic pure: kernel ignores n_writes - # via IS_DYNAMIC DCE, but the wrapper needs a - # valid (1,) tensor pointer. Pass the pre-allocated - # zero tensor to avoid any in-capture alloc. - extra_kwargs["_n_writes_dev"] = _n_writes_dev_pure + # n_writes is now always passed as the slim wrapper's + # required `n_writes=` kwarg above; the legacy + # `_n_writes` (host int) / `_n_writes_dev` (device + # tensor) plumbing is no longer needed. variant_fn( state_work, old_x_work, @@ -3107,61 +3119,56 @@ def _run_incr( ) parts = [] + # Every tag knob emits its value or "auto" when unset (i.e. the + # bench-level value is None, meaning the wrapper resolves from + # the _DEFAULT_TUNING table per-cell). Uniform key set across + # all cells keeps the JSONL keys stable and prevents collisions + # between "I didn't set this" and "I explicitly set this to 0". + def _val(v): + return "auto" if v is None else v + # When tied (not _any_split), emit the shared single-value tag # (M=8 etc). When split, emit explicit Mw / Mnw tags so cells # with the same shared value but different per-main values get # unique JSON keys. def _emit_split(name_w, name_nw, val_w, val_nw): if val_w is None and val_nw is None: + parts.append(f"{name_w[:-1]}=auto") # tied form, both auto return if not _any_split or val_w == val_nw: - parts.append(f"{name_w[:-1]}={val_w}") # strip the 'w' suffix + parts.append(f"{name_w[:-1]}={_val(val_w)}") else: - parts.append(f"{name_w}={val_w}") - parts.append(f"{name_nw}={val_nw}") + parts.append(f"{name_w}={_val(val_w)}") + parts.append(f"{name_nw}={_val(val_nw)}") _emit_split("Mw", "Mnw", block_size_m_w, block_size_m_nw) _emit_split("Ww", "Wnw", num_warps_w, num_warps_nw) _emit_split("Sw", "Snw", num_stages_w, num_stages_nw) - if precompute_num_warps is not None: - parts.append(f"pW={precompute_num_warps}") - if precompute_num_stages is not None: - parts.append(f"pS={precompute_num_stages}") - if heads_per_block is not None: - parts.append(f"H={heads_per_block}") - if maxnreg is not None: - parts.append(f"R={maxnreg}") - if num_ctas is not None: - parts.append(f"CT={num_ctas}") + parts.append(f"pW={_val(precompute_num_warps)}") + parts.append(f"pS={_val(precompute_num_stages)}") + parts.append(f"H={_val(heads_per_block)}") + parts.append(f"R={_val(maxnreg)}") + parts.append(f"CT={_val(num_ctas)}") # Persistent-only knobs (only meaningful when MODE=persistent_main; # printed unconditionally so output rows are uniformly comparable # across modes when the user passed these sweeps). _emit_split("CPSw", "CPSnw", cta_per_sm_w, cta_per_sm_nw) _emit_split("LSw", "LSnw", num_loop_stages_w, num_loop_stages_nw) - if flatten is not None: - parts.append(f"FL={flatten}") - if warp_specialize is not None: - parts.append(f"WS={warp_specialize}") + parts.append(f"FL={_val(flatten)}") + parts.append(f"WS={_val(warp_specialize)}") # TMA sweep tags. Four wrapper-level flags map to three # kernel-level constexprs (rect-load and replay-nowrite-load # share `USE_TMA_LOAD_NOWRITE`, picked by the wrapper based on # RECTANGLE). TMARL specifically gates the rectangle path's # state load; TMANL specifically gates the replay-style - # nowrite path's state load. Distinct because their measured - # perf profiles differ (see CHECKPOINTING_DESIGN.md item #17: - # rect TMA is "not a win" while replay-nowrite TMA is the - # biggest measured win at int8 b>=64). - if use_tma_rect_load is not None: - parts.append(f"TMARL={use_tma_rect_load}") # rect path load - if use_tma_replay_write_load is not None: - parts.append(f"TMAWL={use_tma_replay_write_load}") # replay-write load - if use_tma_replay_nowrite_load is not None: - parts.append(f"TMANL={use_tma_replay_nowrite_load}") # replay-NOWRITE load (NOT rect) - if use_tma_replay_write_store is not None: - parts.append(f"TMAWS={use_tma_replay_write_store}") # replay-write store + # nowrite path's state load. + parts.append(f"TMARL={_val(use_tma_rect_load)}") + parts.append(f"TMAWL={_val(use_tma_replay_write_load)}") + parts.append(f"TMANL={_val(use_tma_replay_nowrite_load)}") + parts.append(f"TMAWS={_val(use_tma_replay_write_store)}") parts.append(f"SR={1 if use_philox else 0}") - parts.append(f"RECT={1 if rectangle_for_nowrite else 0}") + parts.append(f"RECT={'auto' if rectangle_for_nowrite is None else (1 if rectangle_for_nowrite else 0)}") parts.append(f"WC={1 if write_checkpoint else 0}") - parts.append(f"MODE={mode}") + parts.append(f"MODE={_val(mode)}") parts.append(f"SORT={1 if sort_slots else 0}") parts.append(f"REVN={1 if reverse_nowrite else 0}") parts.append(f"HSORT={1 if hardcode_sort else 0}") @@ -4495,12 +4502,13 @@ def _parse_args() -> argparse.Namespace: parser.add_argument( "--rectangle-for-nowrite", type=str, - default="0", + default=None, help="Comma-separated 0/1 values: 0 = replay-style nowrite kernel, " "1 = dedicated rectangle nowrite kernel. Sweep both with '0,1' to " "compare in one invocation. Silently no-op for write cells (the " "write path always uses replay-style). Only applies to the " - "checkpointing variant.", + "checkpointing variant. When unset (default), the wrapper resolves " + "from the _DEFAULT_TUNING lookup per (batch, dtype, sr) cell.", ) parser.add_argument( "--use-tma-rect-load", @@ -4538,14 +4546,16 @@ def _parse_args() -> argparse.Namespace: parser.add_argument( "--modes", type=str, - default="persistent_dynamic", + default=None, help="Comma-separated dispatch modes to sweep, any of " "{persistent_dynamic, persistent_main}. " "persistent_dynamic = single persistent-CTA kernel that dispatches " "per-slot at runtime based on PNAT. " "persistent_main = persistent-CTA kernel with two halves (write + " "nowrite), requires caller-provided _n_writes and a write-first " - "sorted slot_perm. Both modes ignore --write-modes (per-slot from PNAT).", + "sorted slot_perm. Both modes ignore --write-modes (per-slot from PNAT). " + "When unset (default), the wrapper resolves from the _DEFAULT_TUNING " + "lookup per (batch, dtype, sr) cell.", ) parser.add_argument( "--mix-csv", @@ -4713,12 +4723,20 @@ def _round_iters_to_group(name, val): parser.error(f"--sr-modes value must be RN or SR, got {m!r}") args.sr_modes_list = sr_modes - rect_modes = [v.strip() for v in args.rectangle_for_nowrite.split(",") if v.strip()] - rect_list = [] - for v in rect_modes: - if v not in ("0", "1"): - parser.error(f"--rectangle-for-nowrite value must be 0 or 1, got {v!r}") - rect_list.append(v == "1") + # rectangle_for_nowrite=None means "let the wrapper resolve from + # _DEFAULT_TUNING". Empty/unset argparse default produces [None] in the + # sweep list; the kernel call passes None and the wrapper picks per-cell. + if args.rectangle_for_nowrite is None: + rect_list = [None] + else: + rect_modes = [v.strip() for v in args.rectangle_for_nowrite.split(",") if v.strip()] + rect_list = [] + for v in rect_modes: + if v not in ("0", "1"): + parser.error(f"--rectangle-for-nowrite value must be 0 or 1, got {v!r}") + rect_list.append(v == "1") + if not rect_list: + rect_list = [None] args.rectangle_for_nowrite_list = rect_list sort_modes = [v.strip() for v in (args.sort_slots or "0").split(",") if v.strip()] @@ -4756,16 +4774,21 @@ def _round_iters_to_group(name, val): else: args.write_modes_list = [args.write_checkpoint] - modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] - valid_modes = { - "persistent_main", "persistent_dynamic", - } - for m in modes_raw: - if m not in valid_modes: - parser.error( - f"--modes value must be one of {sorted(valid_modes)}, got {m!r}" - ) - args.modes_list = modes_raw or ["persistent_dynamic"] + # mode=None means "let the wrapper resolve from _DEFAULT_TUNING". Same + # convention as --rectangle-for-nowrite. + if args.modes is None: + args.modes_list = [None] + else: + modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] + valid_modes = { + "persistent_main", "persistent_dynamic", + } + for m in modes_raw: + if m not in valid_modes: + parser.error( + f"--modes value must be one of {sorted(valid_modes)}, got {m!r}" + ) + args.modes_list = modes_raw if modes_raw else [None] return args From d7c1e391a4e73ff36af0bf3d0fd4bac723ee6517 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 19 May 2026 15:19:55 -0700 Subject: [PATCH 56/89] mamba replay benchmark: harden driver cell handling Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/benchmark_replay_selective_state_update.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index f1b9de9c244a..9e1f2ea236ac 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -4008,6 +4008,8 @@ def _phase(label: str) -> None: batch_sizes = [int(x) for x in args.batch_sizes.split(",")] mtp_lengths = [int(x) for x in args.mtp_lengths.split(",")] + if args.mix_csv is not None and len(mtp_lengths) != 1: + sys.exit("--mix-csv requires exactly one --mtp-lengths value") dtype_map = { "bf16": torch.bfloat16, @@ -4019,6 +4021,8 @@ def _phase(label: str) -> None: } state_dtypes = [dtype_map[s] for s in args.state_dtypes.split(",")] act_dtypes = [dtype_map[s] for s in args.act_dtypes.split(",")] + if args.json_output and len(act_dtypes) != 1: + sys.exit("--json-output requires exactly one --act-dtypes value") # Resolve baseline function. if args.baseline == "flashinfer_pr3324": @@ -4157,8 +4161,12 @@ def _phase(label: str) -> None: for mode in modes_list: for rect in rect_list: can_sort = mix_samples_cpu is not None + cell_list_active = bool( + getattr(args, "_cell_list_keys", ()) + ) effective_hsort_list = ( - hsort_list if can_sort else [False] + hsort_list if (can_sort or cell_list_active) + else [False] ) for hardcode_sort in effective_hsort_list: _bench_config( From 4a5a39874c2b53b9158c552eab649e31dc884ae4 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 19 May 2026 16:12:47 -0700 Subject: [PATCH 57/89] mamba replay: fix rectangle precompute write-read race Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 207 +++++++----------- 1 file changed, 82 insertions(+), 125 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 26e72c1587c5..5ce3e1c6fbf8 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -540,75 +540,66 @@ def _rectangle_precompute_impl( is_new_k = (k_new_idx >= 0) & (k_new_idx < T) safe_k_new = tl.where(is_new_k, k_new_idx, 0) - # Loop 1: per-head dt processing. dt → dt_processed → dA_cumsum → - # decay_vec_new (= exp(cumAdt_new)). Stored to write_buf for next step. - # decay_vec_full (= total_decay * decay_vec_new) is finalized in loop 2 - # once total_decay is loaded; loop 1 stores raw decay_vec_new to scratch. - for h_local in range(HEADS_PER_BLOCK): - head_idx = first_head + h_local - - dt_base = dt_ptr + pid_b * stride_dt_batch + head_idx * stride_dt_head - dt = tl.load(dt_base + offs_t * stride_dt_T, mask=t_mask, other=0.0).to(tl.float32) - if HAS_DT_BIAS: - dt_bias = tl.load(dt_bias_ptr + head_idx * stride_dt_bias_head).to(tl.float32) - dt = dt + dt_bias - if DT_SOFTPLUS: - dt = softplus(dt) - - A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) - dA_cumsum = tl.cumsum(A * dt, axis=0) - - # Cross-step continuity for old_dA_cumsum: rectangle precompute runs - # only on the nowrite path (write_buf == buf_active, write_offset == PNAT). - # Add the running tail from buf_active[head_idx, PNAT-1] so the buffer - # holds one continuous cumsum across back-to-back nowrites. PNAT is - # scalar/uniform, use scalar if to short-circuit the load at PNAT=0. - if prev_num_accepted_tokens == 0: - prev_total = 0.0 - else: - last_cumsum_ptr = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_active * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T - ) - prev_total = tl.load(last_cumsum_ptr).to(tl.float32) - - # Store dt and dA_cumsum to write_buf at [write_offset, write_offset+T) - # for next step's replay/rectangle use. - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - tl.store( - old_dt_base + (write_offset + offs_t) * stride_old_dt_T, - dt, - mask=t_mask, - ) + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h + + # Precompute this step's (H, T) dt and continuous dA_cumsum tiles. + # Keep them in registers for the rectangle path below; reloading from + # global memory after storing would create a same-kernel write/read race. + dt_addrs = ( + dt_ptr + + pid_b * stride_dt_batch + + heads_block[:, None] * stride_dt_head + + offs_t[None, :] * stride_dt_T + ) + dt_new = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) + + if HAS_DT_BIAS: + dt_bias_heads = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) + dt_new = dt_new + dt_bias_heads[:, None] + if DT_SOFTPLUS: + dt_new = softplus(dt_new) + + A_heads = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) + dA_cumsum_step = tl.cumsum(A_heads[:, None] * dt_new, axis=1) - old_dA_cumsum_base = ( + if prev_num_accepted_tokens == 0: + dA_cumsum_prefix = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) + else: + dA_cumsum_prefix_ptrs = ( old_dA_cumsum_ptr + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - tl.store( - old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, - dA_cumsum + prev_total, - mask=t_mask, + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T ) + dA_cumsum_prefix = tl.load(dA_cumsum_prefix_ptrs).to(tl.float32) - # ---- Hoisted: cache-only loads independent of conv1d ---- - # old_B (group-level, BLOCK_K × BLOCK_DSTATE = ~8KB tile) and the - # decay_vec_full per-head pre-compute (which writes to DRAM and doesn't - # need cross-gdc_wait variables) are issued BEFORE gdc_wait so their - # HBM latency overlaps with conv1d. Per-head factor_dt/exp_diff stay - # below gdc_wait — they need cross-iteration spans, which Triton can't - # express without a DRAM round-trip; the per-head LOADS in the post- - # wait loop are small and cheap, so leave them. + old_dt_write_addrs = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block[:, None] * stride_old_dt_head + + (write_offset + offs_t)[None, :] * stride_old_dt_T + ) + tl.store(old_dt_write_addrs, dt_new, mask=t_mask[None, :]) + + old_dA_cumsum_write_addrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block[:, None] * stride_old_dA_cumsum_head + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T + ) + tl.store( + old_dA_cumsum_write_addrs, + dA_cumsum_step + dA_cumsum_prefix[:, None], + mask=t_mask[None, :], + ) + + # ---- Work independent of conv1d ---- + # Load historical cache and build combo_block before gdc_wait so this + # work can overlap the upstream conv1d latency. group_idx = first_head // nheads_ngroups_ratio # Group-level: old B from active buffer at [0, PNAT) of the K-axis. @@ -626,14 +617,10 @@ def _rectangle_precompute_impl( other=0.0, ) - # Pre-wait: vectorized across HEADS_PER_BLOCK heads. Compute decay_vec_full - # (H, T) and combo = factor_dt * exp_diff (H, T, K). Store decay_vec_full; - # combo_block stays in registers across gdc_wait — used directly post-wait - # to compute rect_CB_scaled without a global memory roundtrip. - offs_h = tl.arange(0, HEADS_PER_BLOCK) - heads_block = first_head + offs_h # (H,) + # combo_block stays in registers across gdc_wait and is used directly + # after the wait to compute rect_CB_scaled. - # Per-head bases (H,) — broadcast with offs_k or offs_t for 2D loads. + # Per-head read bases (H,) - broadcast with offs_k for 2D loads. old_dt_read_h = ( old_dt_ptr + cache_batch_idx * stride_old_dt_cache @@ -646,18 +633,6 @@ def _rectangle_precompute_impl( + buf_active * stride_old_dA_cumsum_dbuf + heads_block * stride_old_dA_cumsum_head ) - old_dt_write_h = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + heads_block * stride_old_dt_head - ) - old_dA_cumsum_write_h = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + heads_block * stride_old_dA_cumsum_head - ) # (H, K) loads at [0, PNAT) — old data from previous step. hk_mask = is_old_k[None, :] # (1, K) @@ -669,34 +644,29 @@ def _rectangle_precompute_impl( old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, mask=hk_mask, other=0.0, ).to(tl.float32) - # (H, T) loads at [PNAT, PNAT+T) — this step's dA_cumsum_new from loop 1. - # With the cross-step continuity fix in loop 1, the values stored at - # [PNAT, PNAT+T) are the continuous cumsum (prefix + per-step new - # cumsum) — i.e., continuous_cumsum[PNAT..PNAT+T-1] in global indexing. + # Use loop-1 registers for this step's newly appended tokens. These are + # exactly the values stored above at [PNAT, PNAT+T); reloading them here + # would read bytes this same kernel just wrote, which Triton does not + # guarantee to fence. ht_mask = t_mask[None, :] # (1, T) - dA_cumsum_new = tl.load( - old_dA_cumsum_write_h[:, None] - + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T, - mask=ht_mask, other=0.0, - ).to(tl.float32) - # (H, K) loads at PNAT-shifted positions for new tokens. - hkn_mask = is_new_k[None, :] - dt_at_kn = tl.load( - old_dt_write_h[:, None] - + (write_offset + safe_k_new)[None, :] * stride_old_dt_T, - mask=hkn_mask, other=0.0, - ).to(tl.float32) - dA_cumsum_at_kn = tl.load( - old_dA_cumsum_write_h[:, None] - + (write_offset + safe_k_new)[None, :] * stride_old_dA_cumsum_T, - mask=hkn_mask, other=0.0, - ).to(tl.float32) + dA_cumsum_new = dA_cumsum_step + dA_cumsum_prefix[:, None] # (H, T) - # decay_vec_full[t] = exp(continuous_cumsum[PNAT+t]) — directly the - # continuous value now stored at buffer position write_offset+t. Was - # decomposed as total_decay * exp(per_step_new[t]) when the buffer held - # per-step (non-continuous) cumsum; with the continuity fix the value - # IS continuous_cumsum[PNAT+t] so no decomposition is needed. + new_token_gather_idx = tl.broadcast_to( + safe_k_new[None, :], (HEADS_PER_BLOCK, BLOCK_SIZE_K) + ) + dt_at_kn = tl.where( + is_new_k[None, :], + tl.gather(dt_new, new_token_gather_idx, axis=1), + 0.0, + ) # (H, K) + dA_cumsum_at_kn = tl.where( + is_new_k[None, :], + tl.gather(dA_cumsum_new, new_token_gather_idx, axis=1), + 0.0, + ) # (H, K) + + # The write buffer stores continuous cumsum, so decay_vec_full[t] is + # exp(continuous_cumsum[PNAT+t]) directly. decay_vec_full_block = tl.exp(dA_cumsum_new) # (H, T) decay_vec_addrs = ( decay_vec_ptr @@ -706,22 +676,9 @@ def _rectangle_precompute_impl( ) # (H, T) tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) - # combo_block = factor_dt * exp_diff — (H, T, K). Stays in registers - # across gdc_wait. With continuous cumsum in the buffer, s_k for any k - # (old or new) is simply -continuous_cumsum[k]; exp_diff[t, k] then - # equals exp(continuous_cumsum[PNAT+t] - continuous_cumsum[k]) — the - # decay weight for token k's contribution to the output at position - # PNAT+t. No need to subtract any "total" — the dA_cumsum_new[t] term - # already carries the full prefix. - # - # Numerical note: pre-fix this kernel computed `total - old_dA[k]` - # (small-minus-small) then summed `+ dA_cumsum_new[t]` (also small, - # per-step). Post-fix `s_k = -old_dA[k]` is large-magnitude positive - # and `dA_cumsum_new[t]` is large-magnitude negative; their sum - # cancels back to the same small value. Cancellation error is bounded - # by ulp(max_magnitude) ≈ 2^-23 · |continuous_cumsum| — negligible - # for max_window ≤ ~1024. Still one exp on the sum (not two muls of - # exps), so no overflow regression vs the original formulation. + # combo_block[t, k] = dt[k] * exp(cumsum[t] - cumsum[k]). + # Old and new K positions share the same formula because both halves use + # continuous cumsum values. factor_dt = tl.where(is_old_k[None, :], old_dt_all, dt_at_kn) # (H, K) s_k = tl.where( is_old_k[None, :], From da8f6e4784e2032c2101612f00ae73327b3210fc Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 19 May 2026 17:09:16 -0700 Subject: [PATCH 58/89] mamba replay: handle rectangle fallback and pad slots Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 138 +++++++++--------- 1 file changed, 72 insertions(+), 66 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 5ce3e1c6fbf8..2bd4465a0fa3 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -651,19 +651,36 @@ def _rectangle_precompute_impl( ht_mask = t_mask[None, :] # (1, T) dA_cumsum_new = dA_cumsum_step + dA_cumsum_prefix[:, None] # (H, T) - new_token_gather_idx = tl.broadcast_to( - safe_k_new[None, :], (HEADS_PER_BLOCK, BLOCK_SIZE_K) - ) - dt_at_kn = tl.where( - is_new_k[None, :], - tl.gather(dt_new, new_token_gather_idx, axis=1), - 0.0, - ) # (H, K) - dA_cumsum_at_kn = tl.where( - is_new_k[None, :], - tl.gather(dA_cumsum_new, new_token_gather_idx, axis=1), - 0.0, - ) # (H, K) + # Production uses matching padded T/K sizes, where tl.gather is the cheap + # path. Some tests use a larger padded K than T; Triton rejects that gather + # axis mismatch, so use a slower one-hot sum fallback for those cases. + if BLOCK_SIZE_T == BLOCK_SIZE_K: + new_token_gather_idx = tl.broadcast_to( + safe_k_new[None, :], (HEADS_PER_BLOCK, BLOCK_SIZE_K) + ) + dt_at_kn = tl.where( + is_new_k[None, :], + tl.gather(dt_new, new_token_gather_idx, axis=1), + 0.0, + ) # (H, K) + dA_cumsum_at_kn = tl.where( + is_new_k[None, :], + tl.gather(dA_cumsum_new, new_token_gather_idx, axis=1), + 0.0, + ) # (H, K) + else: + new_token_selector = ( + (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) + & is_new_k[None, :] + ) + dt_at_kn = tl.sum( + tl.where(new_token_selector[None, :, :], dt_new[:, :, None], 0.0), + axis=1, + ) # (H, K) + dA_cumsum_at_kn = tl.sum( + tl.where(new_token_selector[None, :, :], dA_cumsum_new[:, :, None], 0.0), + axis=1, + ) # (H, K) # The write buffer stores continuous cumsum, so decay_vec_full[t] is # exp(continuous_cumsum[PNAT+t]) directly. @@ -1941,12 +1958,12 @@ def _persistent_main_kernel( ).to(tl.int32) is_pad = cache_batch_idx == pad_slot_id - # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle - # impl. `replay_work_items` carries the cache slot, PNAT and active - # buffer for persistent_main; persistent_dynamic resolves those once - # here from the existing tensors. - if RECTANGLE: - if not is_pad: + if not is_pad: + # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle + # impl. `replay_work_items` carries the cache slot, PNAT and active + # buffer for persistent_main; persistent_dynamic resolves those once + # here from the existing tensors. + if RECTANGLE: if IS_DYNAMIC: is_w = (pnat + T) > MAX_REPLAY_BUFFER_LENGTH else: @@ -1992,17 +2009,10 @@ def _persistent_main_kernel( BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) - True, # WRITE_CHECKPOINT_IS_CONSTEXPR — force inner to use WRITE_CHECKPOINT constexpr - # 3 TMA flags: write-load fires here (we're in the - # is_write branch), nowrite-load is dead (no slot - # reaches it), store fires (write path). + True, # WRITE_CHECKPOINT_IS_CONSTEXPR USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, ) else: - # Rectangle nowrite: pass state_ptr (raw, always) + - # state_tma_descriptor (the single unified descriptor — - # same memory replay paths use). Rect impl gates use - # of the descriptor via its USE_TMA_LOAD constexpr. _persistent_rectangle_impl( pid_m, pid_b, pid_h, cache_batch_idx, active_buf, pnat, @@ -2026,45 +2036,38 @@ def _persistent_main_kernel( LAUNCH_WITH_PDL, QUANT_MAX, USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle ) - # else: pad slot — skip both impls (both would early-return anyway) - else: - # No rectangle path — single _persistent_main_impl call covers - # both write and nowrite slots via WRITE_CHECKPOINT constexpr (non-dynamic) or - # runtime is_write (IS_DYNAMIC=True). Pass all 3 TMA flags; - # impl picks USE_TMA_LOAD_WRITE vs USE_TMA_LOAD_NOWRITE based on - # its computed is_write — constexpr-folds when is_write is - # constexpr (non-dyn), runtime branch when IS_DYNAMIC=True. - _persistent_main_impl( - pid_m, pid_b, pid_h, - cache_batch_idx, active_buf, pnat, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - rand_seed_ptr, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, - stride_old_B_group, stride_old_B_dstate, - stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, - stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - WRITE_CHECKPOINT, IS_DYNAMIC, - False, # WRITE_CHECKPOINT_IS_CONSTEXPR=False — RECT=0 has both write/nowrite slots in one call - USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, - ) + else: + _persistent_main_impl( + pid_m, pid_b, pid_h, + cache_batch_idx, active_buf, pnat, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + rand_seed_ptr, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, + WRITE_CHECKPOINT, IS_DYNAMIC, + False, # WRITE_CHECKPOINT_IS_CONSTEXPR + USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + ) # ============================================================================ @@ -2622,6 +2625,9 @@ def replay_selective_state_update( BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) # Rectangle K-axis bound = window (max_window). Computed unconditionally # so the launch sites can refer to it; only used on the rectangle path. + # If this differs from BLOCK_SIZE_T, rectangle precompute uses a slower + # one-hot fallback because tl.gather requires matching padded axis sizes. + # Production uses matching padded T/K; mismatches are for tests/debugging. BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) # Allocate precomputed intermediates (per-call, not cached). Always From 5acc881586d29985ccfb8e34b748018e95db8f73 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 01:14:15 -0700 Subject: [PATCH 59/89] mamba replay: benchmark cache slot load knob Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../modules/mamba/replay_selective_state_update.py | 13 ++++++++++--- .../benchmark_replay_selective_state_update.py | 12 ++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 2bd4465a0fa3..643a9f8a4f27 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -1892,6 +1892,7 @@ def _persistent_main_kernel( USE_TMA_LOAD_WRITE: tl.constexpr = False, USE_TMA_LOAD_NOWRITE: tl.constexpr = False, USE_TMA_STORE: tl.constexpr = False, + USE_REPLAY_CACHE_SLOT: tl.constexpr = True, ): # PDL signal: fire once at kernel entry (not per work unit). if LAUNCH_DEPENDENT_KERNELS: @@ -1949,9 +1950,12 @@ def _persistent_main_kernel( pid_b = tl.load( work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH ) - cache_batch_idx = tl.load( - work_item_base + _REPLAY_WORK_CACHE_SLOT - ).to(tl.int64) + if USE_REPLAY_CACHE_SLOT: + cache_batch_idx = tl.load( + work_item_base + _REPLAY_WORK_CACHE_SLOT + ).to(tl.int64) + else: + cache_batch_idx = work_item_idx.to(tl.int64) pnat = tl.load(work_item_base + _REPLAY_WORK_PNAT) active_buf = tl.load( work_item_base + _REPLAY_WORK_CACHE_BUF_IDX @@ -2304,6 +2308,7 @@ def replay_selective_state_update( _use_tma_replay_write_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=True _use_tma_replay_write_store: bool | None = None, # SSM state store when WRITE_CHECKPOINT=True _use_tma_replay_nowrite_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=False + _use_replay_cache_slot: bool = True, # Persistent-mode tuning kwargs (consulted for both pd and pm; pd uses # _cta_per_sm / _num_loop_stages, pm uses the _write/_nowrite splits): # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally @@ -3013,6 +3018,7 @@ def launch_persistent_main(write_checkpoint: bool, and not write_checkpoint ), USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), + USE_REPLAY_CACHE_SLOT=bool(_use_replay_cache_slot), num_warps=_nw, **({"num_stages": _ns} if _ns else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), @@ -3087,6 +3093,7 @@ def launch_persistent_dynamic_main(n_writes_tensor: torch.Tensor, _use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load ), USE_TMA_STORE=bool(_use_tma_replay_write_store), + USE_REPLAY_CACHE_SLOT=bool(_use_replay_cache_slot), num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 9e1f2ea236ac..583ea28e20df 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -3240,6 +3240,7 @@ def _run_incr( extra_kwargs["_flatten"] = bool(flatten) if warp_specialize is not None: extra_kwargs["_warp_specialize"] = bool(warp_specialize) + extra_kwargs["_use_replay_cache_slot"] = bool(args.use_cache_slot) replay_selective_state_update( state_work, @@ -3340,6 +3341,8 @@ def _emit_split(name_w, name_nw, val_w, val_nw): parts.append(f"RECT={'auto' if rectangle_for_nowrite is None else (1 if rectangle_for_nowrite else 0)}") parts.append(f"MODE={_val(mode)}") parts.append(f"HSORT={1 if hardcode_sort else 0}") + if not args.use_cache_slot: + parts.append("CSLOT=0") sweep_suffix = (" " + ",".join(parts)) if parts else "" sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") @@ -3404,6 +3407,7 @@ def _emit_split(name_w, name_nw, val_w, val_nw): bool(use_philox), bool(rectangle_for_nowrite), bool(hardcode_sort), + bool(args.use_cache_slot), scenario_pre_iter is not None, expected_K, ) @@ -4555,6 +4559,14 @@ def _parse_args() -> argparse.Namespace: help="External PDL: conv1d launches dependents, precompute waits. " "Only relevant with --with-conv1d. --no-external-pdl disables.", ) + parser.add_argument( + "--use-cache-slot", + action=argparse.BooleanOptionalAction, + default=True, + help="Use the cache-slot field from replay_work_items in persistent_main. " + "--no-use-cache-slot keeps the old identity-cache-slot shortcut for " + "diagnostic comparisons only.", + ) parser.add_argument( "--heads-per-block", type=str, From 18a353e33e85e8f7bfc06f101753c33b5acc1a96 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 01:27:28 -0700 Subject: [PATCH 60/89] mamba replay: specialize rectangle gather fallback Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../modules/mamba/replay_selective_state_update.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 643a9f8a4f27..b682a5c016c8 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -506,6 +506,7 @@ def _rectangle_precompute_impl( BLOCK_SIZE_K: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, HEADS_PER_BLOCK: tl.constexpr, + USE_GATHER_FOR_NEW_TOKENS: tl.constexpr, ): pid_b = tl.program_id(axis=0) pid_hg = tl.program_id(axis=1) @@ -653,8 +654,10 @@ def _rectangle_precompute_impl( # Production uses matching padded T/K sizes, where tl.gather is the cheap # path. Some tests use a larger padded K than T; Triton rejects that gather - # axis mismatch, so use a slower one-hot sum fallback for those cases. - if BLOCK_SIZE_T == BLOCK_SIZE_K: + # axis mismatch, so use a slower one-hot sum fallback for those cases. The + # wrapper passes this as an explicit constexpr so fast-path compilations do + # not carry the fallback branch. + if USE_GATHER_FOR_NEW_TOKENS: new_token_gather_idx = tl.broadcast_to( safe_k_new[None, :], (HEADS_PER_BLOCK, BLOCK_SIZE_K) ) @@ -867,6 +870,7 @@ def _dynamic_precompute_kernel( HEADS_PER_BLOCK: tl.constexpr, # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. RECTANGLE: tl.constexpr, + RECTANGLE_USE_GATHER: tl.constexpr, ): # Hoisted PDL signal: fire as the first thing every program does. if LAUNCH_DEPENDENT_KERNELS: @@ -1008,6 +1012,7 @@ def _dynamic_precompute_kernel( BLOCK_SIZE_K, LAUNCH_WITH_PDL, HEADS_PER_BLOCK, + RECTANGLE_USE_GATHER, ) @@ -2634,6 +2639,7 @@ def replay_selective_state_update( # one-hot fallback because tl.gather requires matching padded axis sizes. # Production uses matching padded T/K; mismatches are for tests/debugging. BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) + rectangle_use_gather = BLOCK_SIZE_T == BLOCK_SIZE_K # Allocate precomputed intermediates (per-call, not cached). Always # allocate (T, K) — the largest layout that any path uses. Replay-style @@ -2906,6 +2912,7 @@ def launch_dynamic_precompute(rectangle: bool): LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, HEADS_PER_BLOCK=heads_per_block, RECTANGLE=rectangle, + RECTANGLE_USE_GATHER=rectangle_use_gather, num_warps=precompute_num_warps, **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), launch_pdl=launch_with_pdl, From b319eac9ed223e17c06ee163f04c9aaa58923ad1 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 01:45:06 -0700 Subject: [PATCH 61/89] mamba replay: remove pad-slot guard and hsort default Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 207 +++++++++--------- ...benchmark_replay_selective_state_update.py | 69 ++++-- 2 files changed, 145 insertions(+), 131 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index b682a5c016c8..af61e8442508 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -23,7 +23,6 @@ import triton import triton.language as tl -from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID from tensorrt_llm._utils import get_sm_version from .mamba2_metadata import (REPLAY_WORK_CACHE_BUF_IDX, @@ -181,7 +180,6 @@ def _replay_precompute_impl( # on no-replay-write steps). prev_num_accepted_tokens_ptr, state_batch_indices_ptr, - pad_slot_id, # Dimensions T: tl.constexpr, dstate: tl.constexpr, @@ -251,8 +249,6 @@ def _replay_precompute_impl( # Resolve cache index for writes if HAS_CACHE_BATCH_INDICES: cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return else: cache_batch_idx = pid_b.to(tl.int64) @@ -450,7 +446,6 @@ def _rectangle_precompute_impl( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, - pad_slot_id, # Dimensions T: tl.constexpr, MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # rectangle K-axis bound @@ -514,8 +509,6 @@ def _rectangle_precompute_impl( if HAS_CACHE_BATCH_INDICES: cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return else: cache_batch_idx = pid_b.to(tl.int64) @@ -811,7 +804,6 @@ def _dynamic_precompute_kernel( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, - pad_slot_id, # Dimensions T: tl.constexpr, MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, @@ -879,8 +871,6 @@ def _dynamic_precompute_kernel( pid_b = tl.program_id(axis=0) if HAS_CACHE_BATCH_INDICES: cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return else: cache_batch_idx = pid_b.to(tl.int64) @@ -905,7 +895,6 @@ def _dynamic_precompute_kernel( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, - pad_slot_id, T, dstate, nheads_ngroups_ratio, @@ -966,7 +955,6 @@ def _dynamic_precompute_kernel( cache_buf_idx_ptr, prev_num_accepted_tokens_ptr, state_batch_indices_ptr, - pad_slot_id, T, MAX_REPLAY_BUFFER_LENGTH, dstate, @@ -1774,7 +1762,6 @@ def _persistent_main_kernel( state_batch_indices_ptr, replay_work_items_ptr, rand_seed_ptr, - pad_slot_id, # Persistent-loop work-distribution scalars. Caller pre-sorts the batch # write-first; the kernel uses (n_writes, batch_total, WRITE_CHECKPOINT) # to derive its own slot range. Write half processes [0, n_writes), @@ -1941,111 +1928,58 @@ def _persistent_main_kernel( pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) work_item_idx = pid_b_local + slot_lo if IS_DYNAMIC: - pid_b = work_item_idx if HAS_CACHE_BATCH_INDICES: + pid_b = work_item_idx cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - is_pad = cache_batch_idx == pad_slot_id else: + pid_b = work_item_idx cache_batch_idx = pid_b.to(tl.int64) - is_pad = False active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) pnat = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) else: work_item_base = replay_work_items_ptr + work_item_idx * _REPLAY_WORK_ITEM_WIDTH - pid_b = tl.load( - work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH - ) if USE_REPLAY_CACHE_SLOT: + pid_b = tl.load( + work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH + ) cache_batch_idx = tl.load( work_item_base + _REPLAY_WORK_CACHE_SLOT ).to(tl.int64) + pnat = tl.load(work_item_base + _REPLAY_WORK_PNAT) + active_buf = tl.load( + work_item_base + _REPLAY_WORK_CACHE_BUF_IDX + ).to(tl.int32) else: + pid_b = tl.load( + work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH + ) cache_batch_idx = work_item_idx.to(tl.int64) - pnat = tl.load(work_item_base + _REPLAY_WORK_PNAT) - active_buf = tl.load( - work_item_base + _REPLAY_WORK_CACHE_BUF_IDX - ).to(tl.int32) - is_pad = cache_batch_idx == pad_slot_id - - if not is_pad: - # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle - # impl. `replay_work_items` carries the cache slot, PNAT and active - # buffer for persistent_main; persistent_dynamic resolves those once - # here from the existing tensors. - if RECTANGLE: - if IS_DYNAMIC: - is_w = (pnat + T) > MAX_REPLAY_BUFFER_LENGTH - else: - is_w = WRITE_CHECKPOINT - if is_w: - # Pass WRITE_CHECKPOINT=True constexpr to specialize this - # impl call for the write path. Under IS_DYNAMIC=True, the - # kernel-level WRITE_CHECKPOINT is False (launcher default), - # but the OUTER is_w branch we are inside narrows the - # runtime path to writes-only, so we override to True here - # so the impl's constexpr-gated `if is_write:` blocks DCE - # to the write-only codegen. Under IS_DYNAMIC=False - # (persistent_main), the kernel-level WRITE_CHECKPOINT is - # itself True for this half (write half launches with - # WRITE_CHECKPOINT=True), and the outer is_w = WRITE_CHECKPOINT = True - # constexpr-folds; passing literal True here is consistent - # and constexpr-equivalent. - _persistent_main_impl( - pid_m, pid_b, pid_h, - cache_batch_idx, active_buf, pnat, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - rand_seed_ptr, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, - stride_old_B_group, stride_old_B_dstate, - stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, - stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) - True, # WRITE_CHECKPOINT_IS_CONSTEXPR - USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, - ) - else: - _persistent_rectangle_impl( - pid_m, pid_b, pid_h, - cache_batch_idx, active_buf, pnat, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, - LAUNCH_WITH_PDL, QUANT_MAX, - USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle - ) + pnat = tl.load(work_item_base + _REPLAY_WORK_PNAT) + active_buf = tl.load( + work_item_base + _REPLAY_WORK_CACHE_BUF_IDX + ).to(tl.int32) + # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle + # impl. `replay_work_items` carries the cache slot, PNAT and active + # buffer for persistent_main; persistent_dynamic resolves those once + # here from the existing tensors. + if RECTANGLE: + if IS_DYNAMIC: + is_w = (pnat + T) > MAX_REPLAY_BUFFER_LENGTH else: + is_w = WRITE_CHECKPOINT + if is_w: + # Pass WRITE_CHECKPOINT=True constexpr to specialize this + # impl call for the write path. Under IS_DYNAMIC=True, the + # kernel-level WRITE_CHECKPOINT is False (launcher default), + # but the OUTER is_w branch we are inside narrows the + # runtime path to writes-only, so we override to True here + # so the impl's constexpr-gated `if is_write:` blocks DCE + # to the write-only codegen. Under IS_DYNAMIC=False + # (persistent_main), the kernel-level WRITE_CHECKPOINT is + # itself True for this half (write half launches with + # WRITE_CHECKPOINT=True), and the outer is_w = WRITE_CHECKPOINT = True + # constexpr-folds; passing literal True here is consistent + # and constexpr-equivalent. _persistent_main_impl( pid_m, pid_b, pid_h, cache_batch_idx, active_buf, pnat, @@ -2073,10 +2007,66 @@ def _persistent_main_kernel( BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - WRITE_CHECKPOINT, IS_DYNAMIC, - False, # WRITE_CHECKPOINT_IS_CONSTEXPR + True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) + True, # WRITE_CHECKPOINT_IS_CONSTEXPR USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, ) + else: + _persistent_rectangle_impl( + pid_m, pid_b, pid_h, + cache_batch_idx, active_buf, pnat, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, + LAUNCH_WITH_PDL, QUANT_MAX, + USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle + ) + else: + _persistent_main_impl( + pid_m, pid_b, pid_h, + cache_batch_idx, active_buf, pnat, + state_ptr, state_tma_descriptor, state_scales_ptr, + old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, + x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, + cb_scaled_ptr, decay_vec_ptr, + rand_seed_ptr, + T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, + stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, + stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, + stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, + stride_old_B_group, stride_old_B_dstate, + stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, + stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, + stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, + stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, + stride_D_head, stride_D_dim, + stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, + stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, + stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, + stride_dv_batch, stride_dv_head, stride_dv_t, + BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, + WRITE_CHECKPOINT, IS_DYNAMIC, + False, # WRITE_CHECKPOINT_IS_CONSTEXPR + USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + ) # ============================================================================ @@ -2272,7 +2262,6 @@ def replay_selective_state_update( dt_bias: torch.Tensor | None = None, dt_softplus: bool = False, state_batch_indices: torch.Tensor | None = None, - pad_slot_id: int = PAD_SLOT_ID, rand_seed: torch.Tensor | None = None, philox_rounds: int = 10, state_scales: torch.Tensor | None = None, @@ -2890,7 +2879,7 @@ def launch_dynamic_precompute(rectangle: bool): cb_scaled, decay_vec, old_B, old_dt, old_dA_cumsum, cache_buf_idx, prev_num_accepted_tokens, - state_batch_indices, pad_slot_id, + state_batch_indices, T, max_window, dstate, nheads // ngroups, dt.stride(0), dt.stride(1), dt.stride(2), dt_bias.stride(0) if dt_bias is not None else 0, @@ -2982,7 +2971,7 @@ def launch_persistent_main(write_checkpoint: bool, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, cb_scaled, decay_vec, - state_batch_indices, replay_work_items_arg, rand_seed, pad_slot_id, + state_batch_indices, replay_work_items_arg, rand_seed, n_writes, batch, nheads, T, max_window, dim, dstate, nheads // ngroups, state.stride(0), state.stride(1), state.stride(2), state.stride(3), @@ -3058,7 +3047,7 @@ def launch_persistent_dynamic_main(n_writes_tensor: torch.Tensor, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, cb_scaled, decay_vec, - state_batch_indices, replay_work_items_arg, rand_seed, pad_slot_id, + state_batch_indices, replay_work_items_arg, rand_seed, n_writes_tensor, batch, nheads, T, max_window, dim, dstate, nheads // ngroups, state.stride(0), state.stride(1), state.stride(2), state.stride(3), diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 583ea28e20df..4bc21ab36065 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -219,7 +219,7 @@ def _load(mod_name: str, file_name: str): spec.loader.exec_module(mod) return mod - # 1. Package __init__ (defines PAD_SLOT_ID = -1) + # 1. Package __init__ (defines replay work-item constants) _load("", "__init__.py") # 2. softplus helper (used by both kernel modules) _load("softplus", "softplus.py") @@ -2803,10 +2803,14 @@ def _split_or_share(split_csv, shared_values): # kernels read the metadata cold. if mix_samples_cpu is not None: device = state_work.device - # Hardcode-sort: per-iter prev_tokens are CPU-sorted write-first. - # Output is scrambled (we don't permute x/B/C/dt to match) but - # timing is meaningful as a clustering experiment. - src = mix_samples_sorted_cpu if (hardcode_sort and mix_samples_sorted_cpu is not None) else mix_samples_cpu + # Hardcode-sort is a precluster diagnostic: sort PNAT samples before + # building replay_work_items. Production keeps decode rows unsorted + # and only sorts the secondary replay metadata. + src = ( + mix_samples_sorted_cpu + if (hardcode_sort and mix_samples_sorted_cpu is not None) + else mix_samples_cpu + ) samples_gpu = torch.from_numpy(src).to(device=device, dtype=torch.int32) n_writes_per_iter_all, replay_work_items_samples_cpu = ( @@ -2926,8 +2930,13 @@ def _graph_pre_iter(j): if baseline_fn is not None and is_pr3324_baseline: baseline_suffix_parts = [f"SR={int(use_philox)}"] - hsort_is_swept = len(getattr(args, "hardcode_sort_list", [False])) > 1 - if scn["fill"] is None and (hardcode_sort or hsort_is_swept): + hsort_list_for_tags = getattr(args, "hardcode_sort_list", [False]) + hsort_in_cell_list = "HSORT" in getattr(args, "_cell_list_keys", ()) + emit_hsort_tag = ( + hardcode_sort or len(hsort_list_for_tags) > 1 + or hsort_in_cell_list + ) + if scn["fill"] is None and emit_hsort_tag: baseline_suffix_parts.append(f"HSORT={1 if hardcode_sort else 0}") baseline_sweep_suffix = ",".join(baseline_suffix_parts) baseline_key = _build_json_key( @@ -3340,7 +3349,10 @@ def _emit_split(name_w, name_nw, val_w, val_nw): parts.append(f"SR={1 if use_philox else 0}") parts.append(f"RECT={'auto' if rectangle_for_nowrite is None else (1 if rectangle_for_nowrite else 0)}") parts.append(f"MODE={_val(mode)}") - parts.append(f"HSORT={1 if hardcode_sort else 0}") + hsort_list_for_tags = getattr(args, "hardcode_sort_list", [False]) + hsort_in_cell_list = "HSORT" in getattr(args, "_cell_list_keys", ()) + if hardcode_sort or len(hsort_list_for_tags) > 1 or hsort_in_cell_list: + parts.append(f"HSORT={1 if hardcode_sort else 0}") if not args.use_cache_slot: parts.append("CSLOT=0") sweep_suffix = (" " + ",".join(parts)) if parts else "" @@ -3972,9 +3984,10 @@ def _phase(label: str) -> None: # Each entry in the JSON file is a dict of canonical knob keys → values, # using the same names that appear in the sweep_tag (Mw/Mnw, Ww/Wnw, # Sw/Snw, pW, pS, H, R, CT, CPSw/CPSnw, LSw/LSnw, FL, WS, TMARL, - # TMAWL, TMANL, TMAWS, SR, RECT, MODE, HSORT). Each - # cell may also use the tied forms M / W / S / CPS / LS (single value - # applied to both write and nowrite halves). + # TMAWL, TMANL, TMAWS, SR, RECT, MODE). Optional diagnostic HSORT cells + # are accepted for compatibility. Each cell may also use the tied forms + # M / W / S / CPS / LS (single value applied to both write and nowrite + # halves). # # On load we: # - Override the bench's CLI knob args (`args.block_size_m_write`, @@ -4412,8 +4425,9 @@ def _parse_args() -> argparse.Namespace: help="Path to a JSON list of cell dicts (one per cell to time). " "Each dict has canonical knob keys → values: Mw, Mnw, Ww, Wnw, Sw, " "Snw, pW, pS, H, R, CT, CPSw, CPSnw, LSw, LSnw, FL, WS, TMARL, " - "TMAWL, TMANL, TMAWS, SR, RECT, MODE, HSORT (tied forms M / W / S / " - "CPS / LS are also accepted and auto-expanded). " + "TMAWL, TMANL, TMAWS, SR, RECT, MODE; optional diagnostic HSORT is " + "also accepted (tied forms M / W / S / CPS / LS are also accepted " + "and auto-expanded). " "When set, bench's CLI knob ranges are auto-overridden to the " "per-knob union across all cells, and the inner-loop filter skips " "any iteration whose knob-value tuple isn't in the list. All cells " @@ -4565,7 +4579,7 @@ def _parse_args() -> argparse.Namespace: default=True, help="Use the cache-slot field from replay_work_items in persistent_main. " "--no-use-cache-slot keeps the old identity-cache-slot shortcut for " - "diagnostic comparisons only.", + "diagnostic comparisons only and requires --hardcode-sort 1.", ) parser.add_argument( "--heads-per-block", @@ -4729,13 +4743,11 @@ def _parse_args() -> argparse.Namespace: parser.add_argument( "--hardcode-sort", type=str, - default="0", - help="Comma-separated 0/1. When 1, the per-iter prev_tokens " - "samples are pre-sorted write-first OFFLINE (CPU-side) before " - "the timed region — kernel runs unchanged (USE_PERM=False) but " - "the EO gate sees sorted PNAT so early-outs cluster naturally. " - "Output is scrambled (we don't permute x/B/C/dt) but timing is " - "meaningful.", + default=None, + help="Comma-separated 0/1, default 0. Diagnostic only: when 1, " + "per-iter PNAT samples are preclustered write-first before " + "replay_work_items are built. Production-like mixed runs leave PNAT " + "unsorted and sort only replay_work_items.", ) parser.add_argument( "--mix-iters", @@ -4865,12 +4877,25 @@ def _round_iters_to_group(name, val): rect_list = [None] args.rectangle_for_nowrite_list = rect_list - hsort_modes = [v.strip() for v in (args.hardcode_sort or "0").split(",") if v.strip()] + hsort_modes = [ + v.strip() + for v in ( + args.hardcode_sort if args.hardcode_sort is not None else "0" + ).split(",") + if v.strip() + ] hsort_list = [] for v in hsort_modes: if v not in ("0", "1"): parser.error(f"--hardcode-sort value must be 0 or 1, got {v!r}") hsort_list.append(v == "1") + if not hsort_list: + hsort_list = [False] + if not args.use_cache_slot and not all(hsort_list): + parser.error( + "--no-use-cache-slot is a diagnostic shortcut and requires " + "--hardcode-sort 1" + ) args.hardcode_sort_list = hsort_list # mode=None means "let the wrapper resolve from _DEFAULT_TUNING". Same From 97f0b6d125d59577091437027d1bdccbd633f79a Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 02:51:17 -0700 Subject: [PATCH 62/89] mamba replay: add nowrite-first launch knob Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 43 ++++++---- ...benchmark_replay_selective_state_update.py | 78 ++++++++++++++----- 2 files changed, 85 insertions(+), 36 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index af61e8442508..aea8f62fe91b 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -2269,6 +2269,7 @@ def replay_selective_state_update( use_internal_pdl=True, write_checkpoint: bool = True, rectangle_for_nowrite: bool | None = None, + nowrite_first: bool = False, mode: str | None = None, _block_size_m: int | None = None, _num_warps: int | None = None, @@ -2385,6 +2386,8 @@ def replay_selective_state_update( use_internal_pdl: enable internal PDL (precompute → main overlap). Defaults True; override for testing only. Ignored on hardware that doesn't support PDL (sm < 90). + nowrite_first: benchmark/tuning knob for mode="persistent_main". + When true, launch the nowrite half before the write half. _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, _precompute_num_warps, _precompute_num_stages, _heads_per_block, @@ -2408,7 +2411,8 @@ def replay_selective_state_update( # write-first; the n_writes tensor partitions the persistent loop # into the two halves with the right WRITE_CHECKPOINT constexpr # each time. RECTANGLE constexpr (= rectangle_for_nowrite) picks - # rect vs replay for the nowrite half. write_checkpoint is ignored. + # rect vs replay for the nowrite half. nowrite_first controls + # launch order only. write_checkpoint is ignored. # Note: mode-and-knob resolution from the default-tuning table happens # below, after we have `batch` and `nheads`. @@ -3113,26 +3117,35 @@ def launch_persistent_dynamic_main(n_writes_tensor: torch.Tensor, elif mode == "persistent_main": # Persistent-CTA main kernel. One shared dynamic_precompute # (per-slot dispatch via PNAT) feeds two persistent_main - # launches (write half + nowrite half). Both halves ALWAYS - # launch; the kernel's runtime check iterates only the slots - # belonging to its half (write: [0, n_writes), nowrite: - # [n_writes, batch)). + # launches (write half + nowrite half). The first main launch in + # program order signals the second one when internal PDL is on. # # Caller-provided contract: `n_writes` is a (1,) int32 device # tensor (the kernel reads it at runtime, after the precompute); # `replay_work_items` is a (batch, 4) int32 device tensor # pre-sorted write-first. + def launch_nowrite(launch_dependent_kernels: bool): + launch_persistent_main( + write_checkpoint=False, + launch_dependent_kernels=launch_dependent_kernels, + rectangle=rectangle_for_nowrite, + ) + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) - launch_persistent_main( - write_checkpoint=True, - launch_dependent_kernels=True, - rectangle=False, # write always replay-style - ) - launch_persistent_main( - write_checkpoint=False, - launch_dependent_kernels=False, - rectangle=rectangle_for_nowrite, - ) + if nowrite_first: + launch_nowrite(launch_dependent_kernels=True) + launch_persistent_main( + write_checkpoint=True, + launch_dependent_kernels=False, + rectangle=False, # write always replay-style + ) + else: + launch_persistent_main( + write_checkpoint=True, + launch_dependent_kernels=True, + rectangle=False, # write always replay-style + ) + launch_nowrite(launch_dependent_kernels=False) else: raise ValueError( f"mode={mode!r} is not supported. Supported modes: " diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 4bc21ab36065..3384db15d9d9 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -2547,6 +2547,7 @@ def _bench_config( mix_samples_cpu=None, mix_label: str = "", hardcode_sort: bool = False, + nowrite_first: bool = False, mix_samples_sorted_cpu=None, mix_write_frac: float | None = None, warmup_only: bool = False, @@ -2626,6 +2627,8 @@ def _bench_config( mode = _resolve_effective_replay_mode( args, batch, state_dtype, use_philox, mode ) + if mode != "persistent_main" and nowrite_first: + return state_work = state0.clone() state_scales_work = state_scales0.clone() if state_scales0 is not None else None @@ -3228,6 +3231,7 @@ def _run_incr( x_call, B_call, C_call = x, B, C extra_kwargs = {} extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite + extra_kwargs["nowrite_first"] = nowrite_first extra_kwargs["mode"] = mode extra_kwargs["n_writes"] = scenario_n_writes extra_kwargs["replay_work_items"] = replay_work_items_buf @@ -3348,6 +3352,10 @@ def _emit_split(name_w, name_nw, val_w, val_nw): parts.append(f"TMAWS={_val(use_tma_replay_write_store)}") parts.append(f"SR={1 if use_philox else 0}") parts.append(f"RECT={'auto' if rectangle_for_nowrite is None else (1 if rectangle_for_nowrite else 0)}") + nowrite_first_list_for_tags = getattr(args, "nowrite_first_list", [False]) + nowrite_first_in_cell_list = "NWF" in getattr(args, "_cell_list_keys", ()) + if nowrite_first or len(nowrite_first_list_for_tags) > 1 or nowrite_first_in_cell_list: + parts.append(f"NWF={1 if nowrite_first else 0}") parts.append(f"MODE={_val(mode)}") hsort_list_for_tags = getattr(args, "hardcode_sort_list", [False]) hsort_in_cell_list = "HSORT" in getattr(args, "_cell_list_keys", ()) @@ -3418,6 +3426,7 @@ def _emit_split(name_w, name_nw, val_w, val_nw): bool(args.internal_pdl), bool(use_philox), bool(rectangle_for_nowrite), + bool(nowrite_first), bool(hardcode_sort), bool(args.use_cache_slot), scenario_pre_iter is not None, @@ -3712,6 +3721,7 @@ def _submit_result_job( "TMANL": "use_tma_replay_nowrite_load", "TMAWS": "use_tma_replay_write_store", "RECT": "rectangle_for_nowrite", + "NWF": "nowrite_first", "HSORT": "hardcode_sort", # MODE and SR get special handling (string values): # MODE → args.modes (single mode name) @@ -3827,6 +3837,7 @@ def _load_cell_list_into_args(args) -> None: "TMANL": "use_tma_replay_nowrite_load", "TMAWS": "use_tma_replay_write_store", "RECT": "rectangle_for_nowrite", + "NWF": "nowrite_first", "MODE": "mode", "HSORT": "hardcode_sort", "SR": "use_philox", @@ -3984,7 +3995,7 @@ def _phase(label: str) -> None: # Each entry in the JSON file is a dict of canonical knob keys → values, # using the same names that appear in the sweep_tag (Mw/Mnw, Ww/Wnw, # Sw/Snw, pW, pS, H, R, CT, CPSw/CPSnw, LSw/LSnw, FL, WS, TMARL, - # TMAWL, TMANL, TMAWS, SR, RECT, MODE). Optional diagnostic HSORT cells + # TMAWL, TMANL, TMAWS, SR, RECT, NWF, MODE). Optional diagnostic HSORT cells # are accepted for compatibility. Each cell may also use the tied forms # M / W / S / CPS / LS (single value applied to both write and nowrite # halves). @@ -4116,6 +4127,7 @@ def _phase(label: str) -> None: rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) modes_list = getattr(args, "modes_list", ["persistent_dynamic"]) hsort_list = getattr(args, "hardcode_sort_list", [False]) + nowrite_first_list = getattr(args, "nowrite_first_list", [False]) # Pre-load AL distribution for mix mode. mix_al = None @@ -4177,27 +4189,29 @@ def _phase(label: str) -> None: for sr_mode in sr_modes_list: for mode in modes_list: for rect in rect_list: - can_sort = mix_samples_cpu is not None - cell_list_active = bool( - getattr(args, "_cell_list_keys", ()) - ) - effective_hsort_list = ( - hsort_list if (can_sort or cell_list_active) - else [False] - ) - for hardcode_sort in effective_hsort_list: - _bench_config( - args, batch, mtp_len, prev_ks, - state_dtype, act_dtype, baseline_fn, - sr_mode=sr_mode, - rectangle_for_nowrite=rect, - mode=mode, - mix_samples_cpu=mix_samples_cpu, - mix_label=mix_label, - hardcode_sort=hardcode_sort, - mix_samples_sorted_cpu=mix_samples_sorted_cpu, - mix_write_frac=mix_write_frac, + for nowrite_first in nowrite_first_list: + can_sort = mix_samples_cpu is not None + cell_list_active = bool( + getattr(args, "_cell_list_keys", ()) + ) + effective_hsort_list = ( + hsort_list if (can_sort or cell_list_active) + else [False] ) + for hardcode_sort in effective_hsort_list: + _bench_config( + args, batch, mtp_len, prev_ks, + state_dtype, act_dtype, baseline_fn, + sr_mode=sr_mode, + rectangle_for_nowrite=rect, + mode=mode, + mix_samples_cpu=mix_samples_cpu, + mix_label=mix_label, + hardcode_sort=hardcode_sort, + nowrite_first=nowrite_first, + mix_samples_sorted_cpu=mix_samples_sorted_cpu, + mix_write_frac=mix_write_frac, + ) _drain_pending_results(args, force=True) @@ -4657,6 +4671,14 @@ def _parse_args() -> argparse.Namespace: "wrapper resolves from the _DEFAULT_TUNING lookup per (batch, dtype, " "sr) cell.", ) + parser.add_argument( + "--nowrite-first", + type=str, + default="0", + help="Comma-separated 0/1 values for mode=persistent_main launch order. " + "0 launches write before nowrite; 1 launches nowrite before write. " + "Ignored for persistent_dynamic.", + ) parser.add_argument( "--use-tma-rect-load", type=str, @@ -4877,6 +4899,20 @@ def _round_iters_to_group(name, val): rect_list = [None] args.rectangle_for_nowrite_list = rect_list + nowrite_first_modes = [ + v.strip() + for v in (args.nowrite_first if args.nowrite_first is not None else "0").split(",") + if v.strip() + ] + nowrite_first_list = [] + for v in nowrite_first_modes: + if v not in ("0", "1"): + parser.error(f"--nowrite-first value must be 0 or 1, got {v!r}") + nowrite_first_list.append(v == "1") + if not nowrite_first_list: + nowrite_first_list = [False] + args.nowrite_first_list = nowrite_first_list + hsort_modes = [ v.strip() for v in ( From f4b106036f3b1046bf412d7c34f2d7de2faaaa23 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 02:52:25 -0700 Subject: [PATCH 63/89] mamba replay: double-buffer old_x to eliminate write-path race MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The replay-write code path had a same-thread write-then-read race on the old_x cache. Inside _persistent_main_impl, the kernel loads old_x[0..PNAT) into registers (line ~1276) for the state += dot(old_x_all, dB_scaled) recurrence, then later stores the fresh x_all into old_x[write_offset..+T) (line ~1466). In the write_checkpoint=True case, write_offset=0 so the store range [0, T) overlaps the load range [0, PNAT) on the SAME buffer (old_x was single-buffered, unlike old_dt/old_dA_cumsum/old_B). Triton's alias analysis does not unify the offset expressions (offs_window vs (write_offset + offs_t)), so the compiler is free to issue the store before the load is fully consumed, corrupting the dot's input and producing per-head cumulative output errors from some t onward — flaky failure rate ~40-45% on persistent_main-write-T55-fp16-16-64-128-1 in tight loops. Fix: double-buffer old_x, matching the existing pattern for old_B / old_dt / old_dA_cumsum. Read from active_buf, write to write_buf (= 1 - active_buf when is_write; = active_buf otherwise — same convention as the others). Eliminates the same-address race entirely. Changes: - Add stride_old_x_dbuf to all 3 kernel signatures (_persistent_main_impl, _persistent_rectangle_impl, _persistent_main_kernel) - Split old_x_base into old_x_read_base / old_x_write_base in both main impls, indexed by active_buf and write_buf respectively - Wrapper: pass 5-stride tuple to all dispatcher sites; update old_x shape assertion and docstring (cache, max_window, nheads, dim) -> (cache, 2, max_window, nheads, dim); max_window now derived from shape[2] - Tests: update old_x allocation in all 5 test functions to add dbuf dim; update per-slot fills to write to the active buffer (matching old_B/ old_dt/old_dA_cumsum); update inspection slicing in test_replay_selective_state_update and test_replay_heads_per_block_multistep to assert against the per-buffer expected values - Bench: same alloc shape update Tests: 1215/1215 pass for the full mamba slim test suite; 1200/1200 pass for the previously-flaky persistent_main-write-T55-fp16 cell across 200 reps x 6 dtypes (vs ~40-45% failure rate without the fix). Perf (fp16/SR, all batches 1..1024, best-(HPB,pW) per cell): - Write path (prev_k=16): consistent +0.4 to +1.8% slower (max +1.79% at b=1024) - the real DB cost from doubled old_x memory footprint and different read/write buffers. - Nowrite path (prev_k=10): mixed direction, mostly ~0% with a few -4 to +3% outliers driven by autotune winners shifting around stride changes; not a real cost. Compared to the alternative false-math-dependency workaround tried earlier, this is roughly half the write-path cost (~+1% vs ~+2%) and is robust to future Triton compiler changes. Cross-module follow-up: mamba_cache_manager.py:1603 still allocates old_x without a dbuf dim - that needs to gain a '2' axis in a coordinated update with this commit (production caller will break otherwise). Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 65 +++++++++++++----- ...benchmark_replay_selective_state_update.py | 4 +- .../test_replay_selective_state_update.py | 66 +++++++++++-------- 3 files changed, 87 insertions(+), 48 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index aea8f62fe91b..2d51ea670bd4 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -1053,8 +1053,9 @@ def _persistent_main_impl( stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - # old_x strides + # old_x strides (double-buffered: cache, dbuf, T, head, dim) stride_old_x_cache, + stride_old_x_dbuf, stride_old_x_T, stride_old_x_head, stride_old_x_dim, @@ -1172,8 +1173,10 @@ def _persistent_main_impl( is_write = WRITE_CHECKPOINT if is_write: write_offset = 0 + write_buf = 1 - active_buf else: write_offset = prev_num_accepted_tokens + write_buf = active_buf offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) @@ -1264,9 +1267,25 @@ def _persistent_main_impl( coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + # Double-buffered old_x: load from active_buf (the "previous step" buffer); + # store at the end writes to write_buf (= 1 - active_buf when is_write, + # else active_buf — same as old_dt / old_dA_cumsum / old_B). Splitting + # the read and write into separate dbuf slots when is_write eliminates + # the same-thread store-vs-load race we previously saw on old_x. + old_x_read_base = ( + old_x_ptr + + cache_batch_idx * stride_old_x_cache + + active_buf * stride_old_x_dbuf + + pid_h * stride_old_x_head + ) + old_x_write_base = ( + old_x_ptr + + cache_batch_idx * stride_old_x_cache + + write_buf * stride_old_x_dbuf + + pid_h * stride_old_x_head + ) old_x_all = tl.load( - old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + old_x_read_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, mask=old_window_mask[:, None] & m_mask[None, :], other=0.0, ) @@ -1438,7 +1457,7 @@ def _persistent_main_impl( other=0.0, ) tl.store( - old_x_base + old_x_write_base + (write_offset + offs_t)[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, x_all, @@ -1522,8 +1541,9 @@ def _persistent_rectangle_impl( stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - # old_x strides + # old_x strides (double-buffered: cache, dbuf, T, head, dim) stride_old_x_cache, + stride_old_x_dbuf, stride_old_x_T, stride_old_x_head, stride_old_x_dim, @@ -1626,7 +1646,17 @@ def _persistent_rectangle_impl( if HAS_Z: z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head + # Rectangle path: nowrite-only. write_buf == active_buf (no flip), so the + # read and write paths target the same dbuf slot. No race: write_offset= + # PNAT puts new tokens at [PNAT, PNAT+T), disjoint from the read range + # [0, PNAT). + old_x_read_base = ( + old_x_ptr + + cache_batch_idx * stride_old_x_cache + + active_buf * stride_old_x_dbuf + + pid_h * stride_old_x_head + ) + old_x_write_base = old_x_read_base if HAS_D: D = tl.load( @@ -1635,7 +1665,7 @@ def _persistent_rectangle_impl( # Hoist: old_x doesn't depend on conv1d/precompute; load before gdc_wait. old_x_load = tl.load( - old_x_base + old_x_read_base + safe_old_k[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, mask=is_old_k[:, None] & m_mask[None, :], @@ -1656,7 +1686,7 @@ def _persistent_rectangle_impl( other=0.0, ) tl.store( - old_x_base + old_x_write_base + offs_k[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, x_K, @@ -1789,8 +1819,9 @@ def _persistent_main_kernel( stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - # old_x strides + # old_x strides (double-buffered: cache, dbuf, T, head, dim) stride_old_x_cache, + stride_old_x_dbuf, stride_old_x_T, stride_old_x_head, stride_old_x_dim, @@ -1991,7 +2022,7 @@ def _persistent_main_kernel( T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_x_cache, stride_old_x_dbuf, stride_old_x_T, stride_old_x_head, stride_old_x_dim, stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, stride_old_B_group, stride_old_B_dstate, stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, @@ -2022,7 +2053,7 @@ def _persistent_main_kernel( T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_x_cache, stride_old_x_dbuf, stride_old_x_T, stride_old_x_head, stride_old_x_dim, stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, stride_D_head, stride_D_dim, @@ -2047,7 +2078,7 @@ def _persistent_main_kernel( T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, + stride_old_x_cache, stride_old_x_dbuf, stride_old_x_T, stride_old_x_head, stride_old_x_dim, stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, stride_old_B_group, stride_old_B_dstate, stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, @@ -2353,7 +2384,7 @@ def replay_selective_state_update( Arguments: state: (cache, nheads, dim, dstate) in-place. After the call, contains the state after replaying prev_num_accepted_tokens old tokens. - old_x: (cache, T, nheads, dim) bf16 — old x cache (single-buffered). + old_x: (cache, 2, T, nheads, dim) bf16 — double-buffered old x cache. old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. @@ -2601,7 +2632,7 @@ def replay_selective_state_update( # `prev_num_accepted_tokens` can be 0..max_window. Window-axis kernel # tiles (BLOCK_SIZE_WINDOW, BLOCK_SIZE_K) are derived independently from # MAX_REPLAY_BUFFER_LENGTH so max_window can exceed BLOCK_SIZE_T freely. - max_window = old_x.shape[1] + max_window = old_x.shape[2] assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" assert x.shape == (batch, T, nheads, dim) @@ -2609,7 +2640,7 @@ def replay_selective_state_update( assert A.shape == (nheads, dim, dstate) assert B.shape == (batch, T, ngroups, dstate) assert C.shape == B.shape - assert old_x.shape == (cache_size, max_window, nheads, dim) + assert old_x.shape == (cache_size, 2, max_window, nheads, dim) assert old_B.shape == (cache_size, 2, max_window, ngroups, dstate) assert old_dt.shape == (cache_size, 2, nheads, max_window) assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_window) @@ -2980,7 +3011,7 @@ def launch_persistent_main(write_checkpoint: bool, T, max_window, dim, dstate, nheads // ngroups, state.stride(0), state.stride(1), state.stride(2), state.stride(3), state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), old_x.stride(4), old_B.stride(0), old_B.stride(1), old_B.stride(2), old_B.stride(3), old_B.stride(4), old_dt.stride(0), old_dt.stride(1), @@ -3056,7 +3087,7 @@ def launch_persistent_dynamic_main(n_writes_tensor: torch.Tensor, T, max_window, dim, dstate, nheads // ngroups, state.stride(0), state.stride(1), state.stride(2), state.stride(3), state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), + old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), old_x.stride(4), old_B.stride(0), old_B.stride(1), old_B.stride(2), old_B.stride(3), old_B.stride(4), old_dt.stride(0), old_dt.stride(1), diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 3384db15d9d9..7ea8524a83c8 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -645,8 +645,8 @@ def _build_tensors( # placeholder/degenerate case where every step is a checkpoint step). # For real replay-style checkpointing, max_window > mtp_len. cache_T = max_window if max_window is not None else mtp_len - # old_x: single-buffered (cache, max_window, nheads, dim) - old_x = torch.randn(batch, cache_T, nheads, head_dim, device=device, dtype=act_dtype) + # old_x: double-buffered (cache, 2, max_window, nheads, dim) + old_x = torch.randn(batch, 2, cache_T, nheads, head_dim, device=device, dtype=act_dtype) # old_B: double-buffered (cache, 2, max_window, ngroups, dstate) old_B = torch.randn(batch, 2, cache_T, ngroups, d_state, device=device, dtype=act_dtype) # old_dt: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous diff --git a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py index 29952b90a5e0..6d74a0fe4903 100644 --- a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py @@ -269,12 +269,12 @@ def test_replay_selective_state_update( ) # Build cache tensors for the replay kernel. - # old_x: (cache, max_window, nheads, dim) bf16 — single-buffered + # old_x: (cache, 2, max_window, nheads, dim) bf16 — double-buffered # old_B: (cache, 2, max_window, ngroups, dstate) bf16 — double-buffered # old_dt: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous # old_dA_cumsum: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous # cache_buf_idx: random 0s and 1s to verify indexing correctness - old_x = torch.zeros(cache_size, max_window, nheads, head_dim, device=device, dtype=dtype) + old_x = torch.randn(cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) old_dA_cumsum = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) @@ -285,7 +285,8 @@ def test_replay_selective_state_update( # values up to max_window are exercised. Inactive buffer has random # garbage to catch indexing bugs. slots = state_batch_indices if paged_cache else slice(None) - old_x[slots, :step1_T] = x1 + # old_x active-buffer fill is done in the per-slot loop below (same pattern + # as old_B/old_dt/old_dA_cumsum since old_x is now double-buffered too). # Compute processed dt and dA_cumsum for step 1 dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) @@ -296,6 +297,7 @@ def test_replay_selective_state_update( for i, slot in enumerate(slot_indices): buf = cache_buf_idx[slot].item() batch_idx = i # maps slot back to the batch index + old_x[slot, buf, :step1_T] = x1[batch_idx] old_B[slot, buf, :step1_T] = B1[batch_idx] old_dt[slot, buf, :, :step1_T] = dt1[batch_idx].T # (step1_T, nheads) → (nheads, step1_T) old_dA_cumsum[slot, buf, :, :step1_T] = dA_cumsum1[batch_idx].T @@ -537,25 +539,31 @@ def test_replay_selective_state_update( active = cache_buf_idx[slot].item() wb = (1 - active) if write_checkpoint else active - # --- old_x (single-buffered): write at [write_offset : +T) of slot --- - written_x = old_x_w[slot, write_offset : write_offset + T] + # --- old_x (double-buffered): write at wb, [write_offset : +T) --- + written_x = old_x_w[slot, wb, write_offset : write_offset + T] torch.testing.assert_close( written_x, x2[batch_idx], rtol=0, atol=0, msg=f"old_x written region wrong at k={k} write={write_checkpoint}", ) - # Untouched ranges of old_x[slot] + # Untouched ranges of old_x[slot, wb] if write_offset > 0: torch.testing.assert_close( - old_x_w[slot, :write_offset], old_x[slot, :write_offset], + old_x_w[slot, wb, :write_offset], old_x[slot, wb, :write_offset], rtol=0, atol=0, msg=f"old_x [0:{write_offset}) modified at k={k} write={write_checkpoint}", ) if write_offset + T < max_window: torch.testing.assert_close( - old_x_w[slot, write_offset + T:], old_x[slot, write_offset + T:], + old_x_w[slot, wb, write_offset + T:], old_x[slot, wb, write_offset + T:], rtol=0, atol=0, msg=f"old_x [{write_offset+T}:) modified at k={k} write={write_checkpoint}", ) + # Other-buffer (= 1-wb) untouched + torch.testing.assert_close( + old_x_w[slot, 1 - wb], old_x[slot, 1 - wb], + rtol=0, atol=0, + msg=f"old_x inactive buffer modified at k={k} write={write_checkpoint}", + ) # --- old_B (double-buffered): write at write_buf, [write_offset:+T) --- torch.testing.assert_close( @@ -706,7 +714,7 @@ def test_replay_selective_state_update_scenarios( disable_state_update=True, ) - old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_x = torch.randn(batch, 2, max_window, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) old_dA_cumsum = torch.randn( @@ -714,11 +722,11 @@ def test_replay_selective_state_update_scenarios( ) cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) - old_x[:, :step1_T] = x1 dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) for i in range(batch): buf = cache_buf_idx[i].item() + old_x[i, buf, :step1_T] = x1[i] old_B[i, buf, :step1_T] = B1[i] old_dt[i, buf, :, :step1_T] = dt1_processed[i].T old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T @@ -873,7 +881,7 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( disable_state_update=True, ) - old_x = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + old_x = torch.randn(batch, 2, max_window, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) old_dA_cumsum = torch.randn( @@ -886,11 +894,11 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( ) torch.testing.assert_close(_n_writes_check, n_writes) - old_x[:, :step1_T] = x1 dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) for i in range(batch): buf = cache_buf_idx[i].item() + old_x[i, buf, :step1_T] = x1[i] old_B[i, buf, :step1_T] = B1[i] old_dt[i, buf, :, :step1_T] = dt1_processed[i].T old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T @@ -1012,8 +1020,8 @@ def test_replay_selective_state_update_philox( ) state0_scales = None - # Cache tensors - old_x = torch.randn(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + # Cache tensors (old_x now double-buffered like the others) + old_x = torch.randn(cache_size, 2, T, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) @@ -1028,8 +1036,8 @@ def test_replay_selective_state_update_philox( prev_tokens = torch.full((cache_size,), T // 2, device=device, dtype=torch.int32) - # max_window is old_x.shape[1] per the wrapper convention; the philox - # test sets old_x = (cache_size, T, ...) so max_window = T here. + # max_window is old_x.shape[2] per the wrapper convention (after dbuf); + # the philox test sets old_x's window axis = T, so max_window = T here. _max_window_philox = T _n_writes_philox, _replay_work_items_philox = _make_replay_work_items( prev_tokens, cache_buf_idx, T, _max_window_philox, batch, @@ -1197,7 +1205,7 @@ def test_philox_rounding_unbiased(state_dtype): batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 ) - old_x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_x = torch.randn(batch, 2, T, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(batch, 2, T, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) old_dA_cumsum = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) @@ -1211,7 +1219,7 @@ def test_philox_rounding_unbiased(state_dtype): prev_tokens = torch.full((batch,), T, device=device, dtype=torch.int32) - # max_window = old_x.shape[1] = T + # max_window = old_x.shape[2] = T (after dbuf at axis 1) _n_writes_unb, _replay_work_items_unb = _make_replay_work_items( prev_tokens, cache_buf_idx, T, T, batch, None, device, ) @@ -1401,12 +1409,12 @@ def test_replay_heads_per_block( disable_state_update=True, ) - # Pre-fill cache buffers. old_x is single-buffer (no dbuf dim); old_B, - # old_dt, old_dA_cumsum are double-buffered. Initialize BOTH buffers - # with controlled random data so "outside write range / other buffer - # unchanged" assertions have well-defined expected values for both. + # Pre-fill cache buffers. All four (old_x, old_B, old_dt, old_dA_cumsum) + # are double-buffered. Initialize BOTH buffers with controlled random + # data so "outside write range / other buffer unchanged" assertions have + # well-defined expected values for both. old_x_init = torch.randn( - cache_size, max_window, nheads, head_dim, device=device, dtype=dtype + cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype ) old_B_init = torch.randn( cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype @@ -1426,9 +1434,9 @@ def test_replay_heads_per_block( dt1_proc = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_proc, dim=1) - old_x_init[:] = x1 # single buffer for slot in range(cache_size): buf = int(cache_buf_idx[slot].item()) + old_x_init[slot, buf] = x1[slot] old_B_init[slot, buf] = B1[slot] old_dt_init[slot, buf] = dt1_proc[slot].T # (nheads, max_window) old_dA_cumsum_init[slot, buf] = dA_cumsum1[slot].T # (nheads, max_window) @@ -1594,16 +1602,16 @@ def test_replay_heads_per_block( write_offset = pnat write_end = write_offset + T - # ----- old_x (single-buffer) ----- + # ----- old_x (double-buffer (cache, 2, max_window, nheads, dim)) ----- expected_old_x_slot = old_x_pre[slot].clone() - expected_old_x_slot[write_offset:write_end] = x2[slot] + expected_old_x_slot[target_buf, write_offset:write_end] = x2[slot] torch.testing.assert_close( old_x_test[slot], expected_old_x_slot, rtol=0, atol=0, msg=( f"old_x slot {slot} (PNAT={pnat}, is_write={is_write}, " - f"write_offset={write_offset}): mismatch " - f"(HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"target_buf={target_buf}, write_offset={write_offset}): " + f"mismatch (HPB={heads_per_block}, T={T}, nheads={nheads}, " f"ngroups={ngroups}, rect={rectangle_nowrite})" ), ) @@ -1788,7 +1796,7 @@ def test_replay_heads_per_block_multistep( ref_outs.append(out_step) test_state = state_init.clone() - old_x = torch.zeros(cache_size, max_window, nheads, head_dim, device=device, dtype=dtype) + old_x = torch.zeros(cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype) old_B = torch.zeros(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.zeros(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) old_dA_cumsum = torch.zeros(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) From 6546272ce40a51030b04bd1c3fccd9ad0a4abbfd Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 03:15:03 -0700 Subject: [PATCH 64/89] mamba replay: double-buffer old_x cache allocation Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py | 8 +++----- .../unittest/_torch/executor/test_mamba_cache_manager.py | 2 +- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 9efda8abd536..66b67b5a2799 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -269,7 +269,7 @@ class SpeculativeState(State): # 0 means temporal saved state is actually the last state, not two back. prev_num_accepted_tokens: torch.Tensor | None = None # (cache,) int — shared across layers cache_buf_idx: torch.Tensor | None = None # (cache,) int32 — shared across layers - old_x: torch.Tensor | None = None # (layers, cache, history, nheads, dim) + old_x: torch.Tensor | None = None # (layers, cache, 2, history, nheads, dim) old_B: torch.Tensor | None = None # (layers, cache, 2, history, ngroups, dstate) # Processed dt: softplus(raw_dt + dt_bias), clamped to dt_limit. old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 @@ -382,9 +382,6 @@ def __init__( self.replay_history_size = max(16, T) # Compact replay cache. - # old_x is single-buffered (written by main kernel after replay). - # old_B, old_dt, old_dA_cumsum are double-buffered (written by - # precompute kernel concurrently with main kernel via PDL). spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( max_batch_size, dtype=int, device=device) spec_kwargs['cache_buf_idx'] = torch.zeros(max_batch_size, @@ -392,6 +389,7 @@ def __init__( device=device) spec_kwargs['old_x'] = torch.zeros(num_local_layers, max_batch_size, + 2, self.replay_history_size, nheads, head_dim, @@ -1662,9 +1660,9 @@ def _setup_replay_buffers(self, spec_config) -> None: self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) - # x is not double-buffered self.old_x = torch.zeros(num_local_mamba_layers, cache_size, + 2, history_size, nheads, head_dim, diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index f7a4cd14e271..91342eeb60bb 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -120,7 +120,7 @@ def test_replay_update_mamba_states_uses_history_window(): mgr = _make_mgr(max_batch_size=4, max_draft_len=5, use_replay_state_update=True) assert mgr.replay_step_width == 6 assert mgr.replay_history_size == 16 - assert mgr.mamba_cache.old_x.shape[2] == 16 + assert mgr.mamba_cache.old_x.shape[3] == 16 assert mgr.mamba_cache.old_B.shape[3] == 16 assert mgr.mamba_cache.old_dt.shape[4] == 16 assert mgr.mamba_cache.old_dA_cumsum.shape[4] == 16 From 2ed22dc0f9dfd16ecccfec251495e26637889640 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 10:36:17 -0700 Subject: [PATCH 65/89] mamba replay: address review cleanup Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../_torch/modules/mamba/mamba2_metadata.py | 6 +- .../_torch/modules/mamba/mamba2_mixer.py | 6 +- .../mamba/replay_selective_state_update.py | 201 +++++++++--------- .../_torch/pyexecutor/mamba_cache_manager.py | 41 ++-- .../executor/test_mamba_cache_manager.py | 2 +- .../modules/mamba/test_mamba2_metadata.py | 11 +- .../test_replay_selective_state_update.py | 63 ++---- 7 files changed, 165 insertions(+), 165 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 677492b74eed..0e387c26efd0 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -252,8 +252,10 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, return self.replay_n_writes.zero_() - prev_num_accepted_tokens, cache_buf_idx, replay_step_width, \ - replay_history_size = replay_metadata + prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens + cache_buf_idx = replay_metadata.cache_buf_idx + replay_step_width = replay_metadata.replay_step_width + replay_history_size = replay_metadata.replay_history_size num_decodes = batch_size - num_contexts self.replay_num_decodes = num_decodes if num_decodes == 0: diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index d6e5a292f18c..3d8dd5760c61 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -514,9 +514,9 @@ def convert_dt(): philox_kwargs['philox_rounds'] = self._philox_rounds if use_replay: - # replay_work_items is sorted write-first and carries - # decode-batch position, cache slot, PNAT, and active - # cache buffer index for persistent-main replay. + # replay_work_items is write-first for persistent_main and + # carries decode-batch position, cache slot, PNAT, and + # active cache buffer index for replay kernels. replay_selective_state_update( ssm_states, layer_cache.old_x, diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 2d51ea670bd4..f7872fffd111 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E501 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -25,10 +26,13 @@ from tensorrt_llm._utils import get_sm_version -from .mamba2_metadata import (REPLAY_WORK_CACHE_BUF_IDX, - REPLAY_WORK_CACHE_SLOT, - REPLAY_WORK_ITEM_WIDTH, REPLAY_WORK_PNAT, - REPLAY_WORK_POSITION_IN_DECODE_BATCH) +from .mamba2_metadata import ( + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_ITEM_WIDTH, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, +) from .softplus import softplus _REPLAY_WORK_POSITION_IN_DECODE_BATCH = tl.constexpr( @@ -263,10 +267,9 @@ def _replay_precompute_impl( # afterward; PNAT_next = PNAT + accepted. [0, PNAT) preserved. # WRITE_CHECKPOINT=True (would overflow): write to STAGING buffer at # [0, T). Caller flips cache_buf_idx afterward; next step's - # active = the one we just wrote. PNAT_next = accepted. Old - # data in the previous active buffer is folded into state via - # the replay update and discarded. This matches today's replay - # kernel behavior exactly. + # active = the one we just wrote. PNAT_next = accepted. The + # previous active-buffer history is folded into state by replay + # before the caller flips to the staging buffer. buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) if write_checkpoint: @@ -876,10 +879,8 @@ def _dynamic_precompute_kernel( pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH - # write_checkpoint is now runtime in replay precompute, so a single - # call site handles both write and nowrite for the replay branch. - # Take rectangle only when RECTANGLE is True AND this slot doesn't - # need write; everything else funnels into replay. + # Replay precompute uses a per-slot write predicate; use rectangle only + # for no-write slots when enabled. if needs_write_runtime or not RECTANGLE: _replay_precompute_impl( dt_ptr, @@ -1267,11 +1268,8 @@ def _persistent_main_impl( coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all - # Double-buffered old_x: load from active_buf (the "previous step" buffer); - # store at the end writes to write_buf (= 1 - active_buf when is_write, - # else active_buf — same as old_dt / old_dA_cumsum / old_B). Splitting - # the read and write into separate dbuf slots when is_write eliminates - # the same-thread store-vs-load race we previously saw on old_x. + # Double buffering keeps old_x replay reads and checkpoint writes in + # disjoint buffers on write steps. old_x_read_base = ( old_x_ptr + cache_batch_idx * stride_old_x_cache @@ -1591,8 +1589,7 @@ def _persistent_rectangle_impl( QUANT_MAX: tl.constexpr, USE_TMA_LOAD: tl.constexpr = False, ): - # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). - write_offset = prev_num_accepted_tokens + # Nowrite-only: new tokens append at [PNAT, PNAT+T). # Rectangle K-axis layout: old at [0, PNAT), new at [PNAT, PNAT+T). @@ -1887,15 +1884,9 @@ def _persistent_main_kernel( QUANT_MAX: tl.constexpr, WRITE_CHECKPOINT: tl.constexpr, LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - # NUM_PERSISTENT: runtime int (not constexpr). Used ONLY as the loop - # stride in `tl.range(pid, total_work, NUM_PERSISTENT, ...)`. Making it - # runtime collapses the cta_per_sm tuning dim from the kernel's compile - # signature: 8 CPS values used to mean 8x recompiles; now they share one - # compiled kernel. Work decomposition (pid_m, pid_b_local, pid_h) does - # NOT depend on NUM_PERSISTENT — it uses constexpr NUM_PID_M_BLOCKS and - # runtime n_slots_local — so loop unrolling and flatten=/num_stages=/ - # warp_specialize= optimizations on `tl.range` operate independently of - # the stride value. + # NUM_PERSISTENT is a runtime loop stride, so CTA-per-SM tuning can vary + # without changing the compiled kernel signature. Work decomposition uses + # constexpr NUM_PID_M_BLOCKS and runtime n_slots_local. NUM_PERSISTENT, NUM_LOOP_STAGES: tl.constexpr, NUM_PID_M_BLOCKS: tl.constexpr, @@ -2425,8 +2416,10 @@ def replay_selective_state_update( _maxnreg, _num_ctas) are benchmark-only overrides; production callers should leave them None to use the heuristic-tuned defaults. """ + sm_version = get_sm_version() + # PDL needs sm >= 90. - if get_sm_version() < 90: + if sm_version < 90: launch_with_pdl = False use_internal_pdl = False @@ -2451,9 +2444,9 @@ def replay_selective_state_update( # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX # instructions (Ada Lovelace introduced them; Hopper/Blackwell carry them). if state.dtype == torch.float8_e4m3fn: - assert get_sm_version() >= 89, ( + assert sm_version >= 89, ( "fp8_e4m3fn state requires SM 89+ (Ada Lovelace / Hopper / Blackwell) " - f"for fp32↔fp8 cvt PTX instructions; current SM is {get_sm_version()}." + f"for fp32↔fp8 cvt PTX instructions; current SM is {sm_version}." ) # PTX cvt.rs.* (stochastic rounding) family lands on Blackwell only. @@ -2462,14 +2455,14 @@ def replay_selective_state_update( # PTX SR instruction needed, runs anywhere. if rand_seed is not None: if state.dtype == torch.float16: - assert get_sm_version() >= 100, ( + assert sm_version >= 100, ( "fp16 stochastic rounding (PTX cvt.rs.f16x2.f32) requires " - f"sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + f"sm_100a (Blackwell B200+); current SM is {sm_version}." ) elif state.dtype == torch.float8_e4m3fn: - assert get_sm_version() >= 100, ( + assert sm_version >= 100, ( "fp8 stochastic rounding (PTX cvt.rs.satfinite.e4m3x4.f32) " - f"requires sm_100a (Blackwell B200+); current SM is {get_sm_version()}." + f"requires sm_100a (Blackwell B200+); current SM is {sm_version}." ) else: assert state.dtype in (torch.int8, torch.int16), ( @@ -2525,15 +2518,10 @@ def replay_selective_state_update( # --- Default-tuning lookup --- # Resolve (mode, knobs) from the table when caller leaves them None. - # Caller-provided kwargs always win. If the caller forces a mode that - # differs from the table's recommendation for this cell, we BRIDGE the - # table's knobs into the forced mode's knob namespace rather than fall - # back to (likely-terrible) kernel defaults: - # table pd → forced pm: copy each unsplit pd knob (M, W, S, CPS, LS) - # to both write and nowrite split knobs. - # table pm → forced pd: take the nowrite split values (Mnw, Wnw, Snw, - # CPSnw, LSnw) as the unsplit knobs. - # Empty table → no-op (caller passes whatever, mode falls back to pd). + # Caller-provided kwargs always win. If the caller forces a different + # mode, translate table knobs into that mode's namespace: persistent_dynamic + # knobs fan out to both persistent_main halves, while persistent_main + # nowrite knobs seed persistent_dynamic. _dt_str = { torch.float32: "fp32", torch.float16: "fp16", @@ -2559,14 +2547,35 @@ def replay_selective_state_update( _num_warps = _num_warps if _num_warps is not None else _table_knobs.get("_num_warps") _num_stages = _num_stages if _num_stages is not None else _table_knobs.get("_num_stages") _heads_per_block = _heads_per_block if _heads_per_block is not None else _table_knobs.get("_heads_per_block") - _precompute_num_warps = _precompute_num_warps if _precompute_num_warps is not None else _table_knobs.get("_precompute_num_warps") - _precompute_num_stages = _precompute_num_stages if _precompute_num_stages is not None else _table_knobs.get("_precompute_num_stages") - _block_size_m_write = _block_size_m_write if _block_size_m_write is not None else _table_knobs.get("_block_size_m_write") - _block_size_m_nowrite = _block_size_m_nowrite if _block_size_m_nowrite is not None else _table_knobs.get("_block_size_m_nowrite") + _precompute_num_warps = ( + _precompute_num_warps + if _precompute_num_warps is not None + else _table_knobs.get("_precompute_num_warps")) + _precompute_num_stages = ( + _precompute_num_stages + if _precompute_num_stages is not None + else _table_knobs.get("_precompute_num_stages")) + _block_size_m_write = ( + _block_size_m_write + if _block_size_m_write is not None + else _table_knobs.get("_block_size_m_write")) + _block_size_m_nowrite = ( + _block_size_m_nowrite + if _block_size_m_nowrite is not None + else _table_knobs.get("_block_size_m_nowrite")) _num_warps_write = _num_warps_write if _num_warps_write is not None else _table_knobs.get("_num_warps_write") - _num_warps_nowrite = _num_warps_nowrite if _num_warps_nowrite is not None else _table_knobs.get("_num_warps_nowrite") - _num_stages_write = _num_stages_write if _num_stages_write is not None else _table_knobs.get("_num_stages_write") - _num_stages_nowrite = _num_stages_nowrite if _num_stages_nowrite is not None else _table_knobs.get("_num_stages_nowrite") + _num_warps_nowrite = ( + _num_warps_nowrite + if _num_warps_nowrite is not None + else _table_knobs.get("_num_warps_nowrite")) + _num_stages_write = ( + _num_stages_write + if _num_stages_write is not None + else _table_knobs.get("_num_stages_write")) + _num_stages_nowrite = ( + _num_stages_nowrite + if _num_stages_nowrite is not None + else _table_knobs.get("_num_stages_nowrite")) _cta_per_sm = _cta_per_sm if _cta_per_sm is not None else _table_knobs.get("_cta_per_sm") _num_loop_stages = _num_loop_stages if _num_loop_stages is not None else _table_knobs.get("_num_loop_stages") # persistent_main uses split write/nowrite tuning knobs. @@ -2609,6 +2618,11 @@ def replay_selective_state_update( _use_tma_replay_write_load = bool(_use_tma_replay_write_load) _use_tma_replay_write_store = bool(_use_tma_replay_write_store) _use_tma_replay_nowrite_load = bool(_use_tma_replay_nowrite_load) + if sm_version < 90: + _use_tma_rect_load = False + _use_tma_replay_write_load = False + _use_tma_replay_write_store = False + _use_tma_replay_nowrite_load = False assert mode in ("persistent_dynamic", "persistent_main"), ( f"unknown mode {mode!r}; expected 'persistent_dynamic' or 'persistent_main'" ) @@ -2626,11 +2640,9 @@ def replay_selective_state_update( ) assert state_scales.device == state.device - # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the - # placeholder degenerate case max_window = T (every step writes replay - # state). For real replay-style history, max_window > T and - # `prev_num_accepted_tokens` can be 0..max_window. Window-axis kernel - # tiles (BLOCK_SIZE_WINDOW, BLOCK_SIZE_K) are derived independently from + # Cache window capacity comes from old_x.shape[2]; it may equal T or be + # larger when retaining replay history. Window-axis kernel tiles + # (BLOCK_SIZE_WINDOW, BLOCK_SIZE_K) are derived independently from # MAX_REPLAY_BUFFER_LENGTH so max_window can exceed BLOCK_SIZE_T freely. max_window = old_x.shape[2] assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" @@ -2787,11 +2799,12 @@ def replay_selective_state_update( if _num_warps is not None: num_warps = _num_warps if _heads_per_block is not None: - # Cap at heads_per_group: HEADS_PER_BLOCK divides the kernel's head - # axis, so a table value larger than the model's heads-per-group - # would overshoot. Protects callers running smaller models than - # the one we tuned against. - heads_per_block = min(_heads_per_block, heads_per_group) + heads_per_block = int(_heads_per_block) + assert heads_per_block > 0, "heads_per_block must be positive" + heads_per_block = min(heads_per_block, heads_per_group) + while (heads_per_group % heads_per_block != 0 + or heads_per_block & (heads_per_block - 1) != 0): + heads_per_block -= 1 if _precompute_num_warps is not None: precompute_num_warps = _precompute_num_warps @@ -2818,6 +2831,9 @@ def replay_selective_state_update( assert heads_per_block <= heads_per_group, ( f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" ) + assert heads_per_group % heads_per_block == 0, ( + f"heads_per_block ({heads_per_block}) must divide heads_per_group ({heads_per_group})" + ) # state_scales pointer + strides: real tensor when quantized, otherwise # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). @@ -2943,18 +2959,12 @@ def launch_dynamic_precompute(rectangle: bool): ) # ---- launch_persistent_main ------------------------------------------ - # Persistent-CTA main kernel. Single launch covers `n_slots` slots - # starting at `slot_offset`. Caller invokes twice: once for the write - # half (slot_offset=0, n_slots=n_writes, write_checkpoint=True) and - # once for the nowrite half (slot_offset=n_writes, - # n_slots=batch-n_writes, write_checkpoint=False). Hard-sort - # contract: caller has pre-sorted slots so [0, n_writes) are writes - # and [n_writes, batch) are nowrites. - - # Resolve persistent-mode bench knobs. Defaults: cta_per_sm = 1 - # (one CTA per SM, matches upstream `_p_matmul_ogs.py`); num_loop_stages - # = 2 (matches in-tree `swiglu` precedent for non-dot persistent loops); - # flatten = True (canonical Triton 3.6 idiom); warp_specialize = False. + # Persistent main launches the write and nowrite halves separately. + # Replay work items are write-first, and the device n_writes tensor + # partitions [0, n_writes) from [n_writes, batch). + + # Resolve persistent-mode tuning knobs. None values fall back to stable + # defaults so direct callers do not need to mirror benchmark search args. _num_sms = torch.cuda.get_device_properties(device).multi_processor_count cta_per_sm_arg = _cta_per_sm if _cta_per_sm else 1 num_persistent_arg = cta_per_sm_arg * _num_sms @@ -2982,26 +2992,27 @@ def launch_persistent_main(write_checkpoint: bool, # `n_writes` is the (1,) int32 device tensor with the write count. # Both halves always launch; an empty half has a zero-length slot # range and the persistent loop does no work. - _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE - _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE - _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE - _cps = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE - _cps = _cps if _cps else 1 - _nls = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE - _nls = _nls if _nls else 2 - _num_persistent = _cps * _num_sms - _num_pid_m_local = (dim + _bsm - 1) // _bsm + block_size_m = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE + launch_num_warps = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE + launch_num_stages = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE + ctas_per_sm = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE + ctas_per_sm = ctas_per_sm if ctas_per_sm else 1 + num_loop_stages = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE + num_loop_stages = num_loop_stages if num_loop_stages else 2 + num_persistent = ctas_per_sm * _num_sms + num_pid_m_local = (dim + block_size_m - 1) // block_size_m # Grid sizing: cap at min(full persistent grid, upper-bound total work). # We use `batch` as the upper bound on slots-per-half — overcounting # by a few CTAs is fine since the kernel derives the exact slot range # from n_writes at runtime. - _total_work_launch = max(1, batch * _num_pid_m_local * nheads) - grid = (min(_num_persistent, _total_work_launch),) - # Per-path TMA descriptor — block_shape[0] must match _bsm. - _desc = (state_tma_descriptor_write if write_checkpoint - else state_tma_descriptor_nowrite) + total_work_launch = max(1, batch * num_pid_m_local * nheads) + grid = (min(num_persistent, total_work_launch),) + # Per-path TMA descriptor — block_shape[0] must match block_size_m. + selected_state_tma_descriptor = ( + state_tma_descriptor_write if write_checkpoint + else state_tma_descriptor_nowrite) _persistent_main_kernel[grid]( - state, _desc, state_scales_arg, old_x, + state, selected_state_tma_descriptor, state_scales_arg, old_x, old_B, old_dt, old_dA_cumsum, prev_num_accepted_tokens, cache_buf_idx, x, C, D, z, out, @@ -3026,14 +3037,14 @@ def launch_persistent_main(write_checkpoint: bool, cb_scaled.stride(0), cb_scaled.stride(1), cb_scaled.stride(2), cb_scaled.stride(3), decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - _bsm, + block_size_m, LAUNCH_WITH_PDL=use_internal_pdl, PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, QUANT_MAX=quant_max, WRITE_CHECKPOINT=write_checkpoint, LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, - NUM_PERSISTENT=_num_persistent, - NUM_LOOP_STAGES=_nls, + NUM_PERSISTENT=num_persistent, + NUM_LOOP_STAGES=num_loop_stages, FLATTEN=flatten_arg, WARP_SPECIALIZE=warp_specialize_arg, IS_DYNAMIC=False, @@ -3050,8 +3061,8 @@ def launch_persistent_main(write_checkpoint: bool, ), USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), USE_REPLAY_CACHE_SLOT=bool(_use_replay_cache_slot), - num_warps=_nw, - **({"num_stages": _ns} if _ns else {}), + num_warps=launch_num_warps, + **({"num_stages": launch_num_stages} if launch_num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), **({"maxnreg": _maxnreg} if _maxnreg else {}), launch_pdl=use_internal_pdl, @@ -3070,8 +3081,8 @@ def launch_persistent_dynamic_main(n_writes_tensor: torch.Tensor, # Grid sizing: cap at total_work (= batch * num_pid_m * nheads) for # the dynamic case (full-batch coverage); see launch_persistent_main # comment for correctness rationale. - _total_work_launch = max(1, batch * _num_pid_m * nheads) - grid = (min(num_persistent_arg, _total_work_launch),) + total_work_launch = max(1, batch * _num_pid_m * nheads) + grid = (min(num_persistent_arg, total_work_launch),) # Persistent-dynamic kernel uses a single BLOCK_SIZE_M (same as the # wrapper's BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE tied convention), so # the write-side descriptor matches. Both write and nowrite slots diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 66b67b5a2799..c184d24da97d 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -17,7 +17,7 @@ import os from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Union import torch @@ -63,6 +63,14 @@ def use_cpp_mamba_cache_manager() -> bool: return os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' +class ReplayStateUpdateMetadata(NamedTuple): + """Shared tensors and fixed sizes for replay state updates.""" + prev_num_accepted_tokens: torch.Tensor + cache_buf_idx: torch.Tensor + replay_step_width: int + replay_history_size: int + + class BaseMambaCacheManager(ABC): """Abstract interface for accessing mamba/recurrent state caches.""" @@ -76,7 +84,7 @@ def get_state_indices(self, *args, **kwargs) -> torch.Tensor: def get_replay_state_update_metadata( self - ) -> Optional[tuple[torch.Tensor, torch.Tensor, int, int]]: + ) -> Optional[ReplayStateUpdateMetadata]: """Return replay metadata tensors and fixed replay sizes.""" return None @@ -609,7 +617,7 @@ def use_replay_state_update(self) -> bool: def get_replay_state_update_metadata( self - ) -> Optional[tuple[torch.Tensor, torch.Tensor, int, int]]: + ) -> Optional[ReplayStateUpdateMetadata]: if (not self._use_replay_state_update or not isinstance(self.mamba_cache, self.SpeculativeState) or self.mamba_cache.prev_num_accepted_tokens is None @@ -617,9 +625,12 @@ def get_replay_state_update_metadata( or self.replay_step_width is None or self.replay_history_size is None): return None - return (self.mamba_cache.prev_num_accepted_tokens, - self.mamba_cache.cache_buf_idx, self.replay_step_width, - self.replay_history_size) + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=( + self.mamba_cache.prev_num_accepted_tokens), + cache_buf_idx=self.mamba_cache.cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) def shutdown(self): """Release tensor memory.""" @@ -814,7 +825,7 @@ def use_replay_state_update(self) -> bool: def get_replay_state_update_metadata( self - ) -> Optional[tuple[torch.Tensor, torch.Tensor, int, int]]: + ) -> Optional[ReplayStateUpdateMetadata]: get_metadata = getattr(self._impl, 'get_replay_state_update_metadata', None) if get_metadata is None: @@ -1358,13 +1369,14 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests): # block (prefix-cache hit or block recycled across requests) may carry # stale prev_num_accepted_tokens / cache_buf_idx values from a prior # owner; the replay kernel reads these on the first decode step. - if self._use_replay_state_update and self.prev_num_accepted_tokens is not None: + if (self._use_replay_state_update + and self.prev_num_accepted_tokens is not None + and self.cache_buf_idx is not None): num_contexts = len(scheduled_batch.context_requests) if num_contexts > 0: ctx_slots = self.cuda_state_indices[:num_contexts].long() self.prev_num_accepted_tokens[ctx_slots] = 0 - # don't care which half of doulbe-buffer is using - # self.cache_buf_idx[ctx_slots] = 0 + self.cache_buf_idx[ctx_slots] = 0 def prepare_resources(self, scheduled_batch: ScheduledRequests): super().prepare_resources(scheduled_batch) @@ -1698,15 +1710,18 @@ def use_replay_state_update(self) -> bool: def get_replay_state_update_metadata( self - ) -> Optional[tuple[torch.Tensor, torch.Tensor, int, int]]: + ) -> Optional[ReplayStateUpdateMetadata]: if (not self._use_replay_state_update or self.prev_num_accepted_tokens is None or self.cache_buf_idx is None or self.replay_step_width is None or self.replay_history_size is None): return None - return (self.prev_num_accepted_tokens, self.cache_buf_idx, - self.replay_step_width, self.replay_history_size) + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=self.prev_num_accepted_tokens, + cache_buf_idx=self.cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) def get_mamba_ssm_cache_dtype(self) -> torch.dtype: return self.ssm_state_dtype diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 91342eeb60bb..bd140c45b90e 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -116,7 +116,7 @@ def _fake(rid): @skip_no_cuda def test_replay_update_mamba_states_uses_history_window(): - """Replay path appends PNAT until the layer kernels checkpointed.""" + """Replay path accumulates PNAT until layer kernels write a checkpoint.""" mgr = _make_mgr(max_batch_size=4, max_draft_len=5, use_replay_state_update=True) assert mgr.replay_step_width == 6 assert mgr.replay_history_size == 16 diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 9e88699be875..25859b8dfa4b 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -20,14 +20,15 @@ import torch from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( - Mamba2Metadata, REPLAY_WORK_CACHE_BUF_IDX, REPLAY_WORK_CACHE_SLOT, REPLAY_WORK_PNAT, REPLAY_WORK_POSITION_IN_DECODE_BATCH, + Mamba2Metadata, cu_seqlens_to_chunk_indices_offsets, cu_seqlens_to_chunk_indices_offsets_triton, ) +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ReplayStateUpdateMetadata skip_no_cuda = pytest.mark.skipif( not torch.cuda.is_available(), @@ -109,8 +110,12 @@ def get_state_indices(self, request_ids, is_padding): return self.state_indices[:len(request_ids)] def get_replay_state_update_metadata(self): - return (self.prev_num_accepted_tokens, self.cache_buf_idx, 6, - 16) + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=self.prev_num_accepted_tokens, + cache_buf_idx=self.cache_buf_idx, + replay_step_width=6, + replay_history_size=16, + ) metadata = Mamba2Metadata(max_batch_size=5, chunk_size=8) seq_lens = torch.tensor([2, 7, 7, 7, 7], dtype=torch.int) diff --git a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py index 6d74a0fe4903..030fe31e86dd 100644 --- a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py @@ -86,6 +86,9 @@ def _make_replay_work_items(prev_tokens, cache_buf_idx, T, max_window, batch, (16, 64, 128, 1), # TP=8 production config (32, 64, 128, 2), # TP=4, ngroups>1 (more heads than B/C groups) ] +_HEADS_PER_BLOCK_CONFIGS = _CONFIGS + [ + (6, 64, 128, 2), # heads_per_group=3 exercises HPB divisor fallback +] # Quantized state dtypes and their representable-magnitude limits (== QUANT_MAX # in the kernel). fp8_e4m3fn cells require SM 89+ for the fp32↔fp8 cvt PTX @@ -622,13 +625,8 @@ def test_replay_selective_state_update( ("all_write", [12, 13, 14, 15], None, False), # All-nowrite: every slot fits in the window. ("all_nowrite", [3, 4, 5, 6], None, False), - # Mixed PNATs with HAND-CODED work-item order. The auto-computed order - # for these PNATs would be [2, 3, 0, 1] — same as the explicit - # value — so this scenario is functionally redundant with the - # auto variant ONLY if `_make_replay_work_items` (the test - # helper) is itself correct. Keeping a hand-coded copy guards - # against a buggy helper: the kernel still gets a known-good order - # and the scenario would still pass even if the helper regressed. + # Mixed PNATs with hand-coded work-item order, independent of the + # _make_replay_work_items test helper. ("mixed_explicit", [3, 10, 12, 16], [2, 3, 0, 1], False), ("mixed_explicit_rect", [3, 10, 12, 16], [2, 3, 0, 1], True), # Mixed PNATs with AUTO-COMPUTED work-item order via `_make_replay_work_items` @@ -655,11 +653,7 @@ def test_replay_selective_state_update_scenarios( - all_write / all_nowrite: every slot on one branch — verifies the empty-half early-return on pm and the all-uniform per-slot dispatch on pd. - - mixed_explicit: hand-coded work-item order, bypassing the test's - `_make_replay_work_items` helper. Guards against a buggy helper: - if the auto-computation regressed, the auto scenarios would still - pass with the broken value, but this one runs against a known-good - order and would still detect the kernel-side issue. + - mixed_explicit: hand-coded work-item order independent of the helper. - mixed_auto: unsorted PNATs, work-item order auto-computed via the helper — the production-shaped flow. @@ -1290,10 +1284,8 @@ def test_philox_rounding_unbiased(state_dtype): # * int16: residual std ~1e-4 → SE ~9e-8 (very tight bound) # * int8: residual std ~3e-2 → SE ~2e-5 # * fp8: residual std ~1e-1 → SE ~9e-5 (loosest, magnitude-driven) - # The previous fixed-1e-5 threshold was below SE for int8/fp8 and would - # always fail by chance. Note the |sr|<|det| fallback was also dropped: - # on Gaussian (symmetric) inputs RN's bias is ~0 by symmetry, so SR vs RN - # is just two unbiased estimators racing — unreliable as a unbias test. + # A fixed absolute threshold is below SE for int8/fp8. Gaussian inputs + # also make RN nearly unbiased, so |sr| < |det| is not reliable here. se_sr = stochastic_std / (num_nonzero ** 0.5) K = 4 assert abs(stochastic_mean) < K * se_sr, ( @@ -1305,22 +1297,9 @@ def test_philox_rounding_unbiased(state_dtype): ) -# HEADS_PER_BLOCK > 1 test. The default heuristic only picks HPB > 1 at large -# total_heads (>= 256-512), which the main test with batch=2 never reaches. -# This test overrides _heads_per_block to exercise multi-head precompute tiles. -# -# Beyond OUTPUT and STATE checks, this also asserts the kernel's WRITE -# CONTRACT on every cache buffer (old_x, old_B, old_dt, old_dA_cumsum) -# across a sweep of PNATs covering nowrite (PNAT=0,1,T,max_window-T-1, -# max_window-T) and write (PNAT=max_window-T+1, max_window-1) paths. The -# old_dA_cumsum check is the one that originally hid the -# continuous-across-nowrite bug — direct verification prevents regression. -# Both `rectangle_for_nowrite` arms are exercised explicitly (don't rely on -# tuning). -# -# Configs: (nheads=16, ngroups=1) and (nheads=32, ngroups=2) both have -# heads_per_group=16. The heuristic caps HPB at min(2|4, hpg), so HPB=2, 4. -@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +# HEADS_PER_BLOCK coverage for multi-head precompute tiles and non-power +# heads_per_group, where the wrapper must keep each tile within one B/C group. +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _HEADS_PER_BLOCK_CONFIGS) @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) @pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) @@ -1355,13 +1334,6 @@ def test_replay_heads_per_block( device = "cuda" dtype = torch.bfloat16 - if nheads % heads_per_block != 0: - pytest.skip(f"nheads ({nheads}) not divisible by heads_per_block ({heads_per_block})") - if heads_per_block > nheads // ngroups: - pytest.skip( - f"heads_per_block ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})" - ) - torch.manual_seed(42) A_base = -torch.rand(nheads, device=device) - 0.5 @@ -1486,7 +1458,7 @@ def test_replay_heads_per_block( dt2_proc = F.softplus(dt2_base.float() + dt_bias_base.float()[None, None, :]) dA_cumsum2_step = torch.cumsum(A_base.float()[None, None, :] * dt2_proc, dim=1) - # Reference state walk identical to the original test. + # Build reference by replaying old history, then this step's tokens. ref_state_f32 = state0.float().clone() for slot in range(batch): if pnat_list[slot] > 0: @@ -1545,7 +1517,7 @@ def test_replay_heads_per_block( _heads_per_block=heads_per_block, ) - # ---------------- Output + state checks (existing coverage) ------------- + # ---------------- Output + state checks --------------------------------- torch.testing.assert_close( test_out, ref_out, @@ -1651,7 +1623,7 @@ def test_replay_heads_per_block( # WRITE: per-step cumsum starting from 0 (fresh staging buf). # NOWRITE: continuous — cumsum offset by the prefix value at # old_dA_cumsum_pre[slot, active_buf, head, PNAT-1] (or 0 if PNAT=0). - # This is the direct regression assertion for the bug just fixed. + # Verifies dA_cumsum continuity across no-write appends. expected_old_dAcs_slot = old_dA_cumsum_pre[slot].clone() step_cumsum = dA_cumsum2_step[slot].T # (nheads, T) if is_write: @@ -1688,7 +1660,7 @@ def test_replay_heads_per_block( # `rectangle_nowrite` forces the rectangle vs non-rectangle nowrite path # (via `mode="persistent_main"` + `rectangle_for_nowrite=…` kwargs) # rather than relying on the tuning table's mode pick. -@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _HEADS_PER_BLOCK_CONFIGS) @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("T", [6, 16], ids=["T6", "T16"]) @pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) @@ -1711,11 +1683,6 @@ def test_replay_heads_per_block_multistep( dtype = torch.bfloat16 n_steps = 8 - if nheads % heads_per_block != 0: - pytest.skip(f"nheads ({nheads}) not divisible by HPB ({heads_per_block})") - if heads_per_block > nheads // ngroups: - pytest.skip(f"HPB ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})") - torch.manual_seed(42) A_base = -torch.rand(nheads, device=device) - 0.5 From c1d4c5a543ea85e1da61605ac1b561e1792e68dd Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 12:06:16 -0700 Subject: [PATCH 66/89] mamba benchmark: restore FlashInfer replay baseline Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 4 +- .../flashinfer_checkpointing_ssu_pr3324.py | 330 +++ .../csrc/checkpointing_ssu.cu | 546 ++++ .../checkpointing_ssu_customize_config.jinja | 38 + .../csrc/checkpointing_ssu_jit_binding.cu | 53 + .../csrc/checkpointing_ssu_kernel_inst.cu | 12 + .../include/flashinfer/exception.h | 106 + .../flashinfer/mamba/checkpointing_ssu.cuh | 154 ++ .../include/flashinfer/mamba/common.cuh | 208 ++ .../include/flashinfer/mamba/conversion.cuh | 474 ++++ .../mamba/kernel_checkpointing_ssu.cuh | 1117 +++++++++ .../mamba/kernel_checkpointing_ssu_8bit.cuh | 1495 +++++++++++ .../mamba/kernel_checkpointing_ssu_common.cuh | 1671 +++++++++++++ .../mamba/launch_checkpointing_ssu.cuh | 181 ++ .../flashinfer/mamba/ssu_mtp_common.cuh | 144 ++ .../include/flashinfer/utils.cuh | 539 ++++ .../include/flashinfer/vec_dtypes.cuh | 2194 +++++++++++++++++ 17 files changed, 9265 insertions(+), 1 deletion(-) create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_customize_config.jinja create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh create mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 7ea8524a83c8..9dbaa8eaf73c 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -2979,9 +2979,11 @@ def _run_pr3324_baseline(): ) else: x_call, B_call, C_call = x, B, C + # PR3324 predates double-buffered old_x. The benchmark keeps + # cache_buf_idx at zero, so buffer 0 is the active baseline view. baseline_fn( state_work, - old_x_work, + old_x_work[:, 0], old_B_work, old_dt_work, old_dA_cumsum_work, diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py new file mode 100644 index 000000000000..86f52d35cfda --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py @@ -0,0 +1,330 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Benchmark-only wrapper for FlashInfer PR 3324 checkpointing SSU. + +The vendored CUDA/C++ sources in ``flashinfer_checkpointing_ssu_pr3324/`` +come from https://github.com/flashinfer-ai/flashinfer/pull/3324. +""" + +from __future__ import annotations + +import functools +import os +from pathlib import Path +from typing import Optional + +import jinja2 +import torch + +from flashinfer.compilation_context import CompilationContext +from flashinfer.jit import env as jit_env +from flashinfer.jit.core import JitSpec, gen_jit_spec +from flashinfer.jit.utils import write_if_different + +_ROOT = Path(__file__).resolve().parent / "flashinfer_checkpointing_ssu_pr3324" +_CSRC_DIR = _ROOT / "csrc" +_INCLUDE_DIR = _ROOT / "include" +_SUPPORTED_MAJORS = [8, 9, 10, 11, 12] + +_DTYPE_MAP = { + torch.float16: "half", + torch.bfloat16: "nv_bfloat16", + torch.float32: "float", + torch.int8: "int8_t", + torch.int16: "int16_t", + torch.int32: "int32_t", + torch.int64: "int64_t", + torch.float8_e4m3fn: "__nv_fp8_e4m3", +} + +_FILENAME_SAFE_DTYPE_MAP = { + torch.float16: "f16", + torch.bfloat16: "bf16", + torch.float32: "f32", + torch.int8: "i8", + torch.int16: "i16", + torch.int32: "i32", + torch.int64: "i64", + torch.float8_e4m3fn: "e4m3", +} + + +def _arch_flags() -> list[str]: + context = CompilationContext() + return context.get_nvcc_flags_list(supported_major_versions=_SUPPORTED_MAJORS) + + +def _uri( + state_dtype: torch.dtype, + input_dtype: torch.dtype, + dt_dtype: torch.dtype, + weight_dtype: torch.dtype, + matrix_a_dtype: torch.dtype, + state_index_dtype: torch.dtype, + state_scale_dtype: Optional[torch.dtype], + dim: int, + dstate: int, + npredicted: int, + max_window: int, + heads_per_group: int, + philox_rounds: int, + enable_pdl: bool, +) -> str: + dtype = _FILENAME_SAFE_DTYPE_MAP + uri = ( + "trtllm_pr3324_checkpointing_ssu_v4_" + f"s_{dtype[state_dtype]}_i_{dtype[input_dtype]}_dt_{dtype[dt_dtype]}_" + f"w_{dtype[weight_dtype]}_a_{dtype[matrix_a_dtype]}_" + f"si_{dtype[state_index_dtype]}_d_{dim}_ds_{dstate}_" + f"np_{npredicted}_mw_{max_window}_hpg_{heads_per_group}" + ) + if state_scale_dtype is not None: + uri += f"_sc_{dtype[state_scale_dtype]}" + if philox_rounds > 0: + uri += f"_pr_{philox_rounds}" + if enable_pdl: + uri += "_pdl" + return uri + + +def _gen_module( + state_dtype: torch.dtype, + input_dtype: torch.dtype, + dt_dtype: torch.dtype, + weight_dtype: torch.dtype, + matrix_a_dtype: torch.dtype, + state_index_dtype: torch.dtype, + state_scale_dtype: Optional[torch.dtype], + dim: int, + dstate: int, + npredicted: int, + max_window: int, + heads_per_group: int, + philox_rounds: int, + enable_pdl: bool, +) -> JitSpec: + uri = _uri( + state_dtype, + input_dtype, + dt_dtype, + weight_dtype, + matrix_a_dtype, + state_index_dtype, + state_scale_dtype, + dim, + dstate, + npredicted, + max_window, + heads_per_group, + philox_rounds, + enable_pdl, + ) + gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri + os.makedirs(gen_directory, exist_ok=True) + + with open(_CSRC_DIR / "checkpointing_ssu_customize_config.jinja") as file: + config_template = jinja2.Template(file.read()) + + state_scale_type = ( + _DTYPE_MAP[state_scale_dtype] if state_scale_dtype is not None else "void" + ) + config = config_template.render( + state_dtype=_DTYPE_MAP[state_dtype], + input_dtype=_DTYPE_MAP[input_dtype], + dt_dtype=_DTYPE_MAP[dt_dtype], + weight_dtype=_DTYPE_MAP[weight_dtype], + matrixA_dtype=_DTYPE_MAP[matrix_a_dtype], + stateIndex_dtype=_DTYPE_MAP[state_index_dtype], + state_scale_type=state_scale_type, + dim=dim, + dstate=dstate, + npredicted=npredicted, + max_window=max_window, + heads_per_group=heads_per_group, + philox_rounds=philox_rounds, + enable_pdl="true" if enable_pdl else "false", + ) + write_if_different(gen_directory / "checkpointing_ssu_config.inc", config) + + source_paths = [] + for filename in ( + "checkpointing_ssu.cu", + "checkpointing_ssu_kernel_inst.cu", + "checkpointing_ssu_jit_binding.cu", + ): + source_path = _CSRC_DIR / filename + dest_path = gen_directory / filename + source_paths.append(dest_path) + with open(source_path) as file: + write_if_different(dest_path, file.read()) + + return gen_jit_spec( + uri, + source_paths, + extra_cuda_cflags=_arch_flags(), + extra_include_paths=[_INCLUDE_DIR], + ) + + +@functools.cache +def _get_module( + state_dtype: torch.dtype, + input_dtype: torch.dtype, + dt_dtype: torch.dtype, + weight_dtype: torch.dtype, + matrix_a_dtype: torch.dtype, + state_index_dtype: torch.dtype, + state_scale_dtype: Optional[torch.dtype], + dim: int, + dstate: int, + npredicted: int, + max_window: int, + heads_per_group: int, + philox_rounds: int, + enable_pdl: bool, +): + return _gen_module( + state_dtype, + input_dtype, + dt_dtype, + weight_dtype, + matrix_a_dtype, + state_index_dtype, + state_scale_dtype, + dim, + dstate, + npredicted, + max_window, + heads_per_group, + philox_rounds, + enable_pdl, + ).build_and_load() + + +def checkpointing_ssu( + state: torch.Tensor, + old_x: torch.Tensor, + old_B: torch.Tensor, + old_dt: torch.Tensor, + old_cumAdt: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + out: torch.Tensor, + D: Optional[torch.Tensor] = None, + z: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + dt_softplus: bool = False, + state_batch_indices: Optional[torch.Tensor] = None, + pad_slot_id: int = -1, + state_scale: Optional[torch.Tensor] = None, + rand_seed: Optional[torch.Tensor] = None, + philox_rounds: int = 10, + d_split: Optional[int] = None, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, + enable_pdl: bool = False, +) -> torch.Tensor: + quantized_state_dtypes = (torch.int8, torch.float8_e4m3fn) + if state.dtype in quantized_state_dtypes: + if state_scale is None: + raise ValueError(f"state dtype {state.dtype} requires state_scale") + elif state_scale is not None: + raise ValueError( + f"state_scale must be None for non-quantized state dtype {state.dtype}" + ) + if cu_seqlens is not None: + npredicted = max_seqlen if max_seqlen is not None else old_x.size(1) + else: + if max_seqlen is not None: + raise ValueError("max_seqlen is only valid with cu_seqlens") + npredicted = x.size(1) + + max_window = old_x.size(1) + if max_window > 16: + raise ValueError(f"PR3324 checkpointing SSU supports max_window <= 16, got {max_window}") + if npredicted > max_window: + raise ValueError(f"npredicted ({npredicted}) must be <= max_window ({max_window})") + + if d_split is None: + d_split = 1 + if state.dtype in quantized_state_dtypes and d_split != 1: + raise ValueError(f"8-bit state requires d_split=1, got {d_split}") + + state_index_dtype = ( + state_batch_indices.dtype if state_batch_indices is not None else torch.int32 + ) + nheads = state.size(1) + ngroups = B.size(-2) + if nheads % ngroups != 0: + raise ValueError(f"nheads ({nheads}) must be divisible by ngroups ({ngroups})") + heads_per_group = nheads // ngroups + + if rand_seed is None: + philox_rounds = 0 + elif philox_rounds <= 0: + raise ValueError(f"philox_rounds must be > 0 with rand_seed, got {philox_rounds}") + + weight_dtype = ( + D.dtype if D is not None else (dt_bias.dtype if dt_bias is not None else dt.dtype) + ) + + module = _get_module( + state.dtype, + x.dtype, + dt.dtype, + weight_dtype, + A.dtype, + state_index_dtype, + state_scale.dtype if state_scale is not None else None, + state.size(2), + state.size(3), + npredicted, + max_window, + heads_per_group, + philox_rounds, + enable_pdl, + ) + module.checkpointing_ssu( + state, + x, + dt, + A, + B, + C, + out, + old_x, + old_B, + old_dt, + old_cumAdt, + cache_buf_idx, + prev_num_accepted_tokens, + D, + z, + dt_bias, + dt_softplus, + state_batch_indices, + pad_slot_id, + state_scale, + rand_seed, + d_split, + cu_seqlens, + ) + return out diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu new file mode 100644 index 000000000000..df0e7e189274 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu @@ -0,0 +1,546 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * 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. + */ +// clang-format off +// config.inc MUST come before the header: it defines DIM, DSTATE, NPREDICTED, +// MAX_WINDOW constexprs that the header's function templates rely on. +#include "checkpointing_ssu_config.inc" +#include +#include +// clang-format on +#include "tvm_ffi_utils.h" + +using namespace flashinfer; +using tvm::ffi::Optional; + +namespace flashinfer::mamba::checkpointing { + +void checkpointing_ssu( + TensorView state, // (state_cache_size, nheads, dim, dstate) + TensorView x, // (batch, NPREDICTED, nheads, dim) / (1, total_tokens, nheads, dim) under varlen + TensorView dt, // (batch, NPREDICTED, nheads, dim) tie_hdim / (1, total_tokens, nheads, dim) + TensorView A, // (nheads, dim, dstate) tie_hdim + TensorView B, // (batch, NPREDICTED, ngroups, dstate) / (1, total_tokens, ngroups, dstate) + TensorView C, // same as B + TensorView output, // same layout as x + // Cache tensors + TensorView old_x, // (state_cache_size, MAX_WINDOW, nheads, dim) + TensorView old_B, // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) + TensorView old_dt, // (state_cache_size, 2, nheads, MAX_WINDOW) f32 + TensorView old_cumAdt, // (state_cache_size, 2, nheads, MAX_WINDOW) f32 + TensorView cache_buf_idx, // (state_cache_size,) int32 + TensorView prev_num_accepted, // (state_cache_size,) int32 + // Optional tensors + Optional D, // (nheads, dim) + Optional z, // same layout as x + Optional dt_bias, // (nheads, dim) tie_hdim + bool dt_softplus, + Optional state_batch_indices, // (batch,) int32 + int64_t pad_slot_id, + Optional state_scale, // (state_cache_size, nheads, dim) f32 + Optional rand_seed, // single int64 + int64_t d_split, // v12 §59: per-head DIM split factor (1, 2, or 4) + Optional cu_seqlens) { // (batch+1,) int32, varlen mode + + bool const is_varlen = cu_seqlens.has_value(); + + // ── Extract dimensions ── + auto const state_cache_size = state.size(0); + auto const nheads = state.size(1); + auto const dim = state.size(2); + auto const dstate = state.size(3); + auto const max_window = old_x.size(1); + auto const ngroups = B.size(2); + + // In non-varlen mode, batch = x.size(0) and npredicted = x.size(1) (the + // 4D batched layout). In varlen, the JIT compile-time NPREDICTED is the + // max seq_len the caller commits to; the wrapper stamped it into the JIT + // URI from `max_seqlen` and we read it back as the constexpr `NPREDICTED` + // (validated against runtime cu_seqlens on the host side below). `batch` + // = number of sequences = `cu_seqlens.size(0) - 1`. + int64_t batch; + int64_t npredicted; + if (is_varlen) { + auto const& cs = cu_seqlens.value(); + CHECK_CUDA(cs); + CHECK_DIM(1, cs); + FLASHINFER_CHECK( + cs.size(0) >= 2, + "cu_seqlens must have shape (batch+1,) with batch >= 1, got size(0)=", cs.size(0)); + FLASHINFER_CHECK(cs.dtype().code == kDLInt && cs.dtype().bits == 32, + "cu_seqlens must be int32"); + CHECK_CONTIGUOUS(cs); + batch = cs.size(0) - 1; + npredicted = NPREDICTED; // JIT-stamped — wrapper ensures max(seq_lens) <= NPREDICTED. + } else { + batch = x.size(0); + npredicted = x.size(1); + } + + // ── JIT compile-time / runtime cross-check ── + // NPREDICTED and MAX_WINDOW are JIT compile-time constants stamped by the + // wrapper. In non-varlen NPREDICTED = x.shape[1]; in varlen NPREDICTED = + // user-supplied `max_seqlen` (upper bound on every cu_seqlens diff). + FLASHINFER_CHECK(npredicted == NPREDICTED, is_varlen ? "max_seqlen=" : "x.size(1)=", npredicted, + " must equal JIT NPREDICTED=", NPREDICTED); + FLASHINFER_CHECK(max_window == MAX_WINDOW, "old_x.size(1)=", max_window, + " must equal JIT MAX_WINDOW=", MAX_WINDOW); + FLASHINFER_CHECK(npredicted <= max_window, "npredicted=", npredicted, + " must be <= max_window=", max_window); + + // ── Validate state ── + CHECK_CUDA(state); + CHECK_DIM(4, state); + { + auto s = state.strides(); + auto sz = state.sizes(); + FLASHINFER_CHECK(s[3] == 1, "state dim 3 (dstate) must have stride 1, got ", s[3]); + FLASHINFER_CHECK(s[2] == sz[3], "state dim 2 (dim) must be contiguous with dim 3, got stride ", + s[2], " expected ", sz[3]); + FLASHINFER_CHECK(s[1] == sz[2] * sz[3], + "state dim 1 (nheads) must be contiguous with dim 2, got stride ", s[1], + " expected ", sz[2] * sz[3]); + } + + // ── Validate x ── + // Non-varlen: shape (batch, NPREDICTED, nheads, dim). + // Varlen : shape (1, total_tokens, nheads, dim) — batch axis collapsed, + // token axis is the outer iteration. The kernel reads x via + // `bos * x_stride_token + …` so x_stride_token is the per-token + // stride in either layout (= nheads*dim for contig). + CHECK_CUDA(x); + CHECK_DIM(4, x); + if (is_varlen) { + FLASHINFER_CHECK(x.size(0) == 1, "varlen: x.size(0)=", x.size(0), " must be 1"); + } else { + FLASHINFER_CHECK(x.size(0) == batch, "x.size(0)=", x.size(0), " must equal batch=", batch); + FLASHINFER_CHECK(x.size(1) == npredicted, "x.size(1)=", x.size(1), + " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(x.size(2) == nheads, "x.size(2)=", x.size(2), " must equal nheads=", nheads); + FLASHINFER_CHECK(x.size(3) == dim, "x.size(3)=", x.size(3), " must equal dim=", dim); + CHECK_LAST_DIM_CONTIGUOUS(x); + FLASHINFER_CHECK(x.stride(2) == dim, "x.stride(2)=", x.stride(2), " must equal dim=", dim, + " ((nheads, dim) must be contiguous)"); + + // In varlen, all per-token tensors share the flattened token axis — use + // x.size(1) as the canonical total_tokens and cross-check the others below. + int64_t const total_tokens = is_varlen ? x.size(1) : 0; + + // ── Validate dt ── + CHECK_CUDA(dt); + CHECK_DIM(4, dt); + if (is_varlen) { + FLASHINFER_CHECK(dt.size(0) == 1, "varlen: dt.size(0)=", dt.size(0), " must be 1"); + FLASHINFER_CHECK(dt.size(1) == total_tokens, "varlen: dt.size(1)=", dt.size(1), + " must equal x.size(1)=", total_tokens); + } else { + FLASHINFER_CHECK(dt.size(0) == batch, "dt.size(0)=", dt.size(0), " must equal batch=", batch); + FLASHINFER_CHECK(dt.size(1) == npredicted, "dt.size(1)=", dt.size(1), + " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(dt.size(2) == nheads, "dt.size(2)=", dt.size(2), " must equal nheads=", nheads); + FLASHINFER_CHECK(dt.size(3) == dim, "dt.size(3)=", dt.size(3), " must equal dim=", dim); + FLASHINFER_CHECK(dt.stride(2) == 1, "dt.stride(2) must be 1 (tie_hdim), got ", dt.stride(2)); + FLASHINFER_CHECK(dt.stride(3) == 0, "dt.stride(3) must be 0 (tie_hdim), got ", dt.stride(3)); + + // ── Validate A: (nheads, dim, dstate) tie_hdim ── + CHECK_CUDA(A); + CHECK_DIM(3, A); + FLASHINFER_CHECK(A.size(0) == nheads, "A.size(0)=", A.size(0), " must equal nheads=", nheads); + FLASHINFER_CHECK(A.size(1) == dim, "A.size(1)=", A.size(1), " must equal dim=", dim); + FLASHINFER_CHECK(A.size(2) == dstate, "A.size(2)=", A.size(2), " must equal dstate=", dstate); + FLASHINFER_CHECK(A.stride(0) == 1, "A.stride(0) must be 1, got ", A.stride(0)); + FLASHINFER_CHECK(A.stride(1) == 0, "A.stride(1) must be 0 (tie_hdim), got ", A.stride(1)); + FLASHINFER_CHECK(A.stride(2) == 0, "A.stride(2) must be 0 (tie_hdim), got ", A.stride(2)); + + // ── Validate B ── + CHECK_CUDA(B); + CHECK_DIM(4, B); + if (is_varlen) { + FLASHINFER_CHECK(B.size(0) == 1, "varlen: B.size(0)=", B.size(0), " must be 1"); + FLASHINFER_CHECK(B.size(1) == total_tokens, "varlen: B.size(1)=", B.size(1), + " must equal x.size(1)=", total_tokens); + } else { + FLASHINFER_CHECK(B.size(0) == batch, "B.size(0)=", B.size(0), " must equal batch=", batch); + FLASHINFER_CHECK(B.size(1) == npredicted, "B.size(1)=", B.size(1), + " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(B.size(3) == dstate, "B.size(3)=", B.size(3), " must equal dstate=", dstate); + CHECK_LAST_DIM_CONTIGUOUS(B); + FLASHINFER_CHECK(B.stride(2) == dstate, "B.stride(2)=", B.stride(2), + " must equal dstate=", dstate, " ((ngroups, dstate) must be contiguous)"); + FLASHINFER_CHECK(nheads % ngroups == 0, "nheads=", nheads, + " must be divisible by ngroups=", ngroups); + + // ── Validate C ── + CHECK_CUDA(C); + CHECK_DIM(4, C); + if (is_varlen) { + FLASHINFER_CHECK(C.size(0) == 1, "varlen: C.size(0)=", C.size(0), " must be 1"); + FLASHINFER_CHECK(C.size(1) == total_tokens, "varlen: C.size(1)=", C.size(1), + " must equal x.size(1)=", total_tokens); + } else { + FLASHINFER_CHECK(C.size(0) == batch, "C.size(0)=", C.size(0), " must equal batch=", batch); + FLASHINFER_CHECK(C.size(1) == npredicted, "C.size(1)=", C.size(1), + " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(C.size(2) == ngroups, "C.size(2)=", C.size(2), " must equal ngroups=", ngroups); + FLASHINFER_CHECK(C.size(3) == dstate, "C.size(3)=", C.size(3), " must equal dstate=", dstate); + CHECK_LAST_DIM_CONTIGUOUS(C); + FLASHINFER_CHECK(C.stride(2) == dstate, "C.stride(2)=", C.stride(2), + " must equal dstate=", dstate, " ((ngroups, dstate) must be contiguous)"); + + // ── Validate output ── + CHECK_CUDA(output); + CHECK_DIM(4, output); + if (is_varlen) { + FLASHINFER_CHECK(output.size(0) == 1, "varlen: output.size(0)=", output.size(0), " must be 1"); + FLASHINFER_CHECK(output.size(1) == total_tokens, "varlen: output.size(1)=", output.size(1), + " must equal x.size(1)=", total_tokens); + } else { + FLASHINFER_CHECK(output.size(0) == batch, "output.size(0)=", output.size(0), + " must equal batch=", batch); + FLASHINFER_CHECK(output.size(1) == npredicted, "output.size(1)=", output.size(1), + " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(output.size(2) == nheads, "output.size(2)=", output.size(2), + " must equal nheads=", nheads); + FLASHINFER_CHECK(output.size(3) == dim, "output.size(3)=", output.size(3), + " must equal dim=", dim); + CHECK_LAST_DIM_CONTIGUOUS(output); + FLASHINFER_CHECK(output.stride(2) == dim, "output.stride(2)=", output.stride(2), + " must equal dim=", dim, " ((nheads, dim) must be contiguous)"); + + // ── Validate cache tensors ── + // old_x: kernel uses `head * DIM + d_tile_off` → (nheads, dim) contig. + CHECK_CUDA(old_x); + CHECK_DIM(4, old_x); // (state_cache_size, MAX_WINDOW, nheads, dim) + FLASHINFER_CHECK(old_x.size(0) == state_cache_size, "old_x.size(0)=", old_x.size(0), + " must equal state_cache_size=", state_cache_size); + FLASHINFER_CHECK(old_x.size(1) == max_window, "old_x.size(1)=", old_x.size(1), + " must equal max_window=", max_window); + FLASHINFER_CHECK(old_x.size(2) == nheads, "old_x.size(2)=", old_x.size(2), + " must equal nheads=", nheads); + FLASHINFER_CHECK(old_x.size(3) == dim, "old_x.size(3)=", old_x.size(3), " must equal dim=", dim); + CHECK_LAST_DIM_CONTIGUOUS(old_x); + FLASHINFER_CHECK(old_x.stride(2) == dim, "old_x.stride(2)=", old_x.stride(2), + " must equal dim=", dim, " ((nheads, dim) must be contiguous)"); + + // old_B: kernel uses `group_idx * DSTATE` → (ngroups, dstate) contig. + CHECK_CUDA(old_B); + CHECK_DIM(5, old_B); // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) + FLASHINFER_CHECK(old_B.size(0) == state_cache_size, "old_B.size(0)=", old_B.size(0), + " must equal state_cache_size=", state_cache_size); + FLASHINFER_CHECK(old_B.size(1) == 2, "old_B.size(1) must be 2 (double-buffered), got ", + old_B.size(1)); + FLASHINFER_CHECK(old_B.size(2) == max_window, "old_B.size(2)=", old_B.size(2), + " must equal max_window=", max_window); + FLASHINFER_CHECK(old_B.size(3) == ngroups, "old_B.size(3)=", old_B.size(3), + " must equal ngroups=", ngroups); + FLASHINFER_CHECK(old_B.size(4) == dstate, "old_B.size(4)=", old_B.size(4), + " must equal dstate=", dstate); + CHECK_LAST_DIM_CONTIGUOUS(old_B); + FLASHINFER_CHECK(old_B.stride(3) == dstate, "old_B.stride(3)=", old_B.stride(3), + " must equal dstate=", dstate, " ((ngroups, dstate) must be contiguous)"); + + // old_dt: kernel only assumes last-dim contig (head row). + CHECK_CUDA(old_dt); + CHECK_DIM(4, old_dt); // (state_cache_size, 2, nheads, MAX_WINDOW) + FLASHINFER_CHECK(old_dt.size(0) == state_cache_size, "old_dt.size(0)=", old_dt.size(0), + " must equal state_cache_size=", state_cache_size); + FLASHINFER_CHECK(old_dt.size(1) == 2, "old_dt.size(1) must be 2, got ", old_dt.size(1)); + FLASHINFER_CHECK(old_dt.size(2) == nheads, "old_dt.size(2)=", old_dt.size(2), + " must equal nheads=", nheads); + FLASHINFER_CHECK(old_dt.size(3) == max_window, "old_dt.size(3)=", old_dt.size(3), + " must equal max_window=", max_window); + CHECK_LAST_DIM_CONTIGUOUS(old_dt); + + // old_cumAdt: same as old_dt. + CHECK_CUDA(old_cumAdt); + CHECK_DIM(4, old_cumAdt); // (state_cache_size, 2, nheads, MAX_WINDOW) + FLASHINFER_CHECK(old_cumAdt.size(0) == state_cache_size, + "old_cumAdt.size(0)=", old_cumAdt.size(0), + " must equal state_cache_size=", state_cache_size); + FLASHINFER_CHECK(old_cumAdt.size(1) == 2, "old_cumAdt.size(1) must be 2, got ", + old_cumAdt.size(1)); + FLASHINFER_CHECK(old_cumAdt.size(2) == nheads, "old_cumAdt.size(2)=", old_cumAdt.size(2), + " must equal nheads=", nheads); + FLASHINFER_CHECK(old_cumAdt.size(3) == max_window, "old_cumAdt.size(3)=", old_cumAdt.size(3), + " must equal max_window=", max_window); + CHECK_LAST_DIM_CONTIGUOUS(old_cumAdt); + + CHECK_CUDA(cache_buf_idx); + CHECK_DIM(1, cache_buf_idx); + FLASHINFER_CHECK(cache_buf_idx.size(0) == state_cache_size, + "cache_buf_idx.size(0)=", cache_buf_idx.size(0), + " must equal state_cache_size=", state_cache_size); + CHECK_CONTIGUOUS(cache_buf_idx); + + CHECK_CUDA(prev_num_accepted); + CHECK_DIM(1, prev_num_accepted); + FLASHINFER_CHECK(prev_num_accepted.size(0) == state_cache_size, + "prev_num_accepted.size(0)=", prev_num_accepted.size(0), + " must equal state_cache_size=", state_cache_size); + CHECK_CONTIGUOUS(prev_num_accepted); + + // ── Validate optional D ── + if (D.has_value()) { + auto& Dv = D.value(); + CHECK_CUDA(Dv); + CHECK_DIM(2, Dv); + FLASHINFER_CHECK(Dv.size(0) == nheads, "D.size(0)=", Dv.size(0), " must equal nheads=", nheads); + FLASHINFER_CHECK(Dv.size(1) == dim, "D.size(1)=", Dv.size(1), " must equal dim=", dim); + FLASHINFER_CHECK(Dv.stride(0) == 1, "D.stride(0) must be 1 (tie_hdim), got ", Dv.stride(0)); + FLASHINFER_CHECK(Dv.stride(1) == 0, "D.stride(1) must be 0 (tie_hdim), got ", Dv.stride(1)); + } + + // ── Validate optional dt_bias ── + if (dt_bias.has_value()) { + auto& db = dt_bias.value(); + CHECK_CUDA(db); + CHECK_DIM(2, db); + FLASHINFER_CHECK(db.size(0) == nheads, "dt_bias.size(0)=", db.size(0), + " must equal nheads=", nheads); + FLASHINFER_CHECK(db.size(1) == dim, "dt_bias.size(1)=", db.size(1), " must equal dim=", dim); + FLASHINFER_CHECK(db.stride(0) == 1, "dt_bias.stride(0) must be 1 (tie_hdim), got ", + db.stride(0)); + FLASHINFER_CHECK(db.stride(1) == 0, "dt_bias.stride(1) must be 0 (tie_hdim), got ", + db.stride(1)); + } + + // ── Validate optional z: same layout/contig rules as x ── + if (z.has_value()) { + auto& zv = z.value(); + CHECK_CUDA(zv); + CHECK_DIM(4, zv); + if (is_varlen) { + FLASHINFER_CHECK(zv.size(0) == 1, "varlen: z.size(0)=", zv.size(0), " must be 1"); + FLASHINFER_CHECK(zv.size(1) == total_tokens, "varlen: z.size(1)=", zv.size(1), + " must equal x.size(1)=", total_tokens); + } else { + FLASHINFER_CHECK(zv.size(0) == batch, "z.size(0)=", zv.size(0), " must equal batch=", batch); + FLASHINFER_CHECK(zv.size(1) == npredicted, "z.size(1)=", zv.size(1), + " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(zv.size(2) == nheads, "z.size(2)=", zv.size(2), " must equal nheads=", nheads); + FLASHINFER_CHECK(zv.size(3) == dim, "z.size(3)=", zv.size(3), " must equal dim=", dim); + CHECK_LAST_DIM_CONTIGUOUS(zv); + FLASHINFER_CHECK(zv.stride(2) == dim, "z.stride(2)=", zv.stride(2), " must equal dim=", dim, + " ((nheads, dim) must be contiguous)"); + } + + // ── Validate optional state_batch_indices ── + if (state_batch_indices.has_value()) { + auto& sbi = state_batch_indices.value(); + CHECK_CUDA(sbi); + CHECK_DIM(1, sbi); + FLASHINFER_CHECK(sbi.size(0) == batch, "state_batch_indices.size(0)=", sbi.size(0), + " must equal batch=", batch); + CHECK_CONTIGUOUS(sbi); + } + + // ── Validate optional state_scale: (state_cache_size, nheads, dim) ── + // Inner two dims (nheads, dim) must be contiguous; only batch stride is + // parameterized in the params struct. + if (state_scale.has_value()) { + auto const& ss = state_scale.value(); + CHECK_CUDA(ss); + CHECK_DIM(3, ss); + FLASHINFER_CHECK(ss.size(0) == state_cache_size, "state_scale.size(0)=", ss.size(0), + " must equal state_cache_size=", state_cache_size); + FLASHINFER_CHECK(ss.size(1) == nheads, "state_scale.size(1)=", ss.size(1), + " must equal nheads=", nheads); + FLASHINFER_CHECK(ss.size(2) == dim, "state_scale.size(2)=", ss.size(2), + " must equal dim=", dim); + FLASHINFER_CHECK(ss.stride(2) == 1, "state_scale.stride(2) must be 1, got ", ss.stride(2)); + FLASHINFER_CHECK(ss.stride(1) == dim, "state_scale.stride(1)=", ss.stride(1), + " must equal dim=", dim, " ((nheads, dim) must be contiguous)"); + } + + // ── Dtype consistency ── + // input_dtype = x.dtype; all activation tensors (B, C, output, z, old_x, + // old_B) and the state cache's "input-side" mirrors must match it. + // weight_dtype = D.dtype = dt_bias.dtype (kernel template sees one + // weight_t for both). + // Cache scalar tensors have fixed dtypes hardcoded in the kernel. + { + auto input_dtype = x.dtype(); + FLASHINFER_CHECK(B.dtype() == input_dtype, "B.dtype must match x.dtype"); + FLASHINFER_CHECK(C.dtype() == input_dtype, "C.dtype must match x.dtype"); + FLASHINFER_CHECK(output.dtype() == input_dtype, "output.dtype must match x.dtype"); + FLASHINFER_CHECK(old_x.dtype() == input_dtype, "old_x.dtype must match x.dtype"); + FLASHINFER_CHECK(old_B.dtype() == input_dtype, "old_B.dtype must match x.dtype"); + if (z.has_value()) { + FLASHINFER_CHECK(z.value().dtype() == input_dtype, "z.dtype must match x.dtype"); + } + if (D.has_value() && dt_bias.has_value()) { + FLASHINFER_CHECK(D.value().dtype() == dt_bias.value().dtype(), + "D.dtype must equal dt_bias.dtype (kernel uses a single weight_t)"); + } + // old_dt / old_cumAdt are produced by this same kernel in f32 and + // consumed back in f32 on the next call. + FLASHINFER_CHECK(old_dt.dtype().code == kDLFloat && old_dt.dtype().bits == 32, + "old_dt must be float32"); + FLASHINFER_CHECK(old_cumAdt.dtype().code == kDLFloat && old_cumAdt.dtype().bits == 32, + "old_cumAdt must be float32"); + // Index tensors used by the kernel as int32 scalars. + FLASHINFER_CHECK(cache_buf_idx.dtype().code == kDLInt && cache_buf_idx.dtype().bits == 32, + "cache_buf_idx must be int32"); + FLASHINFER_CHECK( + prev_num_accepted.dtype().code == kDLInt && prev_num_accepted.dtype().bits == 32, + "prev_num_accepted must be int32"); + if (state_batch_indices.has_value()) { + auto sbi_dt = state_batch_indices.value().dtype(); + FLASHINFER_CHECK(sbi_dt.code == kDLInt && (sbi_dt.bits == 32 || sbi_dt.bits == 64), + "state_batch_indices must be int32 or int64"); + } + if (state_scale.has_value()) { + auto ss_dt = state_scale.value().dtype(); + FLASHINFER_CHECK(ss_dt.code == kDLFloat && ss_dt.bits == 32, "state_scale must be float32"); + } + // Quantized state dtypes (int8, fp8_e4m3fn, ...) require a state_scale + // tensor; non-quantized dtypes must not pass one. Mirrors the Python + // wrapper assertion and matches the kernel's compile-time + // `state_scale_t == void` gating. + { + auto sd = state.dtype(); + bool const is_int8 = (sd.code == kDLInt && sd.bits == 8); + bool const is_fp8 = (sd.code == kDLFloat8_e4m3fn && sd.bits == 8); + bool const is_quantized_state = is_int8 || is_fp8; + if (is_quantized_state) { + FLASHINFER_CHECK(state_scale.has_value(), + "Quantized state.dtype (int8/fp8_e4m3fn) requires a state_scale tensor " + "of shape (state_cache_size, nheads, dim) and dtype float32"); + // The 8-bit replay path uses Layout<_4, _1> (M-shard per warp) which + // needs per-warp M = D_PER_CTA / 4 >= 16 (m16n8 atom M). This forces + // D_PER_CTA >= 64, i.e. d_split == 1. + FLASHINFER_CHECK( + d_split == 1, + "Quantized state.dtype (int8/fp8_e4m3fn) requires d_split=1 (got d_split=", d_split, + "); the M-shard-per-warp replay layout needs D_PER_CTA / 4 >= 16."); + } else { + FLASHINFER_CHECK(!state_scale.has_value(), + "state_scale must be None for non-quantized state.dtype " + "(allowed quantized dtypes: {int8, fp8_e4m3fn})"); + } + } + } + + // ── Populate params ── + CheckpointingSsuParams p; + + // ── Validate d_split (v12 §59) ── + // Allowed for v12: {1, 2}. d_split=4 deferred to v12.x (needs warp-count + // restructure — output MMA `_1×4` layout requires D_PER_CTA ≥ 32). + FLASHINFER_CHECK(d_split == 1 || d_split == 2, "d_split=", d_split, + " must be one of {1, 2} (d_split=4 is deferred to v12.x)"); + FLASHINFER_CHECK(dim % d_split == 0, "dim=", dim, " must be divisible by d_split=", d_split); + FLASHINFER_CHECK(dim / d_split >= 32, "d_split=", d_split, " gives D_PER_CTA=", dim / d_split, + " < 32 (output MMA m16n8 atom floor with _1×4 warp layout)"); + + p.batch = batch; + p.nheads = nheads; + p.dim = dim; + p.dstate = dstate; + p.ngroups = ngroups; + p.state_cache_size = state_cache_size; + p.npredicted = npredicted; + p.max_window = max_window; + p.pad_slot_id = pad_slot_id; + p.d_split = static_cast(d_split); + p.dt_softplus = dt_softplus; + + // Pointers + p.state = state.data_ptr(); + p.x = const_cast(x.data_ptr()); + p.dt = const_cast(dt.data_ptr()); + p.A = const_cast(A.data_ptr()); + p.B = const_cast(B.data_ptr()); + p.C = const_cast(C.data_ptr()); + p.output = output.data_ptr(); + + p.old_x = old_x.data_ptr(); + p.old_B = const_cast(old_B.data_ptr()); + p.old_dt = const_cast(old_dt.data_ptr()); + p.old_cumAdt = const_cast(old_cumAdt.data_ptr()); + p.cache_buf_idx = const_cast(cache_buf_idx.data_ptr()); + p.prev_num_accepted = const_cast(prev_num_accepted.data_ptr()); + + if (D.has_value()) p.D = const_cast(D.value().data_ptr()); + if (z.has_value()) { + p.z = const_cast(z.value().data_ptr()); + // Same seq-dim selection as the rest of the batch-side tensors below. + p.z_stride_seq = z.value().stride(is_varlen ? 1 : 0); + p.z_stride_token = z.value().stride(1); + } + if (dt_bias.has_value()) p.dt_bias = const_cast(dt_bias.value().data_ptr()); + if (state_batch_indices.has_value()) + p.state_batch_indices = const_cast(state_batch_indices.value().data_ptr()); + if (is_varlen) { + p.cu_seqlens = const_cast(cu_seqlens.value().data_ptr()); + } + if (state_scale.has_value()) { + p.state_scale = state_scale.value().data_ptr(); + p.state_scale_stride_seq = state_scale.value().stride(0); + } + if (rand_seed.has_value()) { + auto const& rs = rand_seed.value(); + CHECK_CUDA(rs); + FLASHINFER_CHECK(rs.numel() == 1, "rand_seed must be single-element, got numel=", rs.numel()); + FLASHINFER_CHECK(rs.dtype().code == kDLInt && rs.dtype().bits == 64, "rand_seed must be int64"); + p.rand_seed = static_cast(rs.data_ptr()); + } + + // Strides + p.state_stride_seq = state.stride(0); + + // `*_stride_seq` is the outer iteration stride. Non-varlen iterates over + // dim 0 (per-batch), varlen iterates over dim 1 (per-token) — sequences + // are packed into a single batch in the (1, total_tokens, ...) layout. + // The kernel uses one formula `seq * *_stride_seq` for both modes. + int const seq_dim = is_varlen ? 1 : 0; + p.x_stride_seq = x.stride(seq_dim); + p.x_stride_token = x.stride(1); + p.dt_stride_seq = dt.stride(seq_dim); + p.dt_stride_token = dt.stride(1); + p.B_stride_seq = B.stride(seq_dim); + p.B_stride_token = B.stride(1); + p.C_stride_seq = C.stride(seq_dim); + p.C_stride_token = C.stride(1); + p.out_stride_seq = output.stride(seq_dim); + p.out_stride_token = output.stride(1); + + p.old_x_stride_seq = old_x.stride(0); + p.old_x_stride_token = old_x.stride(1); + p.old_B_stride_seq = old_B.stride(0); + p.old_B_stride_dbuf = old_B.stride(1); + p.old_B_stride_token = old_B.stride(2); + p.old_dt_stride_seq = old_dt.stride(0); + p.old_dt_stride_dbuf = old_dt.stride(1); + p.old_dt_stride_head = old_dt.stride(2); + p.old_cumAdt_stride_seq = old_cumAdt.stride(0); + p.old_cumAdt_stride_dbuf = old_cumAdt.stride(1); + p.old_cumAdt_stride_head = old_cumAdt.stride(2); + + // Launch + ffi::CUDADeviceGuard device_guard(state.device().device_id); + const cudaStream_t stream = get_stream(state.device()); + + launchCheckpointingSsu( + p, stream); +} + +} // namespace flashinfer::mamba::checkpointing diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_customize_config.jinja b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_customize_config.jinja new file mode 100644 index 000000000000..61bf35cf17ca --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_customize_config.jinja @@ -0,0 +1,38 @@ +#pragma once +#include +#include +#include +#include + +using state_t = {{ state_dtype }}; +using input_t = {{ input_dtype }}; +// dt accepted in its native dtype (e.g. bf16) — converted to f32 internally. +// Eliminates the need for a separate dtype conversion kernel launch on the host side. +using dt_t = {{ dt_dtype }}; +using weight_t = {{ weight_dtype }}; +using matrixA_t = {{ matrixA_dtype }}; +using stateIndex_t = {{ stateIndex_dtype }}; +// Type for block-scale decode factors (e.g. float, __half). +// void = no scaling (state_t is used as-is). +using state_scale_t = {{ state_scale_type }}; + +constexpr int DIM = {{ dim }}; +constexpr int DSTATE = {{ dstate }}; +constexpr int NPREDICTED = {{ npredicted }}; +constexpr int MAX_WINDOW = {{ max_window }}; +// nheads / ngroups — JIT-stamped so the kernel compiles only one +// HEADS_PER_GROUP specialization per .so (was 7 via `dispatchRatio`). +// The wrapper computes this from the runtime tensors and selects the +// matching JIT URI; the launcher reads HEADS_PER_GROUP directly without +// a runtime dispatch step. +constexpr int HEADS_PER_GROUP = {{ heads_per_group }}; +// Philox PRNG rounds for stochastic rounding of fp16 state stores. +// 0 = no stochastic rounding; typical value = 10. +constexpr int PHILOX_ROUNDS = {{ philox_rounds }}; +// Programmatic Dependent Launch. When true, the kernel emits the +// griddepcontrol.{wait,launch_dependents} PTX and the load is split around +// `gdc_wait` for cache-load-during-wait overlap. When false, a single-pass +// load_data path is used (no PDL barriers) — matches v21.0 register profile +// and load order. Stamped as a JIT URI key so each (enable_pdl) value +// compiles its own .so; no runtime branch in the kernel binary. +constexpr bool ENABLE_PDL = {{ enable_pdl }}; diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu new file mode 100644 index 000000000000..89034be5c9cc --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu @@ -0,0 +1,53 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * 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. + */ +#include "tvm_ffi_utils.h" + +using tvm::ffi::Optional; + +namespace flashinfer::mamba::checkpointing { + +void checkpointing_ssu( + TensorView state, // (cache, nheads, dim, dstate) + TensorView x, // 4D (batch, T, nheads, dim) or 4D (1, total_tokens, nheads, dim) under varlen + TensorView + dt, // (batch, T, nheads, dim) tie_hdim / (1, total_tokens, nheads, dim) under varlen + TensorView A, // (nheads, dim, dstate) tie_hdim + TensorView B, // (batch, T, ngroups, dstate) / (1, total_tokens, ngroups, dstate) under varlen + TensorView C, // same as B + TensorView output, // same layout as x + // Cache tensors + TensorView old_x, // (cache, T, nheads, dim) + TensorView old_B, // (cache, 2, T, ngroups, dstate) + TensorView old_dt, // (cache, 2, nheads, T) f32 + TensorView old_cumAdt, // (cache, 2, nheads, T) f32 + TensorView cache_buf_idx, // (cache,) int32 + TensorView prev_num_accepted, // (cache,) int32 + // Optional tensors + Optional D, // (nheads, dim) + Optional z, // same layout as x + Optional dt_bias, // (nheads, dim) tie_hdim + bool dt_softplus, + Optional state_batch_indices, // (batch,) int32 + int64_t pad_slot_id, + Optional state_scale, // (cache, nheads, dim) f32 + Optional rand_seed, // single int64 + int64_t d_split, // v12 §59: per-head DIM split factor (1, 2, or 4) + Optional cu_seqlens); // (batch+1,) int32 — varlen mode + +} // namespace flashinfer::mamba::checkpointing + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(checkpointing_ssu, + flashinfer::mamba::checkpointing::checkpointing_ssu); diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu new file mode 100644 index 000000000000..07eaa5e597a5 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu @@ -0,0 +1,12 @@ +// clang-format off +#include "checkpointing_ssu_config.inc" +#include +#include +// clang-format on + +namespace flashinfer::mamba::checkpointing { + +template void launchCheckpointingSsu(CheckpointingSsuParams&, cudaStream_t); + +} // namespace flashinfer::mamba::checkpointing diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h new file mode 100644 index 000000000000..aaaa2b5b3e51 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h @@ -0,0 +1,106 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * 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. + */ +#ifndef FLASHINFER_EXCEPTION_H_ +#define FLASHINFER_EXCEPTION_H_ + +#include +#include +#include + +#define FLASHINFER_ERROR(message) throw flashinfer::Error(__FUNCTION__, __FILE__, __LINE__, message) + +// Base case for empty arguments +inline void write_to_stream(std::ostringstream& oss) { + // No-op for empty arguments +} + +template +void write_to_stream(std::ostringstream& oss, T&& val) { + oss << std::forward(val); +} + +template +void write_to_stream(std::ostringstream& oss, T&& val, Args&&... args) { + oss << std::forward(val) << " "; + write_to_stream(oss, std::forward(args)...); +} + +// Helper macro to handle empty __VA_ARGS__ +#define FLASHINFER_CHECK_IMPL(condition, message) \ + if (!(condition)) { \ + FLASHINFER_ERROR(message); \ + } + +// Main macro that handles both cases +#define FLASHINFER_CHECK(condition, ...) \ + do { \ + if (!(condition)) { \ + std::ostringstream oss; \ + write_to_stream(oss, ##__VA_ARGS__); \ + std::string msg = oss.str(); \ + if (msg.empty()) { \ + msg = "Check failed: " #condition; \ + } \ + FLASHINFER_ERROR(msg); \ + } \ + } while (0) + +// Warning macro +#define FLASHINFER_WARN(...) \ + do { \ + std::ostringstream oss; \ + write_to_stream(oss, ##__VA_ARGS__); \ + std::string msg = oss.str(); \ + if (msg.empty()) { \ + msg = "Warning triggered"; \ + } \ + flashinfer::Warning(__FUNCTION__, __FILE__, __LINE__, msg).emit(); \ + } while (0) + +namespace flashinfer { +class Error : public std::exception { + private: + std::string message_; + + public: + Error(const std::string& func, const std::string& file, int line, const std::string& message) { + std::ostringstream oss; + oss << "Error in function '" << func << "' " + << "at " << file << ":" << line << ": " << message; + message_ = oss.str(); + } + + virtual const char* what() const noexcept override { return message_.c_str(); } +}; + +class Warning { + private: + std::string message_; + + public: + Warning(const std::string& func, const std::string& file, int line, const std::string& message) { + std::ostringstream oss; + oss << "Warning in function '" << func << "' " + << "at " << file << ":" << line << ": " << message; + message_ = oss.str(); + } + + void emit() const { std::cerr << message_ << std::endl; } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_EXCEPTION_H_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh new file mode 100644 index 000000000000..d763d05135c6 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh @@ -0,0 +1,154 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * 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. + */ +#ifndef FLASHINFER_MAMBA_CHECKPOINTING_SSU_CUH_ +#define FLASHINFER_MAMBA_CHECKPOINTING_SSU_CUH_ + +#include + +namespace flashinfer::mamba::checkpointing { + +struct CheckpointingSsuParams { + uint32_t batch{}, nheads{}, dim{}, dstate{}, ngroups{}; + uint32_t state_cache_size{}; + uint32_t npredicted{}; + uint32_t max_window{}; + int32_t pad_slot_id{-1}; + + // v12 §59: per-head DIM split factor. Must be one of {1, 2, 4}. The host + // launcher dispatches to a kernel template specialized on this value; the + // kernel cross-checks via assert(params.d_split == D_SPLIT). + int32_t d_split{1}; + + bool dt_softplus{false}; + + // Note: Programmatic Dependent Launch is JIT-stamped via the `ENABLE_PDL` + // constexpr (see checkpointing_ssu_customize_config.jinja). Each .so has + // its PDL mode baked in; no runtime field needed. + + // ── Tensor pointers ── + void* __restrict__ state{nullptr}; // (state_cache_size, nheads, dim, dstate) + void* __restrict__ x{nullptr}; // (batch, NPREDICTED, nheads, dim) + void* __restrict__ dt{nullptr}; // (batch, NPREDICTED, nheads, dim) tie_hdim + void* __restrict__ A{nullptr}; // (nheads, dim, dstate) tie_hdim + void* __restrict__ B{nullptr}; // (batch, NPREDICTED, ngroups, dstate) + void* __restrict__ C{nullptr}; // (batch, NPREDICTED, ngroups, dstate) + void* __restrict__ D{nullptr}; // (nheads, dim), optional + void* __restrict__ z{nullptr}; // (batch, NPREDICTED, nheads, dim), optional + void* __restrict__ dt_bias{nullptr}; // (nheads, dim) tie_hdim, optional + void* __restrict__ output{nullptr}; // (batch, NPREDICTED, nheads, dim) + + // ── Cache tensors for incremental replay ── + void* __restrict__ old_x{nullptr}; // (state_cache_size, MAX_WINDOW, nheads, dim) single-buffered + void* __restrict__ old_B{ + nullptr}; // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) double-buffered + void* __restrict__ old_dt{ + nullptr}; // (state_cache_size, 2, nheads, MAX_WINDOW) double-buffered, f32 + void* __restrict__ old_cumAdt{ + nullptr}; // (state_cache_size, 2, nheads, MAX_WINDOW) double-buffered, f32 + void* __restrict__ cache_buf_idx{nullptr}; // (state_cache_size,) int32 + void* __restrict__ prev_num_accepted{nullptr}; // (state_cache_size,) int32 + + // ── Index tensors ── + void* __restrict__ state_batch_indices{nullptr}; // (batch,) optional + + // ── Varlen (v20): packed inputs ── + // When non-null, `x/dt/B/C/z/out` are laid out as + // `(1, total_tokens, nheads, dim)` / `(1, total_tokens, ngroups, dstate)` + // and `cu_seqlens[i]` gives the token-axis base of sequence i. + // `seq_len_i = cu_seqlens[i+1] - cu_seqlens[i]`. Kernel dispatch on + // `cu_seqlens != nullptr` selects a `VARLEN=true` template. + // + // The `*_stride_seq` fields below already encode the outer iteration + // stride for both modes — the wrapper sets them to: + // non-varlen: `tensor.stride(0)` (per-batch) + // varlen : `tensor.stride(1)` (per-token, since sequences are packed + // into a single batch of total_tokens) + // so the kernel uses one formula `seq * *_stride_seq` regardless of mode. + void* __restrict__ cu_seqlens{nullptr}; // (batch+1,) int32, optional + + // ── Block-scale decode factors for quantized state ── + void* __restrict__ state_scale{nullptr}; // float32: (state_cache_size, nheads, dim) + + // ── Philox PRNG seed for stochastic rounding ── + const int64_t* rand_seed{nullptr}; + + // ── Strides ── + // state: (state_cache_size, nheads, dim, dstate) — inner 3 dims contiguous + int64_t state_stride_seq{}; + + // For the six batch-side tensors (x, dt, B, C, out, z), `*_stride_seq` + // is the outer iteration stride — per-batch in non-varlen, per-token in + // varlen. `*_stride_token` is the inner per-row (T-axis) stride, same + // in both modes. + + // x: (batch, NPREDICTED, nheads, dim) [non-varlen] / (1, total_tokens, nheads, dim) [varlen] + int64_t x_stride_seq{}; + int64_t x_stride_token{}; + + // dt: (batch, NPREDICTED, nheads, dim) — tie_hdim (stride_dim=0) + int64_t dt_stride_seq{}; + int64_t dt_stride_token{}; + + // B: (batch, NPREDICTED, ngroups, dstate) + int64_t B_stride_seq{}; + int64_t B_stride_token{}; + + // C: (batch, NPREDICTED, ngroups, dstate) + int64_t C_stride_seq{}; + int64_t C_stride_token{}; + + // output: (batch, NPREDICTED, nheads, dim) + int64_t out_stride_seq{}; + int64_t out_stride_token{}; + + // z: (batch, NPREDICTED, nheads, dim) + int64_t z_stride_seq{}; + int64_t z_stride_token{}; + + // old_x: (state_cache_size, MAX_WINDOW, nheads, dim) — single-buffered + int64_t old_x_stride_seq{}; + int64_t old_x_stride_token{}; + + // old_B: (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) — double-buffered + int64_t old_B_stride_seq{}; + int64_t old_B_stride_dbuf{}; + int64_t old_B_stride_token{}; + + // old_dt: (state_cache_size, 2, nheads, MAX_WINDOW) — double-buffered, MAX_WINDOW contiguous + int64_t old_dt_stride_seq{}; + int64_t old_dt_stride_dbuf{}; + int64_t old_dt_stride_head{}; + + // old_cumAdt: (state_cache_size, 2, nheads, MAX_WINDOW) — double-buffered, MAX_WINDOW contiguous + int64_t old_cumAdt_stride_seq{}; + int64_t old_cumAdt_stride_dbuf{}; + int64_t old_cumAdt_stride_head{}; + + // state_scale: (state_cache_size, nheads, dim) + int64_t state_scale_stride_seq{}; +}; + +// Forward declaration — defined in kernel_checkpointing_ssu.cuh. +// `launchCheckpointingSsu` is the public dispatcher: it reads +// `params.d_split` and routes to the matching `launchCheckpointingSsuImpl` +// specialization (v12 §59). Caller side stays single-entry. +template +void launchCheckpointingSsu(CheckpointingSsuParams& params, cudaStream_t stream); + +} // namespace flashinfer::mamba::checkpointing + +#endif // FLASHINFER_MAMBA_CHECKPOINTING_SSU_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh new file mode 100644 index 000000000000..c36883f523cb --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh @@ -0,0 +1,208 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * 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. + */ +#ifndef FLASHINFER_MAMBA_COMMON_CUH_ +#define FLASHINFER_MAMBA_COMMON_CUH_ + +#include + +#include +#include +#include + +#include + +namespace flashinfer::mamba { + +constexpr unsigned warpSize = 32; + +// ============================================================================= +// Common types and utilities +// ============================================================================= + +// Largest power of 2 that divides v (i.e. v & -v). Returns 1 when v == 0. +inline constexpr unsigned largestPow2Divisor(unsigned v) { return v ? (v & (~v + 1)) : 1; } + +// Simple packed vector type for loading N elements of type T. +// Alignment is the largest power-of-2 factor of the total byte size, +// so it is always valid even when N * sizeof(T) is not a power of 2 (e.g. 3 × 2 = 6). +template +struct alignas(largestPow2Divisor(N * sizeof(T))) PackedAligned { + T val[N]; + static constexpr int count = N; + using dtype = T; +}; + +template +__device__ __forceinline__ auto make_zeros() -> load_t { + load_t ret{}; +#pragma unroll + for (int i = 0; i < ret.count; i++) + ret.val[i] = typename load_t::dtype{}; // default initialization + return ret; +}; + +// Computes the vector load size that ensures full warp utilization. +// Avoids cases like: dstate=64, load_t = sizeof(float4)/sizeof(f16), warpsize=32 (32 * 8 > 64) +// in which case a part of the warp would be idle. +template +inline constexpr auto getVectorLoadSizeForFullUtilization() -> unsigned { + static_assert(sizeof(float4) >= sizeof(T)); + constexpr unsigned maxHardwareLoadSize = sizeof(float4) / sizeof(T); + constexpr unsigned maxLogicalLoadSize = (unsigned)DSTATE / warpSize; + return maxHardwareLoadSize < maxLogicalLoadSize ? maxHardwareLoadSize : maxLogicalLoadSize; +} + +__device__ __forceinline__ float warpReduceSum(float val) { + for (int s = warpSize / 2; s > 0; s /= 2) { + val += __shfl_down_sync(UINT32_MAX, val, s); + } + return val; +} + +__device__ __forceinline__ float warpReduceMax(float val) { + for (int s = warpSize / 2; s > 0; s /= 2) { + val = max(val, __shfl_down_sync(UINT32_MAX, val, s)); + } + return val; +} + +__forceinline__ __device__ float softplus(float x) { return __logf(1.f + __expf(x)); } + +__device__ __forceinline__ float thresholded_softplus(float dt_value) { + constexpr float threshold = 20.f; + return (dt_value <= threshold) ? softplus(dt_value) : dt_value; +} + +// ============================================================================= +// Dispatch helpers +// ============================================================================= + +// Format an integer_sequence as a comma-separated string for error messages +template +std::string format_sequence(std::integer_sequence) { + std::ostringstream oss; + bool first = true; + ((oss << (first ? (first = false, "") : ", ") << Values), ...); + return oss.str(); +} + +// Helper function to dispatch dim and dstate with a kernel launcher +template +void dispatchDimDstate(ParamsType& params, std::integer_sequence dims_seq, + std::integer_sequence dstates_seq, + KernelLauncher&& launcher) { + auto dispatch_dstate = [&]() { + auto try_dstate = [&]() { + if (params.dstate == DSTATE) { + launcher.template operator()(); + return true; + } + return false; + }; + bool dispatched = (try_dstate.template operator()() || ...); + FLASHINFER_CHECK(dispatched, "Unsupported dstate value: ", params.dstate, + ".\nSupported values: ", format_sequence(dstates_seq)); + }; + + auto try_dim = [&]() { + if (params.dim == DIM) { + dispatch_dstate.template operator()(); + return true; + } + return false; + }; + + bool dim_dispatched = (try_dim.template operator()() || ...); + FLASHINFER_CHECK(dim_dispatched, "Unsupported dim value: ", params.dim, + ".\nSupported values: ", format_sequence(dims_seq)); +} + +// Helper function to dispatch ratio with a kernel launcher +template +void dispatchRatio(ParamsType& params, std::integer_sequence ratios_seq, + KernelLauncher&& launcher) { + auto try_ratio = [&]() { + if (params.nheads / params.ngroups == RATIO) { + launcher.template operator()(); + return true; + } + return false; + }; + + bool ratio_dispatched = (try_ratio.template operator()() || ...); + FLASHINFER_CHECK(ratio_dispatched, + "Unsupported nheads/ngroups ratio: ", params.nheads / params.ngroups, + ".\nSupported values: ", format_sequence(ratios_seq)); +} + +// Helper function to dispatch dim, dstate, and ntokens_mtp with a kernel launcher +// Reuses dispatchDimDstate by wrapping the launcher to add token dispatch +template +void dispatchDimDstateTokens(ParamsType& params, + std::integer_sequence dims_seq, + std::integer_sequence dstates_seq, + std::integer_sequence tokens_seq, + KernelLauncher&& launcher) { + // Wrap the launcher to add token dispatch as the innermost level + auto dim_dstate_launcher = [&]() { + auto try_tokens = [&]() { + if (params.ntokens_mtp == TOKENS_MTP) { + launcher.template operator()(); + return true; + } + return false; + }; + bool dispatched = (try_tokens.template operator()() || ...); + FLASHINFER_CHECK(dispatched, "Unsupported ntokens_mtp value: ", params.ntokens_mtp, + ".\nSupported values: ", format_sequence(tokens_seq)); + }; + + dispatchDimDstate(params, dims_seq, dstates_seq, dim_dstate_launcher); +} + +// ============================================================================= +// Alignment checks +// ============================================================================= + +// Check alignment for common input variables (x, z, B, C) +// Works for both STP (SelectiveStateUpdateParams) and MTP (SelectiveStateMTPParams) +template +void check_ptr_alignment_input_vars(const ParamsType& params) { + using load_input_t = PackedAligned; + FLASHINFER_CHECK(reinterpret_cast(params.x) % sizeof(load_input_t) == 0, + "x pointer must be aligned to ", sizeof(load_input_t), " bytes"); + FLASHINFER_CHECK((params.x_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, + "x batch stride must be aligned to ", sizeof(load_input_t), " bytes"); + if (params.z) { + FLASHINFER_CHECK(reinterpret_cast(params.z) % sizeof(load_input_t) == 0, + "z pointer must be aligned to ", sizeof(load_input_t), " bytes"); + FLASHINFER_CHECK((params.z_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, + "z batch stride must be aligned to ", sizeof(load_input_t), " bytes"); + } + FLASHINFER_CHECK(reinterpret_cast(params.B) % sizeof(load_input_t) == 0, + "B pointer must be aligned to ", sizeof(load_input_t), " bytes"); + FLASHINFER_CHECK(reinterpret_cast(params.C) % sizeof(load_input_t) == 0, + "C pointer must be aligned to ", sizeof(load_input_t), " bytes"); + FLASHINFER_CHECK((params.B_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, + "B batch stride must be aligned to ", sizeof(load_input_t), " bytes"); + FLASHINFER_CHECK((params.C_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, + "C batch stride must be aligned to ", sizeof(load_input_t), " bytes"); +} + +} // namespace flashinfer::mamba + +#endif // FLASHINFER_MAMBA_COMMON_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh new file mode 100644 index 000000000000..d8b46058b4a7 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh @@ -0,0 +1,474 @@ +#pragma once +#include +#include +#include +#ifdef FLASHINFER_ENABLE_BF16 +#include +#endif + +namespace flashinfer::mamba::conversion { + +inline __device__ float toFloat(float f) { return f; } + +inline __device__ float toFloat(__half h) { return __half2float(h); } + +#ifdef FLASHINFER_ENABLE_BF16 +inline __device__ float toFloat(__nv_bfloat16 val) { return __bfloat162float(val); } +#endif + +// No accuracy loss: int8_t / int16_t range fits exactly in float32 (24-bit +// mantissa represents all integers up to 2^24 = 16M exactly). +inline __device__ float toFloat(int8_t val) { return static_cast(val); } +inline __device__ float toFloat(int16_t val) { return static_cast(val); } + +// fp8 e4m3 → fp32. Goes via __half (cuda_fp8 library has the implicit +// conversion that compiles to `cvt.rn.f16.e4m3` PTX on sm_89+), then +// __half2float for the final step. No direct fp8→fp32 PTX op exists. +inline __device__ float toFloat(__nv_fp8_e4m3 val) { + return __half2float(static_cast<__half>(val)); +} + +// Packed 2-element conversion: convert a packed pair to float2. +// Uses native packed intrinsics for bf16/fp16 (fewer PRMT/SHF instructions). +inline __device__ float2 toFloat2(float2 packed) { return packed; } + +inline __device__ float2 toFloat2(__half2 packed) { return __half22float2(packed); } + +// Pointer-based overloads: read two consecutive elements and convert to float2. +// Dispatches to the packed intrinsic for bf16/fp16 via the overloads above. +inline __device__ float2 toFloat2(float const* ptr) { return {ptr[0], ptr[1]}; } + +inline __device__ float2 toFloat2(__half const* ptr) { + return toFloat2(*reinterpret_cast<__half2 const*>(ptr)); +} + +#ifdef FLASHINFER_ENABLE_BF16 +// inline __device__ float2 toFloat2(__nv_bfloat162 packed) { return __bfloat1622float2(packed); } + +inline __device__ float2 toFloat2(__nv_bfloat162 packed) { + // bf16 is the upper 16 bits of f32 — shift/mask is cheaper than PRMT byte permutation. + // NOTE: this ignores denormals + uint32_t bits = reinterpret_cast(packed); + float2 out; + out.x = __uint_as_float(bits << 16); // low bf16 → upper 16 bits of f32 + out.y = __uint_as_float(bits & 0xFFFF0000u); // high bf16 already in upper 16 bits + return out; +} + +inline __device__ float2 toFloat2(__nv_bfloat16 const* ptr) { + return toFloat2(*reinterpret_cast<__nv_bfloat162 const*>(ptr)); +} + +// Paired f32 → bf16 conversion: pack two f32 values into __nv_bfloat162. +// Uses native cvt.rn.bf16x2.f32 — single instruction, round-to-nearest-even. +inline __device__ __nv_bfloat162 fromFloat2(float2 val) { + uint32_t result; + asm("cvt.rn.bf16x2.f32 %0, %1, %2;\n" : "=r"(result) : "f"(val.y), "f"(val.x)); + return reinterpret_cast<__nv_bfloat162 const&>(result); +} + +#endif + +inline __device__ float2 toFloat2(int8_t const* ptr) { return {toFloat(ptr[0]), toFloat(ptr[1])}; } +inline __device__ float2 toFloat2(int16_t const* ptr) { return {toFloat(ptr[0]), toFloat(ptr[1])}; } + +inline __device__ void convertAndStore(float* output, float input) { *output = input; } + +inline __device__ void convertAndStore(__half* output, float input) { + *output = __float2half(input); +} + +#ifdef FLASHINFER_ENABLE_BF16 +inline __device__ void convertAndStore(__nv_bfloat16* output, float input) { + *output = __float2bfloat16(input); +} +#endif + +inline __device__ void convertAndStore(int16_t* output, float input) { + // Symmetric clip: [-max, max] (not [-max-1, max]) so that negation is safe. + // Matches Triton reference which clips to [-32767, 32767] before storing. + constexpr float int16_max = static_cast(std::numeric_limits::max()); + input = fminf(fmaxf(input, -int16_max), int16_max); + *output = static_cast(__float2int_rn(input)); +} + +// ============================================================================= +// Philox-4x32 PRNG (matches Triton's tl.randint) +// ============================================================================= + +// Generates four pseudorandom uint32s from (seed, offset) using the Philox-4x32 algorithm. +// Produces bit-identical output to Triton's tl.randint4x(seed, offset, n_rounds). +// The offset is int64 and split across Philox c0 (low 32 bits) and c1 (high +// 32 bits) — matches the i64 path of `randint4x` in triton/language/random.py. +// Provides 2^64 unique counter values per seed, avoiding collisions in large +// caches where `cache_slot * stride` exceeds 2^32. +// All four outputs (c0..c3) are independent and uniformly distributed. +template +__device__ __forceinline__ void philox_randint4x(int64_t seed, int64_t offset, uint32_t& r0, + uint32_t& r1, uint32_t& r2, uint32_t& r3) { + constexpr uint32_t PHILOX_KEY_A = 0x9E3779B9u; + constexpr uint32_t PHILOX_KEY_B = 0xBB67AE85u; + constexpr uint32_t PHILOX_ROUND_A = 0xD2511F53u; + constexpr uint32_t PHILOX_ROUND_B = 0xCD9E8D57u; + + uint32_t k0 = static_cast(static_cast(seed)); + uint32_t k1 = static_cast(static_cast(seed) >> 32); + uint64_t uoffset = static_cast(offset); + uint32_t c0 = static_cast(uoffset); + uint32_t c1 = static_cast(uoffset >> 32); + uint32_t c2 = 0, c3 = 0; + +#pragma unroll + for (int i = 0; i < n_rounds; i++) { + uint32_t _c0 = c0, _c2 = c2; + c0 = __umulhi(PHILOX_ROUND_B, _c2) ^ c1 ^ k0; + c2 = __umulhi(PHILOX_ROUND_A, _c0) ^ c3 ^ k1; + c1 = PHILOX_ROUND_B * _c2; + c3 = PHILOX_ROUND_A * _c0; + k0 += PHILOX_KEY_A; + k1 += PHILOX_KEY_B; + } + r0 = c0; + r1 = c1; + r2 = c2; + r3 = c3; +} + +// Generates a pseudorandom uint32 from (seed, offset) using the Philox-4x32 algorithm. +// Produces bit-identical output to Triton's tl.randint(seed, offset, n_rounds). +// The offset is int64 (low/high split across Philox c0/c1) — see +// philox_randint4x for the full rationale. +// NOTE: This discards 3 of the 4 Philox outputs. For better throughput, use +// philox_randint4x to get all 4 outputs from a single Philox invocation. +template +__device__ __forceinline__ uint32_t philox_randint(int64_t seed, int64_t offset) { + uint32_t r0, r1, r2, r3; + philox_randint4x(seed, offset, r0, r1, r2, r3); + return r0; +} + +// ============================================================================= +// Stochastic rounding: fp32 → fp16 +// ============================================================================= + +// Software stochastic rounding: convert one fp32 value to fp16 using 13 random bits. +// Adds random noise at the sub-fp16-mantissa position, then truncates. +// rand13: 13-bit random value in bits [12:0]. +__device__ __forceinline__ uint16_t cvt_rs_f16_sw(float x, uint32_t rand13) { + uint32_t bits = __float_as_uint(x); + uint32_t sign = bits & 0x80000000u; + uint32_t abs_bits = bits & 0x7FFFFFFFu; + + // fp32 has 23 mantissa bits, fp16 has 10. The 13 LSBs are the remainder. + // Add 13-bit random noise at bits [12:0]. Carry into bit 13 → round up. + abs_bits += (rand13 & 0x1FFFu); + + // Convert to fp16 by truncation. + uint32_t f32_exp = (abs_bits >> 23) & 0xFFu; + uint32_t f32_mantissa = abs_bits & 0x7FFFFFu; + + uint16_t f16_bits; + if (f32_exp == 0xFF) { + f16_bits = (f32_mantissa != 0) ? 0x7E00u : 0x7C00u; // NaN or Inf + } else if (f32_exp > 142) { // 127 + 15 = 142 → overflow to Inf + f16_bits = 0x7C00u; + } else if (f32_exp < 113) { // 127 - 14 = 113 → underflow to zero + f16_bits = 0; + } else { + uint16_t f16_exp = static_cast(f32_exp - 112); // rebias: 127→15 + uint16_t f16_mantissa = static_cast(f32_mantissa >> 13); + f16_bits = (f16_exp << 10) | f16_mantissa; + } + + return static_cast(sign >> 16) | f16_bits; +} + +// Forward declaration (defined below, after cvt_rs_f16x2_f32). +__device__ __forceinline__ uint32_t cvt_rs_f16x2_f32(float a, float b, uint32_t rbits); + +// Stochastic rounding: convert one fp32 value to fp16 using 13 random bits. +// On sm_100a+: uses PTX cvt.rs.f16x2.f32 with a dummy zero second input. +// On other archs: software emulation. +__device__ __forceinline__ __half cvt_rs_f16_f32(float x, uint32_t rand13) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL) + // Pack rand13 into rbits[12:0] (for PTX operand b → low half → our x). + // High half gets zero noise for the dummy input. + uint32_t rbits = rand13 & 0x1FFFu; + uint32_t packed = cvt_rs_f16x2_f32(x, 0.0f, rbits); + return __ushort_as_half(static_cast(packed & 0xFFFFu)); +#else + return __ushort_as_half(cvt_rs_f16_sw(x, rand13)); +#endif +} + +// Stochastic rounding: convert two fp32 values to packed fp16x2 using random bits. +// On sm_100a+: uses PTX cvt.rs.f16x2.f32 instruction. +// On other archs: software emulation matching the hardware behavior. +// +// rbits layout (from PTX docs): +// bits [28:16] = 13 random bits for PTX operand "a" (→ d[31:16], high half) +// bits [12:0] = 13 random bits for PTX operand "b" (→ d[15:0], low half) +// bits [31:29] and [15:13] = unused (zero) +// from: https://docs.nvidia.com/cuda/parallel-thread-execution/#cvt-rs-rbits-layout-f16 +// +// Our asm maps: %1→C++ a→PTX b→d[15:0], %2→C++ b→PTX a→d[31:16] +// So: C++ a uses rbits[12:0], C++ b uses rbits[28:16]. +__device__ __forceinline__ uint32_t cvt_rs_f16x2_f32(float a, float b, uint32_t rbits) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL) + uint32_t packed; + asm("cvt.rs.f16x2.f32 %0, %2, %1, %3;" + : "=r"(packed) + : "r"(__float_as_uint(a)), "r"(__float_as_uint(b)), "r"(rbits)); + return packed; +#else + uint32_t rand_a = rbits & 0x1FFFu; // bits [12:0] → C++ a (PTX b → low half) + uint32_t rand_b = (rbits >> 16) & 0x1FFFu; // bits [28:16] → C++ b (PTX a → high half) + uint16_t a_fp16 = __half_as_ushort(cvt_rs_f16_f32(a, rand_a)); + uint16_t b_fp16 = __half_as_ushort(cvt_rs_f16_f32(b, rand_b)); + return static_cast(a_fp16) | (static_cast(b_fp16) << 16); +#endif +} + +// Stochastic rounding store: generates Philox random bits and converts fp32 → fp16 in one call. +// PHILOX_ROUNDS: number of Philox rounds (compile-time), must be > 0. +// seed: Philox seed (from params.rand_seed). +// offset: unique per-element offset (e.g. d * DSTATE + i) for deterministic randomness. +template +inline __device__ void convertSRAndStore(__half* output, float input, int64_t seed, + uint32_t offset) { + uint32_t rand = philox_randint(seed, offset); + *output = cvt_rs_f16_f32(input, rand & 0x1FFFu); +} + +// ============================================================================= +// Stochastic rounding: fp32 → fp8 (e4m3) +// ============================================================================= + +// Bit-reverse 16 bits. HW gives both elements of a pair 16 bits of independent +// SR randomness while consuming a single 16-bit chunk per pair: one element uses +// the chunk straight, the other uses bitrev16 of it. Two bitrev16'd halves of a +// uniformly random 16-bit value remain uniformly distributed and are statistically +// independent, so each element gets full 16-bit unbiased SR while sharing one +// 16-bit register with its pair-mate. +// +// Implementation: PTX `brev.b32` (CUDA intrinsic `__brev`) is a single-ALU-op +// 32-bit reverse; we right-shift by 16 to land the 16 LSBs of input in the 16 +// LSBs of output. Total: 2 SASS instructions (brev + shf/shr), vs the prior +// software 4-step mask/shift/OR chain (~12-16 inst). +__device__ __forceinline__ uint32_t bitrev16(uint32_t b) { return __brev(b) >> 16; } + +// Software stochastic rounding: convert one fp32 value to e4m3 (FN, satfinite) using 16 random +// bits. +// +// Algorithm: place `rand16` at the top of the discarded mantissa range, then truncate. +// shift_truncate = 20 for normal binade (unbiased >= -6) +// = 14 - unbiased for subnormal/underflow +// contribution = rand16 << (shift_truncate - 16) +// total = mant24 + contribution (in uint64 to avoid overflow) +// int_part = total >> shift_truncate +// Then re-encode int_part as e4m3, handling subnormal→normal transitions and saturation. +// +// Saturation (satfinite): +// |x| > 448 → ±448 (max finite e4m3 = 0x7E) +// ±Inf → ±448 +// NaN → canonical NaN with sign preserved (S|1111|111 = 0x7F or 0xFF) +// +// Verified bitwise against HW (cvt.rs.satfinite.e4m3x4.f32 on sm_100a) across 22528 +// inputs spanning subnormal, normal, and saturation regions during the SR +// reverse-engineering effort — see .plans/e4m3_stochastic_rounding.md. +__device__ __forceinline__ uint8_t cvt_rs_e4m3_sw(float x, uint32_t rand16) { + uint32_t bits = __float_as_uint(x); + uint32_t sign = (bits >> 31) & 1u; + uint32_t abs_bits = bits & 0x7FFFFFFFu; + uint32_t f32_exp = (abs_bits >> 23) & 0xFFu; + uint32_t f32_mant = abs_bits & 0x7FFFFFu; + + // NaN / Inf + if (f32_exp == 0xFFu) { + if (f32_mant != 0) { + return static_cast(0x7Fu | (sign << 7)); // canonical e4m3 NaN + } else { + return static_cast(0x7Eu | (sign << 7)); // Inf → ±max finite + } + } + + // fp32 zero / denormal → e4m3 zero (with sign). + if (f32_exp == 0u) { + return static_cast(sign << 7); + } + + int unbiased = static_cast(f32_exp) - 127; + uint64_t mant24 = 0x800000u | f32_mant; // implicit-1 + mantissa, 24-bit + int shift_truncate = (unbiased >= -6) ? 20 : (14 - unbiased); + int rand_shift = shift_truncate - 16; + uint64_t rand_contrib; + if (rand_shift < 0) { + // Defensive: shift_truncate < 16 shouldn't happen for valid normal/subnormal e4m3. + rand_contrib = static_cast(rand16 & 0xFFFFu) >> (-rand_shift); + } else if (rand_shift < 56) { + rand_contrib = static_cast(rand16 & 0xFFFFu) << rand_shift; + } else { + rand_contrib = 0; + } + uint64_t total = mant24 + rand_contrib; + // Guard the shift: shift_truncate reaches 140 for tiny-normal fp32 + // (unbiased < -49), which is UB for a uint64_t shift. Mathematically the + // result is 0 there (the value is far below e4m3's smallest subnormal), + // so flush to int_part = 0. The downstream subnormal branch rounds it to ±0. + uint32_t int_part = (shift_truncate >= 64) ? 0u : static_cast(total >> shift_truncate); + + if (unbiased >= -6) { + // Started in normal binade. int_part ∈ [8, 15] normally; can overflow to 16+ if rand + // bumped the exponent. + int e4m3_exp = unbiased + 7; + while (int_part >= 16u) { + int_part >>= 1; + e4m3_exp += 1; + } + if (e4m3_exp > 15 || (e4m3_exp == 15 && (int_part & 0x7u) == 7u)) { + return static_cast(0x7Eu | (sign << 7)); + } + return static_cast((sign << 7) | (e4m3_exp << 3) | (int_part & 0x7u)); + } else { + // Started in subnormal/underflow. int_part: + // 0 → zero + // 1..7 → subnormal e4m3 + // 8..15 → smallest normal binade (e4m3_exp = 1) + // 16+ → higher normal binades (rare; only if rand pushed up multiple binades) + if (int_part == 0u) { + return static_cast(sign << 7); + } + if (int_part <= 7u) { + return static_cast((sign << 7) | int_part); + } + int e4m3_exp = 1; + while (int_part >= 16u) { + int_part >>= 1; + e4m3_exp += 1; + } + if (e4m3_exp > 15 || (e4m3_exp == 15 && (int_part & 0x7u) == 7u)) { + return static_cast(0x7Eu | (sign << 7)); + } + return static_cast((sign << 7) | (e4m3_exp << 3) | (int_part & 0x7u)); + } +} + +// Stochastic rounding: convert four fp32 values to packed fp8x4 e4m3 using random bits. +// On sm_100a+: uses PTX cvt.rs.satfinite.e4m3x4.f32 (combined stochastic-round + saturate). +// On other archs: software fallback via cvt_rs_e4m3_sw. +// +// Output layout (low byte first): +// packed[ 7: 0] = e4m3(a) +// packed[15: 8] = e4m3(b) +// packed[23:16] = e4m3(c) +// packed[31:24] = e4m3(d) +// +// rbits layout (per PTX docs + empirical HW oracle, see .plans/e4m3_stochastic_rounding.md): +// bits [31:16] = pair rbits for PTX operands a/b (high two outputs) +// bits [15: 0] = pair rbits for PTX operands e/f (low two outputs) +// Each PAIR shares its 16-bit chunk: HW uses the chunk straight for the "even" +// PTX operand (b, f → low byte of pair output) and bitrev16 of the chunk for +// the "odd" operand (a, e → high byte of pair output). Both elements get the +// full 16 bits of independent SR randomness this way. +// our `a` (PTX f, → byte 0) uses rbits[15: 0] +// our `b` (PTX e, → byte 1) uses bitrev16(rbits[15: 0]) +// our `c` (PTX b, → byte 2) uses rbits[31:16] +// our `d` (PTX a, → byte 3) uses bitrev16(rbits[31:16]) +// +// PTX syntax `cvt.rs.satfinite.e4m3x4.f32 d, {a3, a2, a1, a0}, rbits` writes +// e4m3(a_i) into byte i of d. We want byte 0 = e4m3(a), so the source-vector +// ordering is {d, c, b, a} = {%4, %3, %2, %1}. See PTX ISA: +// https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt +__device__ __forceinline__ uint32_t cvt_rs_e4m3x4_f32(float a, float b, float c, float d, + uint32_t rbits) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL) + uint32_t packed; + asm("cvt.rs.satfinite.e4m3x4.f32 %0, {%4, %3, %2, %1}, %5;" + : "=r"(packed) + : "r"(__float_as_uint(a)), "r"(__float_as_uint(b)), "r"(__float_as_uint(c)), + "r"(__float_as_uint(d)), "r"(rbits)); + return packed; +#else + uint32_t low_chunk = rbits & 0xFFFFu; + uint32_t high_chunk = (rbits >> 16) & 0xFFFFu; + uint8_t pa = cvt_rs_e4m3_sw(a, low_chunk); // PTX f → byte 0 + uint8_t pb = cvt_rs_e4m3_sw(b, bitrev16(low_chunk)); // PTX e → byte 1 + uint8_t pc = cvt_rs_e4m3_sw(c, high_chunk); // PTX b → byte 2 + uint8_t pd = cvt_rs_e4m3_sw(d, bitrev16(high_chunk)); // PTX a → byte 3 + return static_cast(pa) | (static_cast(pb) << 8) | + (static_cast(pc) << 16) | (static_cast(pd) << 24); +#endif +} + +// ============================================================================= +// Round-to-nearest-even + saturate: fp32 → int8 +// ============================================================================= + +// cvt.rni.sat.s8.f32: single PTX instruction on sm_80+, replaces +// the F2I.S32 + VIMNMX(min 127) + VIMNMX(max -127) chain. +// Saturates to [-128, 127]. Callers using encode_scale = 127/amax +// guarantee |input| ≤ 127.0, so -128 is never produced. +__device__ __forceinline__ int8_t cvt_rni_sat_s8(float x) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + int32_t result; + asm("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(result) : "f"(x)); + return static_cast(result); +#else + return static_cast(max(-128, min(127, __float2int_rn(x)))); +#endif +} + +// ============================================================================= +// Stochastic rounding + saturate: fp32 → int8 +// ============================================================================= + +// Software SR for int8: add uniform noise in [0, 1) then floor. +// Matches Triton's `floor(scaled_value + rand01)` where +// `rand01 = (rand & 0x00FFFFFF) * (1.0 / (1 << 24))`. +// Saturates to [-127, 127] (symmetric, matching encode_scale = 127/amax). +__device__ __forceinline__ int8_t cvt_rs_sat_s8(float x, uint32_t rand_bits) { + float const rand01 = + static_cast(rand_bits & 0x00FFFFFFu) * (1.0f / static_cast(1 << 24)); + // `__float2int_rd` (round toward -infinity) fuses `floorf` + `__float2int_rz` + // into a single `cvt.rmi.s32.f32` SASS instruction, saving one FRND per call. + int32_t const clamped = max(-127, min(127, __float2int_rd(x + rand01))); + return static_cast(clamped); +} + +// Stochastic rounding: convert four fp32 values to packed s8x4 using a single +// 32-bit random integer. Analogous to cvt_rs_e4m3x4_f32: 16-bit chunks are +// reused via bitrev16 so each output gets 16 bits of independent randomness +// while consuming only one shared 16-bit chunk per pair (two bitrev16'd halves +// of a uniform 16-bit value remain uniformly distributed and statistically +// independent). +// +// 16-bit entropy per element is far more than int8 SR requires: the rounding +// decision compares against a fractional residual with at most ~7 bits of +// meaningful precision for int8, so no quality loss vs the 24-bit scalar path. +// +// Amortization: 1 random u32 → 4 SR int8s. A single Philox call (4 u32s) +// covers 16 int8 conversions, a 4× reduction in PRNG cost vs the scalar +// cvt_rs_sat_s8 path. +__device__ __forceinline__ uint32_t cvt_rs_sat_s8x4_f32(float a, float b, float c, float d, + uint32_t rbits) { + uint32_t const low_chunk = rbits & 0xFFFFu; + uint32_t const high_chunk = (rbits >> 16) & 0xFFFFu; + constexpr float kInv16 = 1.0f / static_cast(1u << 16); + + float const r_a = static_cast(low_chunk) * kInv16; + float const r_b = static_cast(bitrev16(low_chunk)) * kInv16; + float const r_c = static_cast(high_chunk) * kInv16; + float const r_d = static_cast(bitrev16(high_chunk)) * kInv16; + + // `__float2int_rd` (round toward -infinity) emits a single `cvt.rmi.s32.f32` + // SASS op, fusing the `floorf` + `__float2int_rz` chain into one instruction. + int32_t const pa = max(-127, min(127, __float2int_rd(a + r_a))); + int32_t const pb = max(-127, min(127, __float2int_rd(b + r_b))); + int32_t const pc = max(-127, min(127, __float2int_rd(c + r_c))); + int32_t const pd = max(-127, min(127, __float2int_rd(d + r_d))); + + return (static_cast(pa) & 0xFFu) | ((static_cast(pb) & 0xFFu) << 8) | + ((static_cast(pc) & 0xFFu) << 16) | ((static_cast(pd) & 0xFFu) << 24); +} + +} // namespace flashinfer::mamba::conversion diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh new file mode 100644 index 000000000000..c12ae8b5841b --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh @@ -0,0 +1,1117 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * 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. + */ +#ifndef FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_CUH_ +#define FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_CUH_ + +// Incremental SSU kernel — matmul-based, tensor-core MMA (single path). +// Single CTA per (batch, head). Grid: (batch, nheads). +// 4 warps per CTA, 128 threads total. +// +// Single __syncthreads() via per-warp data ownership. Every smem +// read before the final barrier is served by data the same warp loaded. +// No mbarriers, no cross-warp visibility for the first half of the kernel. +// +// Phase 0 (per-warp cp.async, no cross-warp sync): +// State: each warp loads own DIM slice (rows [16W : 16W+16]). +// B, C: redundant on W0, W1 (both do 2-warp CB). +// old_B: redundant on all 4 warps (each warp's replay needs full DSTATE). +// old_x: redundant on all 4 warps. +// x: W2 only (Phase-2 read — covered by the single syncthreads). +// z: W3 only (Phase-2 read — covered by the single syncthreads). +// Scalars + cumAdt: redundant on each warp's first NPREDICTED/MAX_WINDOW lanes. +// Each warp: __pipeline_commit → __pipeline_wait_prior(0) → __syncwarp. +// +// Phase 1 (runs with *no* barrier; CB ‖ replay parallelism preserved): +// - store_old_B hoisted here — W0,W1 only (they hold valid smem.B). +// - Warps 0,1: compute_CB_scaled_2warp (bf16 HMMA → swizzled smem.CB_scaled). +// - All warps: replay_state_mma (HMMA; state in smem updated in-place, +// each warp touches only its own DIM rows). +// +// __syncthreads() ← THE ONE. Provides cross-warp visibility of: +// CB_scaled (W0,W1→all), x (W2→all), z (W3→all). +// +// Phase 2: compute_and_store_output +// = (C @ state^T) * decay + CB_scaled @ x + D*x, z-gate → direct gmem STG. +// State writeback hoisted inside the orchestrator once matmul 3 has +// finished consuming smem.state. +// +// Phase 3: old_x / old_dt / old_cumAdt cache writes. + +#include "kernel_checkpointing_ssu_common.cuh" + +namespace flashinfer::mamba::checkpointing { + +// ============================================================================= +// Shared memory layout +// ============================================================================= +// smem holds the state in its native dtype (`state_t`). cp.async pulls the +// native dtype straight into smem (with the matching `SmemSwizzle`), +// and the conversion to `MMA_prop::operand_t` happens on the register read inside +// add_init_out / replay_state_mma. +// `D_PER_CTA` is the per-CTA D dimension after D-split. +// For D_SPLIT = 1 (default), D_PER_CTA == DIM (per-head DIM). At D_SPLIT > 1 +// the storage is sliced: each CTA owns a contiguous D_PER_CTA-row slice of +// the head's D axis. Buffers that aren't D-owned (B, C, old_B, scalars) are +// unaffected. +// +// Note on D_SMEM_COLS: the Swizzle<3,3,3> atom for bf16 is (8, 64), so +// `make_swizzled_layout_rc` requires col counts to be multiples of 64. When +// D_PER_CTA < 64 (e.g. D_SPLIT=2 → D_PER_CTA=32) we pad the D-owned buffer +// cols up to the swizzle atom width. The cp.async only fills the first +// D_PER_CTA cols; the padded tail is unused but keeps the swizzle layout +// well-formed. Cost: 1 KB per [NPREDICTED_PAD_MMA_M, 64] buffer at D_PER_CTA=32. +template +struct CheckpointingSsuStorage { + // Re-export the two T/W axis sizes so helpers that only see SmemT can + // recover them. + static constexpr int NPREDICTED = NPREDICTED_; + static constexpr int MAX_WINDOW = MAX_WINDOW_; + // Swizzle atom width for input_t (= 64 cols for 2-byte types). + static constexpr int D_SMEM_COLS = next_multiple_of::ATOM_COLS>(D_PER_CTA); + // M-dim of the output MMAs (C, x, z, CB_scaled): always m16-tiled (keyed + // off NPREDICTED, the new-tokens count). + static constexpr int NPREDICTED_PAD_MMA_M = next_multiple_of(NPREDICTED); + // N-dim of the precompute-CB MMA (matmul-1: C @ B^T). B's row count is + // the matmul N-axis → padded to MMA::N=8. When NPREDICTED ≤ 8 only warp + // 0 has valid B rows; warp 1 zero-fills its CB slice. + static constexpr int NPREDICTED_PAD_MMA_N = next_multiple_of(NPREDICTED); + // K-dim of the replay MMA (matmul-2: old_x^T @ dB_scaled). Padded to + // the small atom's K (== the LDSM unit for 2-byte elements). When + // MAX_WINDOW ≤ MMA::K_SMALL=8, replay picks the small atom (1 K-tile, + // smaller smem, +1 CTA/SM occupancy); otherwise the big atom. Assumes + // MAX_WINDOW ≤ MMA::K_BIG (asserted in the wrapper). + static constexpr int MAX_WINDOW_PAD_MMA_K = next_multiple_of(MAX_WINDOW); + // Row count for buffers padded only to the input-type swizzle atom's row + // extent (8 for 2-byte, 4 for 4-byte) — used by C and z, which alias the + // second m-tile back onto the first via `make_aliased_swizzled_layout_rc`. + // Keyed off NPREDICTED. + static constexpr int NPREDICTED_SWIZZLE_R = + next_multiple_of::ATOM_ROWS>(NPREDICTED); + + // All 2D smem buffers below are stored as flat 1D arrays — the actual + // physical layout is determined by `make_swizzled_layout_rc<...>` at each + // access site, which scrambles (row, col) → physical offset via the + // Swizzle XOR. Declaring them as `T[ROWS][COLS]` would falsely suggest a + // row-major C-array layout that nobody ever uses; the only thing that + // matters here is total byte count and 16-byte alignment. + + // CB_scaled — logical (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE) Swizzle<3,3,3>. + // CB_ROW_STRIDE pads each row to one bank cycle (128 B = 32 banks × 4 B) + // worth of `input_t` so LDSM reads in matmul-4's A operand are + // conflict-free. Equals the swizzle atom's col extent for `input_t` + // (64 for 2-byte, 32 for 4-byte). Logical CB matrix is + // (NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M); trailing cols are padding. + static constexpr int CB_ROW_STRIDE = SmemSwizzle::ATOM_COLS; + alignas(16) input_t CB_scaled[NPREDICTED_PAD_MMA_M * CB_ROW_STRIDE]; + + // B — logical (NPREDICTED_PAD_MMA_N, DSTATE). Row count is matmul-1's + // N-axis (since matmul-1 = C @ B^T). Padding rows inside [NPREDICTED, + // NPREDICTED_PAD_MMA_N) contain garbage — valid output uses only + // [0, NPREDICTED). Warp-1 of compute_CB_scaled_2warp reads rows ≥ 8 of a + // 16-row view; those reads spill into C/old_B smem but are masked to 0 by + // the (j < NPREDICTED) CB-store predicate since j ≥ 8 ≥ NPREDICTED when + // NPREDICTED_PAD_MMA_N == 8. + alignas(16) input_t B[NPREDICTED_PAD_MMA_N * DSTATE]; + + // C — physical (next_multiple_of(NPREDICTED), DSTATE). Padded + // only to the swizzle atom's row extent (8 for 2-byte, 4 for 4-byte), not + // to MMA_prop::M=16. cp.async writes to this exact extent (CShape's first + // dim shrunk to match — see load_data). The MMA still views it as + // NPREDICTED_PAD_MMA_M=16 rows via `make_aliased_swizzled_layout_rc`, + // which aliases the second m-tile back onto the first via stride-0 + // row-tile mode. Garbage feeds output rows ≥ NPREDICTED — predicated + // out at gmem store. Saves up to 2 KB of smem at NPREDICTED ≤ ATOM_ROWS, + // no-op when NPREDICTED > ATOM_ROWS. + alignas(16) input_t C[NPREDICTED_SWIZZLE_R * DSTATE]; + + // x — logical (NPREDICTED_PAD_MMA_M, D_SMEM_COLS). Cols padded to + // D_SMEM_COLS for swizzle atom alignment; cp.async only fills cols + // [0, D_PER_CTA), the tail is unused. + alignas(16) input_t x[NPREDICTED_PAD_MMA_M * D_SMEM_COLS]; + + // z — physical (next_multiple_of(NPREDICTED), D_SMEM_COLS). + // Padded only to the swizzle atom's row extent (8 for 2-byte, 4 for + // 4-byte), not to MMA_prop::M=16. z is never an MMA operand — the + // z-gating epilogue reads it via `partition_C` of the m16n8 c-frag, so + // the MMA still views it as NPREDICTED_PAD_MMA_M=16 rows via + // `make_aliased_swizzled_layout_rc`, which aliases the second m-tile back + // onto the first via stride-0 row-tile mode. Garbage feeds output rows + // ≥ NPREDICTED — predicated out at gmem store. Saves up to 1 KB of smem + // at NPREDICTED ≤ ATOM_ROWS, no-op when NPREDICTED > ATOM_ROWS. + alignas(16) input_t z[NPREDICTED_SWIZZLE_R * D_SMEM_COLS]; + + // Old cache data loaded in Phase 0 (consumed in Phase 1 replay). + // old_x — logical (MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS); ldmatrix.trans + // feeds replay MMA A-operand (only the first D_PER_CTA cols are valid + // data). + alignas(16) input_t old_x[MAX_WINDOW_PAD_MMA_K * D_SMEM_COLS]; + + // old_B — logical (MAX_WINDOW_PAD_MMA_K, DSTATE) Swizzle<3,3,3>. Replay + // MMA reads via ldmatrix.trans (LDSM_T) + register scaling. Padding + // rows zero-filled via cp.async ZFILL. + alignas(16) input_t old_B[MAX_WINDOW_PAD_MMA_K * DSTATE]; + + float old_dt[MAX_WINDOW]; + float old_cumAdt[MAX_WINDOW]; + + // Processed dt for new tokens (Phase 1a uses this for CB_scaled + cumAdt) + float dt_proc[NPREDICTED]; + + // Cumulative A*dt — computed once by warp 0, read by all warps after sync + float cumAdt[NPREDICTED]; + + // state — logical (D_PER_CTA, DSTATE) in `state_t` (native dtype). The + // MMA path reinterprets 2-byte state as bf16 for LDSM; f32 state is loaded + // via UniversalCopy and converted to bf16 in registers inside + // add_init_out. + alignas(16) state_t state[D_PER_CTA * DSTATE]; +}; + +// ============================================================================= +// Stochastic-round one fp32 pair to a packed f16x2 u32 with amortized philox +// refresh. rand_idx[4] is mutated in place every 4th call (when pair_idx & 3 +// == 0): a single philox_randint4x feeds 4 consecutive cvt_rs calls, then +// gets refreshed. Each refresh uses a per-lane unique `philox_off` so the +// generated randints don't collide across threads. Triton bit-equality is +// intentionally given up here; unbiasedness still holds since each pair's +// cvt_rs gets its own dedicated 32-bit randint. +// ============================================================================= +template +__device__ __forceinline__ uint32_t +stochastic_round_pair_with_philox_refresh(float a, float b, int pair_idx, int64_t rand_seed, + int64_t philox_off, uint32_t (&rand_idx)[4]) { + int const rand_pos = pair_idx & 3; + if (rand_pos == 0) { + conversion::philox_randint4x(rand_seed, philox_off, rand_idx[0], rand_idx[1], + rand_idx[2], rand_idx[3]); + } + return conversion::cvt_rs_f16x2_f32(a, b, rand_idx[rand_pos]); +} + +// ============================================================================= +// Cross-pass shfl_xor + STG.64 state writeback. +// +// Given two passes' worth of post-cvt_rs packed u32s buffered in `my_packed` +// (pass-0 in [0][:], pass-1 in [1][:]), exchange via shfl_xor across lane^1 +// neighbors so that all 32 lanes can issue ONE STG.64 each per pair iter: +// - even lane k stores PASS n0 (cols (k%4)*2..(k%4)*2+3 of warp's n0 slice) +// - odd lane k stores PASS n1 (cols (k%4)*2-2..(k%4)*2+1 of warp's n1 slice) +// +// Halves the STG instruction count vs per-pass writeback: 1 STG.64 per pair +// iter covers BOTH passes' data via cross-lane participation. +// ============================================================================= +template +__device__ __forceinline__ void exchange_ntile_state_store_global( + state_t* __restrict__ state_w_base, int np, int lane, + uint32_t const (&my_packed)[2][PAIRS_PER_PASS], IdPart const& id_part) { + using namespace cute; + static_assert(sizeof(state_t) == 2, + "exchange_ntile_state_store_global requires 2-byte state_t for STG.64 alignment"); + int const n_base_p0 = np * N_PER_PASS; + int const n_base_p1 = (np + 1) * N_PER_PASS; +#pragma unroll + for (int p = 0; p < PAIRS_PER_PASS; ++p) { + int const i = p * 2; + // xor mask = 1 swaps neighbor lanes: lane 0 <-> lane 1, lane 2 <-> lane 3, ... + uint32_t const peer_p0 = __shfl_xor_sync(constants::MASK_ALL_LANES, my_packed[0][p], 1); + uint32_t const peer_p1 = __shfl_xor_sync(constants::MASK_ALL_LANES, my_packed[1][p], 1); + + int const row = get<0>(id_part(i)); + int const col_p0 = get<1>(id_part(i)) + n_base_p0; + int const col_p1 = get<1>(id_part(i)) + n_base_p1; + + uint64_t combined; + int32_t gmem_off; + if ((lane & 1) == 0) { + // Even lane: store PASS n0 — my (lower col) in low, peer in high. + combined = static_cast(my_packed[0][p]) | + (static_cast(peer_p0) << constants::num_bits_uint32); + gmem_off = row * DSTATE + col_p0; + } else { + // Odd lane: store PASS n1 — peer (lower col) in low, my in high. + // STG addr = gmem[row*DSTATE + (peer's col base)] = col_p1 - 2. + combined = static_cast(peer_p1) | + (static_cast(my_packed[1][p]) << constants::num_bits_uint32); + gmem_off = row * DSTATE + (col_p1 - 2); + } + *reinterpret_cast(&state_w_base[gmem_off]) = combined; + } +} + +// ============================================================================= +// Phase 1b: Replay — tensor-core MMA path (matmul 2: state recurrence). +// state[D, dstate] = state * total_decay + old_x^T @ (coeff * old_B) +// All 128 threads cooperate. +// +// Warps along N=DSTATE: +// TiledMMA uses Layout<_1, _4> — per pass covers (M=DIM, N=4×MMA_prop::N=32). +// Each warp owns: full M (DIM/16 m-atoms) and one n-atom of 8 cols. +// Why: A is small (DIM × K), B is bigger (DSTATE × K). M-split (`_4×1`) +// redundantly loaded full B from each warp (4× × 4 KB = 16 KB). N-split +// (`_1×4`) instead redundantly loads full A (4× × 2 KB = 8 KB) and reads +// B disjointly across warps — net smem read drops 18 KB → 12 KB per replay +// (~33%) at K_BIG. Also unlocks D-split D_PER_CTA < 64. +// ============================================================================= +// state_w_base (f16+philox path): pre-offset gmem pointer to this CTA's owned +// [D_PER_CTA, DSTATE] state slice (params.state + cache_slot * +// state_stride_seq + head * DIM*DSTATE + d_tile * D_PER_CTA*DSTATE). +// Computed in the kernel preamble. Combining base + offset into one i64 +// pointer drops the cross-iter live-range cost from 4 regs (state_w ptr + +// state_gmem_off) to 2 regs (just the base), and the per-pair STG.32 uses an +// i32 element offset inside the chunk. Use this instead of separately +// holding params.state-ptr and state_gmem_off. +template +__device__ __forceinline__ void replay_state_mma(SmemT& smem, CheckpointingSsuParams const& params, + int warp, int lane, int prev_k, int d_tile, + int64_t state_ptr_offset, state_t* state_w_base, + int64_t rand_seed, bool must_checkpoint) { + using namespace cute; + static_assert(sizeof(input_t) == 2, "replay_state_mma requires 2-byte input type"); + static_assert(D_PER_CTA % 16 == 0, "D_PER_CTA must be divisible by 16 (m16n8 atom)"); + static_assert(D_PER_CTA >= 16, "D_PER_CTA must be at least 16"); + + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; // 8 or 16 + int const tid = warp * warpSize + lane; + + // Atom K matches the cache-window tile (MAX_WINDOW_PAD_MMA_K). + // K == MMA_prop::K_BIG (16) → m16n8k16 + x4/x2 ldmatrix.trans + // K == MMA_prop::K_SMALL (8) → m16n8k8 + x2/x1 ldmatrix.trans + using MmaAtomType = std::conditional_t; + using LdsmA = std::conditional_t; + using LdsmB = std::conditional_t; + + // 4 warps along N=DSTATE; each warp covers full M (D_PER_CTA/16 m-atoms). + auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(tid); + + // Per-pass output tile is (D_PER_CTA, N_PER_PASS). N_PER_PASS = 4 warps × n8 = 32 cols. + constexpr int N_PER_PASS = 4 * MMA_prop::N; + static_assert(DSTATE % N_PER_PASS == 0, + "DSTATE must be divisible by 4 * MMA_prop::N for _1x4 warp layout"); + constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; + + float total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + float total_decay = (prev_k > 0) ? __expf(total_cumAdt) : 1.f; + + // ── A operand: old_x [MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS] Swizzle<3,3,3>, transposed + // view [M=D_SMEM_COLS, K=MAX_WINDOW_PAD_MMA_K]. D_SMEM_COLS may be padded above + // D_PER_CTA when D_PER_CTA < swizzle atom; local_tile to D_PER_CTA + // restricts the LDSM to the valid sub-tile. Each warp loads the FULL M (4× + // redundant across warps). See header comment for traffic accounting. ── + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + auto layout_A_full = + make_swizzled_layout_rc_transpose(); + Tensor smem_A_full = make_tensor( + make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); + Tensor smem_A = local_tile(smem_A_full, make_shape(Int{}, Int{}), + make_coord(_0{}, _0{})); + + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); + Tensor frag_A = thr_mma.partition_fragment_A(make_tensor( + (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); + Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); + + cute::copy(s2r_A, smem_A_s2r, frag_A_view); + // old_x is input_t == MMA_prop::operand_t (bf16) — no conversion needed. + + // ── B operand: old_B [MAX_WINDOW_PAD_MMA_K, DSTATE] swizzled, transposed view + // [N=DSTATE, K=MAX_WINDOW_PAD_MMA_K]. Per pass loads N_PER_PASS=32 cols across + // 4 warps; partition_S splits — each warp gets its disjoint 8-col slice. ── + auto layout_B = make_swizzled_layout_rc_transpose(); + Tensor smem_B_full = make_tensor( + make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B); + + auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B = s2r_B.get_slice(tid); + + // ── State: per-CTA swizzle layout [D_PER_CTA, DSTATE]. ── + auto layout_state_swz = make_swizzled_layout_rc(); + state_t* state_base = reinterpret_cast(smem.state); + + // ── Per-pass identity for (row, col) coords ── + // partition_C of an identity tensor of the per-pass output shape gives this + // thread's (row, col) at every C-frag position, including warp-N offset. + // Frag size per thread = (M_atoms=D_PER_CTA/16) × (N_atoms_per_warp=1) × 4 elts. + auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_tile); + // Linear order from CuTe's column-major partition_C with m16n8 atom: + // i=0,1: same row (= row_lo of M-atom 0), adjacent cols (col_off, col_off+1) + // i=2,3: same row (= row_hi of M-atom 0), adjacent cols + // i=4,5: same row (= row_lo of M-atom 1) + // ... (V index 0..3 inside each m16n8, then M-atoms in M-major order) + // Pair load at (i, i+1) covers two consecutive bf16 elts → one 32-bit LDS. + + // Precompute dB coefficients once — depend only on K (lane), not on N. + constexpr int LANES_PER_N_COL = warpSize / MMA_prop::N; // = 4 for m16n8k_ + constexpr int DB_COEFFS_PER_LANE = MAX_WINDOW_PAD_MMA_K / LANES_PER_N_COL; + float dB_coeff[DB_COEFFS_PER_LANE]; + precompute_dB_coeff(dB_coeff, smem, total_cumAdt, prev_k, lane); + + using pair_t = Pair; + + // Philox state amortized across 4 consecutive pair conversions: each call + // returns 4 randints, all 4 get consumed before the next refresh (vs. 1-of-4 + // in the Triton-bit-equal layout — see writeback loop below). Compile-time + // pair_idx (n-loop and i-loop both unrolled) keeps `rand_idx[pair_idx & 3]` + // as a known register access — no local-memory spill. + constexpr bool kPhiloxF16 = (PHILOX_ROUNDS > 0) && std::is_same_v; + [[maybe_unused]] uint32_t rand_idx[4]; + // state_w_base is the pre-combined (params.state + state_gmem_off) base + // pointer — see the function header. No separate state_w / state_gmem_off + // alive in this scope. + + // ── Vectorized state writeback (cross-pass STG.64 fusion) ────────── + // smem always gets nearest-even f32→state_t (consumed by matmul 3 — must + // match Triton's f32→bf16 path as closely as possible). Gmem cache, when + // PHILOX_ROUNDS > 0 and state_t == __half, gets PTX cvt.rs.f16x2.f32 + // stochastic rounding direct from registers via cross-pass STG.64; the + // smem→gmem `store_state` is gated off in compute_and_store_output. + // + // Cross-pass STG fusion: do PASS n0 and PASS n1 back-to-back, buffering + // the post-cvt_rs packed u32s of n0 across n1's HMMA + cvt_rs. Then issue + // ONE STG.64 instruction per pair iter, all 32 lanes active: + // - even lane stores PASS n0 data at the warp's n0 column slice + // - odd lane stores PASS n1 data at the warp's n1 column slice + // Halves the STG instruction count vs per-pass writeback (16 STG.64/thread + // per 2 passes vs 16 + 16 = 32 STG.64/thread previously — same byte volume). + // + // Randint amortization: rand_idx[4] refreshed every 4 pairs; each pair's + // cvt_rs uses one of the 4 randints. Triton bit-equality is intentionally + // given up; unbiasedness still holds. + constexpr int PAIRS_PER_PASS = D_PER_CTA / 8; // = (D_PER_CTA/16) × 2 row-pair iters + static_assert(NUM_N_PASSES % 2 == 0, "Cross-pass STG fusion requires even NUM_N_PASSES"); + +#pragma unroll + for (int np = 0; np < NUM_N_PASSES; np += 2) { + // Buffer of post-cvt_rs packed u32s for both passes (philox path only). + [[maybe_unused]] uint32_t my_packed[2][PAIRS_PER_PASS]; + +#pragma unroll + for (int local_n = 0; local_n < 2; ++local_n) { + int const n = np + local_n; + int const n_base = n * N_PER_PASS; + + // ── Allocate per-pass C-frag (4 × M_atoms fp32 elts/thread) ── + Tensor frag_h = thr_mma.partition_fragment_C( + make_tensor((float*)0x0, make_shape(Int{}, Int{}))); + + // ── Load state × total_decay into frag_h. ── +#pragma unroll + for (int i = 0; i < size(frag_h); i += 2) { + int const row = get<0>(id_part(i)); + int const col = get<1>(id_part(i)) + n_base; + int const off = layout_state_swz(row, col); + pair_t const p = *reinterpret_cast(&state_base[off]); + frag_h(i) = toFloat(p[cute::Int<0>{}]) * total_decay; + frag_h(i + 1) = toFloat(p[cute::Int<1>{}]) * total_decay; + } + + // ── LDSM.T per-pass B (per warp = 1 atom of 8 cols of N) ── + Tensor smem_B_n = + local_tile(smem_B_full, make_tile(Int{}, Int{}), + make_coord(n, _0{})); + auto smem_B_s2r_n = s2r_thr_B.partition_S(smem_B_n); + + Tensor frag_B = thr_mma.partition_fragment_B(make_tensor( + (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); + auto frag_B_view = s2r_thr_B.retile_D(frag_B); + + cute::copy(s2r_B, smem_B_s2r_n, frag_B_view); + + compute_dB_scaling(frag_B, dB_coeff); + + // ── HMMA: frag_h += frag_A @ frag_B ── + cute::gemm(tiled_mma, frag_h, frag_A, frag_B, frag_h); + + // ── Smem write (always) + cvt_rs into my_packed (philox path) ── +#pragma unroll + for (int i = 0; i < size(frag_h); i += 2) { + int const row = get<0>(id_part(i)); + int const col = get<1>(id_part(i)) + n_base; + int const off = layout_state_swz(row, col); + + // Smem write — always nearest-even (output's matmul 3 reads this). + pair_t const q = pack_float2(make_float2(frag_h(i), frag_h(i + 1))); + *reinterpret_cast(&state_base[off]) = q; + + if constexpr (kPhiloxF16) { + static_assert(sizeof(state_t) == 2, "STG.64 cooperative path requires 2-byte state_t"); + int const pair_idx = n * PAIRS_PER_PASS + i / 2; + // Per-lane philox_off is unique per (thread, refresh group) — each + // pair gets its own randint bits. Always computed; only consumed + // by the refresh branch inside the helper. + int64_t const philox_off = + state_ptr_offset + (int64_t)(d_tile * D_PER_CTA + row) * DSTATE + col; + // Buffer the SR'd packed u32 — store happens after BOTH passes. + my_packed[local_n][i / 2] = stochastic_round_pair_with_philox_refresh( + frag_h(i), frag_h(i + 1), pair_idx, rand_seed, philox_off, rand_idx); + } + } + } + + // ── Cross-pass STG.64: all 32 lanes active. ───────────────────────── + // m16n8 lane layout: lane k → row k/4, cols (k%4)*2..(k%4)*2+1. Lanes + // (2k, 2k+1) hold adjacent col-pairs of the same row. After shfl_xor, + // the even/odd lane each has a 4-col contiguous block (in different + // bit-orders). Even lane STG.64s the n0-pass block at its own col + // base; odd lane STG.64s the n1-pass block at the peer's (lower) col + // — both 8-byte aligned for state_t = f16. + // Runtime-gated on must_checkpoint: non-checkpoint steps skip the gmem + // STGs entirely (state HBM remains the prior checkpoint). The cvt_rs + // SR + philox refresh above still ran — only the STGs are elided — + // because skipping them would require routing must_checkpoint into the + // pair_idx amortization logic, which lives across the n-loop. + if constexpr (kPhiloxF16) { + if (must_checkpoint) { + exchange_ntile_state_store_global( + state_w_base, np, lane, my_packed, id_part); + } + } + } +} + +// ── Orchestrator: compute_and_store_output ───────────────────────────── +// out = (C @ state^T) * decay + CB_scaled @ x + D*x, then z-gate. +// All operations on register-resident frag_y — no smem round-trip. +// Result converted f32 → input_t in registers and stored directly to gmem +// via partition_C of the global output tensor (like CUTLASS sgemm_sm80 epilogue). +template +__device__ __forceinline__ void compute_and_store_output(SmemT& smem, + CheckpointingSsuParams const& params, + int warp, int lane, int d_tile, + int64_t out_seq_base, int head, + int64_t cache_slot, float D_val, + bool must_checkpoint, int seq_len) { + using namespace cute; + static_assert(sizeof(input_t) == 2, "compute_and_store_output requires 2-byte input type"); + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + int const tid = warp * warpSize + lane; + + // ── TiledMMA: 128 threads, covers [16, 32] output per step ── + auto tiled_mma = + make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(tid); + + // ── Swizzled smem views ── + // When D_PER_CTA < swizzle atom (= 64 for bf16), the underlying + // smem buffer is padded to D_SMEM_COLS so the swizzle layout is well-formed. + // Per-pass MMA loops only iterate D_PER_CTA / N_TILE tiles → never touch + // the padded tail. + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + + // x: swizzled [NPREDICTED_PAD_MMA_M, D_SMEM_COLS] + auto layout_x_swz = make_swizzled_layout_rc(); + Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), + layout_x_swz); + auto layout_x_trans_swz = + make_swizzled_layout_rc_transpose(); + Tensor smem_x_trans = make_tensor( + make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); + + // z: aliased swizzled [NPREDICTED_PAD_MMA_M, D_SMEM_COLS] — physical buffer + // is only next_multiple_of(NPREDICTED) rows tall; second m-tile + // aliases first. Ghost rows feed predicated-out output rows. + auto layout_z_swz = + make_aliased_swizzled_layout_rc(); + Tensor smem_z = + make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); + + // ── S2R copies ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B = s2r_B.get_slice(tid); + auto s2r_B_trans = + make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); + + // ── Load CB_scaled A operand from smem (precomputed by warps 0,1 between syncs) ── + // Row stride matches the buffer's padded width (one swizzle atom of `input_t`). + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + auto layout_cb_swz = + make_swizzled_layout_rc(); + Tensor smem_CB = make_tensor( + make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); + auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); + Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); + auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); + cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); + + // Decay broadcast: cumAdt[t] → [NPREDICTED_PAD_MMA_M, N_TILE] with stride-0 on N. + constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); + Tensor decay_bcast = make_tensor( + make_smem_ptr(smem.cumAdt), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); + Tensor decay_part = thr_mma.partition_C(decay_bcast); + + // ── Gmem output: partition_C for direct register → gmem store ── + auto* __restrict__ output_ptr = reinterpret_cast(params.output); + // out_base lands on this CTA's D-slice within the head. + int64_t const out_base = out_seq_base + (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; + + // Row predicate for padding. The epilogue store loop iterates i in steps + // of 2 and only consults pred(0) and pred(2) — m16n8k16 C-frag per thread + // has 4 elts at rows {t/4, t/4, t/4+8, t/4+8}, so there are only 2 unique + // row predicates. Compute them once and skip the 4-wide pred tensor. + auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_tile); + bool const pred_row_lo = get<0>(id_part(0)) < seq_len; + bool const pred_row_hi = get<0>(id_part(2)) < seq_len; + + // Number of output N-tiles per pass = D_PER_CTA / N_TILE. + // D_SPLIT=1, D_PER_CTA=64, N_TILE=32 → NUM_N_TILES = 2 (current behavior). + // D_SPLIT=2, D_PER_CTA=32 → NUM_N_TILES = 1 (uses _n1 variant). + constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; + static_assert(NUM_N_TILES == 1 || NUM_N_TILES == 2, + "Output epilogue supports NUM_N_TILES = D_PER_CTA / N_TILE in {1, 2}"); + + // ── Epilogue lambda (defined once; called per N-tile from each branch) ── + auto epilogue = [&](auto& frag_y, int n) { + // Decay: frag_y *= exp(cumAdt[t]) +#pragma unroll + for (int i = 0; i < size(frag_y); ++i) { + frag_y(i) *= __expf(decay_part(i)); + } + + // frag_y += CB_scaled @ x (CB from smem LDSM, x from smem via ldmatrix.trans) + add_cb_x( + frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); + + // frag_y += D * x[t, d] + add_D_skip(frag_y, smem_x, thr_mma, D_val, n); + + // frag_y *= z * sigmoid(z) + compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); + + // Store frag_y directly to gmem (register → gmem, no smem round-trip). + auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), + make_layout(make_shape(Int{}, Int{}), + make_stride(params.out_stride_token, _1{}))); + auto gOut_part = thr_mma.partition_C(gOut_tile); + // Vectorized pair store: elements i and i+1 are same-row, consecutive columns + // in the m16n8k16 partition_C layout, so &gOut_part(i+1) == &gOut_part(i) + 1. + // Address is naturally aligned to sizeof(Pair) since MMA column + // index = (lane%4)*2 → even. pack_float2 dispatches to the native packed + // cvt (e.g. cvt.rn.bf16x2.f32 for bf16) — one instruction for the pair. +#pragma unroll + for (int i = 0; i < size(frag_y); i += 2) { + // Bit 1 of i toggles between the two row groups of the m16n8k16 + // C-frag: i∈{0,1} → row t/4, i∈{2,3} → row t/4+8 (repeats per M-atom). + bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; + if (pred_i) { + *reinterpret_cast*>(&gOut_part(i)) = + pack_float2(make_float2(frag_y(i), frag_y(i + 1))); + } + } + }; + + // Skip the smem→gmem state copy when philox+f16: `replay_state_mma` + // already did the gmem store with stochastic rounding direct from registers. + constexpr bool kSkipSmemToGmemState = (PHILOX_ROUNDS > 0) && std::is_same_v; + + // ── Matmul 3 + store_state + epilogue, dispatching on NUM_N_TILES ── + // (NumNTiles is deduced from the variadic frag_y... pack in `add_init_out`.) + if constexpr (NUM_N_TILES == 2) { + Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); + Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); + add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0, + frag_y_1); + // State writeback hoisted here — after matmul 3 has finished consuming + // smem.state, before matmul 4 which reads only smem.x / smem.CB_scaled + // / smem.z. STGs fire-and-forget alongside the epilogue (matmul 4 + + // D*x + z-gate + output STG). Runtime-gated on must_checkpoint: + // non-checkpoint steps leave the prior state HBM intact (saving + // bandwidth — that's the perf win of the checkpointing design). + if constexpr (!kSkipSmemToGmemState) { + if (must_checkpoint) { + store_state(smem, params, warp, lane, d_tile, + head, cache_slot); + } + } + epilogue(frag_y_0, 0); + epilogue(frag_y_1, 1); + } else { // NUM_N_TILES == 1 (D_SPLIT = 2 path) + Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); + add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0); + // No sync needed before store_state: the post-replay __syncthreads() + // in the kernel already established cross-warp visibility of replay's + // writes to smem.state, and nothing after that point writes to it + // (add_init_out is read-only on smem.state). + if constexpr (!kSkipSmemToGmemState) { + if (must_checkpoint) { + store_state(smem, params, warp, lane, d_tile, + head, cache_slot); + } + } + epilogue(frag_y_0, 0); + } +} + +// ── Orchestrator: compute_no_write_output (must_checkpoint == false path) ── +// Skips the replay matmul entirely. smem.state still holds s_0 after Phase 0, +// so matmul-3 via add_init_out computes u^T = C @ s_0^T directly. +// +// y[t, d] = β(t) · u[t, d] +// + Σ_{j +__device__ __forceinline__ void compute_no_write_output( + SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int prev_k, int d_tile, + int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) { + using namespace cute; + static_assert(sizeof(input_t) == 2, "compute_no_write_output requires 2-byte input type"); + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + int const tid = warp * warpSize + lane; + + // ── TiledMMA for matmul-3 + matmul-4-new (K=NPREDICTED_PAD_MMA_M=16 fits K_BIG). ── + auto tiled_mma = + make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(tid); + + // ── TiledMMA for matmul-4-old: K = MAX_WINDOW_PAD_MMA_K ∈ {8, 16} → atom dispatch. ── + using MmaAtomOld = std::conditional_t; + using LdsmAOld = std::conditional_t; + using LdsmBOld = std::conditional_t; + auto tiled_mma_old = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_old = tiled_mma_old.get_slice(tid); + + // ── Swizzled smem views ── + auto layout_x_swz = make_swizzled_layout_rc(); + Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), + layout_x_swz); + auto layout_x_trans_swz = + make_swizzled_layout_rc_transpose(); + Tensor smem_x_trans = make_tensor( + make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); + + auto layout_old_x_trans_swz = + make_swizzled_layout_rc_transpose(); + Tensor smem_old_x_trans = + make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), + layout_old_x_trans_swz); + + auto layout_z_swz = + make_aliased_swizzled_layout_rc(); + Tensor smem_z = + make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); + + // ── S2R copies (matmul-4-new) ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + auto s2r_B_trans = + make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); + + // ── S2R copies (matmul-4-old, K-dispatched atoms) ── + auto s2r_A_old = make_tiled_copy_A(Copy_Atom{}, tiled_mma_old); + auto s2r_thr_A_old = s2r_A_old.get_slice(tid); + auto s2r_B_old_trans = + make_tiled_copy_B(Copy_Atom{}, tiled_mma_old); + auto s2r_thr_B_old_trans = s2r_B_old_trans.get_slice(tid); + + // ── Load CB_scaled A operand (cols [0, NPREDICTED_PAD_MMA_M)) ── + auto layout_cb_swz = + make_swizzled_layout_rc(); + Tensor smem_CB = make_tensor( + make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); + auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); + Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); + auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); + cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); + + // ── Load CB_old A operand (cols [NPREDICTED_PAD_MMA_M, +MAX_WINDOW_PAD_MMA_K)) ── + // Use the full physical (T_pad, CB_ROW_STRIDE) padded swizzle view — byte- + // compatible with both the CB_scaled (T_pad, T_pad, CB_ROW_STRIDE) write + // layout and compute_CB_old_2warp's wide write layout (inner offset + // r*CB_ROW_STRIDE + c is identical across the three views). + auto layout_cb_full = make_swizzled_layout_rc(); + Tensor smem_CB_full = make_tensor( + make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); + Tensor smem_CB_old = + local_tile(smem_CB_full, make_tile(Int{}, Int{}), + make_coord(_0{}, NPREDICTED_PAD_MMA_M / MAX_WINDOW_PAD_MMA_K)); + auto smem_CB_old_s2r = s2r_thr_A_old.partition_S(smem_CB_old); + Tensor frag_CB_old_A = thr_mma_old.partition_fragment_A(smem_CB_old); + auto frag_CB_old_A_view = s2r_thr_A_old.retile_D(frag_CB_old_A); + cute::copy(s2r_A_old, smem_CB_old_s2r, frag_CB_old_A_view); + + // ── Decay broadcast: cumAdt[t] (per-T scalar) with stride-0 on N. ── + constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); + Tensor decay_bcast = make_tensor( + make_smem_ptr(smem.cumAdt), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); + Tensor decay_part = thr_mma.partition_C(decay_bcast); + + // ── β extra factor: exp(total_old_cumAdt) — uniform constant across (t, d) ── + float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + float const beta_extra = __expf(total_old_cumAdt); + + // ── Gmem output base ── + auto* __restrict__ output_ptr = reinterpret_cast(params.output); + int64_t const out_base = out_seq_base + (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; + + // ── Row predicate (same pattern as compute_and_store_output) ── + auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_tile); + bool const pred_row_lo = get<0>(id_part(0)) < seq_len; + bool const pred_row_hi = get<0>(id_part(2)) < seq_len; + + constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; + static_assert(NUM_N_TILES == 1 || NUM_N_TILES == 2, + "compute_no_write_output supports NUM_N_TILES in {1, 2}"); + + // ── Epilogue per N-tile ── + auto epilogue = [&](auto& frag_y, int n) { + // 1. β-scale: frag_y(t, d) *= exp(total_old_cumAdt + cumAdt[t]). + // Matmul-3 produced u^T = C @ s_0^T; this scales the u term by β + // BEFORE matmul-4 adds the CB·x and CB_old·old_x contributions. +#pragma unroll + for (int i = 0; i < size(frag_y); ++i) { + frag_y(i) *= beta_extra * __expf(decay_part(i)); + } + + // 2. frag_y += CB_scaled @ x (matmul-4 over new tokens). + add_cb_x( + frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); + + // 3. frag_y += CB_old @ old_x (matmul-4 over old tokens — NEW). + add_cb_old_x( + frag_y, frag_CB_old_A, smem_old_x_trans, s2r_B_old_trans, s2r_thr_B_old_trans, thr_mma_old, + tiled_mma_old, n); + + // 4. frag_y += D · x[t, d]. + add_D_skip(frag_y, smem_x, thr_mma, D_val, n); + + // 5. frag_y *= z · sigmoid(z). + compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); + + // 6. Store frag_y → gmem via partition_C (same pattern as compute_and_store_output). + auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), + make_layout(make_shape(Int{}, Int{}), + make_stride(params.out_stride_token, _1{}))); + auto gOut_part = thr_mma.partition_C(gOut_tile); +#pragma unroll + for (int i = 0; i < size(frag_y); i += 2) { + bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; + if (pred_i) { + *reinterpret_cast*>(&gOut_part(i)) = + pack_float2(make_float2(frag_y(i), frag_y(i + 1))); + } + } + }; + + // ── Matmul-3: frag_y = C @ s_0^T (smem.state retains s_0 since replay skipped) ── + if constexpr (NUM_N_TILES == 2) { + Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); + Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); + add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0, + frag_y_1); + epilogue(frag_y_0, 0); + epilogue(frag_y_1, 1); + } else { // NUM_N_TILES == 1 (D_SPLIT = 2 path) + Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); + add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0); + epilogue(frag_y_0, 0); + } +} + +// ── Per-path dispatchers (called from checkpointing_ssu_kernel) ── +// ssu_checkpoint: replay → sync → output (today's body). +// ssu_nocheckpoint: sync → no-write output (skips replay). +template +__device__ __forceinline__ void ssu_checkpoint(SmemT& smem, CheckpointingSsuParams const& params, + int warp, int lane, int prev_k, int d_tile, + int64_t out_seq_base, int head, int64_t cache_slot, + float D_val, int seq_len) { + // ── DO NOT HOIST `rand_seed` ── see kernel preamble for the perf rationale. + int64_t const rand_seed = (PHILOX_ROUNDS > 0) ? *params.rand_seed : 0; + // `state_ptr_offset` is int64 — matches Triton's `base_rand = + // cache_batch_idx * stride_state_batch + ...` (cache_batch_idx is .to(int64)). + // Full 64 bits flow through `philox_randint4x`, which splits low/high + // across Philox c0/c1. No collision risk at large serving cache sizes. + int64_t const state_ptr_offset = + cache_slot * params.state_stride_seq + (int64_t)head * DIM * DSTATE; + state_t* const state_w_base = reinterpret_cast(params.state) + + cache_slot * params.state_stride_seq + + (int64_t)head * DIM * DSTATE + (int64_t)d_tile * D_PER_CTA * DSTATE; + replay_state_mma( + smem, params, warp, lane, prev_k, d_tile, state_ptr_offset, state_w_base, rand_seed, + /*must_checkpoint=*/true); + + __syncthreads(); + + compute_and_store_output(smem, params, warp, lane, d_tile, out_seq_base, head, + cache_slot, D_val, /*must_checkpoint=*/true, seq_len); +} + +template +__device__ __forceinline__ void ssu_nocheckpoint(SmemT& smem, CheckpointingSsuParams const& params, + int warp, int lane, int prev_k, int d_tile, + int64_t out_seq_base, int head, int64_t cache_slot, + float D_val, int seq_len) { + // Sync makes warps 0,1's CB_scaled writes and warps 2,3's CB_old writes + // visible to all warps before matmul-4 reads CB_scaled + CB_old. Also + // covers smem.x (warp 2-loaded) and smem.z (warp 3-loaded) for Phase 2. + __syncthreads(); + + compute_no_write_output(smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, + cache_slot, D_val, seq_len); +} + +// ============================================================================= +// Kernel +// ============================================================================= +template +__global__ void checkpointing_ssu_kernel(CheckpointingSsuParams params) { + // Per-head DIM is sharded across `D_SPLIT` CTAs (D_PER_CTA each). + static_assert(DIM % D_SPLIT == 0, "DIM must be divisible by D_SPLIT"); + constexpr int D_PER_CTA = DIM / D_SPLIT; + static_assert(D_PER_CTA >= 32, + "D_PER_CTA must be >= 32 (output MMA m16n8 with _1×4 warp layout). " + "D_SPLIT=4 (D_PER_CTA=16) needs warp-count restructure."); + static_assert(NPREDICTED <= MAX_WINDOW, + "NPREDICTED must be <= MAX_WINDOW (new tokens must fit in cache)"); + static_assert(MAX_WINDOW <= MMA_prop::K_BIG, + "MAX_WINDOW must be <= MMA::K_BIG=16 (single replay K-tile assumption)"); + // Cross-check: host launcher must dispatch the template specialization + // matching the runtime params.d_split it stamped into the struct. + assert(params.d_split == D_SPLIT); + using SmemT = + CheckpointingSsuStorage; + extern __shared__ __align__(128) char smem_buf[]; + auto& smem = *reinterpret_cast(smem_buf); + + // Grid layout (D_SPLIT, batch, nheads). + int const d_tile = blockIdx.x; + int const seq = blockIdx.y; + int const head = blockIdx.z; + int const lane = threadIdx.x; + int const warp = threadIdx.y; + int const group_idx = head / HEADS_PER_GROUP; + + // ── Resolve cache slot ── + auto const* __restrict__ sbi = reinterpret_cast(params.state_batch_indices); + int64_t const cache_slot = sbi ? static_cast(sbi[seq]) : seq; + if (cache_slot == params.pad_slot_id) return; + + // ── Double-buffer index ── + auto const* __restrict__ buf_idx_ptr = reinterpret_cast(params.cache_buf_idx); + int const buf_read = __ldg(&buf_idx_ptr[cache_slot]); + + // ── prev_num_accepted_tokens ── + auto const* __restrict__ prev_ptr = reinterpret_cast(params.prev_num_accepted); + int const prev_k = prev_ptr[cache_slot]; + + // ── Varlen vs non-varlen prologue. See checkpointing_ssu_kernel_8bit for + // the rationale: `seq_len` flows downstream as a constexpr-foldable + // NPREDICTED in non-varlen, runtime int in varlen. + // + // Uniform gmem-base formula: `outer * *_stride_seq` where + // non-varlen: outer = seq (= blockIdx.y), stride_seq = x.stride(0). + // varlen : outer = cu_seqlens[seq], stride_seq = x.stride(1). + // The wrapper picks the right stride_seq value; the kernel only branches + // on whether to load cu_seqlens. + int seq_len; + int64_t outer; + if constexpr (VARLEN) { + auto const* __restrict__ cu_seqlens = reinterpret_cast(params.cu_seqlens); + // Two LDG.E.32 (not one LDG.E.64): cu_seqlens is only 4-byte aligned + // at `&cu_seqlens[seq]` when seq is odd, and PTX + // `ld.global.v2.b32` faults on a 4-byte-aligned address. ptxas emits + // the two scalar loads back-to-back; latency is hidden against the + // following ALU work. + int const bos = __ldg(&cu_seqlens[seq]); + int const eos = __ldg(&cu_seqlens[seq + 1]); + seq_len = eos - bos; + if (seq_len <= 0) return; + outer = (int64_t)bos; + } else { + seq_len = NPREDICTED; + outer = (int64_t)seq; + } + // x/B/C bases are computed inside `load_post_pdl_wait_data` from `outer` + // so the products don't get pinned in registers across `gdc_wait` (asm + // volatile blocks rematerialization; cost was ~6 extra regs). dt/z bases + // are only consumed pre-wait, and out_base only post-replay — fine to + // precompute. + int64_t const dt_seq_base = outer * params.dt_stride_seq + head; + int64_t const z_seq_base = outer * params.z_stride_seq; + int64_t const out_seq_base = outer * params.out_stride_seq; + + // ── Per-CTA implicit checkpoint criterion ── + // When the new tokens would overflow the cache buffer, we must checkpoint: + // replay [0, prev_k) into state, write state to HBM, write the new tokens + // to the **staging** buffer (1 - buf_read) at offset 0. Otherwise, we + // append the new tokens to the **active** buffer (buf_read) at offset + // prev_k and skip the state HBM write entirely. Cache writes always + // happen — only their target buffer + offset depends on must_checkpoint. + bool const must_checkpoint = (prev_k + seq_len > MAX_WINDOW); + int const buf_write = must_checkpoint ? (1 - buf_read) : buf_read; + int const write_offset = must_checkpoint ? 0 : prev_k; + + // ── Load A (scalar, tie_hdim), dt_bias, and D (hoisted to hide gmem latency) ── + auto const* __restrict__ A_ptr = reinterpret_cast(params.A); + auto const* __restrict__ dt_bias_ptr = reinterpret_cast(params.dt_bias); + auto const* __restrict__ D_ptr = reinterpret_cast(params.D); + float const A_val = toFloat(A_ptr[head]); + float const dt_bias_val = dt_bias_ptr ? toFloat(dt_bias_ptr[head]) : 0.f; + float const D_val = D_ptr ? toFloat(D_ptr[head]) : 0.f; + + // ════════════════════════════════════════════════════════════════════════ + // Phase 0: Load all data into smem (per-warp ownership) + // ════════════════════════════════════════════════════════════════════════ + // Two-phase load around the PDL barrier: + // 1. Issue cp.async for cache (state, old_B, old_x) and in_proj-derived + // data (z); run scalar LDGs (old_dt, old_cumAdt, dt → dt_proc) and the + // cumAdt warp scan. None of these depend on conv1d, so they overlap + // with the upstream's tail. + // 2. `gdc_wait()` — wait for the upstream conv1d to signal (no-op when + // the kernel isn't launched with the PDL attribute). + // 3. Issue cp.async for conv1d outputs (x, B, C), then __pipeline_commit + // + __pipeline_wait_prior(0) + __syncwarp drains BOTH halves' cp.async + // (they share the per-thread async group). + // + // Each warp sees its own cp.async via __syncwarp. Cross-warp visibility + // is established by the post-replay __syncthreads below — replay reads of + // state are safe because (a) replay's frag_h initial load sees only the + // current warp's lane positions, and (b) the actual _1×4 cross-warp + // dependency is on writes that haven't happened yet at this point. + // ENABLE_PDL is JIT-stamped (see checkpointing_ssu_customize_config.jinja). + // `if constexpr` keeps only the chosen branch in the binary — no register + // pressure leak from the unused path. + if constexpr (ENABLE_PDL) { + load_pre_pdl_wait_data(smem, params, lane, warp, d_tile, head, group_idx, cache_slot, + buf_read, A_val, dt_bias_val, dt_seq_base, z_seq_base, + seq_len); + gdc_wait(); + load_post_pdl_wait_data( + smem, params, lane, warp, d_tile, head, group_idx, outer, seq_len); + } else { + load_data( + smem, params, lane, warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, + outer, seq_len); + } + + // old_B writeback hoisted ahead of Phase 1. Source (smem.B) is consumed + // only by Phase 1a CB; the STGs fire-and-forget onto the memory subsystem + // and complete in parallel with all subsequent compute. Only W0, W1 hold + // valid smem.B at this point (they're the ones that cp.async'd B). Gate + // accordingly — store halves its thread count but B is small (4 KB) so + // still cheap. old_B is D-independent (per-group, full DSTATE) — only + // d_tile == 0 writes; other d_tiles would emit identical payloads. + if (d_tile == 0 && warp < 2) { + store_old_B( + smem, params, warp, lane, head, group_idx, cache_slot, buf_write, write_offset, seq_len); + } + + // CB precompute (4-warp split): warps 0,1 compute CB_scaled (new tokens); + // warps 2,3 compute CB_old (old tokens) in the no-write path only. Both + // halves write to disjoint col ranges of the same swizzled smem.CB_scaled + // buffer. In the checkpoint path, warps 2,3 stay idle here and pick up + // work below in `ssu_checkpoint`'s replay matmul. + if (warp < 2) { + compute_CB_scaled_2warp(smem, warp, lane, seq_len); + } else if (!must_checkpoint) { + compute_CB_old_2warp(smem, warp, lane, prev_k, + seq_len); + } + + // ════════════════════════════════════════════════════════════════════════ + // Phase 1b + 2: Per-path dispatch + // ════════════════════════════════════════════════════════════════════════ + // Checkpoint path: replay + sync + compute_and_store_output (today's body). + // No-write path : sync + compute_no_write_output (skips replay; matmul-3 + // reads s_0 directly from smem.state, matmul-4 extends with + // the CB_old @ old_x contribution over [0, prev_k)). + // must_checkpoint is uniform across the CTA (derived from broadcast prev_k + // + compile-time NPREDICTED + MAX_WINDOW), so both branches contain a + // __syncthreads and divergence is balanced. + if (must_checkpoint) { + ssu_checkpoint( + smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); + } else { + ssu_nocheckpoint( + smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); + } + + // ── PDL: signal downstream that `output` is written. The cache writes + // below target tensors that only the next SSU step reads, not the + // immediate downstream kernel, so we can signal before issuing them. + if constexpr (ENABLE_PDL) { + gdc_launch_dependents(); + } + + // ── Phase 3: Store to global memory ── + // (old_B hoisted to pre-Phase-1; state hoisted into compute_and_store_output.) + + // Cache writes — old_x uses all warps (vectorized), dt/cumAdt one warp each. + // Each writes the new NPREDICTED tokens at gmem offset `write_offset` into + // buffer `buf_write` (computed above from must_checkpoint). + store_old_x(smem, params, warp, lane, d_tile, head, + cache_slot, write_offset, seq_len); + // dt_proc / cumAdt are D-independent — only d_tile == 0 writes. + if (d_tile == 0 && warp == 0 && lane < seq_len) { + auto* __restrict__ old_dt_w = reinterpret_cast(params.old_dt); + int64_t const dt_w_base = cache_slot * params.old_dt_stride_seq + + buf_write * params.old_dt_stride_dbuf + + head * params.old_dt_stride_head; + old_dt_w[dt_w_base + write_offset + lane] = smem.dt_proc[lane]; + } + if (d_tile == 0 && warp == 1 && lane < seq_len) { + auto* __restrict__ old_cumAdt_w = reinterpret_cast(params.old_cumAdt); + int64_t const ca_w_base = cache_slot * params.old_cumAdt_stride_seq + + buf_write * params.old_cumAdt_stride_dbuf + + head * params.old_cumAdt_stride_head; + old_cumAdt_w[ca_w_base + write_offset + lane] = smem.cumAdt[lane]; + } +} + +} // namespace flashinfer::mamba::checkpointing + +#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh new file mode 100644 index 000000000000..6733919d477b --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh @@ -0,0 +1,1495 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * 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. + */ +#ifndef FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_8BIT_CUH_ +#define FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_8BIT_CUH_ + +// 8-bit (int8, future e4m3) kernel path for the incremental SSU kernel: +// storage, replay, encode, output, and __global__ kernel. + +#include "kernel_checkpointing_ssu_common.cuh" + +namespace flashinfer::mamba::checkpointing { + +// ============================================================================= +// 8-bit chain-rewrite storage (sibling of CheckpointingSsuStorage) +// ============================================================================= +// Used by `checkpointing_ssu_kernel_8bit` for int8 and fp8 (e4m3) state. +// Differs from the generic `CheckpointingSsuStorage`: +// 1. No `new_state` staging buffer — matmul-3 chains state's fp32 C-frag +// directly into the next mma's A-operand in registers (à la +// `convert_layout_acc_Aregs`), so no smem round-trip is needed. +// 2. Adds an `output_transpose` buffer used to flip the (D, T) output frag +// back to (T, D) before the gmem STG. Matmul-3/4 in the chain path +// compute init_out^T[D, T] (M=D), so the per-warp M-shard frags must be +// transposed via smem before storing into the (T, D) gmem layout. +// +// All shared Phase 0/1 buffers (CB_scaled, B, C, x, z, old_x, old_B, scalars, +// state) are byte-for-byte identical to the generic struct — Phase 0/1 +// helpers (`compute_CB_scaled_2warp`, B/C/x/z loaders, etc.) are templated on +// `SmemT` and read these by name, so they work unchanged. +template +struct CheckpointingSsuStorage8bit { + using state_t = state_t_; + static_assert(sizeof(state_t) == 1, "CheckpointingSsuStorage8bit requires a 1-byte state_t"); + + static constexpr int NPREDICTED = NPREDICTED_; + static constexpr int MAX_WINDOW = MAX_WINDOW_; + static constexpr int D_SMEM_COLS = next_multiple_of::ATOM_COLS>(D_PER_CTA); + static constexpr int NPREDICTED_PAD_MMA_M = next_multiple_of(NPREDICTED); + static constexpr int NPREDICTED_PAD_MMA_N = next_multiple_of(NPREDICTED); + static constexpr int MAX_WINDOW_PAD_MMA_K = next_multiple_of(MAX_WINDOW); + static constexpr int NPREDICTED_SWIZZLE_R = + next_multiple_of::ATOM_ROWS>(NPREDICTED); + static constexpr int CB_ROW_STRIDE = SmemSwizzle::ATOM_COLS; + + // Shared Phase 0/1 buffers (same shape/swizzle as `CheckpointingSsuStorage`). + alignas(16) input_t CB_scaled[NPREDICTED_PAD_MMA_M * CB_ROW_STRIDE]; + alignas(16) input_t B[NPREDICTED_PAD_MMA_N * DSTATE]; + alignas(16) input_t C[NPREDICTED_SWIZZLE_R * DSTATE]; + alignas(16) input_t x[NPREDICTED_PAD_MMA_M * D_SMEM_COLS]; + alignas(16) input_t z[NPREDICTED_SWIZZLE_R * D_SMEM_COLS]; + alignas(16) input_t old_x[MAX_WINDOW_PAD_MMA_K * D_SMEM_COLS]; + alignas(16) input_t old_B[MAX_WINDOW_PAD_MMA_K * DSTATE]; + + float old_dt[MAX_WINDOW]; + float old_cumAdt[MAX_WINDOW]; + float dt_proc[NPREDICTED]; + float cumAdt[NPREDICTED]; + // decay[t] = exp(cumAdt[t]) — precomputed at Phase 0 alongside cumAdt so the + // output decay broadcast in `compute_output_8bit` is a plain LDS instead of a + // per-element __expf. Each of the 4 warps redundantly writes the same values + // (same pattern as cumAdt); cross-warp visibility comes via the kernel's + // existing __syncthreads before compute_output_8bit. + float decay[NPREDICTED]; + + // state — int8 input, only LDS'd in the single replay pass. After replay + // completes its dequant + matmul into the C-frag, smem.state is dead — so + // `output_transpose` could in principle alias it (8 KB int8 vs 2 KB bf16 + // overlap easily), but for clarity we keep them separate; alias is a + // Phase-4 micro-optimization. + alignas(16) state_t state[D_PER_CTA * DSTATE]; + + // output_transpose — physical (NPREDICTED_PAD_MMA_M, OUTPUT_TRANSPOSE_ROW_STRIDE) + // input_t scratch buffer with PADDED row stride for bank-conflict-free per-thread + // STS + 16-byte-aligned cooperative LDS.128. Used by `compute_output_int8` to + // flip the per-warp `frag_y_DxT[D, T]` register layout into `(T, D)` gmem order. + // + // Row stride: D_PER_CTA + 8 = 72 bf16 elts = 144 bytes. The 8-elt (16-byte) pad + // gives: + // - 144 % 16 == 0 → LDS.128 / STG.128 stays 16-byte aligned across all rows. + // - 144 / 4 % 32 == 4 → adjacent t-rows shift bank assignment by 4 banks. + // For the m16n8 partition_C STS pattern (per-elt: 4 lanes write at fixed d, + // t ∈ {0, 2, 4, 6} → banks {0, 4, 8, 12} on the padded layout — all distinct, + // no conflicts), the padded layout cuts STS bank conflicts from ~63% of + // wavefronts (NCU v16.0) down to 0%. + // Volume: 16 × 72 × 2 B = 2.25 KB (vs unswizzled 2 KB; +256 B). + static constexpr int OUTPUT_TRANSPOSE_ROW_STRIDE = D_PER_CTA + 8; + alignas(16) input_t output_transpose[NPREDICTED_PAD_MMA_M * OUTPUT_TRANSPOSE_ROW_STRIDE]; +}; + +// ============================================================================= +// State-dtype dispatch helpers +// ============================================================================= +// `state_t` is one of: `int8_t` (symmetric int8, ±127) or `__nv_fp8_e4m3` (fp8 +// e4m3, ±448). Both are 1-byte storage; the kernel's smem layout is identical. +// Differences live in: (a) the RN encode primitive, (b) the QUANT_MAX clip +// bound, and (c) packing/unpacking the byte from a u16 `Pair`. + +template +__device__ __forceinline__ uint8_t state_byte_of(state_t v) { + if constexpr (std::is_same_v) { + return static_cast(static_cast(v)); + } else { + static_assert(std::is_same_v, + "8-bit state_t must be int8_t or __nv_fp8_e4m3"); + return reinterpret_cast<__nv_fp8_storage_t const&>(v); + } +} + +// fp32 → state_t with RN + saturate. Single-element scalar — the kernel's +// smem layout writes pairs as u16, so per-element conversion fits the +// per-thread fragment topology directly. +template +__device__ __forceinline__ state_t encode_rn_8bit(float x) { + if constexpr (std::is_same_v) { + return conversion::cvt_rni_sat_s8(x); + } else { + static_assert(std::is_same_v, + "8-bit state_t must be int8_t or __nv_fp8_e4m3"); + // cuda_fp8 ctor compiles to `cvt.rn.satfinite.e4m3.f32` on sm_89+. + return __nv_fp8_e4m3(x); + } +} + +// Per-state-dtype symmetric clip / encode-scale denominator. +// int8: ±127 (matches Triton reference, leaves -128 unused) +// fp8_e4m3fn: ±448 (max finite e4m3 value) +template +__device__ __forceinline__ constexpr float quant_max_8bit() { + if constexpr (std::is_same_v) { + return 127.0f; + } else { + static_assert(std::is_same_v, + "8-bit state_t must be int8_t or __nv_fp8_e4m3"); + return 448.0f; + } +} + +// SM80 m16n8k16 C-frag → A-frag layout reshape for chained mma (state → +// matmul-3 in the int8 chain rewrite). +// +// Pattern mirrors the SM90 helper at attention/hopper/utils.cuh:103, but: +// - SM80 m16n8 C-frag inner per-thread layout is rank-2 ((col_pair=2, +// row_pair=2)) — there's no inner "N/8" stride mode like SM90. +// - We instead `logical_divide` the *outer* MMA_N axis by 2: each pair of +// m16n8 N-atoms (= 16 cols of the producing mma's N) becomes one K=16 +// atom of the chained m16n8k16 mma's A operand. +// +// Lane-element mapping (verified by hand on the m16n8k16 PTX layout): +// C-frag at (cp, rp, mma_n=2k+kh) maps to: row=tid/4+rp*8, col=4*(tid%2)+cp+(2k+kh)*8 +// A-frag at (cp, rp, kh, mma_k=k) maps to: row=tid/4+rp*8, col=4*(tid%2)+cp+8*kh + 16k +// Same element: (2k+kh)*8 + cp == 8*kh + cp + 16k. ✓ +// +// Input layout: ((2, 2), MMA_M, MMA_N) — m16n8 C-frag +// Output layout: ((2, 2, 2), MMA_M, MMA_N / 2) — m16n8k16 A-frag, +// MMA_K = MMA_N / 2 +template +__forceinline__ __device__ auto convert_layout_acc_Aregs_sm80(Layout acc_layout) { + using namespace cute; + using X = Underscore; + static_assert(decltype(size<0, 0>(acc_layout))::value == 2, + "C-frag inner mode must be (col_pair=2, row_pair=2)"); + static_assert(decltype(size<0, 1>(acc_layout))::value == 2, + "C-frag inner mode must be (col_pair=2, row_pair=2)"); + static_assert(decltype(rank(acc_layout))::value == 3, + "C-frag must be rank-3 ((C0,C1), MMA_M, MMA_N)"); + static_assert(decltype(rank(get<0>(acc_layout)))::value == 2, + "SM80 m16n8 C-frag inner is rank-2 (no inner stride mode like SM90)"); + // logical_divide the outer MMA_N axis by 2 → ((2, 2), MMA_M, (2, MMA_N/2)) + auto l = logical_divide(acc_layout, Shape{}); + return make_layout( + make_layout(get<0, 0>(l), get<0, 1>(l), get<2, 0>(l)), // ((col_pair, row_pair, k_half)) + get<1>(l), // MMA_M + get<2, 1>(l)); // MMA_K = MMA_N / 2 +} + +// ============================================================================= +// Phase 1b: Replay for QUANTIZED state (int8) with RN encoding. +// ============================================================================= +// state[D, dstate] = dequant(state_q, decode_scale) * total_decay +// + old_x^T @ (coeff * old_B) +// +// Layout: per-warp M-shard via TiledMma `Layout<_4, _1>`. Each warp owns +// D_PER_CTA / 4 D-rows × full DSTATE. This makes amax-over-dstate fully +// warp-local (no atomic, no cross-warp __syncthreads), at the cost of +// loading full B (old_B) from smem in every warp (vs partitioning N +// across warps in the bf16/fp16 path). Constraint: per-warp M must equal +// the m16n8 atom M (=16), so D_PER_CTA must be 64 — the wrapper enforces +// d_split == 1 for int8. +// +// Pipeline: +// 1. Replay n-loop: m16n8 matmul, write fp32 frag → smem.new_state. +// 2. STG redistribution + amax + encode pass (warp-local): +// Each warp covers M_PER_WARP = 16 D-rows × 128 cols of new_state. +// Re-tile 32 lanes as 4 row-groups × 8 col-segments. Per round +// r ∈ [0, 4): each lane reads 16 fp32 (4× LDS.128) for one D-row, +// computes a lane-amax (16 fmaxf), `__shfl_xor` over the 8 +// col-lanes (mask 1, 2, 4) for the full-row amax, encodes 16 int8, +// and STG.128's them to gmem. One writer per row stores +// decode_scale = amax/QUANT_MAX to params.state_scale. +// +// matmul-3 reads new_state (fp32) on the same M-shard partition, so no +// cross-warp visibility is needed. The `__syncthreads` after this +// function returns is for dt_proc / cumAdt visibility (Phase 2), not for +// state. + +// ───────────────────────────────────────────────────────────────────────── +// replay_state_mma_8bit_chain: int8-state chain rewrite — PASS 1 only. +// +// Drops bf16 new_state smem buffer entirely; matmul-3 is fused inline with +// replay HMMA on a per-K-pair cadence (1 K-atom of A in flight at a time). +// +// Pipeline (single fused loop over K-pairs): +// - For kpair ∈ [0, NUM_K_PAIRS=8): +// - Replay 2 m16n8 N-atoms → fp32 frag_h × 2 (16 dstate cols of state). +// - Update per-thread amax (fp32, bit-exact). +// - Cast fp32 → bf16, pack into K-atom-sized A frag (`a_kpair`, +// 8 bf16/thread = 4 32-bit regs). Layout matches the m16n8k16 A +// operand directly (`partition_fragment_A` of the chain TiledMma). +// - LDS one K-atom of B from `smem.C[T_pad, kpair*16..+16]`. +// - `cute::gemm` accumulates one K-atom into `frag_y_DxT`: +// `frag_y_DxT[D, T] += new_state[D, kpair*16..+16] +// @ smem.C[T_pad, kpair*16..+16]^T`. +// - Both `a_kpair` and `b_kpair` go out of scope at iter end. +// Post-loop: +// - Warp-local amax reduce (`__shfl_xor` over 4 col-lanes per row pair). +// - Compute `decode_scale = amax/127`, `encode_scale = 127/amax` per row. +// - STG `decode_scale` to gmem (one writer per (cache, head, d_row)). +// - Return `encode_scale_per_row[2]` to the caller — needed by the +// PASS 2 helper (`encode_state_replay_8bit`) which runs *after* +// `compute_output_8bit` so that `frag_y_DxT`'s 8 fp32 regs are dead by +// the time PASS 2's replay-again runs. +// +// Math identity for the chain (why writing `frag_h(j)` to `a_kpair(local_n*4+j)` +// places the bytes in the m16n8k16 A operand's expected position): +// linear(cp, rp, kh, _, mma_k) = cp + 2*rp + 4*kh + 8*mma_k +// = cp + 2*rp + 4*(mma_n%2) + 8*(mma_n/2) +// = cp + 2*rp + 4*mma_n +// = same linear index as the C-frag for the 2 m16n8 N-atoms making up this +// K-pair. No layout helper needed. +// +// No internal __syncthreads — smem.C is redundantly loaded by all 4 warps +// so chain matmul-3 sees each warp's own data without cross-warp sync. +// The caller's single __syncthreads between all replay passes and +// compute_output_8bit provides smem.CB_scaled / smem.x / smem.z visibility. +template +__device__ __forceinline__ void replay_state_mma_8bit_chain( + SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int prev_k, int d_tile, + int64_t cache_slot, int head, bool must_checkpoint, FragYDxT& frag_y_DxT, + float (&encode_scale_per_row_out)[2], float (&total_scale_out)[2]) { + using namespace cute; + static_assert(sizeof(input_t) == 2, "replay_state_mma_8bit_chain requires 2-byte input_t"); + static_assert(sizeof(state_t) == 1, + "replay_state_mma_8bit_chain is for 1-byte state_t (int8/fp8) only"); + static_assert(D_PER_CTA == 64, + "replay_state_mma_8bit_chain requires D_PER_CTA == 64 (M-shard, per-warp M=16)."); + + constexpr int NUM_WARPS = 4; + constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; // 16 + static_assert(M_PER_WARP == MMA_prop::M, "Per-warp M must equal m16n8 atom M (=16)"); + + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + int const tid = warp * warpSize + lane; + + // Atom-K dispatch: K_BIG=16 (default), K_SMALL=8 if MAX_WINDOW ≤ 8. + using MmaAtomReplayType = std::conditional_t; + using LdsmA = std::conditional_t; + using LdsmB = std::conditional_t; + + // Replay TiledMma: M-shard, 4 warps along M, 1 along N. Output is + // ((2,2), 1, NUM_N_PASSES) per thread of fp32 (or bf16 view for new_state). + auto tiled_mma_replay = + make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_replay = tiled_mma_replay.get_slice(tid); + + // Chain TiledMma: m16n8k16 (always K_BIG=16 since K=DSTATE/16 atoms ≥ 1), + // same M-shard layout as replay. M_per_warp=16 (1 m-atom), + // N=NPREDICTED_PAD_MMA_M (T_pad, ≤ 16 = up to 2 n-atoms per warp), K=DSTATE. + auto tiled_mma_chain = + make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_chain = tiled_mma_chain.get_slice(tid); + + constexpr int N_PER_PASS = MMA_prop::N; // 8 + constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; // 16 + constexpr int FRAG_SIZE = 4; + constexpr int D_ROWS_PER_THREAD = 2; + constexpr float QUANT_MAX = quant_max_8bit(); + + float const total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + float const total_decay = (prev_k > 0) ? __expf(total_cumAdt) : 1.f; + + int const lane_d = lane / 4; + int const warp_d_base = warp * M_PER_WARP; + + // ── Per-row decode_scale for state init. + auto const* __restrict__ state_scale_ptr = reinterpret_cast(params.state_scale); + int64_t const state_scale_base = cache_slot * params.state_scale_stride_seq + + (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; + float decode_scale_in[D_ROWS_PER_THREAD]; + decode_scale_in[0] = state_scale_ptr[state_scale_base + warp_d_base + lane_d]; + decode_scale_in[1] = state_scale_ptr[state_scale_base + warp_d_base + lane_d + 8]; + float total_scale[D_ROWS_PER_THREAD]; + total_scale[0] = decode_scale_in[0] * total_decay; + total_scale[1] = decode_scale_in[1] * total_decay; + + // ── A operand (replay): old_x [MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS] → LDSM_T. + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + auto layout_A_full = + make_swizzled_layout_rc_transpose(); + Tensor smem_A_full = make_tensor( + make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); + Tensor smem_A = local_tile(smem_A_full, make_shape(Int{}, Int{}), + make_coord(_0{}, _0{})); + + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma_replay); + auto s2r_thr_A = s2r_A.get_slice(tid); + Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); + Tensor frag_A_replay = thr_mma_replay.partition_fragment_A(make_tensor( + (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); + Tensor frag_A_replay_view = s2r_thr_A.retile_D(frag_A_replay); + cute::copy(s2r_A, smem_A_s2r, frag_A_replay_view); + + // ── Bake dB coefficients into frag_A once (8 scale ops), replacing 16× + // per-N-pass compute_dB_scaling on frag_B (64 scale ops). + // dB coefficients c[k] baked into frag_A once, replacing per-N-pass B scaling. + apply_dA_coeff(frag_A_replay, smem, total_cumAdt, prev_k, lane); + + // ── B operand (replay): old_B per-pass. + auto layout_B_replay = make_swizzled_layout_rc_transpose(); + Tensor smem_B_full = make_tensor( + make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B_replay); + auto s2r_B_replay = make_tiled_copy_B(Copy_Atom{}, tiled_mma_replay); + auto s2r_thr_B_replay = s2r_B_replay.get_slice(tid); + + // ── State: 1-byte input pointer + manual swizzle offsets (read in BOTH passes). + // Drop bf16 new_state staging — replay's fp32 frag flows directly into the + // register-resident `new_state` tensor below. + state_t* state_base = reinterpret_cast(smem.state); + + // Manual swizzle offsets for m16n8 C-fragment layout (1-byte Swizzle<3,4,3>). + // off = row * 128 + (col ^ ((row & 7) << 4)). + // row_hi = row_lo + 8; (row+8)&7 == row&7 ⇒ off_hi = off_lo + 1024. + // Fragment col within each N_PER_PASS=8 tile: (lane % 4) * 2. + int const row_lo = warp_d_base + lane_d; + int const frag_col_base = (lane & 3) << 1; + int const state_base_lo = row_lo << 7; // row_lo * DSTATE + int const state_xor = (row_lo & 7) << 4; + + float per_thread_amax[D_ROWS_PER_THREAD] = {0.f, 0.f}; + + // No __syncthreads here — smem.C is redundantly loaded by all 4 warps + // (each warp sees its own cp.async via __syncwarp in load_data). Cross-warp + // visibility for smem.CB_scaled / smem.x / smem.z is established by the + // caller's __syncthreads between this function and compute_output_8bit. + + // ── smem.C view + B-operand TiledCopy for chain matmul-3 (hoisted before + // the loop; same view per K-pair, B sliced per K-atom inside the loop). + auto layout_C_swz = + make_aliased_swizzled_layout_rc(); + Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), + layout_C_swz); + auto s2r_B_chain = + make_tiled_copy_B(Copy_Atom{}, tiled_mma_chain); + auto s2r_thr_B_chain = s2r_B_chain.get_slice(tid); + + constexpr int NUM_K_PAIRS = NUM_N_PASSES / 2; // 8 for DSTATE=128 + static_assert(NUM_N_PASSES % 2 == 0, "Per-K-pair fusion requires NUM_N_PASSES to be even"); + static_assert(MMA_prop::K_BIG == 16, "Chain mma assumes m16n8k16 K-atom = 16"); + + // ════════════════════════════════════════════════════════════════════════ + // PASS 1 — fused replay + chain matmul-3 (per-K-pair): + // For each kpair ∈ [0, NUM_K_PAIRS): + // - Run 2 replay HMMAs (N-passes 2*kpair, 2*kpair+1) → fp32 frag_h × 2. + // - Update per-thread amax (bit-exact fp32). + // - Pack each pair's 4 fp32 → 4 bf16 into a tiny K-atom-sized A frag + // (`a_kpair` shape ((2,2,2), 1, 1) of bf16 = 8 elts/thread = 4 + // 32-bit regs). Linear positions [local_n*4 .. local_n*4+3] + // within `a_kpair` map to the m16n8k16 A operand's (kh=local_n) + // slice — proven by the linear-index identity in the deleted + // `new_state`-tensor comment above. + // - LDS one K-atom of B (smem.C[T_pad, kpair*16..+16]) into a + // similarly small `b_kpair` frag (4 32-bit regs / thread). + // - `cute::gemm` accumulates one K-atom into `frag_y_DxT`. + // - Both `a_kpair` and `b_kpair` go out of scope at iter end → the + // compiler frees those ~8 32-bit regs/thread for the next iter. + // Net: register footprint drops from the 32 regs of the old register- + // resident `new_state` array (held across the whole loop) to ~8 regs in + // flight. Frees ~24 regs/thread → potentially +1-2 blocks/SM occupancy. + // ════════════════════════════════════════════════════════════════════════ +#pragma unroll + for (int kpair = 0; kpair < NUM_K_PAIRS; ++kpair) { + // K-atom-sized A frag for chain matmul-3 (filled across the 2 N-passes). + Tensor a_kpair = thr_mma_chain.partition_fragment_A(make_tensor( + (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); + static_assert(decltype(size(a_kpair))::value == 8, + "a_kpair must hold 1 m16n8k16 K-atom of A = 8 bf16/thread"); + +#pragma unroll + for (int local_n = 0; local_n < 2; ++local_n) { + int const n = kpair * 2 + local_n; + int const n_base = n * N_PER_PASS; + + Tensor frag_h = thr_mma_replay.partition_fragment_C( + make_tensor((float*)0x0, make_shape(Int{}, Int{}))); + static_assert(decltype(size(frag_h))::value == FRAG_SIZE, + "FRAG_SIZE must match the partitioned C-fragment size"); + + // Zero-init accumulator — MMA from scratch, state added after. + clear(frag_h); + + // Replay B operand load. + Tensor smem_B_n = + local_tile(smem_B_full, make_tile(Int{}, Int{}), + make_coord(n, _0{})); + auto smem_B_s2r_n = s2r_thr_B_replay.partition_S(smem_B_n); + Tensor frag_B_replay = thr_mma_replay.partition_fragment_B(make_tensor( + (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); + auto frag_B_replay_view = s2r_thr_B_replay.retile_D(frag_B_replay); + cute::copy(s2r_B_replay, smem_B_s2r_n, frag_B_replay_view); + + // Replay HMMA: frag_h = frag_A_scaled @ frag_B (c[k] baked into A). + cute::gemm(tiled_mma_replay, frag_h, frag_A_replay, frag_B_replay, frag_h); + + { + int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); + Pair const p0 = *reinterpret_cast const*>(&state_base[off_lo]); + frag_h(0) += toFloat(p0[Int<0>{}]) * total_scale[0]; + frag_h(1) += toFloat(p0[Int<1>{}]) * total_scale[0]; + Pair const p1 = + *reinterpret_cast const*>(&state_base[off_lo + 1024]); + frag_h(2) += toFloat(p1[Int<0>{}]) * total_scale[1]; + frag_h(3) += toFloat(p1[Int<1>{}]) * total_scale[1]; + } + + // Update amax (fp32, bit-exact) AND pack 4 fp32 → 4 bf16 into a_kpair + // at offset local_n*4 (matches A-frag's (kh=local_n) slice). +#pragma unroll + for (int i = 0; i < FRAG_SIZE; i += 2) { + int const d_idx = i / 2; + float const a0 = fabsf(frag_h(i)); + float const a1 = fabsf(frag_h(i + 1)); + per_thread_amax[d_idx] = fmaxf(per_thread_amax[d_idx], fmaxf(a0, a1)); + + Pair const q = + pack_float2(make_float2(frag_h(i), frag_h(i + 1))); + *reinterpret_cast*>(&a_kpair(local_n * FRAG_SIZE + i)) = q; + } + } + + // ── B operand for chain matmul-3 K-atom: smem.C[T_pad, kpair*16..+16] ── + Tensor smem_C_k = + local_tile(smem_C, make_tile(Int{}, Int{}), + make_coord(_0{}, kpair)); + auto smem_C_k_s2r = s2r_thr_B_chain.partition_S(smem_C_k); + Tensor b_kpair = thr_mma_chain.partition_fragment_B( + make_tensor((MMA_prop::operand_t*)0x0, + make_shape(Int{}, Int{}))); + auto b_kpair_view = s2r_thr_B_chain.retile_D(b_kpair); + cute::copy(s2r_B_chain, smem_C_k_s2r, b_kpair_view); + + // Single K-atom chain matmul-3: frag_y_DxT += a_kpair @ b_kpair + // frag_y_DxT (pre-zeroed by caller) accumulates across all 8 K-atoms. + cute::gemm(tiled_mma_chain, frag_y_DxT, a_kpair, b_kpair, frag_y_DxT); + } + + // ── Warp-local amax reduce (Layout<_4,_1> → fully warp-local; no atomics). +#pragma unroll + for (int i = 0; i < D_ROWS_PER_THREAD; ++i) { + per_thread_amax[i] = fmaxf(per_thread_amax[i], + __shfl_xor_sync(constants::MASK_ALL_LANES, per_thread_amax[i], 1)); + per_thread_amax[i] = fmaxf(per_thread_amax[i], + __shfl_xor_sync(constants::MASK_ALL_LANES, per_thread_amax[i], 2)); + } + + // ── encode scale (Triton fall-through for amax==0). + // decode_scale = 1 / encode_scale (mathematically: decode = amax/QUANT_MAX, + // encode = QUANT_MAX/amax, so decode = 1/encode; and when amax==0 both fall + // through to 1.f → 1/1 == 1). Computed inline at the STG below — keeping + // only `encode_scale_per_row` in regs saves 2 fp32 regs across PASS 2. + float encode_scale_per_row[D_ROWS_PER_THREAD]; +#pragma unroll + for (int i = 0; i < D_ROWS_PER_THREAD; ++i) { + float const a = per_thread_amax[i]; + encode_scale_per_row[i] = (a == 0.f) ? 1.f : (QUANT_MAX / a); + } + + // ── STG decode_scale (one writer per (cache, head, d_row)). + if (must_checkpoint && (lane & 3) == 0) { + auto* __restrict__ state_scale_w = reinterpret_cast(params.state_scale); +#pragma unroll + for (int i = 0; i < D_ROWS_PER_THREAD; ++i) { + int const d_row_in_atom = lane_d + (i & 1) * 8; + int const d_row = warp_d_base + d_row_in_atom; + state_scale_w[state_scale_base + d_row] = 1.f / encode_scale_per_row[i]; + } + } + + // Hand `encode_scale_per_row` AND `total_scale` (= OLD decode_scale_in × + // total_decay) to the caller so PASS 2 (encode replay-again) can: + // - dequantize the OLD int8 state with the OLD decode_scale (NOT the NEW + // one we just STG'd above — re-reading params.state_scale in PASS 2 + // would pick up the new value and corrupt the encode), and + // - encode the NEW state with the right encode_scale = 127 / amax. + encode_scale_per_row_out[0] = encode_scale_per_row[0]; + encode_scale_per_row_out[1] = encode_scale_per_row[1]; + total_scale_out[0] = total_scale[0]; + total_scale_out[1] = total_scale[1]; +} + +// ───────────────────────────────────────────────────────────────────────── +// encode_state_replay_8bit: PASS 2 of the int8 chain rewrite. +// +// Re-runs the replay matmul fresh (replay-again), encodes the post-replay +// state fp32 → int8 using `encode_scale_per_row[]` from PASS 1, and STG.16's +// the int8 pairs to gmem. Bit-exact with Triton's fp32-encode path. +// +// Called *after* `compute_output_8bit` so that: +// - `frag_y_DxT`'s 8 fp32 regs are dead (chain matmul-3's accumulator +// was consumed by the output STG). +// - PASS 2's gmem STGs fire alongside `store_old_x` / dt_proc / cumAdt +// writes — all gmem traffic at the kernel tail where there's nothing +// else to do. +// +// The setup (TiledMma, frag_A_replay, smem layouts) is duplicated +// from `replay_state_mma_8bit_chain` — separate stack frame keeps register +// allocation simple and avoids cross-function lifetime tracking. +template +__device__ __forceinline__ void encode_state_replay_8bit( + SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int prev_k, int d_tile, + int64_t cache_slot, int head, float const (&encode_scale_per_row)[2], + float const (&total_scale)[2], int64_t rand_seed, int64_t state_ptr_offset) { + using namespace cute; + static_assert(sizeof(input_t) == 2, "encode_state_replay_8bit requires 2-byte input_t"); + static_assert(sizeof(state_t) == 1, + "encode_state_replay_8bit is for 1-byte state_t (int8/fp8) only"); + static_assert(D_PER_CTA == 64, + "encode_state_replay_8bit requires D_PER_CTA == 64 (M-shard, per-warp M=16)."); + + constexpr int NUM_WARPS = 4; + constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; + static_assert(M_PER_WARP == MMA_prop::M, "Per-warp M must equal m16n8 atom M (=16)"); + + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + int const tid = warp * warpSize + lane; + + using MmaAtomReplayType = std::conditional_t; + using LdsmA = std::conditional_t; + using LdsmB = std::conditional_t; + + auto tiled_mma_replay = + make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_replay = tiled_mma_replay.get_slice(tid); + + constexpr int N_PER_PASS = MMA_prop::N; + constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; + constexpr int FRAG_SIZE = 4; + constexpr int D_ROWS_PER_THREAD = 2; + + float const total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + + // total_scale (= OLD decode_scale_in × total_decay) was computed in PASS 1 + // and is passed in by reference. We MUST NOT re-load decode_scale_in from + // params.state_scale here — by the time PASS 2 runs, PASS 1 has already + // STG'd the NEW decode_scale to that same gmem location. + + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + auto layout_A_full = + make_swizzled_layout_rc_transpose(); + Tensor smem_A_full = make_tensor( + make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); + Tensor smem_A = local_tile(smem_A_full, make_shape(Int{}, Int{}), + make_coord(_0{}, _0{})); + + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma_replay); + auto s2r_thr_A = s2r_A.get_slice(tid); + Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); + Tensor frag_A_replay = thr_mma_replay.partition_fragment_A(make_tensor( + (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); + Tensor frag_A_replay_view = s2r_thr_A.retile_D(frag_A_replay); + cute::copy(s2r_A, smem_A_s2r, frag_A_replay_view); + + // dB coefficients baked into frag_A (same identity as PASS 1). + apply_dA_coeff(frag_A_replay, smem, total_cumAdt, prev_k, lane); + + auto layout_B_replay = make_swizzled_layout_rc_transpose(); + Tensor smem_B_full = make_tensor( + make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B_replay); + auto s2r_B_replay = make_tiled_copy_B(Copy_Atom{}, tiled_mma_replay); + auto s2r_thr_B_replay = s2r_B_replay.get_slice(tid); + + state_t* state_base = reinterpret_cast(smem.state); + + // Manual swizzle offsets (same derivation as replay_state_mma_8bit_chain). + int const lane_d = lane / 4; + int const warp_d_base = warp * M_PER_WARP; + int const row_lo = warp_d_base + lane_d; + int const frag_col_base = (lane & 3) << 1; + int const state_base_lo = row_lo << 7; + int const state_xor = (row_lo & 7) << 4; + + // Philox state for SR — one refresh every 4 n-passes (cvt_rs_sat_s8x4_f32 + // packs 4 int8s per u32 of randomness, so 1 Philox call covers 16 int8s). + [[maybe_unused]] uint32_t rand_idx[4]; + +#pragma unroll + for (int n = 0; n < NUM_N_PASSES; ++n) { + int const n_base = n * N_PER_PASS; + + Tensor frag_h = thr_mma_replay.partition_fragment_C( + make_tensor((float*)0x0, make_shape(Int{}, Int{}))); + + // Zero-init accumulator — MMA from scratch, state added after. + clear(frag_h); + + Tensor smem_B_n = + local_tile(smem_B_full, make_tile(Int{}, Int{}), + make_coord(n, _0{})); + auto smem_B_s2r_n = s2r_thr_B_replay.partition_S(smem_B_n); + Tensor frag_B_replay = thr_mma_replay.partition_fragment_B(make_tensor( + (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); + auto frag_B_replay_view = s2r_thr_B_replay.retile_D(frag_B_replay); + cute::copy(s2r_B_replay, smem_B_s2r_n, frag_B_replay_view); + + // HMMA: frag_h = frag_A_scaled @ frag_B (c[k] baked into A). + cute::gemm(tiled_mma_replay, frag_h, frag_A_replay, frag_B_replay, frag_h); + + { + int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); + Pair const p0 = *reinterpret_cast const*>(&state_base[off_lo]); + frag_h(0) += toFloat(p0[Int<0>{}]) * total_scale[0]; + frag_h(1) += toFloat(p0[Int<1>{}]) * total_scale[0]; + Pair const p1 = *reinterpret_cast const*>(&state_base[off_lo + 1024]); + frag_h(2) += toFloat(p1[Int<0>{}]) * total_scale[1]; + frag_h(3) += toFloat(p1[Int<1>{}]) * total_scale[1]; + } + + // ── Encode + in-place STS to smem.state at cols [n*8, n*8+8) ── + // Overwrites the OLD 8-bit input *for this n-pass's cols only*. The next + // n-pass dequants from a DIFFERENT col band [(n+1)*8, +8) — still OLD — + // so no read-after-write hazard. Each warp writes only to its own + // M-shard rows; cross-warp visibility is established by the + // caller's __syncthreads before the cooperative store_state call. + { + int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); + float const e0 = encode_scale_per_row[0]; + float const e1 = encode_scale_per_row[1]; + + if constexpr (PHILOX_ROUNDS > 0) { + // One Philox4x call yields 4 independent u32s — enough for 4 n-passes + // when each pass packs all 4 of its 8-bit outputs into one u32 via the + // dtype-specific x4 cvt_rs (int8: `cvt_rs_sat_s8x4_f32`, 16-bit + // randomness/elt via bitrev16 trick; fp8 e4m3: `cvt_rs_e4m3x4_f32`, + // native PTX `cvt.rs.satfinite.e4m3x4.f32` on sm_100a+ with SW fallback). + int const rand_pos = n & 3; + if (rand_pos == 0) { + int64_t const philox_off = + state_ptr_offset + (int64_t)row_lo * DSTATE + (frag_col_base + n_base); + conversion::philox_randint4x(rand_seed, philox_off, rand_idx[0], + rand_idx[1], rand_idx[2], rand_idx[3]); + } + // Packed layout: byte 0 = q0_lo, byte 1 = q1_lo (→ row_lo store at off_lo) + // byte 2 = q0_hi, byte 3 = q1_hi (→ row_hi store at off_lo + 1024) + uint32_t packed; + if constexpr (std::is_same_v) { + packed = conversion::cvt_rs_sat_s8x4_f32(frag_h(0) * e0, frag_h(1) * e0, frag_h(2) * e1, + frag_h(3) * e1, rand_idx[rand_pos]); + } else { + static_assert(std::is_same_v, + "8-bit SR supports state_t in {int8_t, __nv_fp8_e4m3}"); + packed = conversion::cvt_rs_e4m3x4_f32(frag_h(0) * e0, frag_h(1) * e0, frag_h(2) * e1, + frag_h(3) * e1, rand_idx[rand_pos]); + } + Pair q_lo, q_hi; + q_lo.raw = static_cast(packed & 0xFFFFu); + q_hi.raw = static_cast(packed >> 16); + *reinterpret_cast*>(&state_base[off_lo]) = q_lo; + *reinterpret_cast*>(&state_base[off_lo + 1024]) = q_hi; + } else { + // d_idx=0: row_lo + state_t const q0_lo = encode_rn_8bit(frag_h(0) * e0); + state_t const q1_lo = encode_rn_8bit(frag_h(1) * e0); + Pair q_lo; + q_lo.raw = static_cast(state_byte_of(q0_lo)) | + (static_cast(state_byte_of(q1_lo)) << 8); + *reinterpret_cast*>(&state_base[off_lo]) = q_lo; + // d_idx=1: row_hi = row_lo + 8, off_hi = off_lo + 1024 + state_t const q0_hi = encode_rn_8bit(frag_h(2) * e1); + state_t const q1_hi = encode_rn_8bit(frag_h(3) * e1); + Pair q_hi; + q_hi.raw = static_cast(state_byte_of(q0_hi)) | + (static_cast(state_byte_of(q1_hi)) << 8); + *reinterpret_cast*>(&state_base[off_lo + 1024]) = q_hi; + } + } + } + + // No __syncthreads or cooperative STG here — the caller's single sync + // provides cross-warp smem.state visibility, then calls store_state. +} + +// ──────────────────────────────────────────────────────────────────────── +// compute_output_8bit: transposed matmul-4 + epilogue + smem-transpose STG +// ──────────────────────────────────────────────────────────────────────── +// Companion to `replay_state_mma_int8_chain`. Consumes the per-warp +// `frag_y_DxT` (shape ((2,2), 1, T_pad/8) of fp32 per thread; M=D-shard, +// N=T_pad) — pre-loaded with init_out^T from chain matmul-3 — and: +// 1. Decay broadcast: frag_y_DxT *= exp(cumAdt[t]) (per T-col, scalar LDS). +// 2. Chain matmul-4 transposed: frag_y_DxT += x^T[D, T] @ CB_scaled^T[T, T] +// A operand: smem.x viewed via x_trans (D, T) → LDSM_N feeds A(M=D, K=T). +// B operand: smem.CB_scaled (T, T) → LDSM_T feeds B(K=T, N=T). +// 3. D*x skip: frag_y_DxT(d, t) += D_val * x[t, d] (scalar LDS per element; +// consecutive frag elts at fixed D, varying T → not pair-loadable). +// 4. z-gate: frag_y_DxT *= z * sigmoid(z) (scalar LDS per element). +// 5. fp32 → input_t pack (in-place register cvt via pack_float2). +// 6. Per-thread STS to smem.output_transpose at (T, D) layout. +// 7. __syncthreads. +// 8. Cooperative STG.128 from smem.output_transpose (T, D) to gmem (T, D). +// +// Cross-warp dependencies (smem.x, smem.z, smem.CB_scaled) are already +// visible because the caller's __syncthreads fires between all replay +// passes and this function. +template +__device__ __forceinline__ void compute_output_8bit(SmemT& smem, + CheckpointingSsuParams const& params, int warp, + int lane, int d_tile, int64_t out_seq_base, + int head, int64_t cache_slot, float D_val, + int seq_len, FragYDxT& frag_y_DxT) { + using namespace cute; + static_assert(sizeof(input_t) == 2, "compute_output_8bit requires 2-byte input_t"); + static_assert(D_PER_CTA == 64, "compute_output_8bit requires D_PER_CTA == 64"); + static_assert(NUM_WARPS == 4, "compute_output_8bit requires 4 warps"); + + int const tid = warp * warpSize + lane; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; // 16 + + // Same TiledMma as replay_state_mma_int8_chain (M-shard, m16n8k16). + auto tiled_mma_chain = + make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_chain = tiled_mma_chain.get_slice(tid); + + // ── Smem views ── + // x_trans: x physically stored at (T, D); transposed view at (D, T). + // Used as the A operand of the chain matmul-4. + auto layout_x_trans = + make_swizzled_layout_rc_transpose(); + Tensor smem_x_trans = make_tensor( + make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans); + Tensor smem_x_trans_tile = + local_tile(smem_x_trans, make_shape(Int{}, Int{}), + make_coord(_0{}, _0{})); + + // x natural (T, D) view — for D-skip + z-gate per-element scalar LDS. + auto layout_x = make_swizzled_layout_rc(); + + // z natural (T, D) view (aliased so padded rows alias valid rows). + auto layout_z = make_aliased_swizzled_layout_rc(); + + // CB_scaled (T, T_pad) within (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE). + auto layout_cb = + make_swizzled_layout_rc(); + Tensor smem_CB = make_tensor( + make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb); + + // ── Per-thread (d, t) coord lookup for epilogue scalar reads + smem-transpose write ── + auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma_chain.partition_C(id_tile); + + // ── 1. Decay broadcast: frag_y(i) *= exp(cumAdt[t]) ── + // Reads `smem.decay[t]` (= exp(cumAdt[t])) precomputed in Phase 0 by + // `compute_cumAdt` — fused EX2 with the cumsum write, replacing ~512 per-CTA + // __expf calls in this inner loop with a single LDS per element. + // For padded T-cols (t >= NPREDICTED), the read returns garbage but the STG + // at the end is predicated on t < NPREDICTED, so the garbage never reaches gmem. +#pragma unroll + for (int i = 0; i < size(frag_y_DxT); ++i) { + int const t = get<1>(id_part(i)); + if (t < seq_len) { + frag_y_DxT(i) *= smem.decay[t]; + } + } + + // ── 2. Chain matmul-4: frag_y_DxT += x^T @ CB^T ── + // A operand: smem.x physically (T, D); transposed view (D, T) used as + // A(M=D, K=T). The transposed view has D-stride=1, T-stride=D — same + // pattern as replay's A from old_x — so use LDSM_T to produce row-major + // A from this column-wise smem source. + auto s2r_A_x = + make_tiled_copy_A(Copy_Atom{}, tiled_mma_chain); + auto s2r_thr_A_x = s2r_A_x.get_slice(tid); + auto smem_x_s2r = s2r_thr_A_x.partition_S(smem_x_trans_tile); + Tensor frag_A_x = thr_mma_chain.partition_fragment_A(make_tensor( + (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); + auto frag_A_x_view = s2r_thr_A_x.retile_D(frag_A_x); + cute::copy(s2r_A_x, smem_x_s2r, frag_A_x_view); + + // B operand for chain matmul-4 = CB^T. smem.CB natural view shape (T, T) + // already has T_inner stride 1 = K-major. Use LDSM_N (no transpose). + auto s2r_B_CB = + make_tiled_copy_B(Copy_Atom{}, tiled_mma_chain); + auto s2r_thr_B_CB = s2r_B_CB.get_slice(tid); + auto smem_CB_s2r = s2r_thr_B_CB.partition_S(smem_CB); + Tensor frag_B_CB = thr_mma_chain.partition_fragment_B( + make_tensor((MMA_prop::operand_t*)0x0, + make_shape(Int{}, Int{}))); + auto frag_B_CB_view = s2r_thr_B_CB.retile_D(frag_B_CB); + cute::copy(s2r_B_CB, smem_CB_s2r, frag_B_CB_view); + + cute::gemm(tiled_mma_chain, frag_y_DxT, frag_A_x, frag_B_CB, frag_y_DxT); + + // ── 3. D*x skip: frag_y(d, t) += D_val * x[t, d] (scalar LDS per element) ── + if (D_val != 0.f) { + auto* __restrict__ smem_x_base = reinterpret_cast(smem.x); +#pragma unroll + for (int i = 0; i < size(frag_y_DxT); ++i) { + int const d = get<0>(id_part(i)); + int const t = get<1>(id_part(i)); + if (t < seq_len) { + int const off = layout_x(t, d); + frag_y_DxT(i) += D_val * toFloat(smem_x_base[off]); + } + } + } + + // ── 4. z-gate: frag_y *= z * sigmoid(z) (scalar LDS per element) ── + if (params.z != nullptr) { + auto* __restrict__ smem_z_base = reinterpret_cast(smem.z); +#pragma unroll + for (int i = 0; i < size(frag_y_DxT); ++i) { + int const d = get<0>(id_part(i)); + int const t = get<1>(id_part(i)); + if (t < seq_len) { + int const off = layout_z(t, d); + float const z = toFloat(smem_z_base[off]); + frag_y_DxT(i) *= z * __fdividef(1.f, (1.f + __expf(-z))); + } + } + } + + // ── 5. Pack fp32 → input_t per element + 6. STS to smem.output_transpose (T, D) ── + // Padded row stride (D_PER_CTA + 8 = 72 bf16 = 144 bytes) gives: + // - 16-byte-aligned LDS.128 / STG.128 across all rows. + // - 4-bank shift per row → m16n8 STS pattern hits {bank 0, 4, 8, 12} for the + // 4 t-rows of an elt → bank-conflict-free (vs 4-way conflict at stride 64). + // See CheckpointingSsuStorage8bit::OUTPUT_TRANSPOSE_ROW_STRIDE for derivation. + constexpr int kSmemRowStride = SmemT::OUTPUT_TRANSPOSE_ROW_STRIDE; // 72 bf16 elts + auto* __restrict__ smem_out_base = reinterpret_cast(smem.output_transpose); +#pragma unroll + for (int i = 0; i < size(frag_y_DxT); ++i) { + int const d = get<0>(id_part(i)); + int const t = get<1>(id_part(i)); + if (t < seq_len) { + // Pack via pack_float2(f, 0.f) and take low elt — emits a single cvt + // (compiler folds the dummy into a no-op for the discarded high half). + smem_out_base[t * kSmemRowStride + d] = + pack_float2(make_float2(frag_y_DxT(i), 0.f))[Int<0>{}]; + } + } + + // ── 7. Warp sync for cross-lane STS→LDS ordering ── + __syncwarp(); + + // ── 8. Warp-local cooperative STG.128: 32 lanes → one warp's 16 D-rows ── + // Each warp's data: 16 D-rows × T_pad=16 cols × 2 B = 512 B. + // Re-tile 32 lanes: (t = lane%16, d_group = lane/16 ∈ {0, 1}) → covers + // T_pad × 2 D-groups = 32 slots, each STG.128 = 8 D-cols × 2 B = 16 B. + // No cross-warp coordination → no __syncthreads. + constexpr int kElsPerSTG = 16 / sizeof(input_t); // 8 bf16 elts per STG.128 + constexpr int kDGroupsPerWarp = M_PER_WARP / kElsPerSTG; // = 16 / 8 = 2 + static_assert(NPREDICTED_PAD_MMA_M * kDGroupsPerWarp == 32, + "warp-local STG re-tile: T_pad × dGroupsPerWarp must equal warpSize"); + + int const stg_t = lane % NPREDICTED_PAD_MMA_M; + int const stg_d_group = lane / NPREDICTED_PAD_MMA_M; + int const warp_d_base = warp * M_PER_WARP; + int const stg_d = warp_d_base + stg_d_group * kElsPerSTG; + + if (stg_t < seq_len) { + int const smem_off = stg_t * kSmemRowStride + stg_d; + + auto* __restrict__ output_ptr = reinterpret_cast(params.output); + int64_t const out_base = out_seq_base + (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; + int64_t const gmem_off = out_base + (int64_t)stg_t * params.out_stride_token + stg_d; + + // 128-bit copy. smem_off * 2 B = (t * 144 + d * 2) is 16-byte aligned + // for any t when d % 8 == 0 (here d_offset_within_warp = 0 or 8). + using Vec = uint4; + *reinterpret_cast(&output_ptr[gmem_off]) = + *reinterpret_cast(&smem_out_base[smem_off]); + } +} + +// ============================================================================= +// add_init_out_8bit: matmul-3 for the no-checkpoint path (N-shard). +// ============================================================================= +// Computes `frag_y[T, D] = C @ dequant(smem.state)^T` in N-shard `Layout<_1,_4>`, +// reading int8/fp8 state directly from smem and dequanting per-element to bf16 +// in registers before the HMMA. Mirrors the bf16 path's `add_init_out` but with +// a custom B-operand loader since the 1-byte state can't use LDSM directly. +// +// Per-thread B-frag layout (m16n8k16, PTX ISA): +// Per-lane 4 bf16 elts at (K, N) = {(2t, gID), (2t+1, gID), (2t+8, gID), (2t+9, gID)} +// where t = lane%4, gID = lane/4. +// In our matmul-3: B = state^T = (DSTATE = K-axis, D_PER_CTA = N-axis). N-shard +// gives each warp `MMA::N = 8` D-cols per N-tile, so this lane's d_row = +// n_tile*N_TILE + warp*8 + lane/4. +// +// Per K-tile: load 4 int8 bytes per lane (2 byte-pair LDS into Pair), +// CAST to bf16 (no per-row scale), pack into B-frag, HMMA into the n-th frag_y. +// `decode_scale[d]` is per-output-col (constant across the K reduction), so +// the caller pulls it OUT of the inner product and applies it post-matmul in +// the β-scale loop: +// y[t, d] = decode_scale[d] · Σ_n C[t, n] · state_byte[d, n] +// This eliminates the per-cell `... * scale` FMUL chain (was the +// long_scoreboard hotspot at line 1066) at the cost of 2 extra FMUL/elt in +// the post-matmul C-frag scale (net 1792× fewer FMUL per warp). +template +__device__ __forceinline__ void add_init_out_8bit(SmemT const& smem, int warp, int lane, + TiledMma const& tiled_mma, ThrMma const& thr_mma, + int tid, FragY&... frag_y) { + using namespace cute; + static_assert(sizeof(state_t) == 1, "add_init_out_8bit requires 1-byte state"); + static_assert(D_PER_CTA == 64, "add_init_out_8bit requires D_PER_CTA == 64 (8-bit D_SPLIT=1)"); + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int K_TILE = MMA_prop::K_BIG; // 16 + constexpr int NUM_K_TILES = DSTATE / K_TILE; // 8 + constexpr int N_TILE = cute::tile_size<1>(TiledMma{}); // 32 (4 warps × MMA::N=8) + constexpr int NUM_N_TILES = sizeof...(FragY); // 2 (D_PER_CTA / N_TILE) + static_assert(NUM_N_TILES * N_TILE == D_PER_CTA, "FragY count must match D_PER_CTA / N_TILE"); + + // ── Per-thread coords ── + int const t = lane & 3; // K-pair index within K-atom + int const lane_d = lane >> 2; // gID = lane/4; selects N-col within atom + int const warp_d_base = warp * MMA_prop::N; // warp's 8-col offset within an N-tile + + // ── A operand (C): swizzled (T_pad, DSTATE), K-tiled per K-loop iter ── + auto layout_C_swz = + make_aliased_swizzled_layout_rc(); + Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), + layout_C_swz); + Tensor smem_C_ktiled = local_tile(smem_C, make_tile(Int{}, Int{}), + make_coord(_0{}, _)); + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + auto smem_A_s2r = s2r_thr_A.partition_S(smem_C_ktiled); + Tensor frag_A = thr_mma.partition_fragment_A(smem_C_ktiled(_, _, _0{})); + auto frag_A_view = s2r_thr_A.retile_D(frag_A); + + // ── B-frag (one m16n8k16 atom of B; 4 bf16 elts/lane) ── + Tensor frag_B = thr_mma.partition_fragment_B( + make_tensor((MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); + static_assert(decltype(size(frag_B))::value == 4, "B-frag must be 4 elts/lane for m16n8k16"); + + // ── Smem state base (1-byte) + Swizzle<3,4,3> XOR formula: + // off = d_row * DSTATE + (K XOR ((d_row & 7) << 4)) + // K within the same swizzle row-group {0..7 mod 8} shares the same XOR mask. ── + state_t const* state_base = reinterpret_cast(smem.state); + + // Pre-clear accumulators (caller doesn't pre-zero — matches bf16 add_init_out). + (clear(frag_y), ...); + + // Parameter-pack indexing via pointer array (same pattern as pipelined_kloop_gemm). + using FragY0 = std::tuple_element_t<0, std::tuple>; + FragY0* frag_y_p[NUM_N_TILES] = {(&frag_y)...}; + + // ── K-loop ── +#pragma unroll + for (int k = 0; k < NUM_K_TILES; ++k) { + int const K_base = k * K_TILE; + + // Load A K-tile via LDSM (shared across all N-tiles within this K-tile). + cute::copy(s2r_A, smem_A_s2r(_, _, _, k), frag_A_view); + + CUTE_UNROLL + for (int n = 0; n < NUM_N_TILES; ++n) { + int const d_row = n * N_TILE + warp_d_base + lane_d; + int const state_base_lo = d_row << 7; // d_row * DSTATE (DSTATE=128 → <<7) + int const state_xor = (d_row & 7) << 4; // Swizzle<3,4,3> + int const off_lo = state_base_lo + ((K_base + (t << 1)) ^ state_xor); + int const off_hi = state_base_lo + ((K_base + (t << 1) + 8) ^ state_xor); + + Pair const p_lo = *reinterpret_cast const*>(&state_base[off_lo]); + Pair const p_hi = *reinterpret_cast const*>(&state_base[off_hi]); + + // Pure int8/fp8 → bf16 cast. decode_scale is applied post-matmul in + // the caller's β-scale loop. + Pair const b_lo = pack_float2( + make_float2(toFloat(p_lo[Int<0>{}]), toFloat(p_lo[Int<1>{}]))); + Pair const b_hi = pack_float2( + make_float2(toFloat(p_hi[Int<0>{}]), toFloat(p_hi[Int<1>{}]))); + + // frag_B(0,1) = K-pair at {K_base+2t, K_base+2t+1}; (2,3) = at {+8, +9}. + *reinterpret_cast*>(&frag_B(0)) = b_lo; + *reinterpret_cast*>(&frag_B(2)) = b_hi; + + cute::gemm(tiled_mma, *frag_y_p[n], frag_A, frag_B, *frag_y_p[n]); + } + } +} + +// ============================================================================= +// compute_no_write_output_8bit — N-shard output for the no-checkpoint path. +// ============================================================================= +// Mirror of the bf16 path's `compute_no_write_output`, but with int8/fp8 state. +// Uses N-shard `Layout<_1,_4>` (best smem traffic) instead of the M-shard chain +// — the M-shard exists only for amax reduction in the checkpoint path, which +// doesn't run here. +// +// y[t, d] = β(t) · u[t, d] (matmul-3 via +// add_init_out_8bit) +// + Σ_{j +__device__ __forceinline__ void compute_no_write_output_8bit( + SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int prev_k, int d_tile, + int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) { + using namespace cute; + static_assert(sizeof(input_t) == 2, "compute_no_write_output_8bit requires 2-byte input_t"); + static_assert(sizeof(state_t) == 1, "compute_no_write_output_8bit is for 1-byte state"); + static_assert(D_PER_CTA == 64, "compute_no_write_output_8bit requires D_PER_CTA == 64"); + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + int const tid = warp * warpSize + lane; + + // ── TiledMMA for matmul-3 + matmul-4-new (m16n8k16) ── + auto tiled_mma = + make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(tid); + + // ── TiledMMA for matmul-4-old: K = MAX_WINDOW_PAD_MMA_K ∈ {8, 16} → atom dispatch ── + using MmaAtomOld = std::conditional_t; + using LdsmAOld = std::conditional_t; + using LdsmBOld = std::conditional_t; + auto tiled_mma_old = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_old = tiled_mma_old.get_slice(tid); + + // ── Swizzled smem views ── + auto layout_x_swz = make_swizzled_layout_rc(); + Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), + layout_x_swz); + auto layout_x_trans_swz = + make_swizzled_layout_rc_transpose(); + Tensor smem_x_trans = make_tensor( + make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); + + auto layout_old_x_trans_swz = + make_swizzled_layout_rc_transpose(); + Tensor smem_old_x_trans = + make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), + layout_old_x_trans_swz); + + auto layout_z_swz = + make_aliased_swizzled_layout_rc(); + Tensor smem_z = + make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); + + // ── S2R copies (matmul-4-new) ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + auto s2r_B_trans = + make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); + + // ── S2R copies (matmul-4-old) ── + auto s2r_A_old = make_tiled_copy_A(Copy_Atom{}, tiled_mma_old); + auto s2r_thr_A_old = s2r_A_old.get_slice(tid); + auto s2r_B_old_trans = + make_tiled_copy_B(Copy_Atom{}, tiled_mma_old); + auto s2r_thr_B_old_trans = s2r_B_old_trans.get_slice(tid); + + // ── Load CB_scaled A operand (cols [0, NPREDICTED_PAD_MMA_M)) ── + auto layout_cb_swz = + make_swizzled_layout_rc(); + Tensor smem_CB = make_tensor( + make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); + auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); + Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); + auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); + cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); + + // ── Load CB_old A operand (cols [NPREDICTED_PAD_MMA_M, +MAX_WINDOW_PAD_MMA_K)) ── + auto layout_cb_full = make_swizzled_layout_rc(); + Tensor smem_CB_full = make_tensor( + make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); + Tensor smem_CB_old = + local_tile(smem_CB_full, make_tile(Int{}, Int{}), + make_coord(_0{}, NPREDICTED_PAD_MMA_M / MAX_WINDOW_PAD_MMA_K)); + auto smem_CB_old_s2r = s2r_thr_A_old.partition_S(smem_CB_old); + Tensor frag_CB_old_A = thr_mma_old.partition_fragment_A(smem_CB_old); + auto frag_CB_old_A_view = s2r_thr_A_old.retile_D(frag_CB_old_A); + cute::copy(s2r_A_old, smem_CB_old_s2r, frag_CB_old_A_view); + + // ── Decay broadcast (per-T scalar, stride-0 on N) ── + constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); + Tensor decay_bcast = make_tensor( + make_smem_ptr(smem.cumAdt), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); + Tensor decay_part = thr_mma.partition_C(decay_bcast); + + float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + float const beta_extra = __expf(total_old_cumAdt); + + // ── Post-matmul-3 state decode_scale (factored OUT of the inner K-product). ── + // y[t, d] = decode_scale[d] · raw_y[t, d] where raw_y = C @ (bf16)(state_byte)^T. + // Per lane, the m16n8 C-frag holds 4 elts spanning 2 unique d-cols (d_lo = 2t, + // d_hi = 2t+1) per N-tile. Indexed by `i & 1` in the epilogue scale loop. + auto const* __restrict__ state_scale_ptr = reinterpret_cast(params.state_scale); + int64_t const state_scale_base = cache_slot * params.state_scale_stride_seq + + (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; + constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; + static_assert(NUM_N_TILES == 2, + "compute_no_write_output_8bit assumes NUM_N_TILES == 2 (D_PER_CTA=64, N_TILE=32)"); + int const t_col = lane & 3; // 2t and 2t+1 are this lane's two C-frag d-cols + float decode_scale[NUM_N_TILES][2]; + CUTE_UNROLL + for (int n = 0; n < NUM_N_TILES; ++n) { + int const d_lo = n * N_TILE + warp * MMA_prop::N + (t_col << 1); + decode_scale[n][0] = state_scale_ptr[state_scale_base + d_lo]; + decode_scale[n][1] = state_scale_ptr[state_scale_base + d_lo + 1]; + } + + // ── Gmem output base ── + auto* __restrict__ output_ptr = reinterpret_cast(params.output); + int64_t const out_base = out_seq_base + (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; + + // ── Row predicate ── + auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_tile); + bool const pred_row_lo = get<0>(id_part(0)) < seq_len; + bool const pred_row_hi = get<0>(id_part(2)) < seq_len; + + auto epilogue = [&](auto& frag_y, int n) { + // β-scale + state decode_scale fused: frag_y(i) *= β · exp(cumAdt[t]) · decode_scale[d]. + // The decode_scale[d] absorbs the per-row state quant factor (was previously + // multiplied into each B-element during dequant). +#pragma unroll + for (int i = 0; i < size(frag_y); ++i) { + int const d_idx = i & 1; // i=0,2 → d_lo; i=1,3 → d_hi + frag_y(i) *= beta_extra * __expf(decay_part(i)) * decode_scale[n][d_idx]; + } + + // matmul-4-new (CB_scaled @ x). + add_cb_x( + frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); + + // matmul-4-old (CB_old @ old_x). + add_cb_old_x( + frag_y, frag_CB_old_A, smem_old_x_trans, s2r_B_old_trans, s2r_thr_B_old_trans, thr_mma_old, + tiled_mma_old, n); + + // D·x. + add_D_skip(frag_y, smem_x, thr_mma, D_val, n); + + // z-gate. + compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); + + // Direct partition_C STG. + auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), + make_layout(make_shape(Int{}, Int{}), + make_stride(params.out_stride_token, _1{}))); + auto gOut_part = thr_mma.partition_C(gOut_tile); +#pragma unroll + for (int i = 0; i < size(frag_y); i += 2) { + bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; + if (pred_i) { + *reinterpret_cast*>(&gOut_part(i)) = + pack_float2(make_float2(frag_y(i), frag_y(i + 1))); + } + } + }; + + // ── Matmul-3: frag_y = C @ (bf16)(state_byte)^T (smem.state retains s_0 since + // replay skipped; decode_scale[d] applied post-matmul in the epilogue). ── + Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); + Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); + add_init_out_8bit(smem, warp, lane, tiled_mma, thr_mma, + tid, frag_y_0, frag_y_1); + epilogue(frag_y_0, 0); + epilogue(frag_y_1, 1); +} + +// ============================================================================= +// Per-path dispatcher: no-checkpoint branch (must_checkpoint == false). +// ============================================================================= +// Sync makes warps 0,1's CB_scaled writes AND warps 2,3's CB_old writes +// visible to all warps before matmul-3 and matmul-4 read smem.{CB_scaled, +// CB_old, x, z}. Matches the bf16 path's `ssu_nocheckpoint`. +template +__device__ __forceinline__ void ssu_nocheckpoint_8bit( + SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int prev_k, int d_tile, + int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) { + __syncthreads(); + compute_no_write_output_8bit(smem, params, warp, lane, prev_k, d_tile, out_seq_base, + head, cache_slot, D_val, seq_len); +} + +// ============================================================================= +// Per-path dispatcher: checkpoint branch (must_checkpoint == true). +// ============================================================================= +// Encapsulates the existing M-shard chain: PASS 1 replay+matmul-3 → frag_y_DxT, +// PASS 2 re-replay+encode (always runs since must_checkpoint==true here), the +// single __syncthreads, cooperative state STG, and the transposed matmul-4 + +// transpose-STG output. +// +// Pulled out of `checkpointing_ssu_kernel_8bit` to mirror the bf16 path's +// `ssu_checkpoint` and make the kernel-body dispatch on +// must_checkpoint readable. +template +__device__ __forceinline__ void ssu_checkpoint_8bit(SmemT& smem, + CheckpointingSsuParams const& params, int warp, + int lane, int prev_k, int d_tile, + int64_t out_seq_base, int head, + int64_t cache_slot, float D_val, int seq_len) { + using namespace cute; + int const tid = warp * warpSize + lane; + + // ── Allocate per-warp frag_y_DxT (chain mma C-frag, fp32) ── + // Layout ((2, 2), MMA_M=1, MMA_N=NPREDICTED_PAD_MMA_M/8) per thread. + // Caller must zero before chain matmul-3 accumulates. + auto tiled_mma_chain = + make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_chain = tiled_mma_chain.get_slice(tid); + auto id_DxT = + make_identity_tensor(make_shape(Int{}, Int{})); + Tensor frag_y_DxT = thr_mma_chain.partition_fragment_C(id_DxT); + cute::clear(frag_y_DxT); + + // ── Phase 1b: replay + amax + chain matmul-3 → frag_y_DxT (init_out^T). + // `encode_scale_per_row[]` is computed at the end of PASS 1 (from the warp- + // reduced amax) and consumed by `encode_state_replay_8bit` further below — + // *after* `compute_output_8bit` consumes `frag_y_DxT` and STGs the output. + float encode_scale_per_row[2]; + float total_scale[2]; + replay_state_mma_8bit_chain( + smem, params, warp, lane, prev_k, d_tile, cache_slot, head, /*must_checkpoint=*/true, + frag_y_DxT, encode_scale_per_row, total_scale); + + // ── Philox seed for stochastic rounding (deferred to reduce register pressure) ── + [[maybe_unused]] int64_t const rand_seed = (PHILOX_ROUNDS > 0) ? *params.rand_seed : 0; + // `state_ptr_offset` is int64 — matches Triton's `base_rand = + // cache_batch_idx * stride_state_batch + ...` (cache_batch_idx is .to(int64)). + // Full 64 bits flow through `philox_randint4x`, which splits low/high + // across Philox c0/c1. No collision risk at large serving cache sizes. + int64_t const state_ptr_offset = + cache_slot * params.state_stride_seq + (int64_t)head * DIM * DSTATE; + + // ── PASS 2 (replay-again): re-run replay HMMA, encode fp32 → int8 to + // smem.state. Runs BEFORE the sync so both replay passes overlap with + // warps 0,1's CB precompute — one fewer __syncthreads in the kernel. + // frag_y_DxT stays live through PASS 2 (extra register pressure accepted). ── + encode_state_replay_8bit( + smem, params, warp, lane, prev_k, d_tile, cache_slot, head, encode_scale_per_row, total_scale, + rand_seed, state_ptr_offset); + + // ── Single sync: cross-warp visibility for smem.CB_scaled (warps 0,1) / + // smem.x (warp 2) / smem.z (warp 3) / smem.state (all warps' M-shards). ── + __syncthreads(); + + // ── Cooperative STG.128 for encoded state (after sync for cross-warp + // smem.state visibility). Fire-and-forget before compute_output_8bit. ── + store_state(smem, params, warp, lane, d_tile, head, + cache_slot); + + // ── Phase 2: transposed matmul-4 + epilogue + smem-transpose STG ── + compute_output_8bit( + smem, params, warp, lane, d_tile, out_seq_base, head, cache_slot, D_val, seq_len, frag_y_DxT); +} + +// ============================================================================= +// Kernel — int8 chain rewrite (separate kernel from the generic path) +// ============================================================================= +// The int8 path uses a fundamentally different output computation: +// 1. M-shard replay (Layout<_4, _1>) — same as v15.4. +// 2. Chained matmul-3: replay's fp32 C-frag → bf16 A-frag in registers via +// `convert_layout_acc_Aregs_sm80` (no smem.new_state staging). +// 3. Transposed matmul-4: x as A (M=D), CB^T as B → output^T(D, T) in regs. +// 4. Smem-transpose + cooperative STG.128 to (T, D) gmem. +// To keep the generic kernel uncluttered (no `if constexpr (sizeof(state_t) == 1)` +// branches), the int8 kernel is a standalone function that calls the new +// helpers (`replay_state_mma_8bit_chain`, `compute_output_8bit`) and uses +// `CheckpointingSsuStorage8bit` for smem. Phase 0/1 helpers (`load_data`, +// `store_old_B`, `compute_CB_scaled_2warp`) are reused verbatim — they only +// touch shared smem fields that both storage structs expose by name. +// +template +__global__ void checkpointing_ssu_kernel_8bit(CheckpointingSsuParams params) { + using namespace cute; + static_assert(sizeof(state_t) == 1, + "checkpointing_ssu_kernel_8bit requires 1-byte state_t (int8 or fp8 e4m3)"); + static_assert(NPREDICTED <= MAX_WINDOW); + static_assert(MAX_WINDOW <= MMA_prop::K_BIG); + // int8 path uses M-shard layout (Layout<_4,_1>): per-warp M = 16 = m16n8 + // atom M. D_PER_CTA must equal DIM (D_SPLIT=1) to give 4×16=64 D-rows/CTA. + // The wrapper enforces d_split == 1 for int8. + constexpr int D_PER_CTA = DIM; + static_assert(D_PER_CTA == 64, "int8 chain kernel requires DIM == 64"); + assert(params.d_split == 1); + + using SmemT = + CheckpointingSsuStorage8bit; + extern __shared__ __align__(128) char smem_buf[]; + auto& smem = *reinterpret_cast(smem_buf); + + // Grid: (1, batch, nheads). D-tile is always 0 for int8 (D_SPLIT=1). + int const d_tile = blockIdx.x; + int const seq = blockIdx.y; + int const head = blockIdx.z; + int const lane = threadIdx.x; + int const warp = threadIdx.y; + int const group_idx = head / HEADS_PER_GROUP; + + // ── Resolve cache slot ── + auto const* __restrict__ sbi = reinterpret_cast(params.state_batch_indices); + int64_t const cache_slot = sbi ? static_cast(sbi[seq]) : seq; + if (cache_slot == params.pad_slot_id) return; + + auto const* __restrict__ buf_idx_ptr = reinterpret_cast(params.cache_buf_idx); + int const buf_read = __ldg(&buf_idx_ptr[cache_slot]); + + auto const* __restrict__ prev_ptr = reinterpret_cast(params.prev_num_accepted); + int const prev_k = prev_ptr[cache_slot]; + + // ── Varlen vs non-varlen prologue. The kernel branches once on the + // VARLEN template; downstream helpers receive `seq_len` (constexpr-foldable + // NPREDICTED in non-varlen, runtime in varlen) and pre-computed per-sequence + // gmem base offsets (`x_seq_base` etc.) — they're varlen-agnostic. + // + // Uniform gmem-base formula: `outer * *_stride_seq` where + // non-varlen: outer = seq (= blockIdx.y), stride_seq = x.stride(0). + // varlen : outer = cu_seqlens[seq], stride_seq = x.stride(1). + // The wrapper picks the right stride_seq value; the kernel only branches + // on whether to load cu_seqlens. + int seq_len; + int64_t outer; + if constexpr (VARLEN) { + auto const* __restrict__ cu_seqlens = reinterpret_cast(params.cu_seqlens); + // Two LDG.E.32 (not one LDG.E.64): cu_seqlens is only 4-byte aligned + // at `&cu_seqlens[seq]` when seq is odd, and PTX + // `ld.global.v2.b32` faults on a 4-byte-aligned address. ptxas emits + // the two scalar loads back-to-back; latency is hidden against the + // following ALU work. + int const bos = __ldg(&cu_seqlens[seq]); + int const eos = __ldg(&cu_seqlens[seq + 1]); + seq_len = eos - bos; + if (seq_len <= 0) return; + outer = (int64_t)bos; + } else { + seq_len = NPREDICTED; + outer = (int64_t)seq; + } + // x/B/C bases computed inside `load_post_pdl_wait_data` from `outer` — + // see generic kernel for rationale (avoid pinning 6 regs across gdc_wait). + int64_t const dt_seq_base = outer * params.dt_stride_seq + head; + int64_t const z_seq_base = outer * params.z_stride_seq; + int64_t const out_seq_base = outer * params.out_stride_seq; + + bool const must_checkpoint = (prev_k + seq_len > MAX_WINDOW); + int const buf_write = must_checkpoint ? (1 - buf_read) : buf_read; + int const write_offset = must_checkpoint ? 0 : prev_k; + + // ── Load scalars (A, dt_bias, D) ── + auto const* __restrict__ A_ptr = reinterpret_cast(params.A); + auto const* __restrict__ dt_bias_ptr = reinterpret_cast(params.dt_bias); + auto const* __restrict__ D_ptr = reinterpret_cast(params.D); + float const A_val = toFloat(A_ptr[head]); + float const dt_bias_val = dt_bias_ptr ? toFloat(dt_bias_ptr[head]) : 0.f; + float const D_val = D_ptr ? toFloat(D_ptr[head]) : 0.f; + + // ── Phase 0: two-phase load around the PDL barrier (see generic kernel + // for the full rationale). Pre-wait: state + old_* cache + in_proj + // outputs (dt, z) + scalar scans. Post-wait: x/B/C from conv1d. ── + // ENABLE_PDL is JIT-stamped; `if constexpr` keeps only one load path in + // the binary (no register pressure leak from the unused path). + if constexpr (ENABLE_PDL) { + load_pre_pdl_wait_data(smem, params, lane, warp, d_tile, head, group_idx, cache_slot, + buf_read, A_val, dt_bias_val, dt_seq_base, z_seq_base, + seq_len); + gdc_wait(); + load_post_pdl_wait_data( + smem, params, lane, warp, d_tile, head, group_idx, outer, seq_len); + } else { + load_data( + smem, params, lane, warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, + outer, seq_len); + } + + // ── store_old_B hoist (warps 0,1 only, d_tile == 0) ── + if (d_tile == 0 && warp < 2) { + store_old_B( + smem, params, warp, lane, head, group_idx, cache_slot, buf_write, write_offset, seq_len); + } + + // ── CB precompute (4-warp split): warps 0,1 compute CB_scaled (new tokens); + // warps 2,3 compute CB_old (old tokens) in the no-write path only. Mirrors + // the bf16 path's dispatch — warps 2,3 stay idle in checkpoint mode and + // pick up work below inside `ssu_checkpoint_8bit`'s replay. ── + if (warp < 2) { + compute_CB_scaled_2warp(smem, warp, lane, seq_len); + } else if (!must_checkpoint) { + compute_CB_old_2warp(smem, warp, lane, prev_k, + seq_len); + } + + // ── Phase 1b + 2: per-path dispatch ── + // Checkpoint: M-shard chain (PASS 1 + PASS 2 + sync + state STG + transposed + // matmul-4 with smem-transpose STG). + // No-write : N-shard matmul-3 from int8/fp8 state + matmul-4-new + matmul-4-old + // + direct partition_C STG (mirrors the bf16 no-write path). + // must_checkpoint is uniform across the CTA — both branches contain a + // __syncthreads so divergence is balanced. + if (must_checkpoint) { + ssu_checkpoint_8bit(smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, + cache_slot, D_val, seq_len); + } else { + ssu_nocheckpoint_8bit(smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, + cache_slot, D_val, seq_len); + } + + // ── PDL: signal downstream that `output` is written. Cache writes below + // target tensors only the next SSU step reads, not the immediate + // downstream kernel — safe to signal first. ── + if constexpr (ENABLE_PDL) { + gdc_launch_dependents(); + } + + // ── Phase 3: cache writes (old_x, dt_proc, cumAdt) ── + store_old_x(smem, params, warp, lane, d_tile, head, + cache_slot, write_offset, seq_len); + if (d_tile == 0 && warp == 0 && lane < seq_len) { + auto* __restrict__ old_dt_w = reinterpret_cast(params.old_dt); + int64_t const dt_w_base = cache_slot * params.old_dt_stride_seq + + buf_write * params.old_dt_stride_dbuf + + head * params.old_dt_stride_head; + old_dt_w[dt_w_base + write_offset + lane] = smem.dt_proc[lane]; + } + if (d_tile == 0 && warp == 1 && lane < seq_len) { + auto* __restrict__ old_cumAdt_w = reinterpret_cast(params.old_cumAdt); + int64_t const ca_w_base = cache_slot * params.old_cumAdt_stride_seq + + buf_write * params.old_cumAdt_stride_dbuf + + head * params.old_cumAdt_stride_head; + old_cumAdt_w[ca_w_base + write_offset + lane] = smem.cumAdt[lane]; + } +} + +} // namespace flashinfer::mamba::checkpointing + +#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_8BIT_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh new file mode 100644 index 000000000000..83f0d2f8bfee --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh @@ -0,0 +1,1671 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * 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. + */ +#ifndef FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_COMMON_CUH_ +#define FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_COMMON_CUH_ + +// Shared infrastructure for the incremental SSU kernel: utilities, loaders, +// stores, MMA helpers, and functions used by both the 2/4-byte (bf16/fp16/fp32) +// and 8-bit (int8, future e4m3) kernel paths. + +#include +#include + +#include +#include + +#include "../utils.cuh" +#include "../vec_dtypes.cuh" +#include "checkpointing_ssu.cuh" +#include "common.cuh" +#include "conversion.cuh" +#include "cute/tensor.hpp" +#include "ssu_mtp_common.cuh" + +namespace flashinfer::mamba::checkpointing { + +using namespace conversion; + +// ldmatrix.b8 (SM100_U8x16_LDSM_T) was tried as a replacement for per-lane +// LDS.16 int8 state loads. It's 5-18% slower (bench v16.7b vs v16.8) because: +// (1) inherent 2-way bank conflicts (16 threads × 16B vs 128B banks), +// (2) state is the accumulator (C-frag), not an A/B operand — layout +// remapping costs 8 shuffles + byte extractions, +// (3) dynamic byte selection via SHF adds 15%+ short_scoreboard stalls. +namespace constants { +constexpr unsigned int MASK_ALL_LANES = 0xFFFFFFFFu; +constexpr unsigned int num_bits_uint32 = 32u; +} // namespace constants + +// ── Programmatic Dependent Launch (PDL) helpers ──────────────────────────── +// `gdc_wait` enforces no gmem access before the upstream PDL-paired kernel +// has signaled. `gdc_launch_dependents` hints the downstream PDL-paired +// kernel to launch early. Both are no-ops on SM<90 and harmless without +// the launch-time `cudaLaunchAttributeProgrammaticStreamSerialization` +// attribute, so the kernel can always emit them; the host-side `enable_pdl` +// toggle is what flips the launch attribute. +__forceinline__ __device__ void gdc_wait() { +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif +} + +__forceinline__ __device__ void gdc_launch_dependents() { +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +// Round x up to the next multiple of Y (Y must be a power of 2). +template +constexpr int next_multiple_of(int x) { + static_assert(Y > 0 && (Y & (Y - 1)) == 0, "Y must be a power of 2"); + return (x + Y - 1) & ~(Y - 1); +} + +// NativeOf::type: scalar T → 2-wide native CUDA vector type. +template +struct NativeOf; +template <> +struct NativeOf { + using type = float2; +}; +template <> +struct NativeOf<__half> { + using type = __half2; +}; +template <> +struct NativeOf<__nv_bfloat16> { + using type = __nv_bfloat162; +}; + +// Pair: thin wrapper over the native 2-wide type, adding compile-time +// `[cute::Int{}]` indexing so index-driven loops stay branchless. Same +// layout as the native type — `=` compiles to one LDS.U32 / STS.U32 and the +// pair stays in one register. +template +struct Pair { + typename NativeOf::type raw; + template + __device__ __forceinline__ auto operator[](cute::Int) const { + static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); + if constexpr (I == 0) + return raw.x; + else + return raw.y; + } +}; + +// Pair: explicit 16-bit packed specialization. A struct of two +// `int8_t` fields would let the compiler split the pair into two 32-bit +// registers (CUDA registers are 32-bit; sub-word fields are zero-extended +// per access). By backing the pair with a single 16-bit word and +// extracting via shift+cast, we keep the two elements in one register +// throughout the load → unpack → cast pipeline. +template <> +struct Pair { + uint16_t raw; // [bits 7:0] = element 0, [bits 15:8] = element 1 + template + __device__ __forceinline__ int8_t operator[](cute::Int) const { + static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); + if constexpr (I == 0) + return static_cast(raw & 0xFFu); + else + return static_cast(raw >> 8); + } +}; + +// Pair<__nv_fp8_e4m3>: same single-u16 backing as `Pair` — fp8 e4m3 +// is also a 1-byte storage type, so the load → unpack → cast pipeline runs +// through one 16-bit register. +template <> +struct Pair<__nv_fp8_e4m3> { + uint16_t raw; + template + __device__ __forceinline__ __nv_fp8_e4m3 operator[](cute::Int) const { + static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); + __nv_fp8_storage_t const byte = (I == 0) ? static_cast<__nv_fp8_storage_t>(raw & 0xFFu) + : static_cast<__nv_fp8_storage_t>(raw >> 8); + return reinterpret_cast<__nv_fp8_e4m3 const&>(byte); + } +}; + +// pack_float2: float2 → Pair, using packed hardware cvt when available. +template +__device__ __forceinline__ Pair pack_float2(float2 val); +template <> +__device__ __forceinline__ Pair pack_float2(float2 val) { + return {val}; +} +template <> +__device__ __forceinline__ Pair<__half> pack_float2<__half>(float2 val) { + return {__float22half2_rn(val)}; +} +template <> +__device__ __forceinline__ Pair<__nv_bfloat16> pack_float2<__nv_bfloat16>(float2 val) { + return {conversion::fromFloat2(val)}; +} + +// ============================================================================= +// cp.async copy atoms +// ============================================================================= +// 128-bit vector loads, shared by every gmem→smem copy in the kernel. The +// ldmatrix unit has the same vector width so `vec_bytes` is derived from the +// atom's source-register type and reused as the LDSM vector width. +struct Copy_prop { + using Atom = cute::SM80_CP_ASYNC_CACHEALWAYS; + using AtomZFill = cute::SM80_CP_ASYNC_CACHEALWAYS_ZFILL; + static constexpr int vec_bytes = sizeof(std::remove_extent_t); +}; + +// ============================================================================= +// MMA constants +// ============================================================================= +// All MMA-related atom types, dtype, and dims for this kernel, grouped so the +// header doesn't sprinkle loose aliases. The replay step chooses between the +// k=8 and k=16 atoms at compile time (MAX_WINDOW ≤ 8 picks K8 for smaller smem, +// +1 CTA/SM); dims are pulled from MMA_Traits so they stay in sync with the +// atom choice (e.g. m16n8k32 for int8 would just need AtomK16/K8 swapped). +struct MMA_prop { + using AtomK16 = cute::SM80_16x8x16_F32BF16BF16F32_TN; + using AtomK8 = cute::SM80_16x8x8_F32BF16BF16F32_TN; + // Operand dtype — matches the bf16 input of the atoms above. + using operand_t = __nv_bfloat16; + + static constexpr int M = cute::size<0>(typename cute::MMA_Traits::Shape_MNK{}); + static constexpr int N = cute::size<1>(typename cute::MMA_Traits::Shape_MNK{}); + static constexpr int K_BIG = cute::size<2>(typename cute::MMA_Traits::Shape_MNK{}); + static constexpr int K_SMALL = cute::size<2>(typename cute::MMA_Traits::Shape_MNK{}); +}; + +// ============================================================================= +// Swizzled smem layout for mma.sync operands (row-major). +// ============================================================================= +// The swizzle picks the `M` parameter to make each ldmatrix / cp.async atom +// exactly 16 bytes of contiguous element data (one 128-bit vector), and keeps +// B = S = 3 so that each 8-row block XORs row↔column bits to stay +// bank-conflict-free on the 128-byte bank cycle. +// +// sizeof(T) Swizzle atom rows × cols row bytes +// 2B Swizzle<3, 3, 3> 8 × 64 128 +// 4B Swizzle<3, 2, 3> 8 × 32 128 +// 1B Swizzle<3, 4, 3> 8 × 128 128 +// +// The MMA operand element type dictates the smem buffer element type, which in +// turn dictates the swizzle — so every call site passes its own element type. +constexpr int log2_pow2(int x) { + int r = 0; + while (x > 1) { + x >>= 1; + ++r; + } + return r; +} + +template +struct SmemSwizzle { + static_assert(Copy_prop::vec_bytes % sizeof(Elem) == 0, + "element size must divide LDSM atom (16 bytes)"); + static constexpr int ELEMS_PER_ATOM = Copy_prop::vec_bytes / sizeof(Elem); + using type = cute::Swizzle<3, log2_pow2(ELEMS_PER_ATOM), 3>; + static constexpr int ATOM_ROWS = 1 << type::num_bits; + static constexpr int ATOM_COLS = 1 << (type::num_base + type::num_shft); +}; + +// Default (ROW_STRIDE == COLS): tile the swizzle atom into a (ROWS, COLS) +// physical extent — the canonical CuTe pattern. +// Padded (ROW_STRIDE > COLS): logical (ROWS, COLS) view with the row stride +// inflated to ROW_STRIDE. Used when COLS doesn't tile cleanly with the +// swizzle atom's col extent (e.g. CB_scaled: logical 16 cols, atom 64) but +// we want the atom-aligned bank pattern. The extra cols-per-row are not +// "wasted padding" — the swizzle XOR scatters logical cells across the full +// ROW_STRIDE, so the physical extent is what the bijection actually needs. +template +__device__ __forceinline__ auto make_swizzled_layout_rc() { + using namespace cute; + using S = SmemSwizzle; + static_assert(ROWS % S::ATOM_ROWS == 0, "ROWS must be a multiple of the swizzle atom rows"); + static_assert(ROW_STRIDE % S::ATOM_COLS == 0, + "ROW_STRIDE must be a multiple of the swizzle atom cols"); + static_assert(ROW_STRIDE >= COLS, "ROW_STRIDE must be at least COLS"); + if constexpr (ROW_STRIDE == COLS) { + auto atom = composition(typename S::type{}, + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, _1{}))); + return tile_to_shape(atom, make_shape(Int{}, Int{})); + } else { + return composition(typename S::type{}, make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, _1{}))); + } +} + +// Aliased-row swizzled smem layout: logical (LOGICAL_ROWS, COLS) view over a +// physical buffer sized to next_multiple_of(VALID_ROWS) rows. +// When VALID_ROWS ≤ ATOM_ROWS the physical buffer is just one row-atom tall +// (e.g. 8 rows for bf16) but the MMA still wants to address LOGICAL_ROWS=16 +// rows (m16n8k16's M). We achieve the alias with a stride-0 outer mode on +// the row-tile axis: row r ∈ [0, LOGICAL_ROWS) maps to physical row +// (r mod PHYS_ROWS), col c maps unchanged. The first m-tile carries the +// real C data; the second m-tile reads the same bytes, feeds garbage into +// MMA accumulator rows ≥ VALID_ROWS, predicated out at gmem store. +// +// When VALID_ROWS > ATOM_ROWS (VALID_ROWS > 8 for 2-byte), PHYS_ROWS == LOGICAL_ROWS +// and the alias factor collapses to 1 — this then degenerates to the same +// layout that `make_swizzled_layout_rc` produces. +template +__device__ __forceinline__ auto make_aliased_swizzled_layout_rc() { + using namespace cute; + using S = SmemSwizzle; + static_assert(LOGICAL_ROWS % S::ATOM_ROWS == 0, + "LOGICAL_ROWS must be a multiple of the swizzle atom rows"); + static_assert(COLS % S::ATOM_COLS == 0, "COLS must be a multiple of the swizzle atom cols"); + constexpr int PHYS_ROWS = next_multiple_of(VALID_ROWS); + constexpr int LOG_M_TILES = LOGICAL_ROWS / S::ATOM_ROWS; + constexpr int PHYS_M_TILES = PHYS_ROWS / S::ATOM_ROWS; + static_assert(LOG_M_TILES % PHYS_M_TILES == 0, + "LOGICAL_ROWS must be a multiple of PHYS_ROWS for clean alias"); + constexpr int ALIAS = LOG_M_TILES / PHYS_M_TILES; + constexpr int N_TILES = COLS / S::ATOM_COLS; + auto atom = composition(typename S::type{}, + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, _1{}))); + // Outer layout (in atom-units): row-tile mode = (PHYS_M_TILES, ALIAS) strides + // (1, 0); col-tile mode = N_TILES stride PHYS_M_TILES. blocked_product + // scales these by the atom cosize (= ATOM_ROWS * ATOM_COLS). + auto outer = + make_layout(make_shape(make_shape(Int{}, Int{}), Int{}), + make_stride(make_stride(_1{}, _0{}), Int{})); + return blocked_product(atom, outer); +} + +// Transposed swizzled smem layout: maps (col, row) → same physical offset as +// make_swizzled_layout_rc maps (row, col). Enables bank-conflict-free ldmatrix.trans +// reads on data stored with make_swizzled_layout_rc. +// Built by swapping modes of the original inner layout (before swizzle), which +// guarantees correct cross-atom offsets when both dimensions have multiple atoms. +template +__device__ __forceinline__ auto make_swizzled_layout_rc_transpose() { + using namespace cute; + using S = SmemSwizzle; + static_assert(ROWS % S::ATOM_ROWS == 0, "ROWS must be a multiple of the swizzle atom rows"); + static_assert(COLS % S::ATOM_COLS == 0, "COLS must be a multiple of the swizzle atom cols"); + // Build the inner (un-swizzled) tiled layout for the original (ROWS, COLS) layout + auto inner = tile_to_shape(make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, _1{})), + make_shape(Int{}, Int{})); + // Swap modes to get true transpose: result(c, r) == original(r, c) + auto inner_T = make_layout(get<1>(inner), get<0>(inner)); + return composition(typename S::type{}, inner_T); +} + +// ============================================================================= +// B/C/x/z load helper: single-warp cp.async into swizzled smem. +// ============================================================================= + +// Generic swizzled cp.async with ZFILL for padding rows — the six [ROWS_PAD, +// COLS] single-warp loaders (B, old_B, x, z, old_x, C) all collapse into this. +// Gmem tile is [ROWS_PAD, COLS] with runtime row stride; rows >= VALID_ROWS +// are zero-filled in smem without touching gmem (cp.async.ca.ZFILL). Thread +// layout Shape<_4,_8>×val Shape<_1,_8> = 32 threads × 16B each = one warp +// covers 4 rows × 64 cols per step. +// +// Template args are all compile-time so the ZFILL predicate constant-folds; +// the caller pre-offsets `gmem_src` by the tile base, keeping 64-bit pointer +// math at the callsite. `SmemShape` is a CuTe static shape, e.g. +// `cute::Shape, cute::Int<128>>` — (rows_pad, cols_pad). The +// shape travels as a single type so later we can pad cols too (e.g. DSTATE=96 +// rounded up to a full bank cycle) without growing the parameter list. +// `valid_rows_rt` is a runtime bound that overrides the compile-time +// `VALID_ROWS` template parameter — used by the varlen (v20) path to tighten +// the predicate from `< NPREDICTED` to `< seq_len`. When the caller omits it, +// the default is the constexpr `VALID_ROWS` so non-varlen call sites fold to +// the same SASS as before. +template +__device__ __forceinline__ void load_tile_async(input_t* __restrict__ smem_dst, + input_t const* __restrict__ gmem_src, + int gmem_row_stride, int lane, + int valid_rows_rt = VALID_ROWS) { + using namespace cute; + constexpr int ROWS_PAD = size<0>(SmemShape{}); + constexpr int VALID_COLS = size<1>(SmemShape{}); + // Smem cols are padded up to the swizzle atom width. We always use the + // wide thread layout (4 thread-rows × 8 thread-cols × 1×8 val = 4 rows × + // 64 cols/pass for bf16) so each thread-row covers all 8 vec-cols of the + // Swizzle atom — the design contract that makes cp.async writes + // bank-conflict-free (each row consumes one full bank cycle, rows + // serialize across cycles). A "narrow" layout (½-atom-width per row) + // would force adjacent rows to compete for the same 16 banks, costing + // ~3-4× replay (observed as 12-way LDGSTS conflicts in d_split=2 ncu). + // When VALID_COLS < SMEM_COLS (D_SPLIT > 1 path), cp.async ZFILL drops + // the predicated-out cols as zeros without touching gmem — same mechanism + // as ZFILL'ing rows ≥ VALID_ROWS. The padded smem cells are unused. + constexpr int SMEM_COLS = next_multiple_of::ATOM_COLS>(VALID_COLS); + Tensor s_full = + make_tensor(make_smem_ptr(smem_dst), make_swizzled_layout_rc()); + Tensor g_full = make_tensor(make_gmem_ptr(gmem_src), + make_layout(make_shape(Int{}, Int{}), + make_stride(gmem_row_stride, Int<1>{}))); + + constexpr int VAL_COLS_PER_THREAD = Copy_prop::vec_bytes / sizeof(input_t); + static_assert(SMEM_COLS % VAL_COLS_PER_THREAD == 0, + "SMEM_COLS must be divisible by VAL_COLS_PER_THREAD"); + using ThrLayout = Layout, Stride<_8, _1>>; + static_assert(size<1>(ThrLayout{}) * VAL_COLS_PER_THREAD == SmemSwizzle::ATOM_COLS, + "wide thread layout must cover one full swizzle atom width per row"); + auto g2s = make_tiled_copy(Copy_Atom{}, ThrLayout{}, + Layout>>{}); + auto thr = g2s.get_slice(lane); + + auto id = make_identity_tensor(make_shape(Int{}, Int{})); + auto thr_id = thr.partition_S(id); + auto pred = make_tensor(shape(thr_id)); + CUTE_UNROLL + for (int i = 0; i < size(pred); ++i) { + pred(i) = (get<0>(thr_id(i)) < valid_rows_rt) && (get<1>(thr_id(i)) < VALID_COLS); + } + copy_if(g2s, pred, thr.partition_S(g_full), thr.partition_D(s_full)); +} + +// State load — D_SPLIT-conditional dispatch: +// +// D_SPLIT == 1: per-warp partition (warp W loads rows +// [W*DIM/4 : (W+1)*DIM/4)). Large coalesced gmem +// reads per warp. Tests pass without an extra CTA-wide barrier +// because the post-replay __syncthreads covers the eventual +// cross-warp state reads. +// +// D_SPLIT >= 2: 128-thread cooperative load. Required because at +// D_PER_CTA = 16 / 4 = 4 D-rows per warp the per-warp layout no +// longer divides cleanly into the (4, 8) thread-tile atom that +// `Copy_prop::Atom` expects. Cooperative load works for any +// D_PER_CTA ∈ {DIM, DIM/2, DIM/4} that's a multiple of 16. +// +// Both variants write through `make_swizzled_layout_rc` +// followed by a `local_tile` to the (D_PER_CTA, DSTATE) slice this CTA +// owns — the swizzle outer-stride is invariant across D_SPLIT. +template +__device__ __forceinline__ void load_state_per_warp(SmemT& smem, + state_t const* __restrict__ state_ptr, + int64_t state_base, int warp, int lane) { + using namespace cute; + static_assert(NUM_WARPS == 4, "Expected 4 warps"); + static_assert(D_PER_CTA % NUM_WARPS == 0, "D_PER_CTA must be divisible by NUM_WARPS"); + constexpr int DIM_PER_WARP = D_PER_CTA / NUM_WARPS; + + // Single-local_tile path — swizzle layout sized to this CTA's + // D_PER_CTA slice; one local_tile splits it directly per-warp. + Tensor sState_full = make_tensor(make_smem_ptr(reinterpret_cast(smem.state)), + make_swizzled_layout_rc()); + Tensor gState_full = make_tensor(make_gmem_ptr(state_ptr + state_base), + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + + Tensor sState = local_tile(sState_full, make_shape(Int{}, Int{}), + make_coord(warp, _0{})); + Tensor gState = local_tile(gState_full, make_shape(Int{}, Int{}), + make_coord(warp, _0{})); + + constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); + auto g2s = + make_tiled_copy(Copy_Atom{}, + Layout, Stride<_8, _1>>{}, Layout>>{}); + auto thr = g2s.get_slice(lane); + copy(g2s, thr.partition_S(gState), thr.partition_D(sState)); +} + +template +__device__ __forceinline__ void load_state_cta(SmemT& smem, state_t const* __restrict__ state_ptr, + int64_t state_base, int tid) { + using namespace cute; + static_assert(NUM_WARPS == 4, "Expected 4 warps"); + + Tensor sState = make_tensor(make_smem_ptr(reinterpret_cast(smem.state)), + make_swizzled_layout_rc()); + Tensor gState = make_tensor(make_gmem_ptr(state_ptr + state_base), + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + + constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); + using ThrLayout = Layout, Stride<_8, _1>>; + constexpr int THR_ROWS = decltype(size<0>(ThrLayout{}))::value; + static_assert(D_PER_CTA % THR_ROWS == 0, + "D_PER_CTA must be divisible by the thread layout's row count"); + auto g2s = make_tiled_copy(Copy_Atom{}, ThrLayout{}, + Layout>>{}); + auto thr = g2s.get_slice(tid); + copy(g2s, thr.partition_S(gState), thr.partition_D(sState)); +} + +// ============================================================================= +// Phase 0: cooperative data load into smem (all warps). +// Compute cumAdt[T] = cumsum(A * dt_proc) → smem. +// Warp-level inclusive prefix sum using Hillis-Steele shuffles. +// Only the first NPREDICTED lanes participate; the rest are idle. +// +// If the storage struct exposes a `decay` field, this also writes +// `decay[lane] = exp(val)` — fuses the EX2 with the cumsum write so the output +// decay broadcast in `compute_output_8bit` becomes a plain LDS (no per-element +// __expf). Detected via SFINAE so the 2/4-byte storage (no decay field) is +// unaffected. +namespace detail { +template +struct has_decay : std::false_type {}; +template +struct has_decay().decay[0])>> : std::true_type {}; +} // namespace detail + +template +__device__ __forceinline__ void compute_cumAdt(SmemT& smem, int lane, float A_val) { + float val = (lane < NPREDICTED) ? A_val * smem.dt_proc[lane] : 0.f; + // Inclusive prefix sum (Hillis-Steele) + for (int offset = 1; offset < NPREDICTED; offset *= 2) { + float other = __shfl_up_sync(constants::MASK_ALL_LANES, val, offset); + if (lane >= offset) val += other; + } + if (lane < NPREDICTED) { + smem.cumAdt[lane] = val; + if constexpr (detail::has_decay::value) { + smem.decay[lane] = __expf(val); + } + } +} + +// Load phase. Split into two halves around the PDL barrier (`gdc_wait`): +// +// load_pre_pdl_wait_data: data NOT produced by the immediate upstream +// kernel (conv1d) — state and old_* are cache from the previous SSU +// step; dt and z are in_proj outputs (in_proj fully completed before +// conv1d began, so they are visible by the time we hit `gdc_wait`). +// Issues cp.async for cache + z, runs the scalar LDGs (old_dt, +// old_cumAdt, dt→dt_proc) and the cumAdt warp scan. No commit/wait — +// the cp.async stays in flight while we `gdc_wait` on conv1d. +// +// load_post_pdl_wait_data: x/B/C cp.async (conv1d outputs — must wait) +// and the single `__pipeline_commit + __pipeline_wait_prior(0) + +// __syncwarp` that drains both halves. Cache cp.async issued in the +// pre-wait half share the per-thread async group with these, so one +// wait_prior(0) covers them all. +// +// Per-warp data ownership (unchanged from the pre-split version): +// state: per-warp contiguous DIM slice (warp W owns rows [16W : 16W+16]). +// B, C: redundant on W0, W1 (both compute 2-warp CB, both need full). +// old_B: redundant on all 4 warps (each warp's replay reads full DSTATE). +// old_x: redundant on all 4 warps (small, ~2 KB — partitioning not worth +// the complication). +// x: W2 only (Phase-2 read, covered by final __syncthreads). +// z: W3 only (Phase-2 read, covered by final __syncthreads). +// scalars (old_dt, old_cumAdt, dt→dt_proc) + cumAdt cumsum: +// redundant on each warp's first NPREDICTED/MAX_WINDOW lanes. Writes +// are idempotent across warps (identical payloads to same slots). +// ============================================================================= +// Per-sequence gmem base offsets (`x_seq_base`, etc.) are computed once in +// the kernel prologue — they encode the "start of this sequence" along the +// outer axis (batch in non-varlen, packed-token in varlen). Helper indexing +// is then uniform `seq_base + inner`. +// +// `seq_len` is the per-sequence new-token count (== NPREDICTED constexpr in +// non-varlen, runtime int in varlen). Used as the cp.async row predicate +// and the dt/scalar lane predicate so trailing rows past `seq_len` ZFILL to +// zero in smem. +template +__device__ __forceinline__ void load_pre_pdl_wait_data( + SmemT& smem, CheckpointingSsuParams const& params, int lane, int warp, int d_tile, int head, + int group_idx, int64_t cache_slot, int buf_read, float A_val, float dt_bias_val, + int64_t dt_seq_base, int64_t z_seq_base, int seq_len) { + constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 + static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); + static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); + + int const d_tile_off = d_tile * D_PER_CTA; + + auto const* __restrict__ z_ptr = reinterpret_cast(params.z); + auto const* __restrict__ old_x_ptr = reinterpret_cast(params.old_x); + auto const* __restrict__ old_B_ptr = reinterpret_cast(params.old_B); + auto const* __restrict__ old_dt_ptr = reinterpret_cast(params.old_dt); + auto const* __restrict__ old_cumAdt_ptr = reinterpret_cast(params.old_cumAdt); + auto const* __restrict__ dt_ptr = reinterpret_cast(params.dt); + + int64_t const ox_base = cache_slot * params.old_x_stride_seq + head * DIM + d_tile_off; + int64_t const oB_base = cache_slot * params.old_B_stride_seq + + buf_read * params.old_B_stride_dbuf + group_idx * DSTATE; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + // ZShape: shrunk to the swizzle atom's row extent (z is read via + // partition_C alias, the physical buffer only needs to be one swizzle + // row-atom tall). + using ZShape = cute::Shape, cute::Int>; + // old_B / old_x's smem row count = replay matmul K-axis = MAX_WINDOW_PAD_MMA_K. + using OldBShape = cute::Shape, cute::Int>; + using OxShape = cute::Shape, cute::Int>; + + // ── State: per-CTA D-slice ([D_PER_CTA, DSTATE]). Dispatch on D_SPLIT + // (= DIM == D_PER_CTA): per-warp coalesced load when one CTA owns the + // full head's D, cooperative 128-thread load when D is sharded. ── + { + auto const* __restrict__ state_ptr = reinterpret_cast(params.state); + int64_t const state_base = cache_slot * params.state_stride_seq + (int64_t)head * DIM * DSTATE + + (int64_t)d_tile_off * DSTATE; + if constexpr (DIM == D_PER_CTA) { + // D_SPLIT=1: per-warp partition (warp w loads contiguous DIM/4 D-rows). + load_state_per_warp(smem, state_ptr, state_base, warp, + lane); + } else { + // D_SPLIT>=2: 128-thread cooperative load (per-warp doesn't divide + // cleanly when D_PER_CTA/4 is too small for the (4,8) thread atom). + int const tid = warp * warpSize + lane; + load_state_cta(smem, state_ptr, state_base, tid); + } + } + + // ── old_B: redundant on all 4 warps (each warp's replay consumes full + // DSTATE). Identical payloads to same smem dest — final bytes + // deterministic. VALID_ROWS = MAX_WINDOW (cache rows). ── + load_tile_async(smem.old_B, old_B_ptr + oB_base, params.old_B_stride_token, + lane); + + // ── old_x: redundant on all 4 warps (small, simpler than partitioning). + // VALID_ROWS = MAX_WINDOW (cache rows). ── + load_tile_async(smem.old_x, old_x_ptr + ox_base, params.old_x_stride_token, + lane); + + // ── z: W3 only (Phase-2 read, final __syncthreads makes it visible). + // Sourced from in_proj — not from conv1d — so safe to issue pre-wait. ── + if (warp == 3 && z_ptr) { + int64_t const z_base = z_seq_base + head * DIM + d_tile_off; + load_tile_async(smem.z, z_ptr + z_base, params.z_stride_token, lane, + seq_len); + } + + // Commit the cache cp.async group BEFORE the caller's `gdc_wait()` so the + // hardware actually issues the gmem→smem transfers while the wait is in + // flight (without commit, the operations sit pending and only kick off + // once the post-wait commit fires — no overlap). Placed immediately after + // the last cp.async (z); the synchronous LDGs + cumAdt scan below are not + // part of any pipeline group and run in parallel with the in-flight + // transfers. The post half issues a second group; `__pipeline_wait_prior(0)` + // there drains both. + __pipeline_commit(); + + // ── Scalar loads + cumAdt cumsum: redundant per warp. + // old_dt / old_cumAdt: load up to MAX_WINDOW lanes (cache scalars). + // dt_proc: load up to NPREDICTED lanes (new-token scalars from in_proj). + // Synchronous LDG + plain smem stores — no cp.async. Writes from 4 + // warps to the same slots are idempotent (same payloads). ── + static_assert(MAX_WINDOW <= warpSize, "MAX_WINDOW must fit in a single warp"); + if (lane < MAX_WINDOW) { + int64_t const dt_rd_base = cache_slot * params.old_dt_stride_seq + + buf_read * params.old_dt_stride_dbuf + + head * params.old_dt_stride_head; + smem.old_dt[lane] = old_dt_ptr[dt_rd_base + lane]; + + int64_t const ca_rd_base = cache_slot * params.old_cumAdt_stride_seq + + buf_read * params.old_cumAdt_stride_dbuf + + head * params.old_cumAdt_stride_head; + smem.old_cumAdt[lane] = old_cumAdt_ptr[ca_rd_base + lane]; + } + // dt → softplus → smem.dt_proc. Under varlen the active lane range is + // `[0, seq_len)`; lanes `[seq_len, NPREDICTED)` are left uninitialized — + // `compute_cumAdt` will scan over them and produce garbage in the + // `cumAdt[seq_len:NPREDICTED]` tail, but every downstream consumer + // (`compute_CB_scaled_2warp` mask, output STG, dt_proc/cumAdt tape writes) + // is gated on `seq_len`, so the garbage never reaches gmem or contaminates + // valid rows. + // + // Per-lane stride along the T-axis is `dt_stride_token` in both layouts + // (4D batch and 1D packed varlen); the caller bakes `head` into + // `dt_seq_base` so the inner indexing is `dt_seq_base + lane * + // dt_stride_token`. + if (lane < seq_len) { + float dt_val = toFloat(dt_ptr[dt_seq_base + (int64_t)lane * params.dt_stride_token]); + dt_val += dt_bias_val; + if (params.dt_softplus) dt_val = thresholded_softplus(dt_val); + smem.dt_proc[lane] = dt_val; + } + // cumAdt = cumsum(A * dt_proc) — warp-local Hillis-Steele shuffle. Each + // of the 4 warps runs the same reduction on identical inputs (dt_proc + // just written above) and writes the same smem.cumAdt slots. + compute_cumAdt(smem, lane, A_val); +} + +// Post-wait half. Issues cp.async for conv1d outputs (x, B, C) and drains +// the per-thread async group (which includes both the cache cp.async issued +// in `load_pre_pdl_wait_data` and these conv1d cp.async). Caller must have +// called `gdc_wait()` between the two halves; otherwise this reads stale +// conv1d data. +// +// Takes `outer` (the per-sequence outer index) rather than pre-multiplied +// `*_seq_base` scalars. Computing `outer * stride_seq` inside this function +// keeps the multipliers transient instead of pinning them across the +// `gdc_wait()` asm-volatile barrier (which the compiler can't reorder around +// and thus can't rematerialize through). Saves ~6 registers vs. pre-computed +// bases. +template +__device__ __forceinline__ void load_post_pdl_wait_data(SmemT& smem, + CheckpointingSsuParams const& params, + int lane, int warp, int d_tile, int head, + int group_idx, int64_t outer, int seq_len) { + constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 + static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); + static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); + + int const d_tile_off = d_tile * D_PER_CTA; + + auto const* __restrict__ B_ptr = reinterpret_cast(params.B); + auto const* __restrict__ C_ptr = reinterpret_cast(params.C); + auto const* __restrict__ x_ptr = reinterpret_cast(params.x); + + int64_t const B_base = outer * params.B_stride_seq + (int64_t)group_idx * DSTATE; + int64_t const C_base = outer * params.C_stride_seq + (int64_t)group_idx * DSTATE; + int64_t const x_base = outer * params.x_stride_seq + head * DIM + d_tile_off; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; + // CShape: first dim shrunk to the swizzle atom's row extent so cp.async + // writes don't spill past the (also shrunk) C smem buffer. When NPREDICTED + // > ATOM_ROWS this falls back to NPREDICTED_PAD_MMA_M. + using CShape = cute::Shape, cute::Int>; + // B's smem row count = matmul-1 N-axis = NPREDICTED_PAD_MMA_N. + using BShape = cute::Shape, cute::Int>; + using XShape = cute::Shape, cute::Int>; + + // ── B: redundant on W0, W1 (both do 2-warp CB compute) ── + if (warp < 2) { + load_tile_async(smem.B, B_ptr + B_base, params.B_stride_token, lane, + seq_len); + } + // ── C: redundant on all 4 warps — chain matmul-3 reads smem.C from every + // warp, so each warp must see its own cp.async without a cross-warp sync. + // Identical payloads to same smem dest (same pattern as old_B / old_x). ── + load_tile_async(smem.C, C_ptr + C_base, params.C_stride_token, lane, seq_len); + + // ── x: W2 only (Phase-2 read, final __syncthreads makes it visible) ── + if (warp == 2) { + load_tile_async(smem.x, x_ptr + x_base, params.x_stride_token, lane, + seq_len); + } + + // Commit the conv1d cp.async group and drain BOTH groups: the cache + // group committed in `load_pre_pdl_wait_data` (pre-`gdc_wait`) and this + // conv1d group. `__pipeline_wait_prior(0)` waits for ≤0 pending groups. + // __syncwarp() provides acquire semantics across the 32 lanes of each + // warp. No cross-warp sync here; the only __syncthreads is after + // CB + replay. + __pipeline_commit(); + __pipeline_wait_prior(0); + __syncwarp(); +} + +// Single-pass load — used when `params.enable_pdl == false`. All cp.async +// (state, B, C, old_B, old_x, x, z) issue together into one async group; +// scalars + cumAdt scan run while the cp.async are in flight; one commit + +// wait_prior(0) + syncwarp drains the whole thing. This restores the v21.0 +// load order: the synchronous LDG-then-STS for old_dt/old_cumAdt/dt benefits +// from overlap with the conv1d cp.async (B/C/x) — which the split (pre + +// gdc_wait + post) form sacrifices since conv1d cp.async only issue after +// the wait. When PDL is paired with an upstream conv1d, the split's +// cache-load-during-wait overlap dominates; when not paired, the split is +// pure overhead (gdc_wait is a no-op, but the cp.async are delayed). +template +__device__ __forceinline__ void load_data(SmemT& smem, CheckpointingSsuParams const& params, + int lane, int warp, int d_tile, int head, int group_idx, + int64_t cache_slot, int buf_read, float A_val, + float dt_bias_val, int64_t outer, int seq_len) { + constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 + static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); + static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); + + int const d_tile_off = d_tile * D_PER_CTA; + + auto const* __restrict__ B_ptr = reinterpret_cast(params.B); + auto const* __restrict__ C_ptr = reinterpret_cast(params.C); + auto const* __restrict__ x_ptr = reinterpret_cast(params.x); + auto const* __restrict__ z_ptr = reinterpret_cast(params.z); + auto const* __restrict__ old_x_ptr = reinterpret_cast(params.old_x); + auto const* __restrict__ old_B_ptr = reinterpret_cast(params.old_B); + auto const* __restrict__ old_dt_ptr = reinterpret_cast(params.old_dt); + auto const* __restrict__ old_cumAdt_ptr = reinterpret_cast(params.old_cumAdt); + auto const* __restrict__ dt_ptr = reinterpret_cast(params.dt); + + int64_t const B_base = outer * params.B_stride_seq + (int64_t)group_idx * DSTATE; + int64_t const C_base = outer * params.C_stride_seq + (int64_t)group_idx * DSTATE; + int64_t const x_base = outer * params.x_stride_seq + head * DIM + d_tile_off; + int64_t const ox_base = cache_slot * params.old_x_stride_seq + head * DIM + d_tile_off; + int64_t const oB_base = cache_slot * params.old_B_stride_seq + + buf_read * params.old_B_stride_dbuf + group_idx * DSTATE; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + using CShape = cute::Shape, cute::Int>; + using BShape = cute::Shape, cute::Int>; + using XShape = cute::Shape, cute::Int>; + using ZShape = cute::Shape, cute::Int>; + using OldBShape = cute::Shape, cute::Int>; + using OxShape = cute::Shape, cute::Int>; + + // ── State: per-CTA D-slice ([D_PER_CTA, DSTATE]) ── + { + auto const* __restrict__ state_ptr = reinterpret_cast(params.state); + int64_t const state_base = cache_slot * params.state_stride_seq + (int64_t)head * DIM * DSTATE + + (int64_t)d_tile_off * DSTATE; + if constexpr (DIM == D_PER_CTA) { + load_state_per_warp(smem, state_ptr, state_base, warp, + lane); + } else { + int const tid = warp * warpSize + lane; + load_state_cta(smem, state_ptr, state_base, tid); + } + } + + if (warp < 2) { + load_tile_async(smem.B, B_ptr + B_base, params.B_stride_token, lane, + seq_len); + } + load_tile_async(smem.C, C_ptr + C_base, params.C_stride_token, lane, seq_len); + + load_tile_async(smem.old_B, old_B_ptr + oB_base, params.old_B_stride_token, + lane); + load_tile_async(smem.old_x, old_x_ptr + ox_base, params.old_x_stride_token, + lane); + + if (warp == 2) { + load_tile_async(smem.x, x_ptr + x_base, params.x_stride_token, lane, + seq_len); + } + if (warp == 3 && z_ptr) { + int64_t const z_base = outer * params.z_stride_seq + head * DIM + d_tile_off; + load_tile_async(smem.z, z_ptr + z_base, params.z_stride_token, lane, + seq_len); + } + + // ── Scalar loads (overlap with cp.async) + cumAdt cumsum ── + static_assert(MAX_WINDOW <= warpSize, "MAX_WINDOW must fit in a single warp"); + if (lane < MAX_WINDOW) { + int64_t const dt_rd_base = cache_slot * params.old_dt_stride_seq + + buf_read * params.old_dt_stride_dbuf + + head * params.old_dt_stride_head; + smem.old_dt[lane] = old_dt_ptr[dt_rd_base + lane]; + + int64_t const ca_rd_base = cache_slot * params.old_cumAdt_stride_seq + + buf_read * params.old_cumAdt_stride_dbuf + + head * params.old_cumAdt_stride_head; + smem.old_cumAdt[lane] = old_cumAdt_ptr[ca_rd_base + lane]; + } + int64_t const dt_seq_base_local = outer * params.dt_stride_seq + head; + if (lane < seq_len) { + float dt_val = toFloat(dt_ptr[dt_seq_base_local + (int64_t)lane * params.dt_stride_token]); + dt_val += dt_bias_val; + if (params.dt_softplus) dt_val = thresholded_softplus(dt_val); + smem.dt_proc[lane] = dt_val; + } + compute_cumAdt(smem, lane, A_val); + + __pipeline_commit(); + __pipeline_wait_prior(0); + __syncwarp(); +} + +// (compute_cumAdt moved above load_pre_pdl_wait_data so it can be called from there) + +// Compute CB_scaled[T,T] = (C @ B^T) * decay * dt_proc * causal_mask. +// Split across 2 warps: warp 0 computes columns 0:8, warp 1 computes columns 8:16. +// Result stored to swizzled smem.CB_scaled (input_t, row stride 64, Swizzle<3,3,3>). +// Called between the two __syncthreads by warps 0 and 1 only. +// `seq_len` is the runtime row/col bound on the (T, T) CB_scaled tile. +// Caller passes `NPREDICTED` (constexpr) for non-varlen — the mask +// `j <= t && t < seq_len && j < seq_len` then folds to today's SASS. +// Varlen passes the per-sequence `seq_len ≤ NPREDICTED`; rows/cols past it +// get zeroed so downstream matmul-4 / chain matmul-3 see zeros there. +template +__device__ __forceinline__ void compute_CB_scaled_2warp(SmemT& smem, int warp, int lane, + int seq_len) { + using namespace cute; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; + // 2-warp output split: each warp owns NPREDICTED_PAD_MMA_M / 2 cols of + // the (NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M) CB tile. Must be a + // multiple of the MMA atom's N for the partition to be atom-aligned + // (currently 8 == MMA_prop::N for NPREDICTED_PAD_MMA_M=16; if M-pad ever + // grows, this still holds as long as M-pad is a multiple of 2 * MMA_prop::N). + constexpr int N_HALF = NPREDICTED_PAD_MMA_M / 2; + static_assert(N_HALF % MMA_prop::N == 0, + "compute_CB_scaled_2warp: NPREDICTED_PAD_MMA_M / 2 must be a multiple of MMA::N"); + + // CB_scaled output tile layout (used by both warp 0 compute and warp 1 + // zero-fill when smem.B has only 8 rows). Row stride matches the buffer's + // padded width (one swizzle atom of `input_t`). + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + auto layout_cb_swz = + make_swizzled_layout_rc(); + + // ── NPREDICTED_PAD_MMA_N == 8: warp 1 has no valid B rows to read. But + // CB_scaled[:, 8:16] must still be zero so matmul-4's K-reduction sees + // zeros for j ≥ NPREDICTED. Do a simple 32-thread zero-fill and return. + if constexpr (NPREDICTED_PAD_MMA_N == 8) { + if (warp == 1) { + auto* __restrict__ cb = reinterpret_cast(smem.CB_scaled); + constexpr int COLS_TO_CLEAR = NPREDICTED_PAD_MMA_M - N_HALF; // 8 +#pragma unroll + for (int i = lane; i < NPREDICTED_PAD_MMA_M * COLS_TO_CLEAR; i += warpSize) { + int const r = i / COLS_TO_CLEAR; + int const c = N_HALF + (i % COLS_TO_CLEAR); + cb[layout_cb_swz(r, c)] = MMA_prop::operand_t(0.f); + } + return; + } + } + + // ── Swizzled smem views ── + // C is padded to NPREDICTED_PAD_MMA_M; B has NPREDICTED_PAD_MMA_N rows. + // Use NPREDICTED_PAD_MMA_N for smem_B so the physical layout matches the + // write layout from load_tile_async — `tile_to_shape` produces different + // outer strides for (8, 128) vs (16, 128). + // Aliased C view: physical buffer is just next_multiple_of(NPREDICTED) + // rows tall but the MMA atom needs M=16; second m-tile aliases first m-tile + // (predicated rows discarded at output store). + auto layout_C = + make_aliased_swizzled_layout_rc(); + auto layout_B = make_swizzled_layout_rc(); + Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C); + Tensor smem_B = make_tensor(make_smem_ptr(reinterpret_cast(smem.B)), layout_B); + + // ── TiledMMA: _1x_1 = 32 threads, one [16, 8] atom ── + auto tiled_mma = + make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(lane); + + // ── K-tile A operand (C): full [NPREDICTED_PAD_MMA_M, K_TILE], shared by both warps ── + constexpr int K_TILE = MMA_prop::K_BIG; + Tensor smem_C_tiled = local_tile(smem_C, make_tile(Int{}, Int{}), + make_coord(_0{}, _)); + + // ── K-tile B operand ── + // NPREDICTED_PAD_MMA_N == 16: warp 0 → N=[0,8), warp 1 → N=[8,16). + // NPREDICTED_PAD_MMA_N == 8 : only warp 0 runs (warp 1 took the early + // exit above), tile at (_0, _). + Tensor smem_B_half = + local_tile(smem_B, make_tile(Int{}, Int{}), make_coord(warp, _)); + + // ── Register fragments ── + Tensor frag_A = thr_mma.partition_fragment_A(smem_C_tiled(_, _, _0{})); + Tensor frag_B = thr_mma.partition_fragment_B(smem_B_half(_, _, _0{})); + + // ── Output accumulator: [NPREDICTED_PAD_MMA_M, N_HALF] f32 ── + auto layout_cb_half = make_layout(make_shape(Int{}, Int{})); + Tensor frag_acc = thr_mma.partition_fragment_C(make_tensor((float*)nullptr, layout_cb_half)); + clear(frag_acc); + + // ── S2R copies ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(lane); + Tensor smem_C_s2r = s2r_thr_A.partition_S(smem_C_tiled); + Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); + + auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B = s2r_B.get_slice(lane); + Tensor smem_B_s2r = s2r_thr_B.partition_S(smem_B_half); + Tensor frag_B_view = s2r_thr_B.retile_D(frag_B); + + // ── Gemm: 8 K-tiles, 1 HMMA each ── + constexpr int NUM_K_TILES = DSTATE / K_TILE; +#pragma unroll + for (int k = 0; k < NUM_K_TILES; ++k) { + cute::copy(s2r_A, smem_C_s2r(_, _, _, k), frag_A_view); + cute::copy(s2r_B, smem_B_s2r(_, _, _, k), frag_B_view); + cute::gemm(tiled_mma, frag_acc, frag_A, frag_B, frag_acc); + } + + // ── Elementwise: decay * dt_proc * causal mask, convert f32 → MMA_prop::operand_t ── + auto id_half = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_half); + + // ── Store to swizzled smem.CB_scaled ── + Tensor smem_CB = make_tensor( + make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); + // Tile into [NPREDICTED_PAD_MMA_M, N_HALF] halves; warp selects its half + Tensor smem_CB_half = local_tile(smem_CB, make_tile(Int{}, Int{}), + make_coord(_0{}, warp)); + Tensor smem_CB_part = thr_mma.partition_C(smem_CB_half); + +#pragma unroll + for (int i = 0; i < size(frag_acc); ++i) { + int t = get<0>(id_part(i)); + int j = warp * N_HALF + get<1>(id_part(i)); + float val; + if (j <= t && t < seq_len && j < seq_len) { + val = frag_acc(i) * __expf(smem.cumAdt[t] - smem.cumAdt[j]) * smem.dt_proc[j]; + } else { + val = 0.f; + } + smem_CB_part(i) = MMA_prop::operand_t(val); + } +} + +// Compute CB_old[t, i] = (C @ old_B^T)[t, i] * exp(cumAdt[t]) * dB_old(i) for +// i ∈ [0, prev_k); 0 otherwise. +// dB_old(i) = exp(total_old_cumAdt − smem.old_cumAdt[i]) * smem.old_dt[i]. +// The per-t factor exp(cumAdt[t]) is baked in here (vs at matmul-4 time) so the +// epilogue's β-scale = exp(total_old_cumAdt) * exp(cumAdt[t]) on init_out +// composes with a single CB_old @ old_x add — no extra elementwise pass. +// Identity (matches Triton's combined-sequence SSU): +// y_old_contrib[t, d] = exp(cumAdt[t]) * Σ_i dB_old(i) * x_old[i, d] * (C[t] · B_old[i]) +// = Σ_i CB_old[t, i] * old_x[i, d]. +// Written into smem.CB_scaled at cols [NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M + +// MAX_WINDOW_PAD_MMA_K). Sibling of compute_CB_scaled_2warp — runs on warps 2, 3 in parallel with +// warps 0, 1 writing the new-token half at cols [0, NPREDICTED_PAD_MMA_M). Uses the no-write +// path's CB_old region of the same swizzled buffer (32 cols total ≤ CB_ROW_STRIDE=64). +// +// 2-warp N-split: each warp owns one m16n8 N-atom (MMA::N=8 cols). +// MAX_WINDOW_PAD_MMA_K == 16: warp 2 → cols [0, 8); warp 3 → cols [8, 16). +// MAX_WINDOW_PAD_MMA_K == 8 : warp 2 covers all 8 cols; warp 3 returns early. +template +__device__ __forceinline__ void compute_CB_old_2warp(SmemT& smem, int warp, int lane, int prev_k, + int seq_len) { + using namespace cute; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + constexpr int N_HALF = MMA_prop::N; // 8 — one m16n8 atom per warp + constexpr int NUM_N_ATOMS = MAX_WINDOW_PAD_MMA_K / N_HALF; + static_assert(MAX_WINDOW_PAD_MMA_K % N_HALF == 0, + "compute_CB_old_2warp: MAX_WINDOW_PAD_MMA_K must be a multiple of MMA::N"); + static_assert(NPREDICTED_PAD_MMA_M + MAX_WINDOW_PAD_MMA_K <= CB_ROW_STRIDE, + "CB_scaled buffer must fit both CB_new (cols [0,T_pad)) and CB_old " + "(cols [T_pad, T_pad+K_old)) within its physical row stride"); + + int const sub_warp = warp - 2; // ∈ {0, 1} + if (sub_warp >= NUM_N_ATOMS) return; + + float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + + // ── Swizzled smem views (A = C, B = old_B; same shapes as the replay path's + // C/old_B reads, so we get cache locality with no extra cp.async). ── + auto layout_C = + make_aliased_swizzled_layout_rc(); + auto layout_old_B = make_swizzled_layout_rc(); + Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C); + Tensor smem_old_B = + make_tensor(make_smem_ptr(reinterpret_cast(smem.old_B)), layout_old_B); + + // ── TiledMMA: 32 threads, single m16n8k16 atom (K-loops over DSTATE) ── + auto tiled_mma = + make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(lane); + + // ── K-tile A operand (C): full M, K-loop dim ── + constexpr int K_TILE = MMA_prop::K_BIG; + Tensor smem_C_tiled = local_tile(smem_C, make_tile(Int{}, Int{}), + make_coord(_0{}, _)); + + // ── K-tile B operand (old_B): warp picks its 8-col N-atom slice ── + Tensor smem_old_B_half = + local_tile(smem_old_B, make_tile(Int{}, Int{}), make_coord(sub_warp, _)); + + // ── Register fragments ── + Tensor frag_A = thr_mma.partition_fragment_A(smem_C_tiled(_, _, _0{})); + Tensor frag_B = thr_mma.partition_fragment_B(smem_old_B_half(_, _, _0{})); + + auto layout_cb_half = make_layout(make_shape(Int{}, Int{})); + Tensor frag_acc = thr_mma.partition_fragment_C(make_tensor((float*)nullptr, layout_cb_half)); + clear(frag_acc); + + // ── S2R copies ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(lane); + Tensor smem_C_s2r = s2r_thr_A.partition_S(smem_C_tiled); + Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); + + auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B = s2r_B.get_slice(lane); + Tensor smem_old_B_s2r = s2r_thr_B.partition_S(smem_old_B_half); + Tensor frag_B_view = s2r_thr_B.retile_D(frag_B); + + // ── GEMM: DSTATE / K_BIG = 8 K-tiles ── + constexpr int NUM_K_TILES = DSTATE / K_TILE; +#pragma unroll + for (int k = 0; k < NUM_K_TILES; ++k) { + cute::copy(s2r_A, smem_C_s2r(_, _, _, k), frag_A_view); + cute::copy(s2r_B, smem_old_B_s2r(_, _, _, k), frag_B_view); + cute::gemm(tiled_mma, frag_acc, frag_A, frag_B, frag_acc); + } + + // ── Identity coords for elementwise / store ── + auto id_half = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_half); + + // ── Store to swizzled smem.CB_scaled at the CB_old region (cols [T_pad, T_pad+K_old)). + // Use the full physical (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE) padded swizzle view. + // Byte-compatible with compute_CB_scaled_2warp's (T_pad, T_pad, CB_ROW_STRIDE) + // padded view: both produce inner offset r*CB_ROW_STRIDE + c, same Swizzle. ── + auto layout_cb_full = make_swizzled_layout_rc(); + Tensor smem_CB = make_tensor( + make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); + // (16, CB_ROW_STRIDE) tiled by (16, N_HALF) → (1, CB_ROW_STRIDE/N_HALF) tiles. + // Coord (0, T_pad/N_HALF + sub_warp) lands inside the CB_old region. + constexpr int CB_OLD_TILE_BASE = NPREDICTED_PAD_MMA_M / N_HALF; + Tensor smem_CB_half = local_tile(smem_CB, make_tile(Int{}, Int{}), + make_coord(_0{}, CB_OLD_TILE_BASE + sub_warp)); + Tensor smem_CB_part = thr_mma.partition_C(smem_CB_half); + +#pragma unroll + for (int i = 0; i < size(frag_acc); ++i) { + int t = get<0>(id_part(i)); + int j = sub_warp * N_HALF + get<1>(id_part(i)); + float val; + if (j < prev_k && t < seq_len) { + val = frag_acc(i) * __expf(smem.cumAdt[t] + total_old_cumAdt - smem.old_cumAdt[j]) * + smem.old_dt[j]; + } else { + val = 0.f; + } + smem_CB_part(i) = MMA_prop::operand_t(val); + } +} + +// ============================================================================= +// Precompute dB scaling coefficients (called once before the N-tile loop). +// Returns DB_COEFFS_PER_LANE floats in coeff[], one per B-fragment element of +// the replay MMA — equal to (replay K) / 4 (each m16n8k* B-frag holds K*8/32 +// elts per lane). +// DB_COEFFS_PER_LANE = 4 for m16n8k16 B fragment +// DB_COEFFS_PER_LANE = 2 for m16n8k8 B fragment +// coeff[i] = 0 when k >= prev_k, embedding the causal mask so the inner +// loop needs no branch. +// +// K-index derivation (row-major TN MMA, lane = tid % 32): +// K_base = (lane % 4) * 2 +// m16n8k16 B frag (4 elts): 0→K_base, 1→K_base+1, 2→K_base+8, 3→K_base+9 +// m16n8k8 B frag (2 elts): 0→K_base, 1→K_base+1 +// ============================================================================= +template +__device__ __forceinline__ void precompute_dB_coeff(float coeff[DB_COEFFS_PER_LANE], + SmemT const& smem, float total_cumAdt, + int prev_k, int lane) { + static_assert(DB_COEFFS_PER_LANE == 2 || DB_COEFFS_PER_LANE == 4, + "DB_COEFFS_PER_LANE must be 2 (k8) or 4 (k16)"); + int const K_base = (lane % 4) * 2; +#pragma unroll + for (int i = 0; i < DB_COEFFS_PER_LANE; ++i) { + // m16n8k_ V-index → K-offset: (V & 1) is the col-pair offset; (V & 2) ? 8 : 0 + // covers the second K-tile inside the K_BIG (k16) atom. + int const k = K_base + (i & 1) + ((i & 2) << 2); + coeff[i] = (k < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[k]) * smem.old_dt[k] : 0.f; + } +} + +// Apply precomputed dB coefficients to frag_B in-place. +// Scales frag_B in-place by per-coefficient multiplier; frag dtype inferred +// from FragB::value_type. +// coeff[i] = 0 encodes both causal mask and zero-fill for k >= prev_k. +// ============================================================================= +template +__device__ __forceinline__ void compute_dB_scaling(FragB& frag_B, + float const coeff[DB_COEFFS_PER_LANE]) { + using namespace cute; + static_assert(size(FragB{}) == DB_COEFFS_PER_LANE, "frag_B size must match DB_COEFFS_PER_LANE"); + using frag_t = typename FragB::value_type; +#pragma unroll + for (int i = 0; i < DB_COEFFS_PER_LANE; ++i) { + frag_B(i) = frag_t(toFloat(frag_B(i)) * coeff[i]); + } +} + +// Scale frag_A by dB coefficients ONCE before the N-pass loop, replacing +// 16 per-N-pass compute_dB_scaling calls on frag_B (64 scale ops → 8 or 4). +// Identity: sum_k A[m,k]*(c[k]*B[k,n]) == sum_k (c[k]*A[m,k])*B[k,n]. +// +// K-index derivation (PTX ISA, m16n8k{8,16} mma.sync, A operand row-major): +// groupID = lane / 4, threadID_in_group = lane % 4 +// K_base = (lane % 4) * 2 (same formula as B operand) +// m16n8k8 A regs: a0 = A[groupID, K_base:K_base+2] +// a1 = A[groupID+8, K_base:K_base+2] +// frag (4 elts): {0,2}→K_base, {1,3}→K_base+1 (2 unique K) +// m16n8k16 A regs: a0..a1 as above, +// a2 = A[groupID, K_base+8:K_base+10] +// a3 = A[groupID+8, K_base+8:K_base+10] +// frag (8 elts): {0,2}→K_base, {1,3}→K_base+1, +// {4,6}→K_base+8, {5,7}→K_base+9 (4 unique K) +// ============================================================================= +template +__device__ __forceinline__ void apply_dA_coeff(FragA& frag_A, SmemT const& smem, float total_cumAdt, + int prev_k, int lane) { + using namespace cute; + constexpr int FRAG_A_SIZE = size(FragA{}); + static_assert((MAX_WINDOW_PAD_MMA_K == 16 && FRAG_A_SIZE == 8) || + (MAX_WINDOW_PAD_MMA_K == 8 && FRAG_A_SIZE == 4), + "apply_dA_coeff: unsupported MMA K / frag_A size combination"); + using frag_t = typename FragA::value_type; + + int const K_base = (lane % 4) * 2; + + if constexpr (MAX_WINDOW_PAD_MMA_K == 8) { + float const c0 = (K_base < prev_k) + ? __expf(total_cumAdt - smem.old_cumAdt[K_base]) * smem.old_dt[K_base] + : 0.f; + float const c1 = (K_base + 1 < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[K_base + 1]) * + smem.old_dt[K_base + 1] + : 0.f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + frag_A(i) = frag_t(toFloat(frag_A(i)) * ((i & 1) ? c1 : c0)); + } + } else { + float c[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + int const k = K_base + (j & 1) + ((j & 2) ? 8 : 0); + c[j] = (k < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[k]) * smem.old_dt[k] : 0.f; + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + int const ci = (i & 1) | ((i & 4) >> 1); + frag_A(i) = frag_t(toFloat(frag_A(i)) * c[ci]); + } + } +} + +// ── CuTe mma.sync output sub-functions ────────────────────────────────────── +// Each operates on a register-resident frag_y accumulator (f32). +// Called from compute_and_store_output's N-tile loop. + +// Convert fragment elements from src_t to MmaT in-place. +// No-op when src_t == MmaT. For the cross-dtype case: reads a src_t pair, +// converts via f32 intermediate, writes an MmaT pair. `pack_float2` +// dispatches to the native packed cvt for the destination type (e.g. +// cvt.rn.bf16x2.f32 for bf16). +template +__device__ __forceinline__ void convert_frag(Frag& frag) { + if constexpr (!std::is_same_v) { +#pragma unroll + for (int i = 0; i < cute::size(frag); i += 2) { + float2 const vals = toFloat2(reinterpret_cast(&frag(i))); + *reinterpret_cast*>(&frag(i)) = pack_float2(vals); + } + } +} + +// State → MMA B operand: dtype-aware TiledCopy. +// 2-byte smem: LDSM (SM75_U32x2_LDSM_N) — vectorized 16-bit ldmatrix. +// 4-byte smem: scalar UniversalCopy; pairs are converted to +// bf16 in registers by `convert_frag` after the load. +template +__device__ __forceinline__ auto make_state_b_s2r(TiledMma const& tm) { + using namespace cute; + if constexpr (sizeof(state_t) == 2) { + return make_tiled_copy_B(Copy_Atom{}, tm); + } else { + static_assert(sizeof(state_t) == 4, "wide state path expects 4-byte smem"); + return make_tiled_copy_B(Copy_Atom, state_t>{}, tm); + } +} + +// Src → dst fragment conversion — a strict superset of the in-place overload +// above: supports narrowing (e.g. f32 → bf16) via a separate src fragment. +// Three paths: +// (1) src_t == dst_t: bit copy via Pair (sidesteps cutlass-wrapper vs +// native dtype mismatches like cutlass::bfloat16_t vs __nv_bfloat16). +// (2) Same width, different dtype (e.g. fp16 → bf16): paired cvt through f32. +// Works in-place when `src` aliases `dst`. +// (3) Different width (e.g. f32 → bf16): paired element load + pack_float2. +template +__device__ __forceinline__ void convert_frag(SrcFrag const& src, DstFrag& dst) { + using namespace cute; + if constexpr (std::is_same_v) { +#pragma unroll + for (int i = 0; i < size(src); i += 2) { + *reinterpret_cast*>(&dst(i)) = *reinterpret_cast const*>(&src(i)); + } + } else if constexpr (sizeof(src_t) == sizeof(dst_t)) { +#pragma unroll + for (int i = 0; i < size(src); i += 2) { + float2 const vals = toFloat2(reinterpret_cast(&src(i))); + *reinterpret_cast*>(&dst(i)) = pack_float2(vals); + } + } else { + static_assert(sizeof(dst_t) == 2, "only narrowing to 2-byte dst supported"); +#pragma unroll + for (int i = 0; i < size(src); i += 2) { + *reinterpret_cast*>(&dst(i)) = + pack_float2(make_float2(src(i), src(i + 1))); + } + } +} + +// 2b. frag_y += CB_scaled @ x (matmul 4, single K-tile) +// CB_scaled A operand loaded from swizzled smem via LDSM (precomputed by warps 0,1). +// x B operand loaded from smem via ldmatrix.trans. +template +__device__ __forceinline__ void add_cb_x(FragY& frag_y, FragCB const& frag_CB, + SmemXTrans const& smem_x_trans, + S2RBTrans const& s2r_B_trans, + S2RThrBTrans const& s2r_thr_B_trans, ThrMma const& thr_mma, + TiledMma const& tiled_mma, int n) { + using namespace cute; + Tensor smem_x_trans_ntile = local_tile( + smem_x_trans, make_tile(Int{}, Int{}), make_coord(n, _0{})); + auto smem_x_trans_s2r = s2r_thr_B_trans.partition_S(smem_x_trans_ntile); + auto frag_B_x = thr_mma.partition_fragment_B( + make_tensor((MmaT*)0x0, make_shape(Int{}, Int{}))); + auto frag_B_x_view = s2r_thr_B_trans.retile_D(frag_B_x); + + cute::copy(s2r_B_trans, smem_x_trans_s2r, frag_B_x_view); + cute::gemm(tiled_mma, frag_y, frag_CB, frag_B_x, frag_y); +} + +// 2c. frag_y += CB_old @ old_x (matmul-4 over old tokens; sibling of add_cb_x). +// CB_old A-operand: pre-loaded by caller (m16n8k_old A-frag). +// old_x B-operand: ldmatrix.trans from smem.old_x viewed transposed. +// K_OLD = MAX_WINDOW_PAD_MMA_K ∈ {8, 16}. Caller's tiled_mma_old uses the +// matching m16n8k_OLD atom (K_BIG=16 or K_SMALL=8). frag_y partitioned by a +// different (K_BIG) tiled_mma is layout-compatible — the m16n8 C-frag shape is +// the same regardless of K. +template +__device__ __forceinline__ void add_cb_old_x(FragY& frag_y, FragCBOld const& frag_CB_old, + SmemOldXTrans const& smem_old_x_trans, + S2RBTransOld const& s2r_B_trans_old, + S2RThrBTransOld const& s2r_thr_B_trans_old, + ThrMmaOld const& thr_mma_old, + TiledMmaOld const& tiled_mma_old, int n) { + using namespace cute; + Tensor smem_old_x_ntile = local_tile( + smem_old_x_trans, make_tile(Int{}, Int{}), make_coord(n, _0{})); + auto smem_old_x_s2r = s2r_thr_B_trans_old.partition_S(smem_old_x_ntile); + auto frag_B_old_x = thr_mma_old.partition_fragment_B( + make_tensor((MmaT*)0x0, make_shape(Int{}, Int{}))); + auto frag_B_old_x_view = s2r_thr_B_trans_old.retile_D(frag_B_old_x); + + cute::copy(s2r_B_trans_old, smem_old_x_s2r, frag_B_old_x_view); + cute::gemm(tiled_mma_old, frag_y, frag_CB_old, frag_B_old_x, frag_y); +} + +// 3b. frag_y += D * x[t, d] (per-thread skip connection via partition_C) +template +__device__ __forceinline__ void add_D_skip(FragY& frag_y, SmemX const& smem_x, + ThrMma const& thr_mma, float D_val, int n) { + using namespace cute; + if (D_val == 0.f) return; + Tensor smem_x_tile = local_tile(smem_x, make_tile(Int{}, Int{}), + make_coord(_0{}, n)); + Tensor x_part = thr_mma.partition_C(smem_x_tile); + // Load pairs of consecutive bf16 elements and convert via paired toFloat2. + // m16n8k16 partition_C places consecutive N-column pairs adjacent in smem. + static_assert(sizeof(input_t) == 2, "vectorized D_skip requires 2-byte input_t"); +#pragma unroll + for (int i = 0; i < size(frag_y); i += 2) { + float2 vals = toFloat2(reinterpret_cast(&x_part(i))); + frag_y(i) += D_val * vals.x; + frag_y(i + 1) += D_val * vals.y; + } +} + +// 4b. frag_y *= z * sigmoid(z) (z-gating via partition_C) +template +__device__ __forceinline__ void compute_z_gating(FragY& frag_y, SmemZ const& smem_z, + ThrMma const& thr_mma, void const* z_ptr, int n) { + using namespace cute; + if (!z_ptr) return; + Tensor smem_z_tile = local_tile(smem_z, make_tile(Int{}, Int{}), + make_coord(_0{}, n)); + Tensor z_part = thr_mma.partition_C(smem_z_tile); +#pragma unroll + for (int i = 0; i < size(frag_y); i += 2) { + float2 const z = toFloat2(reinterpret_cast(&z_part(i))); + frag_y(i) *= z.x * __fdividef(1.f, (1.f + __expf(-z.x))); + frag_y(i + 1) *= z.y * __fdividef(1.f, (1.f + __expf(-z.y))); + } +} + +// ============================================================================= +// Pipelined K-loop GEMM +// ============================================================================= +// Computes frag_y[n] += A @ B[n] for n ∈ [0, NumNTiles), where A is shared +// across N-tiles and B[n] is the n-th N-tile of `smem_B` (sliced inside). +// +// NumStages-deep register pipeline hides LDSM → HMMA latency: at steady state +// slot (k+NumStages-1) is loading while HMMA consumes slot k. ATypeIn → MmaT +// and BTypeIn → MmaT conversions happen in registers between load and consume +// (in-place when widths match — see `convert_frag`). +// +// Used by matmul 3 (init_out += C @ state^T): A = C (shared), B = state. +// NumNTiles = sizeof...(FragY) = D_PER_CTA / N_TILE (1 for D_SPLIT=2, 2 for +// D_SPLIT=1). +template +__device__ __forceinline__ void pipelined_kloop_gemm(TiledMma const& tiled_mma, + ThrMma const& thr_mma, int tid, + SmemAKtiled const& smem_A_ktiled, + SmemB const& smem_B, FragY&... frag_y) { + using namespace cute; + constexpr int NumNTiles = sizeof...(FragY); + static_assert(NumStages >= 2, "NumStages must be >= 2 for pipelining"); + static_assert(NumKTiles >= NumStages - 1, "NumKTiles must be >= NumStages - 1 for full prologue"); + static_assert(NumNTiles >= 1, "NumNTiles must be >= 1"); + + constexpr int N_TILE = cute::tile_size<1>(TiledMma{}); + constexpr int K_TILE = cute::tile_size<2>(TiledMma{}); + + // ── S2R copies ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + auto s2r_B = make_state_b_s2r(tiled_mma); + auto s2r_thr_B = s2r_B.get_slice(tid); + + // ── Tile B by (N, K): shape (N_TILE, K_TILE, N_OUTER, NumKTiles) ── + auto smem_B_tiled = local_tile(smem_B, make_tile(Int{}, Int{}), make_coord(_, _)); + + // ── Partitioned smem (A shared, B per-N-tile) ── + auto smem_A_s2r = s2r_thr_A.partition_S(smem_A_ktiled); + auto sample_smem_B_n = smem_B_tiled(_, _, _0{}, _); + using SmemBS2RType = decltype(s2r_thr_B.partition_S(sample_smem_B_n)); + SmemBS2RType smem_B_s2r[NumNTiles]; + CUTE_UNROLL + for (int n = 0; n < NumNTiles; ++n) { + smem_B_s2r[n] = s2r_thr_B.partition_S(smem_B_tiled(_, _, n, _)); + } + + // ── Fragment / view types ── + using FragA = decltype(thr_mma.partition_fragment_A(smem_A_ktiled(_, _, _0{}))); + using FragB = decltype(thr_mma.partition_fragment_B(sample_smem_B_n(_, _, _0{}))); + using b_view_t = std::conditional_t; + using FragBStg = decltype(make_fragment_like(std::declval())); + using FragAView = decltype(s2r_thr_A.retile_D(std::declval())); + using FragBStgView = decltype(s2r_thr_B.retile_D(std::declval())); + + // ── Multi-stage register fragments ── + // Storage type matches the MMA fragment for A; for B the staging buffer is + // BTypeIn-typed (when narrowing) or MmaT-typed (when widths match — the two + // alias the same registers and `convert_frag` collapses to a bit-copy / + // in-place reinterpret). + FragA frag_A[NumStages]; + FragB frag_B[NumNTiles][NumStages]; + FragBStg frag_B_stg[NumNTiles][NumStages]; + FragAView frag_A_view[NumStages]; + FragBStgView frag_B_stg_view[NumNTiles][NumStages]; + CUTE_UNROLL + for (int s = 0; s < NumStages; ++s) { + frag_A_view[s] = s2r_thr_A.retile_D(frag_A[s]); + CUTE_UNROLL + for (int n = 0; n < NumNTiles; ++n) { + frag_B_stg_view[n][s] = s2r_thr_B.retile_D(frag_B_stg[n][s]); + } + } + + // Pack frag_y into a pointer array for indexed access (replay kernel pattern). + using FragY0 = std::tuple_element_t<0, std::tuple>; + static_assert((std::is_same_v && ...), + "all FragY parameters must be the same type"); + FragY0* frag_y_p[NumNTiles] = {(&frag_y)...}; + + // ── Per-stage operations (slot is constant after #pragma unroll) ── + auto load_one = [&](int k_src, int slot) { + cute::copy(s2r_A, smem_A_s2r(_, _, _, k_src), frag_A_view[slot]); + CUTE_UNROLL + for (int n = 0; n < NumNTiles; ++n) { + cute::copy(s2r_B, smem_B_s2r[n](_, _, _, k_src), frag_B_stg_view[n][slot]); + } + }; + auto convert_one = [&](int slot) { + convert_frag(frag_A[slot]); + CUTE_UNROLL + for (int n = 0; n < NumNTiles; ++n) { + convert_frag(frag_B_stg[n][slot], frag_B[n][slot]); + } + }; + auto compute_one = [&](int slot) { + CUTE_UNROLL + for (int n = 0; n < NumNTiles; ++n) { + cute::gemm(tiled_mma, *frag_y_p[n], frag_A[slot], frag_B[n][slot], *frag_y_p[n]); + } + }; + + // ── Clear accumulators ── + CUTE_UNROLL + for (int n = 0; n < NumNTiles; ++n) clear(*frag_y_p[n]); + + // ── Prologue: load + convert stages 0..NumStages-2 ── + CUTE_UNROLL + for (int s = 0; s < NumStages - 1; ++s) { + load_one(s, s); + convert_one(s); + } + + // ── Main K-loop: load slot (k+NumStages-1) % NumStages, compute slot k % NumStages ── +#pragma unroll + for (int k = 0; k < NumKTiles; ++k) { + int const k_load = k + NumStages - 1; + int const slot_load = k_load % NumStages; + int const slot_compute = k % NumStages; + if (k_load < NumKTiles) load_one(k_load, slot_load); + compute_one(slot_compute); + if (k_load < NumKTiles) convert_one(slot_load); + } +} + +// ── Matmul 3: init_out = C @ state^T ──────────────────────────────────────── +// Thin wrapper: builds the swizzled smem views for C and state, then dispatches +// to `pipelined_kloop_gemm`. NumNTiles = sizeof...(FragY) = D_PER_CTA / N_TILE +// (1 for D_SPLIT=2, 2 for D_SPLIT=1). +// +// On C: aliased view (see compute_CB_scaled_2warp). On state: 2-byte smem is +// reinterpret-cast to MMA_prop::operand_t so the 16-bit LDSM atom matches the view +// (actual element type recovered inside `convert_frag`); ≥4-byte smem keeps +// the native dtype and uses scalar UniversalCopy + register conversion. +template +__device__ __forceinline__ void add_init_out(SmemT const& smem, TiledMma const& tiled_mma, + ThrMma const& thr_mma, int tid, FragY&... frag_y) { + using namespace cute; + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int K_TILE = cute::tile_size<2>(TiledMma{}); + constexpr int NUM_K_TILES = DSTATE / K_TILE; + // Smem source dtype for matmul-3 (generic kernel only; 8-bit state goes + // through the dedicated `checkpointing_ssu_kernel_8bit` path): + // - sizeof(state_t) == 2 (fp16/bf16): LDSM the native 16-bit, view as bf16. + // - sizeof(state_t) == 4 (fp32): scalar UniversalCopy + on-the-fly convert. + static_assert(sizeof(state_t) != 1, + "add_init_out is the 2/4-byte path; 1-byte state goes through " + "compute_output_8bit"); + constexpr bool is_2byte_smem = (sizeof(state_t) == 2); + using state_view_t = std::conditional_t; + using BTypeIn = state_t; + + auto layout_C_swz = + make_aliased_swizzled_layout_rc(); + Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), + layout_C_swz); + Tensor smem_C_ktiled = local_tile(smem_C, make_tile(Int{}, Int{}), + make_coord(_0{}, _)); + + // Swizzle layout matches the dtype of the buffer being viewed. + auto const layout_state_swz = make_swizzled_layout_rc(); + state_view_t const* smem_state_ptr = reinterpret_cast(smem.state); + Tensor smem_state = make_tensor(make_smem_ptr(smem_state_ptr), layout_state_swz); + + pipelined_kloop_gemm<3, NUM_K_TILES, input_t, BTypeIn, MMA_prop::operand_t>( + tiled_mma, thr_mma, tid, smem_C_ktiled, smem_state, frag_y...); +} + +// store_state: vectorized smem → gmem state writeback (128 threads). +// Defined here (rather than alongside the other Phase 3 store helpers +// below) because compute_and_store_output calls it inline — issued right +// after matmul 3 so the STGs fire-and-forget in parallel with matmul 4 + +// epilogue. smem and gmem hold the same dtype now (no on-egress +// conversion) so this is always a direct 128-bit copy. +template +__device__ __forceinline__ void store_state(SmemT& smem, CheckpointingSsuParams const& params, + int warp, int lane, int d_tile, int head, + int64_t cache_slot) { + using namespace cute; + int const flat_tid = warp * warpSize + lane; + auto* __restrict__ state_w = reinterpret_cast(params.state); + // gmem dest = head's full state base + d_tile's row slice. + int64_t const state_base = cache_slot * params.state_stride_seq + (int64_t)head * DIM * DSTATE + + (int64_t)d_tile * D_PER_CTA * DSTATE; + + // ── Per-CTA smem swizzle layout [D_PER_CTA, DSTATE]. ── + auto layout_smem_swz = make_swizzled_layout_rc(); + state_t const* smem_state_base = reinterpret_cast(smem.state); + + Tensor sState = make_tensor(make_smem_ptr(smem_state_base), layout_smem_swz); + Tensor gState = make_tensor(make_gmem_ptr(state_w + state_base), + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + // Each store is 16 bytes — adjust val cols to the dtype. + constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); + auto s2g = + make_tiled_copy(Copy_Atom, state_t>{}, + Layout, Stride<_8, _1>>{}, Layout>>{}); + auto thr = s2g.get_slice(flat_tid); + copy(s2g, thr.partition_S(sState), thr.partition_D(gState)); +} + +// ── Store functions (called from kernel after compute_y + sync) ── +// (store_state moved above compute_and_store_output — used there for +// the state-writeback hoist.) + +template +__device__ __forceinline__ void store_old_x(SmemT& smem, CheckpointingSsuParams const& params, + int warp, int lane, int d_tile, int head, + int64_t cache_slot, int write_offset, int seq_len) { + using namespace cute; + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + int const flat_tid = warp * warpSize + lane; + + auto* __restrict__ old_x_w = reinterpret_cast(params.old_x); + // gmem dest = head's full slot + d_tile's D-slice offset, shifted by + // `write_offset` along the T-axis (must_checkpoint ? 0 : prev_k). + int64_t const ox_w_base = cache_slot * params.old_x_stride_seq + + (int64_t)write_offset * params.old_x_stride_token + head * DIM + + (int64_t)d_tile * D_PER_CTA; + + // Smem and gmem are both viewed at the full atom-padded width D_SMEM_COLS. + // The wide thread layout (16 row × 8 col × 1×8 val = 16 rows × 64 cols/pass + // for bf16) covers one full atom width per thread-row, which is the + // swizzle's bank-conflict-free contract on the LDS side (load-from-smem). + // A narrow layout would (a) waste 64 threads (warps 2, 3 idle) and + // (b) cause LDS bank conflicts on the smem-read side (observed as + // 4-way LDS conflict in d_split=2 ncu). Cols ≥ D_PER_CTA are predicated + // off via copy_if so STG never fires for them — no OOB write into the + // next d_tile / next head's gmem region. + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + auto layout_x_swz = make_swizzled_layout_rc(); + Tensor sX = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_swz); + Tensor gX = make_tensor(make_gmem_ptr(old_x_w + ox_w_base), + make_layout(make_shape(Int{}, Int{}), + make_stride(params.old_x_stride_token, Int<1>{}))); + + using ThrLayoutX = Layout, Stride<_8, _1>>; + auto s2g = make_tiled_copy(Copy_Atom, input_t>{}, ThrLayoutX{}, + Layout>{}); + auto thr_s2g = s2g.get_slice(flat_tid); + + auto tSsX = thr_s2g.partition_S(sX); + auto tSgX = thr_s2g.partition_D(gX); + + // Per-(row, col) predicate: skip rows ≥ NPREDICTED (m-padding) and cols ≥ + // D_PER_CTA (atom-padding past the d_tile's data). + auto cX = make_identity_tensor(make_shape(Int{}, Int{})); + auto tScX = thr_s2g.partition_D(cX); + auto pred = make_tensor(shape(tScX)); + CUTE_UNROLL + for (int i = 0; i < size(pred); ++i) { + pred(i) = (get<0>(tScX(i)) < seq_len) && (get<1>(tScX(i)) < D_PER_CTA); + } + copy_if(s2g, pred, tSsX, tSgX); +} + +// store_old_B runs on W0, W1 only (64 threads). Caller must gate +// with `if (warp < 2)` — these are the warps that hold valid smem.B +// after their own cp.async + wait. Halving the thread count keeps the +// overlap (writeback fires before CB+replay consume smem.B). +// +// Source: smem.B with NPREDICTED_PAD_MMA_N rows. +// Destination: gmem old_B[buf_write][write_offset:write_offset+NPREDICTED, :]. +// The `write_offset` argument shifts the gmem T-axis base — it's added to the +// base pointer below; the per-element predicate masks rows ≥ NPREDICTED. +// +// Thread layout `(8, 8) × (1, 8)` — **atom-aligned** with the Swizzle<3,3,3> +// (8, 64) atom for conflict-free smem reads. Per-tile 8 × 64 covers one +// full atom. For NPREDICTED_PAD_MMA_N=16: iters (2, 2) = 4 tiles, each +// thread owns 2 rows (t/8 and t/8+8) → per-iteration row predicate. For +// =8: iters (1, 2) = 2 tiles, each thread owns 1 row. The per-element +// predicate works for both. +template +__device__ __forceinline__ void store_old_B(SmemT& smem, CheckpointingSsuParams const& params, + int warp, int lane, int head, int group_idx, + int64_t cache_slot, int buf_write, int write_offset, + int seq_len) { + using namespace cute; + if (head % HEADS_PER_GROUP != 0) return; + constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; // matches smem.B row count + // Called only from warps 0, 1 — flat_tid ∈ [0, 64). + int const flat_tid = warp * warpSize + lane; + + auto* __restrict__ old_B_w = reinterpret_cast(params.old_B); + int64_t const oB_base = cache_slot * params.old_B_stride_seq + + buf_write * params.old_B_stride_dbuf + + (int64_t)write_offset * params.old_B_stride_token + group_idx * DSTATE; + + auto layout_B_swz = make_swizzled_layout_rc(); + Tensor sB = make_tensor(make_smem_ptr(reinterpret_cast(smem.B)), layout_B_swz); + Tensor gB = make_tensor(make_gmem_ptr(old_B_w + oB_base), + make_layout(make_shape(Int{}, Int{}), + make_stride(params.old_B_stride_token, Int<1>{}))); + + // 64 threads, (8, 8) × (1, 8) = atom-aligned per-tile (8, 64). + auto s2g = make_tiled_copy(Copy_Atom, input_t>{}, + Layout, Stride<_8, _1>>{}, Layout>{}); + auto thr_s2g = s2g.get_slice(flat_tid); + auto tSsB = thr_s2g.partition_S(sB); + auto tSgB = thr_s2g.partition_D(gB); + + // Fast path: no smem-side row padding AND no varlen-side truncation. + // The runtime `seq_len == NPREDICTED` is a constexpr-foldable compare in + // the non-varlen path (kernel prologue assigns `seq_len = NPREDICTED`), + // so it eliminates at -O3. In varlen with `seq_len == NPREDICTED` it's + // a runtime check that picks the cheaper unpredicated STG. + if constexpr (NPREDICTED == NPREDICTED_PAD_MMA_N) { + if (seq_len == NPREDICTED) { + copy(s2g, tSsB, tSgB); + return; + } + } + // Predicated: either smem rows > NPREDICTED (m-padding) OR varlen with + // seq_len < NPREDICTED. Mask each iter against `seq_len`. + auto cB = make_identity_tensor(make_shape(Int{}, Int{})); + auto tScB = thr_s2g.partition_D(cB); + auto pred = make_tensor(shape(tScB)); + CUTE_UNROLL + for (int i = 0; i < size(pred); ++i) { + pred(i) = get<0>(tScB(i)) < seq_len; + } + copy_if(s2g, pred, tSsB, tSgB); +} + +} // namespace flashinfer::mamba::checkpointing + +#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_COMMON_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh new file mode 100644 index 000000000000..754217d204da --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh @@ -0,0 +1,181 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * 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. + */ +#ifndef FLASHINFER_MAMBA_LAUNCH_CHECKPOINTING_SSU_CUH_ +#define FLASHINFER_MAMBA_LAUNCH_CHECKPOINTING_SSU_CUH_ + +// Launcher functions for the incremental SSU kernel. +// Includes both the bf16/fp16/fp32 and 8-bit kernel headers. + +#include "kernel_checkpointing_ssu.cuh" +#include "kernel_checkpointing_ssu_8bit.cuh" + +namespace flashinfer::mamba::checkpointing { + +// ── Dispatcher ───────────────────────────────────────────────────────────── +// `D_SPLIT` splits each head's DIM axis across `D_SPLIT` CTAs. +// `VARLEN` selects the packed-token gmem layout (cu_seqlens-driven). +// `launchCheckpointingSsuImpl` is the per-(D_SPLIT, VARLEN) specialization; +// `launchCheckpointingSsu` (below) is the runtime dispatcher. +template +void launchCheckpointingSsuImpl(CheckpointingSsuParams& params, cudaStream_t stream) { + constexpr int NUM_WARPS = 4; + + FLASHINFER_CHECK(params.nheads % params.ngroups == 0, "nheads (", params.nheads, + ") must be divisible by ngroups (", params.ngroups, ")"); + + // cp.async.ca with .L2::128B requires 16B-aligned pointers (128-bit / sizeof element). + // The .L2::128B hint further requires the base address to be 128B-aligned for full + // cache line utilization, but the hardware only faults on < 16B alignment. + // All cp.async-loaded operands need 16B alignment; output is also vectorized + // (Pair stores partitioned by m16n8k16 partition_C — base must be at + // least 16B-aligned for the stride math to keep per-thread stores aligned). + FLASHINFER_CHECK_ALIGNMENT(params.B, 16); + FLASHINFER_CHECK_ALIGNMENT(params.C, 16); + FLASHINFER_CHECK_ALIGNMENT(params.x, 16); + FLASHINFER_CHECK_ALIGNMENT(params.state, 16); + FLASHINFER_CHECK_ALIGNMENT(params.old_x, 16); + FLASHINFER_CHECK_ALIGNMENT(params.old_B, 16); + FLASHINFER_CHECK_ALIGNMENT(params.output, 16); + if (params.z != nullptr) { + FLASHINFER_CHECK_ALIGNMENT(params.z, 16); + } + + // Per-CTA D = DIM / D_SPLIT. Smem footprint shrinks for D-owned + // buffers (state, x, z, old_x); non-D buffers (B, C, old_B, scalars) unchanged. + constexpr int D_PER_CTA = DIM / D_SPLIT; + + // HEADS_PER_GROUP is JIT-stamped via the customize_config jinja, so only + // one (nheads / ngroups) specialization gets baked into this .so. The + // wrapper has already validated `nheads / ngroups == HEADS_PER_GROUP` + // before reaching us — the kernel cross-checks with an assert below. + FLASHINFER_CHECK(params.nheads / params.ngroups == HEADS_PER_GROUP, + "nheads/ngroups (=", params.nheads / params.ngroups, + ") must match JIT HEADS_PER_GROUP=", HEADS_PER_GROUP); + // PDL launch attribute. ENABLE_PDL is JIT-stamped (see + // checkpointing_ssu_customize_config.jinja); the kernel's body has its + // PDL PTX gated on the same constexpr via `if constexpr (ENABLE_PDL)`, so + // the .so contains exactly one load path. When ENABLE_PDL is false the + // attribute is set to 0 (effectively no PDL) — cudaLaunchKernelEx is + // used either way per FlashInfer convention (see norm.cuh:135). + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = ENABLE_PDL ? 1 : 0; + + auto launch_kernel = [&]() { + if constexpr (sizeof(state_t) == 1) { + // int8 chain rewrite — uses checkpointing_ssu_kernel_8bit + + // CheckpointingSsuStorage8bit. Only D_SPLIT == 1 is valid (the wrapper + // asserts this); D_SPLIT == 2 still gets template-instantiated by the + // public dispatcher's switch but is unreachable at runtime — gate the + // body with `if constexpr (D_SPLIT == 1)` so that path doesn't launch. + if constexpr (D_SPLIT == 1) { + auto func = + checkpointing_ssu_kernel_8bit; + constexpr size_t smem_size = + sizeof(CheckpointingSsuStorage8bit); + + if constexpr (smem_size > 0) { + FLASHINFER_CUDA_CHECK( + cudaFuncSetAttribute(func, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + + cudaLaunchConfig_t config; + config.gridDim = dim3(D_SPLIT, params.batch, params.nheads); + config.blockDim = dim3(warpSize, NUM_WARPS); + config.dynamicSmemBytes = smem_size; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = 1; + FLASHINFER_CUDA_CHECK(cudaLaunchKernelEx(&config, func, params)); + } else { + FLASHINFER_CHECK(false, + "checkpointing_ssu_kernel_8bit: unsupported D_SPLIT != 1 for 8-bit " + "state_t (got D_SPLIT=", + D_SPLIT, ")"); + } + } else { + // Generic kernel: bf16 / fp16 / fp32 state, supports D_SPLIT ∈ {1, 2}. + auto func = + checkpointing_ssu_kernel; + + constexpr size_t smem_size = sizeof( + CheckpointingSsuStorage); + + if constexpr (smem_size > 0) { + FLASHINFER_CUDA_CHECK( + cudaFuncSetAttribute(func, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + + // Grid is (D_SPLIT, batch, nheads). D-tile is the fastest axis so the + // `D_SPLIT` CTAs of the same head land on adjacent SMs and share L2 + // lines for the redundantly-loaded inputs (C, B, dt, ...). + cudaLaunchConfig_t config; + config.gridDim = dim3(D_SPLIT, params.batch, params.nheads); + config.blockDim = dim3(warpSize, NUM_WARPS); + config.dynamicSmemBytes = smem_size; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = 1; + FLASHINFER_CUDA_CHECK(cudaLaunchKernelEx(&config, func, params)); + } + }; + + launch_kernel(); +} + +// Public dispatcher: routes on `params.d_split` ({1, 2}) and varlen +// (`params.cu_seqlens != nullptr` → VARLEN=true). Each (D_SPLIT, VARLEN) +// pair gets its own template specialization — the JIT URI distinguishes them +// only via `d_split` today, so the same compiled `.so` will hold all four +// specializations after this commit. +template +void launchCheckpointingSsu(CheckpointingSsuParams& params, cudaStream_t stream) { + bool const is_varlen = (params.cu_seqlens != nullptr); + auto launch = [&]() { + launchCheckpointingSsuImpl(params, stream); + }; + auto launch_d_split = [&]() { + if (is_varlen) { + launch.template operator()(); + } else { + launch.template operator()(); + } + }; + switch (params.d_split) { + case 1: + launch_d_split.template operator()<1>(); + break; + case 2: + launch_d_split.template operator()<2>(); + break; + default: + FLASHINFER_CHECK(false, "Unsupported d_split: ", params.d_split, + ". Allowed values: {1, 2}. d_split=4 needs " + "warp-count restructure."); + } +} + +} // namespace flashinfer::mamba::checkpointing + +#endif // FLASHINFER_MAMBA_LAUNCH_CHECKPOINTING_SSU_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh new file mode 100644 index 000000000000..bb9accba44d9 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh @@ -0,0 +1,144 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * 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. + */ + +// Shared definitions for the vertical and horizontal MTP kernels. + +#pragma once + +#include + +#include "conversion.cuh" + +namespace flashinfer::mamba::mtp { + +// Round up to next power of 2 (compile-time). +constexpr int nextPow2(int v) { + v--; + v |= v >> 1; + v |= v >> 2; + v |= v >> 4; + v |= v >> 8; + v |= v >> 16; + return v + 1; +} + +using barrier_t = cuda::barrier; + +enum class WarpRole { kCompute, kTMALoad, kEpilogue }; + +__device__ __forceinline__ WarpRole get_warp_role(int warp) { + if (warp < 12) return WarpRole::kCompute; + if (warp < 15) return WarpRole::kTMALoad; + return WarpRole::kEpilogue; +} + +// XOR-based bank-conflict-free swizzle for horizontal state traversal. +// Operates on flat byte addresses: XORs the bank index with the row (cycle) index. +// cycle_length = row stride in bytes, bank_size = sizeof(uint32_t). +template +__device__ __forceinline__ int xor_swizzle(int address) { + int const cycle = address / cycle_length; + int const delta = address % cycle_length; + int const bank_idx = delta / bank_size; + int const intra_bank = delta % bank_size; + int const new_bank_idx = bank_idx ^ cycle; + return cycle * cycle_length + new_bank_idx * bank_size + intra_bank; +} + +// ── Parity-based barrier helpers (tight spin, no NANOSLEEP) ───────────────── +// More efficient than cuda::barrier::wait() for latency-sensitive pipelines. +// The standard cuda::barrier::wait() adds a NANOSLEEP backoff loop between +// try_wait attempts, which can overshoot and waste cycles. The raw +// mbarrier.try_wait.parity instruction does a tight spin instead. +// See CUDA Programming Guide §4.9.3 "Explicit Phase Tracking". + +__device__ __forceinline__ void arrive_and_wait_parity(barrier_t& bar, uint32_t& parity) { + uint32_t const smem_addr = + static_cast(__cvta_generic_to_shared(cuda::device::barrier_native_handle(bar))); + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(smem_addr) : "memory"); + uint32_t ready = 0; + while (!ready) { + asm volatile( + "{\n" + ".reg .pred p;\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n" + "selp.b32 %0, 1, 0, p;\n" + "}\n" + : "=r"(ready) + : "r"(smem_addr), "r"(parity)); + } + parity ^= 1; +} + +// ── SM100 f32x2 packed SIMD helpers ────────────────────────────────────────── +// On Blackwell (SM100+), {mul,fma}.f32x2 pack two fp32 operations into one +// instruction and issue on the dedicated FMUL2 pipeline, which runs in parallel +// with the regular FMA pipe. This halves instruction count for element-wise +// fp32 math on independent pairs (e.g. adjacent state-vector components). +// On older architectures the fallback is two scalar ops — zero overhead. +// See: https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/simd_sm100.hpp + +__device__ __forceinline__ void mul_f32x2(float2& c, float2 const& a, float2 const& b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + asm("mul.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast(c)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); +#else + c.x = a.x * b.x; + c.y = a.y * b.y; +#endif +} + +__device__ __forceinline__ void fma_f32x2(float2& d, float2 const& a, float2 const& b, + float2 const& c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + asm("fma.rn.f32x2 %0, %1, %2, %3;\n" + : "=l"(reinterpret_cast(d)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), + "l"(reinterpret_cast(c))); +#else + d.x = a.x * b.x + c.x; + d.y = a.y * b.y + c.y; +#endif +} + +// ============================================================================= +// convertAndStoreSRHorizontal — convert a pair of f32 state values to half. +// When PHILOX_ROUNDS > 0: stochastic rounding via f16x2. +// When PHILOX_ROUNDS == 0: plain nearest-even conversion. +// e is the pair-aligned index within the tile (must be even). +// ============================================================================= + +template +__device__ __forceinline__ void convertAndStoreSRHorizontal(state_t& out0, state_t& out1, float s0, + float s1, int64_t rand_seed, + int state_ptr_offset, int dd, int col0, + int e, uint32_t (&rand_ints)[4]) { + using namespace conversion; + if constexpr (PHILOX_ROUNDS > 0) { + if (e % 4 == 0) + philox_randint4x(rand_seed, state_ptr_offset + dd * DSTATE + col0 + e, + rand_ints[0], rand_ints[1], rand_ints[2], rand_ints[3]); + uint32_t packed = cvt_rs_f16x2_f32(s0, s1, rand_ints[e / 2 % 2]); + out0 = __ushort_as_half(static_cast(packed & 0xFFFFu)); + out1 = __ushort_as_half(static_cast(packed >> 16)); + } else { + convertAndStore(&out0, s0); + convertAndStore(&out1, s1); + } +} + +} // namespace flashinfer::mamba::mtp diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh new file mode 100644 index 000000000000..787a6de6d656 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh @@ -0,0 +1,539 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * 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. + */ +#ifndef FLASHINFER_UTILS_CUH_ +#define FLASHINFER_UTILS_CUH_ +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "exception.h" + +#define STR_HELPER(x) #x +#define STR(x) STR_HELPER(x) + +// macro to turn off fp16 qk reduction to reduce binary +#ifndef FLASHINFER_ALWAYS_DISUSE_FP16_QK_REDUCTION +#define FLASHINFER_ALWAYS_DISUSE_FP16_QK_REDUCTION 0 +#endif + +#ifndef NDEBUG +#define FLASHINFER_CUDA_CALL(func, ...) \ + { \ + cudaError_t e = (func); \ + if (e != cudaSuccess) { \ + std::cerr << "CUDA Error: " << cudaGetErrorString(e) << " (" << e << ") " << __FILE__ \ + << ": line " << __LINE__ << " at function " << STR(func) << std::endl; \ + return e; \ + } \ + } +#else +#define FLASHINFER_CUDA_CALL(func, ...) \ + { \ + cudaError_t e = (func); \ + if (e != cudaSuccess) { \ + return e; \ + } \ + } +#endif + +#define FLASHINFER_CUDA_CHECK(func) \ + do { \ + cudaError_t e = (func); \ + FLASHINFER_CHECK(e == cudaSuccess, "CUDA Error: ", cudaGetErrorString(e), " (", int(e), \ + ") at ", __FILE__, ":", __LINE__, " in ", STR(func)); \ + } while (0) + +#define FLASHINFER_CHECK_ALIGNMENT(ptr, size_bytes) \ + FLASHINFER_CHECK(reinterpret_cast(ptr) % (size_bytes) == 0, #ptr, \ + " must be aligned to ", (size_bytes), " bytes, got address ", (uintptr_t)(ptr)) + +#define FLASHINFER_CHECK_TMA_ALIGNED(ptr) FLASHINFER_CHECK_ALIGNMENT(ptr, 128) + +#define DISPATCH_USE_FP16_QK_REDUCTION(use_fp16_qk_reduction, USE_FP16_QK_REDUCTION, ...) \ + if (use_fp16_qk_reduction) { \ + FLASHINFER_ERROR("FP16_QK_REDUCTION disabled at compile time"); \ + } else { \ + constexpr bool USE_FP16_QK_REDUCTION = false; \ + __VA_ARGS__ \ + } + +#define DISPATCH_NUM_MMA_Q(num_mma_q, NUM_MMA_Q, ...) \ + if (num_mma_q == 1) { \ + constexpr size_t NUM_MMA_Q = 1; \ + __VA_ARGS__ \ + } else if (num_mma_q == 2) { \ + constexpr size_t NUM_MMA_Q = 2; \ + __VA_ARGS__ \ + } else { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported num_mma_q: " << num_mma_q; \ + FLASHINFER_ERROR(err_msg.str()); \ + } + +#define DISPATCH_NUM_MMA_KV(max_mma_kv, NUM_MMA_KV, ...) \ + if (max_mma_kv >= 8) { \ + constexpr size_t NUM_MMA_KV = 8; \ + __VA_ARGS__ \ + } else if (max_mma_kv >= 4) { \ + constexpr size_t NUM_MMA_KV = 4; \ + __VA_ARGS__ \ + } else if (max_mma_kv >= 2) { \ + constexpr size_t NUM_MMA_KV = 2; \ + __VA_ARGS__ \ + } else if (max_mma_kv >= 1) { \ + constexpr size_t NUM_MMA_KV = 1; \ + __VA_ARGS__ \ + } else { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported max_mma_kv: " << max_mma_kv; \ + FLASHINFER_ERROR(err_msg.str()); \ + } + +#define DISPATCH_CTA_TILE_Q(cta_tile_q, CTA_TILE_Q, ...) \ + switch (cta_tile_q) { \ + case 128: { \ + constexpr uint32_t CTA_TILE_Q = 128; \ + __VA_ARGS__ \ + break; \ + } \ + case 64: { \ + constexpr uint32_t CTA_TILE_Q = 64; \ + __VA_ARGS__ \ + break; \ + } \ + case 16: { \ + constexpr uint32_t CTA_TILE_Q = 16; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported cta_tile_q: " << cta_tile_q; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_GQA_GROUP_SIZE(group_size, GROUP_SIZE, ...) \ + if (group_size == 1) { \ + constexpr size_t GROUP_SIZE = 1; \ + __VA_ARGS__ \ + } else if (group_size == 2) { \ + constexpr size_t GROUP_SIZE = 2; \ + __VA_ARGS__ \ + } else if (group_size == 3) { \ + constexpr size_t GROUP_SIZE = 3; \ + __VA_ARGS__ \ + } else if (group_size == 4) { \ + constexpr size_t GROUP_SIZE = 4; \ + __VA_ARGS__ \ + } else if (group_size == 8) { \ + constexpr size_t GROUP_SIZE = 8; \ + __VA_ARGS__ \ + } else { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported group_size: " << group_size; \ + FLASHINFER_ERROR(err_msg.str()); \ + } + +#define DISPATCH_MASK_MODE(mask_mode, MASK_MODE, ...) \ + switch (mask_mode) { \ + case MaskMode::kNone: { \ + constexpr MaskMode MASK_MODE = MaskMode::kNone; \ + __VA_ARGS__ \ + break; \ + } \ + case MaskMode::kCausal: { \ + constexpr MaskMode MASK_MODE = MaskMode::kCausal; \ + __VA_ARGS__ \ + break; \ + } \ + case MaskMode::kCustom: { \ + constexpr MaskMode MASK_MODE = MaskMode::kCustom; \ + __VA_ARGS__ \ + break; \ + } \ + case MaskMode::kMultiItemScoring: { \ + constexpr MaskMode MASK_MODE = MaskMode::kMultiItemScoring; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported mask_mode: " << int(mask_mode); \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +// convert head_dim to compile-time constant +#define DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, ...) \ + switch (head_dim) { \ + case 64: { \ + constexpr size_t HEAD_DIM = 64; \ + __VA_ARGS__ \ + break; \ + } \ + case 128: { \ + constexpr size_t HEAD_DIM = 128; \ + __VA_ARGS__ \ + break; \ + } \ + case 256: { \ + constexpr size_t HEAD_DIM = 256; \ + __VA_ARGS__ \ + break; \ + } \ + case 512: { \ + constexpr size_t HEAD_DIM = 512; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported head_dim: " << head_dim; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +// convert interleave to compile-time constant +#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \ + if (interleave) { \ + constexpr bool INTERLEAVE = true; \ + __VA_ARGS__ \ + } else { \ + constexpr bool INTERLEAVE = false; \ + __VA_ARGS__ \ + } + +#define DISPATCH_ROPE_DIM(rope_dim, ROPE_DIM, ...) \ + switch (rope_dim) { \ + case 16: { \ + constexpr uint32_t ROPE_DIM = 16; \ + __VA_ARGS__ \ + break; \ + } \ + case 32: { \ + constexpr uint32_t ROPE_DIM = 32; \ + __VA_ARGS__ \ + break; \ + } \ + case 64: { \ + constexpr uint32_t ROPE_DIM = 64; \ + __VA_ARGS__ \ + break; \ + } \ + case 128: { \ + constexpr uint32_t ROPE_DIM = 128; \ + __VA_ARGS__ \ + break; \ + } \ + case 256: { \ + constexpr uint32_t ROPE_DIM = 256; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported ROPE_DIM: " << rope_dim; \ + err_msg << ". Supported values: 16, 32, 64, 128, 256"; \ + err_msg << " in DISPATCH_ROPE_DIM"; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_POS_ENCODING_MODE(pos_encoding_mode, POS_ENCODING_MODE, ...) \ + switch (pos_encoding_mode) { \ + case PosEncodingMode::kNone: { \ + constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kNone; \ + __VA_ARGS__ \ + break; \ + } \ + case PosEncodingMode::kRoPELlama: { \ + constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kRoPELlama; \ + __VA_ARGS__ \ + break; \ + } \ + case PosEncodingMode::kALiBi: { \ + constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kALiBi; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported pos_encoding_mode: " << int(pos_encoding_mode); \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_ALIGNED_VEC_SIZE(aligned_vec_size, ALIGNED_VEC_SIZE, ...) \ + switch (aligned_vec_size) { \ + case 16: { \ + constexpr size_t ALIGNED_VEC_SIZE = 16; \ + __VA_ARGS__ \ + break; \ + } \ + case 8: { \ + constexpr size_t ALIGNED_VEC_SIZE = 8; \ + __VA_ARGS__ \ + break; \ + } \ + case 4: { \ + constexpr size_t ALIGNED_VEC_SIZE = 4; \ + __VA_ARGS__ \ + break; \ + } \ + case 2: { \ + constexpr size_t ALIGNED_VEC_SIZE = 2; \ + __VA_ARGS__ \ + break; \ + } \ + case 1: { \ + constexpr size_t ALIGNED_VEC_SIZE = 1; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported aligned_vec_size: " << aligned_vec_size; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, ...) \ + if (compute_capacity.first >= 8) { \ + constexpr uint32_t NUM_STAGES_SMEM = 2; \ + __VA_ARGS__ \ + } else { \ + constexpr uint32_t NUM_STAGES_SMEM = 1; \ + __VA_ARGS__ \ + } + +namespace flashinfer { + +template +__forceinline__ __device__ __host__ constexpr T1 ceil_div(const T1 x, const T2 y) noexcept { + return (x + y - 1) / y; +} + +template +__forceinline__ __device__ __host__ constexpr T1 round_up(const T1 x, const T2 y) noexcept { + return ceil_div(x, y) * y; +} + +template +__forceinline__ __device__ __host__ constexpr T1 round_down(const T1 x, const T2 y) noexcept { + return (x / y) * y; +} + +inline std::pair GetCudaComputeCapability() { + int device_id = 0; + cudaGetDevice(&device_id); + int major = 0, minor = 0; + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_id); + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device_id); + return std::make_pair(major, minor); +} + +// This function is thread-safe and cached the sm_count. +// But it will only check the current CUDA device, thus assuming each process handles single GPU. +inline int GetCudaMultiProcessorCount() { + static std::atomic sm_count{0}; + int cached = sm_count.load(std::memory_order_relaxed); + if (cached == 0) { + int device_id; + cudaGetDevice(&device_id); + cudaDeviceProp device_prop; + cudaGetDeviceProperties(&device_prop, device_id); + cached = device_prop.multiProcessorCount; + sm_count.store(cached, std::memory_order_relaxed); + } + return cached; +} + +template +inline void DebugPrintCUDAArray(T* device_ptr, size_t size, std::string prefix = "") { + std::vector host_array(size); + std::cout << prefix; + cudaMemcpy(host_array.data(), device_ptr, size * sizeof(T), cudaMemcpyDeviceToHost); + for (size_t i = 0; i < size; ++i) { + std::cout << host_array[i] << " "; + } + std::cout << std::endl; +} + +inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_dim) { + if (avg_packed_qo_len > 64 && head_dim < 256) { + return 128; + } else { + auto compute_capacity = GetCudaComputeCapability(); + if (compute_capacity.first >= 8) { + // Ampere or newer + if (avg_packed_qo_len > 16) { + // avg_packed_qo_len <= 64 + return 64; + } else { + // avg_packed_qo_len <= 16 + return 16; + } + } else { + // NOTE(Zihao): not enough shared memory on Turing for 1x4 warp layout + return 64; + } + } +} + +inline int UpPowerOfTwo(int x) { + // Returns the smallest power of two greater than or equal to x + if (x <= 0) return 1; + --x; + x |= x >> 1; + x |= x >> 2; + x |= x >> 4; + x |= x >> 8; + x |= x >> 16; + return x + 1; +} + +#define LOOP_SPLIT_MASK(iter, COND1, COND2, ...) \ + { \ + _Pragma("unroll 1") for (; (COND1); (iter) -= 1) { \ + constexpr bool WITH_MASK = true; \ + __VA_ARGS__ \ + } \ + _Pragma("unroll 1") for (; (COND2); (iter) -= 1) { \ + constexpr bool WITH_MASK = false; \ + __VA_ARGS__ \ + } \ + } + +/*! + * \brief Return x - y if x > y, otherwise return 0. + */ +__device__ __forceinline__ uint32_t sub_if_greater_or_zero(uint32_t x, uint32_t y) { + return (x > y) ? x - y : 0U; +} + +// ======================= PTX Memory Utility Functions ======================= +// Non-atomic global memory access with cache streaming hint (cs) +// These are useful for streaming memory access patterns where data is used once + +/*! + * \brief Get the lane ID within a warp (0-31) + */ +__forceinline__ __device__ int get_lane_id() { + int lane_id; + asm("mov.u32 %0, %%laneid;" : "=r"(lane_id)); + return lane_id; +} + +/*! + * \brief Non-atomic global load for short (2 bytes) with cache streaming hint + */ +__forceinline__ __device__ short ld_na_global_s16(const short* addr) { + short val; + asm volatile("ld.global.cs.b16 %0, [%1];" : "=h"(val) : "l"(addr)); + return val; +} + +/*! + * \brief Non-atomic global store for short (2 bytes) with cache streaming hint + */ +__forceinline__ __device__ void st_na_global_s16(short* addr, short val) { + asm volatile("st.global.cs.b16 [%0], %1;" ::"l"(addr), "h"(val)); +} + +/*! + * \brief Non-atomic global load for int (4 bytes) with cache streaming hint + */ +__forceinline__ __device__ int ld_na_global_v1(const int* addr) { + int val; + asm volatile("ld.global.cs.b32 %0, [%1];" : "=r"(val) : "l"(addr)); + return val; +} + +/*! + * \brief Non-atomic global load for int2 (8 bytes) with cache streaming hint + */ +__forceinline__ __device__ int2 ld_na_global_v2(const int2* addr) { + int2 val; + asm volatile("ld.global.cs.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr)); + return val; +} + +/*! + * \brief Non-atomic global store for int (4 bytes) with cache streaming hint + */ +__forceinline__ __device__ void st_na_global_v1(int* addr, int val) { + asm volatile("st.global.cs.b32 [%0], %1;" ::"l"(addr), "r"(val)); +} + +/*! + * \brief Non-atomic global store for int2 (8 bytes) with cache streaming hint + */ +__forceinline__ __device__ void st_na_global_v2(int2* addr, int2 val) { + asm volatile("st.global.cs.v2.b32 [%0], {%1, %2};" ::"l"(addr), "r"(val.x), "r"(val.y)); +} + +/*! + * \brief Prefetch data to L2 cache + */ +template +__forceinline__ __device__ void prefetch_L2(const T* addr) { + asm volatile("prefetch.global.L2 [%0];" ::"l"(addr)); +} + +__device__ __forceinline__ void swap(uint32_t& a, uint32_t& b) { + uint32_t tmp = a; + a = b; + b = tmp; +} + +__device__ __forceinline__ uint32_t dim2_offset(const uint32_t& dim_a, const uint32_t& idx_b, + const uint32_t& idx_a) { + return idx_b * dim_a + idx_a; +} + +__device__ __forceinline__ uint32_t dim3_offset(const uint32_t& dim_b, const uint32_t& dim_a, + const uint32_t& idx_c, const uint32_t& idx_b, + const uint32_t& idx_a) { + return (idx_c * dim_b + idx_b) * dim_a + idx_a; +} + +__device__ __forceinline__ uint32_t dim4_offset(const uint32_t& dim_c, const uint32_t& dim_b, + const uint32_t& dim_a, const uint32_t& idx_d, + const uint32_t& idx_c, const uint32_t& idx_b, + const uint32_t& idx_a) { + return ((idx_d * dim_c + idx_c) * dim_b + idx_b) * dim_a + idx_a; +} + +#define DEFINE_HAS_MEMBER(member) \ + template \ + struct has_##member : std::false_type {}; \ + template \ + struct has_##member().member)>> : std::true_type {}; \ + template \ + inline constexpr bool has_##member##_v = has_##member::value; + +} // namespace flashinfer + +#endif // FLASHINFER_UTILS_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh new file mode 100644 index 000000000000..25c3b6fc60d4 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh @@ -0,0 +1,2194 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * 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. + */ +#ifndef VEC_DTYPES_CUH_ +#define VEC_DTYPES_CUH_ + +#include +#include +#include +#include +#if CUDA_VERSION >= 12080 +#include +#endif +#include + +#include + +namespace flashinfer { + +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 900)) +#define FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED +#endif + +#define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__ + +__device__ __forceinline__ void st_global_release(int4 const& val, int4* addr) { + asm volatile("st.release.global.sys.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), + "r"(val.z), "r"(val.w), "l"(addr)); +} + +__device__ __forceinline__ int4 ld_global_acquire(int4* addr) { + int4 val; + asm volatile("ld.acquire.global.sys.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(addr)); + return val; +} + +__device__ __forceinline__ void st_global_volatile(int4 const& val, int4* addr) { + asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), + "r"(val.z), "r"(val.w), "l"(addr)); +} + +__device__ __forceinline__ int4 ld_global_volatile(int4* addr) { + int4 val; + asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(addr)); + return val; +} + +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 < 120200) && \ + (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)) +// CUDA version < 12.2 and GPU architecture < 80 +FLASHINFER_INLINE __nv_bfloat162 make_bfloat162(const __nv_bfloat16 x, const __nv_bfloat16 y) { + __nv_bfloat162 t; + t.x = x; + t.y = y; + return t; +} + +FLASHINFER_INLINE __nv_bfloat16 __hmul(const __nv_bfloat16 a, const __nv_bfloat16 b) { + __nv_bfloat16 val; + const float fa = __bfloat162float(a); + const float fb = __bfloat162float(b); + // avoid ftz in device code + val = __float2bfloat16(__fmaf_ieee_rn(fa, fb, -0.0f)); + return val; +} + +FLASHINFER_INLINE __nv_bfloat162 __hmul2(const __nv_bfloat162 a, const __nv_bfloat162 b) { + __nv_bfloat162 val; + val.x = __hmul(a.x, b.x); + val.y = __hmul(a.y, b.y); + return val; +} + +FLASHINFER_INLINE __nv_bfloat162 __floats2bfloat162_rn(const float a, const float b) { + __nv_bfloat162 val; + val = __nv_bfloat162(__float2bfloat16_rn(a), __float2bfloat16_rn(b)); + return val; +} + +FLASHINFER_INLINE __nv_bfloat162 __float22bfloat162_rn(const float2 a) { + __nv_bfloat162 val = __floats2bfloat162_rn(a.x, a.y); + return val; +} +FLASHINFER_INLINE float2 __bfloat1622float2(const __nv_bfloat162 a) { + float hi_float; + float lo_float; + lo_float = __internal_bfloat162float(((__nv_bfloat162_raw)a).x); + hi_float = __internal_bfloat162float(((__nv_bfloat162_raw)a).y); + return make_float2(lo_float, hi_float); +} +#endif + +/******************* vec_t type cast *******************/ + +template +struct vec_cast { + template + FLASHINFER_INLINE static void cast(dst_t* dst, const src_t* src) { +#pragma unroll + for (size_t i = 0; i < vec_size; ++i) { + dst[i] = (dst_t)src[i]; + } + } +}; + +template <> +struct vec_cast<__nv_fp8_e4m3, float> { + template + FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, const float* src) { + if constexpr (vec_size == 1) { + dst[0] = __nv_fp8_e4m3(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((__nv_fp8x2_storage_t*)dst)[i] = + __nv_cvt_float2_to_fp8x2(((float2*)src)[i], __NV_SATFINITE, __NV_E4M3); + } + } + } +}; + +template <> +struct vec_cast<__nv_fp8_e5m2, float> { + template + FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, const float* src) { + if constexpr (vec_size == 1) { + dst[0] = __nv_fp8_e5m2(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((__nv_fp8x2_storage_t*)dst)[i] = + __nv_cvt_float2_to_fp8x2(((float2*)src)[i], __NV_SATFINITE, __NV_E5M2); + } + } + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(float* dst, const half* src) { + if constexpr (vec_size == 1) { + dst[0] = (float)src[0]; + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((float2*)dst)[i] = __half22float2(((half2*)src)[i]); + } + } + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(half* dst, const float* src) { + if constexpr (vec_size == 1) { + dst[0] = __float2half(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((half2*)dst)[i] = __float22half2_rn(((float2*)src)[i]); + } + } + } +}; + +template +constexpr FLASHINFER_INLINE int get_exponent_bits() { + if constexpr (std::is_same_v) { + return 4; + } else if constexpr (std::is_same_v) { + return 5; + } else if constexpr (std::is_same_v) { + return 5; + } else if constexpr (std::is_same_v) { + return 8; + } +} + +template +constexpr FLASHINFER_INLINE int get_mantissa_bits() { + if constexpr (std::is_same_v) { + return 3; + } else if constexpr (std::is_same_v) { + return 2; + } else if constexpr (std::is_same_v) { + return 11; + } else if constexpr (std::is_same_v) { + return 7; + } +} + +/*! + * \brief Fallback to software fast dequant implementation if hardware dequantization is not + * available. + * \note Inspired by Marlin's fast dequantization, but here we don't have to permute + * weights order. + * \ref + * https://github.com/vllm-project/vllm/blob/6dffa4b0a6120159ef2fe44d695a46817aff65bc/csrc/quantization/fp8/fp8_marlin.cu#L120 + */ +template +__device__ void fast_dequant_f8f16x4(uint32_t* input, uint2* output) { + uint32_t q = *input; + if constexpr (std::is_same_v && std::is_same_v) { + output->x = __byte_perm(0U, q, 0x5140); + output->y = __byte_perm(0U, q, 0x7362); + } else { + constexpr int FP8_EXPONENT = get_exponent_bits(); + constexpr int FP8_MANTISSA = get_mantissa_bits(); + constexpr int FP16_EXPONENT = get_exponent_bits(); + + constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT; + // Calculate MASK for extracting mantissa and exponent + constexpr int MASK1 = 0x80000000; + constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA); + constexpr int MASK3 = MASK2 & 0x7fffffff; + constexpr int MASK = MASK3 | (MASK3 >> 16); + q = __byte_perm(q, q, 0x1302); + + // Extract and shift FP8 values to FP16 format + uint32_t Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + uint32_t Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT); + + constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); + // Construct and apply exponent bias + if constexpr (std::is_same_v) { + const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); + + // Convert to half2 and apply bias + *(half2*)&(output->x) = __hmul2(*reinterpret_cast(&Out1), bias_reg); + *(half2*)&(output->y) = __hmul2(*reinterpret_cast(&Out2), bias_reg); + } else { + constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; + const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); + // Convert to bfloat162 and apply bias + *(nv_bfloat162*)&(output->x) = + __hmul2(*reinterpret_cast(&Out1), bias_reg); + *(nv_bfloat162*)&(output->y) = + __hmul2(*reinterpret_cast(&Out2), bias_reg); + } + } +} + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp8_e4m3* src) { + if constexpr (vec_size == 1) { + dst[0] = nv_bfloat16(src[0]); + } else if constexpr (vec_size == 2) { + dst[0] = nv_bfloat16(src[0]); + dst[1] = nv_bfloat16(src[1]); + } else { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) { + fast_dequant_f8f16x4<__nv_fp8_e4m3, nv_bfloat16>((uint32_t*)&src[i * 4], + (uint2*)&dst[i * 4]); + } + } + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp8_e5m2* src) { + if constexpr (vec_size == 1) { + dst[0] = nv_bfloat16(src[0]); + } else if constexpr (vec_size == 2) { + dst[0] = nv_bfloat16(src[0]); + dst[1] = nv_bfloat16(src[1]); + } else { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) { + fast_dequant_f8f16x4<__nv_fp8_e5m2, nv_bfloat16>((uint32_t*)&src[i * 4], + (uint2*)&dst[i * 4]); + } + } + } +}; + +template <> +struct vec_cast<__nv_fp8_e4m3, half> { + template + FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, const half* src) { +#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + if constexpr (vec_size == 1) { + dst[0] = __nv_fp8_e4m3(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint16_t y; + uint32_t x = *(uint32_t*)&src[i * 2]; + asm volatile("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(y) : "r"(x)); + *(uint16_t*)&dst[i * 2] = y; + } + } +#else +#pragma unroll + for (size_t i = 0; i < vec_size; ++i) { + dst[i] = __nv_fp8_e4m3(src[i]); + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + } +}; + +template <> +struct vec_cast<__nv_fp8_e5m2, half> { + template + FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, const half* src) { +#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + if constexpr (vec_size == 1) { + dst[0] = __nv_fp8_e5m2(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint16_t y; + uint32_t x = *(uint32_t*)&src[i * 2]; + asm volatile("cvt.rn.satfinite.e5m2x2.f16x2 %0, %1;" : "=h"(y) : "r"(x)); + *(uint16_t*)&dst[i * 2] = y; + } + } +#else +#pragma unroll + for (size_t i = 0; i < vec_size; ++i) { + dst[i] = __nv_fp8_e5m2(src[i]); + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(half* dst, const __nv_fp8_e4m3* src) { +#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + if constexpr (vec_size == 1) { + dst[0] = half(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint32_t y; + uint16_t x = *(uint16_t*)&src[i * 2]; + asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(y) : "h"(x)); + *(uint32_t*)&dst[i * 2] = y; + } + } +#else + if constexpr (vec_size == 1) { + dst[0] = half(src[0]); + } else if constexpr (vec_size == 2) { + dst[0] = half(src[0]); + dst[1] = half(src[1]); + } else { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) { + fast_dequant_f8f16x4<__nv_fp8_e4m3, half>((uint32_t*)&src[i * 4], (uint2*)&dst[i * 4]); + } + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(half* dst, const __nv_fp8_e5m2* src) { +#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + if constexpr (vec_size == 1) { + dst[0] = half(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint32_t y; + uint16_t x = *(uint16_t*)&src[i * 2]; + asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;" : "=r"(y) : "h"(x)); + *(uint32_t*)&dst[i * 2] = y; + } + } +#else + if constexpr (vec_size == 1) { + dst[0] = half(src[0]); + } else if constexpr (vec_size == 2) { + dst[0] = half(src[0]); + dst[1] = half(src[1]); + } else { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) { + fast_dequant_f8f16x4<__nv_fp8_e5m2, half>((uint32_t*)&src[i * 4], (uint2*)&dst[i * 4]); + } + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + } +}; + +#if defined(FLASHINFER_ENABLE_FP4_E2M1) && CUDA_VERSION >= 12080 +// Convert __nv_fp4x2_e2m1 (2 fp4 values per byte) to fp16. +// vec_size counts fp16 output elements; src has stride-2 layout: +// src[0] holds x0,x1 src[1] is padding +// src[2] holds x2,x3 src[3] is padding ... etc. +// Each valid byte encodes 2 fp4 values -> 2 fp16 via cvt.rn.f16x2.e2m1x2. +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(half* dst, const __nv_fp4x2_e2m1* src) { + static_assert(vec_size % 2 == 0, "vec_size must be even for fp4x2 dequantization"); +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint32_t y; + // Valid fp4x2 bytes are at even positions (stride 2); odd positions are padding. + uint32_t b = reinterpret_cast(src)[i * 2]; + asm volatile( + "{\n" + ".reg .b8 fp4_byte;\n" + "mov.b32 {fp4_byte, _, _, _}, %1;\n" + "cvt.rn.f16x2.e2m1x2 %0, fp4_byte;\n" + "}" + : "=r"(y) + : "r"(b)); + reinterpret_cast(dst)[i] = y; + } +#else + // Software LUT fallback for arch < SM100. + // e2m1 encoding: bit[3]=sign, bit[2:0]=magnitude index in {0,0.5,1,1.5,2,3,4,6}. + // Each packed byte holds two fp4 values: bits[3:0]=first, bits[7:4]=second. + constexpr uint16_t lut[16] = { + 0x0000, // +0.0 + 0x3800, // +0.5 + 0x3C00, // +1.0 + 0x3E00, // +1.5 + 0x4000, // +2.0 + 0x4200, // +3.0 + 0x4400, // +4.0 + 0x4600, // +6.0 + 0x8000, // -0.0 + 0xB800, // -0.5 + 0xBC00, // -1.0 + 0xBE00, // -1.5 + 0xC000, // -2.0 + 0xC200, // -3.0 + 0xC400, // -4.0 + 0xC600, // -6.0 + }; +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint8_t b = reinterpret_cast(src)[i * 2]; + reinterpret_cast(dst)[i * 2 + 0] = lut[b & 0x0F]; + reinterpret_cast(dst)[i * 2 + 1] = lut[(b >> 4) & 0x0F]; + } +#endif + } +}; +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp4x2_e2m1* src) { + static_assert(vec_size % 2 == 0, "vec_size must be even for fp4x2 dequantization"); +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint32_t y; + // Valid fp4x2 bytes are at even positions (stride 2); odd positions are padding. + uint32_t b = reinterpret_cast(src)[i * 2]; +#if (defined __CUDACC_VER_MAJOR__) && (defined __CUDACC_VER_MINOR__) && \ + ((__CUDACC_VER_MAJOR__ > 13) || ((__CUDACC_VER_MAJOR__ == 13) && (__CUDACC_VER_MINOR__ >= 2))) + // cvt.rn.bf16x2.e2m1x2 requires CUDA Toolkit >= 13.2 + asm volatile( + "{\n" + ".reg .b8 fp4_byte;\n" + "mov.b32 {fp4_byte, _, _, _}, %1;\n" + "cvt.rn.bf16x2.e2m1x2 %0, fp4_byte;\n" + "}" + : "=r"(y) + : "r"(b)); +#else + // Fallback: convert e2m1 -> fp16 -> bf16 when cvt.rn.bf16x2.e2m1x2 is unavailable + uint32_t fp16x2; + asm volatile( + "{\n" + ".reg .b8 fp4_byte;\n" + "mov.b32 {fp4_byte, _, _, _}, %1;\n" + "cvt.rn.f16x2.e2m1x2 %0, fp4_byte;\n" + "}" + : "=r"(fp16x2) + : "r"(b)); + __half2 h2 = reinterpret_cast<__half2&>(fp16x2); + __nv_bfloat162 bf16x2 = __float22bfloat162_rn(__half22float2(h2)); + y = reinterpret_cast(bf16x2); +#endif + reinterpret_cast(dst)[i] = y; + } +#else + // Software LUT fallback for arch < SM100. + // e2m1 encoding: bit[3]=sign, bit[2:0]=magnitude index in {0,0.5,1,1.5,2,3,4,6}. + // Each packed byte holds two fp4 values: bits[3:0]=first, bits[7:4]=second. + constexpr uint16_t lut[16] = { + 0x0000, // +0.0 + 0x3F00, // +0.5 + 0x3F80, // +1.0 + 0x3FC0, // +1.5 + 0x4000, // +2.0 + 0x4040, // +3.0 + 0x4080, // +4.0 + 0x40C0, // +6.0 + 0x8000, // -0.0 + 0xBF00, // -0.5 + 0xBF80, // -1.0 + 0xBFC0, // -1.5 + 0xC000, // -2.0 + 0xC040, // -3.0 + 0xC080, // -4.0 + 0xC0C0, // -6.0 + }; +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint8_t b = reinterpret_cast(src)[i * 2]; + reinterpret_cast(dst)[i * 2 + 0] = lut[b & 0x0F]; + reinterpret_cast(dst)[i * 2 + 1] = lut[(b >> 4) & 0x0F]; + } +#endif + } +}; + +#endif // FLASHINFER_ENABLE_FP4_E2M1 && CUDA_VERSION >= 12080 + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(float* dst, const nv_bfloat16* src) { + if constexpr (vec_size == 1) { + dst[0] = (float)src[0]; + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((float2*)dst)[i] = __bfloat1622float2(((nv_bfloat162*)src)[i]); + } + } + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const float* src) { + if constexpr (vec_size == 1) { + dst[0] = nv_bfloat16(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((nv_bfloat162*)dst)[i] = __float22bfloat162_rn(((float2*)src)[i]); + } + } + } +}; + +template +struct vec_t { + FLASHINFER_INLINE float_t& operator[](size_t i); + FLASHINFER_INLINE const float_t& operator[](size_t i) const; + FLASHINFER_INLINE void fill(float_t val); + FLASHINFER_INLINE void load(const float_t* ptr); + FLASHINFER_INLINE void store(float_t* ptr) const; + FLASHINFER_INLINE void load_global_acquire(float* addr); + FLASHINFER_INLINE void store_global_release(float* addr) const; + FLASHINFER_INLINE void load_global_volatile(float* addr); + FLASHINFER_INLINE void store_global_volatile(float* addr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src); + template + FLASHINFER_INLINE void cast_load(const T* ptr); + template + FLASHINFER_INLINE void cast_store(T* ptr) const; + FLASHINFER_INLINE static void memcpy(float_t* dst, const float_t* src); + FLASHINFER_INLINE float_t* ptr(); +}; + +template +FLASHINFER_INLINE void cast_from_impl(vec_t& dst, + const vec_t& src) { + vec_cast::cast( + dst.ptr(), const_cast*>(&src)->ptr()); +} + +template +FLASHINFER_INLINE void cast_load_impl(vec_t& dst, + const src_float_t* src_ptr) { + if constexpr (std::is_same_v) { + dst.load(src_ptr); + } else { + vec_t tmp; + tmp.load(src_ptr); + dst.cast_from(tmp); + } +} + +template +FLASHINFER_INLINE void cast_store_impl(tgt_float_t* dst_ptr, + const vec_t& src) { + if constexpr (std::is_same_v) { + src.store(dst_ptr); + } else { + vec_t tmp; + tmp.cast_from(src); + tmp.store(dst_ptr); + } +} + +/******************* vec_t<__nv_fp8_e4m3> *******************/ + +// __nv_fp8_e4m3 x 1 +template <> +struct vec_t<__nv_fp8_e4m3, 1> { + __nv_fp8_e4m3 data; + + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { + return ((const __nv_fp8_e4m3*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::fill(__nv_fp8_e4m3 val) { data = val; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::load(const __nv_fp8_e4m3* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::store(__nv_fp8_e4m3* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::memcpy(__nv_fp8_e4m3* dst, + const __nv_fp8_e4m3* src) { + *dst = *src; +} + +// __nv_fp8_e4m3 x 2 +template <> +struct vec_t<__nv_fp8_e4m3, 2> { + __nv_fp8x2_e4m3 data; + + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { + return ((const __nv_fp8_e4m3*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::fill(__nv_fp8_e4m3 val) { + data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::load(const __nv_fp8_e4m3* ptr) { + data = *((__nv_fp8x2_e4m3*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::store(__nv_fp8_e4m3* ptr) const { + *((__nv_fp8x2_e4m3*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::memcpy(__nv_fp8_e4m3* dst, + const __nv_fp8_e4m3* src) { + *((__nv_fp8x2_e4m3*)dst) = *((__nv_fp8x2_e4m3*)src); +} + +// __nv_fp8_e4m3 x 4 + +template <> +struct vec_t<__nv_fp8_e4m3, 4> { + __nv_fp8x4_e4m3 data; + + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { + return ((const __nv_fp8_e4m3*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::fill(__nv_fp8_e4m3 val) { + data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::load(const __nv_fp8_e4m3* ptr) { + data = *((__nv_fp8x4_e4m3*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::store(__nv_fp8_e4m3* ptr) const { + *((__nv_fp8x4_e4m3*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::memcpy(__nv_fp8_e4m3* dst, + const __nv_fp8_e4m3* src) { + *((__nv_fp8x4_e4m3*)dst) = *((__nv_fp8x4_e4m3*)src); +} + +// __nv_fp8_e4m3 x 8 + +template <> +struct vec_t<__nv_fp8_e4m3, 8> { + uint2 data; + + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { + return ((const __nv_fp8_e4m3*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::fill(__nv_fp8_e4m3 val) { + ((__nv_fp8x4_e4m3*)(&data.x))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*)(&data.y))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::load(const __nv_fp8_e4m3* ptr) { + data = *((uint2*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::store(__nv_fp8_e4m3* ptr) const { + *((uint2*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::memcpy(__nv_fp8_e4m3* dst, + const __nv_fp8_e4m3* src) { + *((uint2*)dst) = *((uint2*)src); +} + +// __nv_fp8_e4m3 x 16 or more +template +struct vec_t<__nv_fp8_e4m3, vec_size> { + static_assert(vec_size % 16 == 0, "Invalid vector size"); + int4 data[vec_size / 16]; + + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)data)[i]; } + FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { + return ((const __nv_fp8_e4m3*)data)[i]; + } + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((__nv_fp8x4_e4m3*)(&(data[i].x)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*)(&(data[i].y)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*)(&(data[i].z)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*)(&(data[i].w)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + } + } + FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e4m3* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + *((int4*)(data + i)) = ld_global_acquire((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_release(__nv_fp8_e4m3* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_release(data[i], (int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e4m3* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ld_global_volatile((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e4m3* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_volatile(data[i], (int4*)(addr + i * 16)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +/******************* vec_t<__nv_fp8_e5m2> *******************/ + +// __nv_fp8_e5m2 x 1 +template <> +struct vec_t<__nv_fp8_e5m2, 1> { + __nv_fp8_e5m2 data; + + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { + return ((const __nv_fp8_e5m2*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::fill(__nv_fp8_e5m2 val) { data = val; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::load(const __nv_fp8_e5m2* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::store(__nv_fp8_e5m2* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::memcpy(__nv_fp8_e5m2* dst, + const __nv_fp8_e5m2* src) { + *dst = *src; +} + +// __nv_fp8_e5m2 x 2 +template <> +struct vec_t<__nv_fp8_e5m2, 2> { + __nv_fp8x2_e5m2 data; + + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { + return ((const __nv_fp8_e5m2*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::fill(__nv_fp8_e5m2 val) { + data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::load(const __nv_fp8_e5m2* ptr) { + data = *((__nv_fp8x2_e5m2*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::store(__nv_fp8_e5m2* ptr) const { + *((__nv_fp8x2_e5m2*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::memcpy(__nv_fp8_e5m2* dst, + const __nv_fp8_e5m2* src) { + *((__nv_fp8x2_e5m2*)dst) = *((__nv_fp8x2_e5m2*)src); +} + +// __nv_fp8_e5m2 x 4 + +template <> +struct vec_t<__nv_fp8_e5m2, 4> { + __nv_fp8x4_e5m2 data; + + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { + return ((const __nv_fp8_e5m2*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::fill(__nv_fp8_e5m2 val) { + data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::load(const __nv_fp8_e5m2* ptr) { + data = *((__nv_fp8x4_e5m2*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::store(__nv_fp8_e5m2* ptr) const { + *((__nv_fp8x4_e5m2*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::memcpy(__nv_fp8_e5m2* dst, + const __nv_fp8_e5m2* src) { + *((__nv_fp8x4_e5m2*)dst) = *((__nv_fp8x4_e5m2*)src); +} + +// __nv_fp8_e5m2 x 8 + +template <> +struct vec_t<__nv_fp8_e5m2, 8> { + uint2 data; + + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { + return ((const __nv_fp8_e5m2*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::fill(__nv_fp8_e5m2 val) { + ((__nv_fp8x4_e5m2*)(&data.x))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*)(&data.y))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::load(const __nv_fp8_e5m2* ptr) { + data = *((uint2*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::store(__nv_fp8_e5m2* ptr) const { + *((uint2*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::memcpy(__nv_fp8_e5m2* dst, + const __nv_fp8_e5m2* src) { + *((uint2*)dst) = *((uint2*)src); +} + +// __nv_fp8_e5m2 x 16 or more + +template +struct vec_t<__nv_fp8_e5m2, vec_size> { + static_assert(vec_size % 16 == 0, "Invalid vector size"); + int4 data[vec_size / 16]; + + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)data)[i]; } + FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { + return ((const __nv_fp8_e5m2*)data)[i]; + } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((__nv_fp8x4_e5m2*)(&(data[i].x)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*)(&(data[i].y)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*)(&(data[i].z)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*)(&(data[i].w)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + } + } + FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void store_global_release(__nv_fp8_e5m2* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_release(data[i], (int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e5m2* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ld_global_acquire((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e5m2* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_volatile(data[i], (int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e5m2* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ld_global_volatile((int4*)(addr + i * 16)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +#if defined(FLASHINFER_ENABLE_FP4_E2M1) && CUDA_VERSION >= 12080 +/******************* vec_t<__nv_fp4_e2m1> *******************/ + +// __nv_fp4_e2m1 x 2 +template <> +struct vec_t<__nv_fp4_e2m1, 2> { + uint8_t data; + // index access is not supported for sub-byte data type + FLASHINFER_INLINE __nv_fp4_e2m1* ptr() { return reinterpret_cast<__nv_fp4_e2m1*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) { + data = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); + } + FLASHINFER_INLINE void load(const __nv_fp4_e2m1* ptr) { data = *((uint8_t*)ptr); } + FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const { *((uint8_t*)ptr) = data; } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, const __nv_fp4_e2m1* src) { + *((uint8_t*)dst) = *((uint8_t*)src); + } +}; + +// __nv_fp4_e2m1 x 4 +template <> +struct vec_t<__nv_fp4_e2m1, 4> { + uint16_t data; + + FLASHINFER_INLINE __nv_fp4_e2m1* ptr() { return reinterpret_cast<__nv_fp4_e2m1*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) { + __nv_fp4x2_storage_t val8 = + (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); + data = (uint16_t(val8) << 8) | uint16_t(val8); + } + FLASHINFER_INLINE void load(const __nv_fp4_e2m1* ptr) { data = *((uint16_t*)ptr); } + FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const { *((uint16_t*)ptr) = data; } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, const __nv_fp4_e2m1* src) { + *((uint16_t*)dst) = *((uint16_t*)src); + } +}; + +// __nv_fp4_e2m1 x 8 +template <> +struct vec_t<__nv_fp4_e2m1, 8> { + uint32_t data; + + FLASHINFER_INLINE __nv_fp4_e2m1* ptr() { return reinterpret_cast<__nv_fp4_e2m1*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) { + __nv_fp4x2_storage_t val8 = + (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); + uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); + data = (uint32_t(val16) << 16) | uint32_t(val16); + } + FLASHINFER_INLINE void load(const __nv_fp4_e2m1* ptr) { data = *((uint32_t*)ptr); } + FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const { *((uint32_t*)ptr) = data; } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, const __nv_fp4_e2m1* src) { + *((uint32_t*)dst) = *((uint32_t*)src); + } +}; + +template <> +struct vec_t<__nv_fp4_e2m1, 16> { + uint2 data; + FLASHINFER_INLINE __nv_fp4_e2m1* ptr() { return reinterpret_cast<__nv_fp4_e2m1*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) { + __nv_fp4x2_storage_t val8 = + (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); + uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); + uint32_t val32 = (uint32_t(val16) << 16) | uint32_t(val16); + data.x = val32; + data.y = val32; + } + FLASHINFER_INLINE void load(const __nv_fp4_e2m1* ptr) { data = *((uint2*)ptr); } + FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const { *((uint2*)ptr) = data; } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, const __nv_fp4_e2m1* src) { + *((uint2*)dst) = *((uint2*)src); + } +}; + +// __nv_fp4_e2m1 x 32 or more +template +struct vec_t<__nv_fp4_e2m1, vec_size> { + static_assert(vec_size % 32 == 0, "Invalid vector size"); + int4 data[vec_size / 32]; + + FLASHINFER_INLINE __nv_fp4_e2m1* ptr() { return reinterpret_cast<__nv_fp4_e2m1*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) { + __nv_fp4x2_storage_t val8 = + (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); + uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); + uint32_t val32 = (uint32_t(val16) << 16) | uint32_t(val16); +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) { + data[i].x = val32; + data[i].y = val32; + data[i].z = val32; + data[i].w = val32; + } + } + FLASHINFER_INLINE void load(const __nv_fp4_e2m1* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void store_global_release(__nv_fp4_e2m1* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) { + st_global_release(*(int4*)&data[i], (int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void load_global_acquire(__nv_fp4_e2m1* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) { + *(int4*)&data[i] = ld_global_acquire((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_volatile(__nv_fp4_e2m1* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) { + st_global_volatile(*(int4*)&data[i], (int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void load_global_volatile(__nv_fp4_e2m1* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) { + *(int4*)&data[i] = ld_global_volatile((int4*)(addr + i * 16)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, const __nv_fp4_e2m1* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +#endif // FLASHINFER_ENABLE_FP4_E2M1 && CUDA_VERSION >= 12080 + +/******************* vec_t *******************/ + +// half x 1 +template <> +struct vec_t { + half data; + + FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; } + FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; } + FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(half val); + FLASHINFER_INLINE void load(const half* ptr); + FLASHINFER_INLINE void store(half* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(half* dst, const half* src); +}; + +FLASHINFER_INLINE void vec_t::fill(half val) { data = val; } + +FLASHINFER_INLINE void vec_t::load(const half* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t::store(half* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t::memcpy(half* dst, const half* src) { *dst = *src; } + +// half x 2 +template <> +struct vec_t { + half2 data; + + FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; } + FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; } + FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(half val); + FLASHINFER_INLINE void load(const half* ptr); + FLASHINFER_INLINE void store(half* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(half* dst, const half* src); +}; + +FLASHINFER_INLINE void vec_t::fill(half val) { data = make_half2(val, val); } + +FLASHINFER_INLINE void vec_t::load(const half* ptr) { data = *((half2*)ptr); } + +FLASHINFER_INLINE void vec_t::store(half* ptr) const { *((half2*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(half* dst, const half* src) { + *((half2*)dst) = *((half2*)src); +} + +// half x 4 + +template <> +struct vec_t { + uint2 data; + + FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; } + FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; } + FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(half val); + FLASHINFER_INLINE void load(const half* ptr); + FLASHINFER_INLINE void store(half* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(half* dst, const half* src); +}; + +FLASHINFER_INLINE void vec_t::fill(half val) { + *(half2*)(&data.x) = make_half2(val, val); + *(half2*)(&data.y) = make_half2(val, val); +} + +FLASHINFER_INLINE void vec_t::load(const half* ptr) { data = *((uint2*)ptr); } + +FLASHINFER_INLINE void vec_t::store(half* ptr) const { *((uint2*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(half* dst, const half* src) { + *((uint2*)dst) = *((uint2*)src); +} + +// half x 8 or more + +template +struct vec_t { + static_assert(vec_size % 8 == 0, "Invalid vector size"); + int4 data[vec_size / 8]; + FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)data)[i]; } + FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)data)[i]; } + FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(half val) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + *(half2*)(&(data[i].x)) = make_half2(val, val); + *(half2*)(&(data[i].y)) = make_half2(val, val); + *(half2*)(&(data[i].z)) = make_half2(val, val); + *(half2*)(&(data[i].w)) = make_half2(val, val); + } + } + FLASHINFER_INLINE void load(const half* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(half* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void load_global_acquire(half* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ld_global_acquire((int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void store_global_release(half* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + st_global_release(data[i], (int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void store_global_volatile(half* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + st_global_volatile(data[i], (int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void load_global_volatile(half* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ld_global_volatile((int4*)(addr + i * 8)); + } + } + + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(half* dst, const half* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +/******************* vec_t *******************/ + +// nv_bfloat16 x 1 +template <> +struct vec_t { + nv_bfloat16 data; + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; } + FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { + return ((const nv_bfloat16*)(&data))[i]; + } + FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(nv_bfloat16 val); + FLASHINFER_INLINE void load(const nv_bfloat16* ptr); + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src); +}; + +FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) { data = val; } + +FLASHINFER_INLINE void vec_t::load(const nv_bfloat16* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { + *dst = *src; +} + +// nv_bfloat16 x 2 +template <> +struct vec_t { + nv_bfloat162 data; + + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; } + FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { + return ((const nv_bfloat16*)(&data))[i]; + } + FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(nv_bfloat16 val); + FLASHINFER_INLINE void load(const nv_bfloat16* ptr); + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src); +}; + +FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) { + data = make_bfloat162(val, val); +} + +FLASHINFER_INLINE void vec_t::load(const nv_bfloat16* ptr) { + data = *((nv_bfloat162*)ptr); +} + +FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const { + *((nv_bfloat162*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { + *((nv_bfloat162*)dst) = *((nv_bfloat162*)src); +} + +// nv_bfloat16 x 4 + +template <> +struct vec_t { + uint2 data; + + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; } + FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { + return ((const nv_bfloat16*)(&data))[i]; + } + FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(nv_bfloat16 val); + FLASHINFER_INLINE void load(const nv_bfloat16* ptr); + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src); +}; + +FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) { + *(nv_bfloat162*)(&data.x) = make_bfloat162(val, val); + *(nv_bfloat162*)(&data.y) = make_bfloat162(val, val); +} + +FLASHINFER_INLINE void vec_t::load(const nv_bfloat16* ptr) { + data = *((uint2*)ptr); +} + +FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const { + *((uint2*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { + *((uint2*)dst) = *((uint2*)src); +} + +// nv_bfloat16 x 8 or more + +template +struct vec_t { + static_assert(vec_size % 8 == 0, "Invalid vector size"); + int4 data[vec_size / 8]; + + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)data)[i]; } + FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { + return ((const nv_bfloat16*)data)[i]; + } + FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(nv_bfloat16 val) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + *(nv_bfloat162*)(&(data[i].x)) = make_bfloat162(val, val); + *(nv_bfloat162*)(&(data[i].y)) = make_bfloat162(val, val); + *(nv_bfloat162*)(&(data[i].z)) = make_bfloat162(val, val); + *(nv_bfloat162*)(&(data[i].w)) = make_bfloat162(val, val); + } + } + FLASHINFER_INLINE void load(const nv_bfloat16* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void store_global_release(nv_bfloat16* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + st_global_release(data[i], (int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void load_global_acquire(nv_bfloat16* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ld_global_acquire((int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void store_global_volatile(nv_bfloat16* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + st_global_volatile(data[i], (int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void load_global_volatile(nv_bfloat16* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ld_global_volatile((int4*)(addr + i * 8)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +/******************* vec_t *******************/ + +// uint8_t x 1 +template <> +struct vec_t { + uint8_t data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } + FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { + return ((const uint8_t*)(&data))[i]; + } + FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(const uint8_t* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) { data = val; } + +FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { *dst = *src; } + +// uint8_t x 2 +template <> +struct vec_t { + uint16_t data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } + FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { + return ((const uint8_t*)(&data))[i]; + } + FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(const uint8_t* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) { + data = (uint16_t(val) << 8) | uint16_t(val); +} + +FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *((uint16_t*)ptr); } + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *((uint16_t*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { + *((uint16_t*)dst) = *((uint16_t*)src); +} + +// uint8_t x 4 + +template <> +struct vec_t { + uint32_t data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } + FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { + return ((const uint8_t*)(&data))[i]; + } + FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(const uint8_t* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) { + data = (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); +} + +FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *((uint32_t*)ptr); } + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *((uint32_t*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { + *((uint32_t*)dst) = *((uint32_t*)src); +} + +// uint8_t x 8 + +template <> +struct vec_t { + uint2 data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } + FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { + return ((const uint8_t*)(&data))[i]; + } + FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(const uint8_t* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) { + uint32_t val32 = + (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); + data.x = val32; + data.y = val32; +} + +FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *((uint2*)ptr); } + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *((uint2*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { + *((uint2*)dst) = *((uint2*)src); +} + +// uint8_t x 16 or more + +template +struct vec_t { + static_assert(vec_size % 16 == 0, "Invalid vector size"); + int4 data[vec_size / 16]; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)data)[i]; } + FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { return ((const uint8_t*)data)[i]; } + FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(uint8_t val) { + uint32_t val32 = + (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i].x = val32; + data[i].y = val32; + data[i].z = val32; + data[i].w = val32; + } + } + FLASHINFER_INLINE void load(const uint8_t* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(uint8_t* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void load_global_acquire(uint8_t* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ld_global_acquire((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_release(uint8_t* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_release(data[i], (int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void load_global_volatile(uint8_t* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ld_global_volatile((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_volatile(uint8_t* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_volatile(data[i], (int4*)(addr + i * 16)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +/******************* vec_t *******************/ + +// float x 1 + +template <> +struct vec_t { + float data; + + FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(&data))[i]; } + FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(&data))[i]; } + FLASHINFER_INLINE float* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(float val); + FLASHINFER_INLINE void load(const float* ptr); + FLASHINFER_INLINE void store(float* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(float* dst, const float* src); +}; + +FLASHINFER_INLINE void vec_t::fill(float val) { data = val; } + +FLASHINFER_INLINE void vec_t::load(const float* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t::store(float* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t::memcpy(float* dst, const float* src) { *dst = *src; } + +// float x 2 + +template <> +struct vec_t { + float2 data; + + FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(&data))[i]; } + FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(&data))[i]; } + FLASHINFER_INLINE float* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(float val); + FLASHINFER_INLINE void load(const float* ptr); + FLASHINFER_INLINE void store(float* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(float* dst, const float* src); +}; + +FLASHINFER_INLINE void vec_t::fill(float val) { data = make_float2(val, val); } + +FLASHINFER_INLINE void vec_t::load(const float* ptr) { data = *((float2*)ptr); } + +FLASHINFER_INLINE void vec_t::store(float* ptr) const { *((float2*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(float* dst, const float* src) { + *((float2*)dst) = *((float2*)src); +} + +// float x 4 or more +template +struct vec_t { + static_assert(vec_size % 4 == 0, "Invalid vector size"); + float4 data[vec_size / 4]; + + FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(data))[i]; } + FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(data))[i]; } + FLASHINFER_INLINE float* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(float val) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + data[i] = make_float4(val, val, val, val); + } + } + FLASHINFER_INLINE void load(const float* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + data[i] = ((float4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(float* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + ((float4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void store_global_release(float* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + st_global_release(*(int4*)(data + i), (int4*)(addr + i * 4)); + } + } + FLASHINFER_INLINE void load_global_acquire(float* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + *((int4*)(data + i)) = ld_global_acquire((int4*)(addr + i * 4)); + } + } + FLASHINFER_INLINE void store_global_volatile(float* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + st_global_volatile(*(int4*)(data + i), (int4*)(addr + i * 4)); + } + } + FLASHINFER_INLINE void load_global_volatile(float* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + *((int4*)(data + i)) = ld_global_volatile((int4*)(addr + i * 4)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(float* dst, const float* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + ((float4*)dst)[i] = ((float4*)src)[i]; + } + } +}; + +template +struct vec2_dtype { + using type = T; +}; + +template <> +struct vec2_dtype { + using type = half2; +}; + +template <> +struct vec2_dtype<__nv_bfloat16> { + using type = __nv_bfloat162; +}; + +template <> +struct vec2_dtype<__nv_fp8_e4m3> { + using type = __nv_fp8x2_e4m3; +}; + +template <> +struct vec2_dtype<__nv_fp8_e5m2> { + using type = __nv_fp8x2_e5m2; +}; + +template +using vec2_dtype_t = typename vec2_dtype::type; + +template +FLASHINFER_INLINE vec2_dtype_t get_vec2_element(vec_t& vec, int i) { + static_assert(VEC_SIZE % 2 == 0, "VEC_SIZE must be a multiple of 2"); + return ((vec2_dtype_t*)&(vec[0]))[i]; +} + +} // namespace flashinfer + +#endif // VEC_DTYPES_CUH_ From 92de174cf621cf469e0eef18ec3b50ff1d42fe9a Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 12:21:07 -0700 Subject: [PATCH 67/89] mamba benchmark: model FlashInfer cache-slot indirection Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../benchmark_replay_selective_state_update.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 9dbaa8eaf73c..e799214ea13a 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -550,6 +550,7 @@ def _build_tensors( x, dt, B, C : (batch, mtp_len, ...) – token inputs for both kernels A, dt_bias, D : SSM parameters (float32, tie_hdim strides) prev_tokens : (batch,) + state_batch_indices : identity cache-slot indirection (batch,) replay_work_items : packed per-slot replay metadata (batch, 4) out_incr : pre-allocated output for replay kernel (batch, mtp_len, nheads, head_dim) out_base : pre-allocated output for baseline kernel (batch, mtp_len, nheads, head_dim) @@ -579,6 +580,7 @@ def _build_tensors( cached["dt_bias"], cached["D"], cached["prev_tokens"][:b], + cached["state_batch_indices"][:b], cached["replay_work_items"][:b], cached["out_incr"][:b], cached["out_base"][:b], @@ -666,14 +668,14 @@ def _build_tensors( # prev_tokens placeholder — overwritten per-run prev_tokens = torch.zeros(batch, device=device, dtype=torch.int32) + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) replay_work_items = torch.empty( batch, REPLAY_WORK_ITEM_WIDTH, device=device, dtype=torch.int32 ) - position_in_decode_batch = torch.arange(batch, device=device, dtype=torch.int32) replay_work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = ( - position_in_decode_batch + state_batch_indices ) - replay_work_items[:, REPLAY_WORK_CACHE_SLOT] = position_in_decode_batch + replay_work_items[:, REPLAY_WORK_CACHE_SLOT] = state_batch_indices replay_work_items[:, REPLAY_WORK_PNAT] = 0 replay_work_items[:, REPLAY_WORK_CACHE_BUF_IDX] = 0 @@ -719,6 +721,7 @@ def _build_tensors( "dt_bias": dt_bias, "D": D, "prev_tokens": prev_tokens, + "state_batch_indices": state_batch_indices, "replay_work_items": replay_work_items, "out_incr": out_incr, "out_base": out_base, @@ -746,6 +749,7 @@ def _build_tensors( dt_bias, D, prev_tokens[:rb], + state_batch_indices[:rb], replay_work_items[:rb], out_incr[:rb], out_base[:rb], @@ -2583,6 +2587,7 @@ def _bench_config( dt_bias, D, prev_tokens, + state_batch_indices, replay_work_items_buf, out_incr, out_base, @@ -2998,7 +3003,9 @@ def _run_pr3324_baseline(): D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, + state_batch_indices=( + state_batch_indices if args.use_cache_slot else None + ), state_scale=state_scales_work, rand_seed=rand_seed, philox_rounds=args.philox_rounds, From ef75adfa697c3852d80a303178c4970ea2f85131 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 12:25:26 -0700 Subject: [PATCH 68/89] mamba benchmark: improve default timing UX Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 45 +++++++++++++------ 1 file changed, 32 insertions(+), 13 deletions(-) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index e799214ea13a..d07198cad7e4 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -788,9 +788,27 @@ def _build_tensors( "_persistent_main", "selective_scan_update", "selective_state_update", + "checkpointing_ssu", "causal_conv1d_update", ) +_SR_SUPPORTED_DTYPES = ( + torch.float16, + torch.int8, + torch.int16, + torch.float8_e4m3fn, +) + + +def _sr_modes_for_dtype(state_dtype: torch.dtype, + requested_modes: list[str]) -> list[str]: + """Return the rounding modes that should run for one state dtype.""" + if state_dtype == torch.float32: + return ["RN"] + if state_dtype not in _SR_SUPPORTED_DTYPES: + return [mode for mode in requested_modes if mode == "RN"] + return requested_modes + def _kernels_per_iter_incremental( mode: str, @@ -2300,7 +2318,7 @@ def _cw(label: str) -> None: prev_ks = _resolve_prev_ks(args, mtp_len) for state_dtype in state_dtypes: for act_dtype in act_dtypes: - for sr_mode in sr_modes_list: + for sr_mode in _sr_modes_for_dtype(state_dtype, sr_modes_list): for mode in modes_list: for rect in rect_list: can_sort = ( @@ -2618,15 +2636,11 @@ def _bench_config( # SR rounding: allow fp16 and the quantized dtypes (int8/int16/fp8). # bf16/fp32 SR is not supported (no PTX path for bf16; fp32 doesn't need - # rounding). When sweeping --sr-modes RN,SR over a mixed dtype set, - # silently skip the SR cell for unsupported dtypes — the RN cell still - # prints, and other dtypes still get their SR row. + # rounding). _sr_modes_for_dtype maps fp32 to RN so production sweeps can + # request SR once while still getting fp32/RN. rand_seed = None - _SR_SUPPORTED = ( - torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn, - ) if use_philox: - if state_dtype not in _SR_SUPPORTED: + if state_dtype not in _SR_SUPPORTED_DTYPES: return rand_seed = torch.randint(0, 2**62, (1,), device="cuda", dtype=torch.int64) mode = _resolve_effective_replay_mode( @@ -4107,6 +4121,7 @@ def _phase(label: str) -> None: # Print header mix_enabled = args.mix_csv is not None or getattr(args, "pmix", False) headline_name = "score_us" if mix_enabled else "median_us" + print(f"conv1d: {'enabled' if args.with_conv1d else 'disabled'}") if mix_enabled: print( "kmix bucket score: mix rows report bucket-weighted score_us; " @@ -4195,7 +4210,7 @@ def _phase(label: str) -> None: for state_dtype in state_dtypes: for act_dtype in act_dtypes: - for sr_mode in sr_modes_list: + for sr_mode in _sr_modes_for_dtype(state_dtype, sr_modes_list): for mode in modes_list: for rect in rect_list: for nowrite_first in nowrite_first_list: @@ -4584,10 +4599,14 @@ def _parse_args() -> argparse.Namespace: ) parser.add_argument( "--with-conv1d", - action="store_true", - help="Include conv1d kernel before replay SSM. " - "Uses realistic L2 flush: cold caches flushed, hot in_proj output " - "kept warm. Measures conv1d → precompute → main span.", + "--with-conv-1d", + action=argparse.BooleanOptionalAction, + default=True, + dest="with_conv1d", + help="Include conv1d kernel before replay SSM. Default on; pass " + "--no-with-conv1d to time state update only. Uses realistic L2 flush: " + "cold caches flushed, hot in_proj output kept warm. Measures " + "conv1d → precompute → main span.", ) parser.add_argument( "--external-pdl", From d17b7fde76dc2c81b5a97a2f255ede87bdabe565 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 14:25:39 -0700 Subject: [PATCH 69/89] remove scratch kernel copies and do precommit fixes Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...aseline_checkpointing_state_update_slim.py | 3182 ----------- ...aseline_checkpointing_state_update_slim.py | 3173 ----------- .../_torch/modules/mamba/mamba2_metadata.py | 26 +- .../mamba/replay_selective_state_update.py | 1740 ++++-- .../_torch/pyexecutor/mamba_cache_manager.py | 15 +- ...benchmark_replay_selective_state_update.py | 996 ++-- .../flashinfer_checkpointing_ssu_pr3324.py | 9 +- .../csrc/checkpointing_ssu.cu | 1016 ++-- .../csrc/checkpointing_ssu_jit_binding.cu | 52 +- .../csrc/checkpointing_ssu_kernel_inst.cu | 10 +- .../include/flashinfer/exception.h | 146 +- .../flashinfer/mamba/checkpointing_ssu.cuh | 247 +- .../include/flashinfer/mamba/common.cuh | 254 +- .../include/flashinfer/mamba/conversion.cuh | 555 +- .../mamba/kernel_checkpointing_ssu.cuh | 1837 ++++--- .../mamba/kernel_checkpointing_ssu_8bit.cuh | 2330 ++++----- .../mamba/kernel_checkpointing_ssu_common.cuh | 2310 ++++---- .../mamba/launch_checkpointing_ssu.cuh | 283 +- .../flashinfer/mamba/ssu_mtp_common.cuh | 152 +- .../include/flashinfer/utils.cuh | 919 ++-- .../include/flashinfer/vec_dtypes.cuh | 4653 ++++++++++------- .../modules/mamba/test_mamba2_metadata.py | 17 +- .../test_replay_selective_state_update.py | 547 +- 23 files changed, 10450 insertions(+), 14019 deletions(-) delete mode 100644 tensorrt_llm/_torch/modules/mamba/_v0_baseline_checkpointing_state_update_slim.py delete mode 100644 tensorrt_llm/_torch/modules/mamba/_v3_baseline_checkpointing_state_update_slim.py diff --git a/tensorrt_llm/_torch/modules/mamba/_v0_baseline_checkpointing_state_update_slim.py b/tensorrt_llm/_torch/modules/mamba/_v0_baseline_checkpointing_state_update_slim.py deleted file mode 100644 index 65c73cdd5aa7..000000000000 --- a/tensorrt_llm/_torch/modules/mamba/_v0_baseline_checkpointing_state_update_slim.py +++ /dev/null @@ -1,3182 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. -# -# Adapted from: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py -# SPDX-FileCopyrightText: Copyright contributors to the sglang project -# -# Copyright (c) 2024, Tri Dao, Albert Gu. -# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py - -import torch -import triton -import triton.language as tl - -from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID -from tensorrt_llm._utils import get_sm_version - -from .softplus import softplus - - -# Lazy global allocator for Triton TMA tensor descriptors. Required by any -# host- or device-built tensor_descriptor; without it Triton raises at first -# launch. See TMA backlog item #17 / scratch experiment notes. -_TMA_ALLOCATOR_SET = False - - -def _ensure_tma_allocator() -> None: - global _TMA_ALLOCATOR_SET - if _TMA_ALLOCATOR_SET: - return - - def _alloc_fn(size, alignment, stream): - # Triton expects an int8 buffer of `size` bytes; alignment is enforced - # by the allocator returning a buffer satisfying it (PyTorch's - # cudaMalloc-backed tensors are 256B-aligned, so we're fine). - return torch.empty(size, device="cuda", dtype=torch.int8) - - triton.set_allocator(_alloc_fn) - _TMA_ALLOCATOR_SET = True - - -@triton.jit -def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: - """Stochastic rounding: fp32 pair → fp16x2 using Philox random bits. - - Uses PTX cvt.rs.f16x2.f32 which rounds each fp32 value to fp16 using - the random bits to break ties, avoiding systematic rounding bias that - accumulates over many decode steps with fp16 state. - - Adapted from flashinfer (Apache-2.0, vLLM/mamba lineage). - """ - return tl.inline_asm_elementwise( - asm="""{ - cvt.rs.f16x2.f32 $0, $2, $1, $3; - }""", - constraints=("=r,r,r,r,r"), - args=(x, rand), - dtype=tl.float16, - is_pure=True, - pack=2, - ) - - -@triton.jit -def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: - """Stochastic rounding: fp32 quad → fp8 e4m3 using Philox random bits. - - Uses PTX cvt.rs.satfinite.e4m3x4.f32 which combines stochastic rounding - and saturating cast in a single op (output is final fp8, no separate - clamp needed). The reversed source-register order {$4,$3,$2,$1} is - load-bearing — PTX packs leftmost source into the high byte but Triton's - pack=4 is little-endian, so the natural {$1,$2,$3,$4} order would - silently shuffle every group of 4 contiguous outputs. - - Requires SM_100a+ (Blackwell B200). Caller must gate at the wrapper - level — this kernel does not check. - - Adapted from vLLM PR #40012 (Apache-2.0). - """ - return tl.inline_asm_elementwise( - asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", - constraints="=r,r,r,r,r,r,r,r,r", - args=(x, rand), - dtype=tl.float8e4nv, - is_pure=True, - pack=4, - ) - - -@triton.jit -def _bitrev32(x: tl.tensor) -> tl.tensor: - return tl.inline_asm_elementwise( - asm="brev.b32 $0, $1;", - constraints="=r,r", - args=(x,), - dtype=tl.uint32, - is_pure=True, - pack=1, - ) - - -@triton.jit -def _stochastic_round_int8_packed( - x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor -) -> tl.tensor: - """Stochastic rounding for int8 using one random uint32 per 4 values.""" - low = rand & 0x0000FFFF - high = (rand >> 16) & 0x0000FFFF - low_rev = _bitrev32(low) >> 16 - high_rev = _bitrev32(high) >> 16 - rand_pos = offs_n & 3 - rand16 = tl.where( - rand_pos == 0, - low, - tl.where(rand_pos == 1, low_rev, tl.where(rand_pos == 2, high, high_rev)), - ) - rand01 = rand16.to(tl.float32) * (1.0 / float(1 << 16)) - return tl.extra.cuda.libdevice.floor(x + rand01) - - -@triton.jit -def _stochastic_round_int16_packed( - x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor -) -> tl.tensor: - """Stochastic rounding for int16 using one random uint32 per 2 values.""" - rand_bits = tl.where((offs_n & 1) == 0, rand, _bitrev32(rand)) - rand01 = (rand_bits & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) - return tl.extra.cuda.libdevice.floor(x + rand01) - - -# Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, -# old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. -# Grid: (batch, nheads // HEADS_PER_BLOCK). - - -@triton.jit() -def _replay_precompute_impl( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, - decay_vec_ptr, - # Cache pointers (both buffers reachable via stride_*_dbuf). This - # kernel writes to either the active (= cache_buf_idx) or inactive - # (= 1 - cache_buf_idx) buffer depending on WRITE_CHECKPOINT — see - # comment block at top of kernel body. - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - # Double-buffer index (per cache slot) — selects this step's "active" - # buffer (= where the historical inputs for this step live). - cache_buf_idx_ptr, - # Per-request accepted-tokens count (already-cached old tokens at - # [0, PNAT) of the active buffer; new tokens this step go after them - # on no-checkpoint steps). - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides: (cache, 2, T, ngroups, dstate) - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, - # Checkpointing flag — selects target buffer + offset for new-token - # cache writes. See "Cache write semantics" block below. - # Runtime (not constexpr): the only WRITE_CHECKPOINT-dependent code in - # this body is the write_buf/write_offset selection, which is plain - # arithmetic — no constexpr-shaped tile or whole-block gate. Letting - # it be runtime lets the dynamic dispatch kernel call us once with the - # per-slot needs_write flag instead of inlining two specializations. - write_checkpoint, -): - pid_b = tl.program_id(axis=0) - pid_hg = tl.program_id(axis=1) # head-group index - first_head = pid_hg * HEADS_PER_BLOCK - - # Resolve cache index for writes - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - - # --- Cache write semantics --- - # cache_buf_idx names this step's "active" buffer — the one with the - # historical inputs at [0, PNAT). The other buffer is "staging". - # - # Where do we write new tokens this step? - # WRITE_CHECKPOINT=False (no overflow): append to ACTIVE buffer at - # offset [PNAT : PNAT+T). Caller does NOT flip cache_buf_idx - # afterward; PNAT_next = PNAT + accepted. [0, PNAT) preserved. - # WRITE_CHECKPOINT=True (would overflow): write to STAGING buffer at - # [0, T). Caller flips cache_buf_idx afterward; next step's - # active = the one we just wrote. PNAT_next = accepted. Old - # data in the previous active buffer is folded into state via - # the replay update and discarded. This matches today's replay - # kernel behavior exactly. - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - if write_checkpoint: - write_buf = 1 - buf_active - write_offset = 0 - else: - write_buf = buf_active - write_offset = prev_num_accepted_tokens - - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - t_mask = offs_t < T - n_mask = offs_n < dstate - - # Causal mask is shared across all heads (depends only on offs_t) - causal_mask = offs_t[:, None] >= offs_t[None, :] - valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] - - # --- Vectorized pre-wait phase across HEADS_PER_BLOCK heads --- - # Compute dt, dA_cumsum, decay_vec as (H, T) tiles. Pre-compute - # scale_combo = decay_matrix * dt[:, None, :] as an (H, T, T) tile that - # stays in registers across gdc_wait — eliminates the post-wait reload - # of dt + dA_cumsum and the per-head loop. - offs_h = tl.arange(0, HEADS_PER_BLOCK) - heads_block = first_head + offs_h # (H,) - - # Load dt (H, T) - dt_addrs = ( - dt_ptr + pid_b * stride_dt_batch - + heads_block[:, None] * stride_dt_head - + offs_t[None, :] * stride_dt_T - ) - dt = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) - if HAS_DT_BIAS: - dt_bias = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) - dt = dt + dt_bias[:, None] - if DT_SOFTPLUS: - dt = softplus(dt) - - A = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) # (H,) - dA_cumsum = tl.cumsum(A[:, None] * dt, axis=1) # (H, T) - decay_vec = tl.exp(dA_cumsum) # (H, T) - - # Cross-step continuity for old_dA_cumsum: when appending to active_buf at - # offset PNAT > 0, the previous step left a running cumsum at [0, PNAT) - # whose tail value lives at active_buf[head, PNAT-1]. Add that tail to - # this step's per-step-restarted cumsum before storing so the buffer - # holds one continuous cumsum across N back-to-back nowrites. Write path - # (write_buf = 1 - buf_active, write_offset = 0) starts fresh, no prefix. - # Both branches are on scalar runtime values (write_checkpoint and PNAT), - # uniform across the block — use scalar if to short-circuit the load. - if write_checkpoint or prev_num_accepted_tokens == 0: - prev_total = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) - else: - last_cumsum_ptrs = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_active * stride_old_dA_cumsum_dbuf - + heads_block * stride_old_dA_cumsum_head - + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T - ) - prev_total = tl.load(last_cumsum_ptrs).to(tl.float32) - - # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of write_buf. - old_dt_addrs = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + heads_block[:, None] * stride_old_dt_head - + (write_offset + offs_t)[None, :] * stride_old_dt_T - ) - tl.store(old_dt_addrs, dt, mask=t_mask[None, :]) - - old_dA_cumsum_addrs = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + heads_block[:, None] * stride_old_dA_cumsum_head - + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T - ) - tl.store(old_dA_cumsum_addrs, dA_cumsum + prev_total[:, None], mask=t_mask[None, :]) - - # decay_vec scratch — always at offs_t. - decay_vec_addrs = ( - decay_vec_ptr + pid_b * stride_dv_batch - + heads_block[:, None] * stride_dv_head - + offs_t[None, :] * stride_dv_t - ) - tl.store(decay_vec_addrs, decay_vec, mask=t_mask[None, :]) - - # scale_combo (H, T, T) = exp(dA_cumsum[h, t1] - dA_cumsum[h, t2]) * dt[h, t2] - # Stays live across gdc_wait — used post-wait to compute CB_scaled. - decay_matrix = tl.exp(dA_cumsum[:, :, None] - dA_cumsum[:, None, :]) # (H, T, T) - scale_combo = decay_matrix * dt[:, None, :] # (H, T, T) - - # --- Wait for upstream kernel (external PDL) before loading B and C --- - # All dt processing above is independent of conv1d outputs. - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- - group_idx = first_head // nheads_ngroups_ratio - C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group - B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group - - C_all = tl.load( - C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - B_all = tl.load( - B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - - # Compute raw CB once — shared across all heads in this block - raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) - - # Store B to cache at [write_offset : write_offset+T) of write_buf. - if first_head % nheads_ngroups_ratio == 0: - old_B_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + write_buf * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - tl.store( - old_B_base - + (write_offset + offs_t)[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - B_all, - mask=t_mask[:, None] & n_mask[None, :], - ) - - # --- Vectorized post-wait phase: scale_combo (H, T, T) is still live in - # registers from pre-wait; multiply by raw_CB (T, T), apply causal mask, - # store as one (H, T, T) tile. --- - CB_scaled_block = tl.where( - valid_mask[None, :, :], - raw_CB[None, :, :] * scale_combo, - 0.0, - ) # (H, T, T) - cb_scaled_addrs = ( - cb_scaled_ptr + pid_b * stride_cb_batch - + heads_block[:, None, None] * stride_cb_head - + offs_t[None, :, None] * stride_cb_t - + offs_t[None, None, :] * stride_cb_j - ) # (H, T, T) - cb_store_mask = ( - (offs_t[None, :, None] < BLOCK_SIZE_T) - & (offs_t[None, None, :] < BLOCK_SIZE_T) - ) - tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) - - -# Replay-style precompute kernel. Thin wrapper around _replay_precompute_impl -# that carries the @triton.heuristics for constexpr derivation; called from -# the Python wrapper on the replay-style path (write or replay-nowrite). -@triton.jit() -def _rectangle_precompute_impl( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle - decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) - # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite - # path: read from buf_active at [0, PNAT), write new tokens at - # [PNAT, PNAT+T) of buf_active (same buffer). - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # rectangle K-axis bound - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides (rectangle: (batch, nheads, T, K)) - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides: (cache, 2, T_max, ngroups, dstate) - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides: (cache, 2, nheads, T_max) - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides: (cache, 2, nheads, T_max) - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, -): - pid_b = tl.program_id(axis=0) - pid_hg = tl.program_id(axis=1) - first_head = pid_hg * HEADS_PER_BLOCK - - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - write_buf = buf_active - write_offset = prev_num_accepted_tokens - - # Static rectangle K-axis layout: old at [0, PNAT) (mask is_old_k); - # new at [MAX-T, MAX) at compile-time shift K_NEW_SHIFT = MAX - T. - # PNAT + T <= MAX is guaranteed on the nowrite path → no overlap. - K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T - - offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis (output rows) - offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis (rectangle input cols) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - t_mask = offs_t < T - n_mask = offs_n < dstate - - # K-axis masks (approach C: runtime PNAT-offset instead of K_NEW_SHIFT) - # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. - is_old_k = offs_k < prev_num_accepted_tokens - safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - prev_num_accepted_tokens - is_new_k = (k_new_idx >= 0) & (k_new_idx < T) - safe_k_new = tl.where(is_new_k, k_new_idx, 0) - - # Loop 1: per-head dt processing. dt → dt_processed → dA_cumsum → - # decay_vec_new (= exp(cumAdt_new)). Stored to write_buf for next step. - # decay_vec_full (= total_decay * decay_vec_new) is finalized in loop 2 - # once total_decay is loaded; loop 1 stores raw decay_vec_new to scratch. - for h_local in range(HEADS_PER_BLOCK): - head_idx = first_head + h_local - - dt_base = dt_ptr + pid_b * stride_dt_batch + head_idx * stride_dt_head - dt = tl.load(dt_base + offs_t * stride_dt_T, mask=t_mask, other=0.0).to(tl.float32) - if HAS_DT_BIAS: - dt_bias = tl.load(dt_bias_ptr + head_idx * stride_dt_bias_head).to(tl.float32) - dt = dt + dt_bias - if DT_SOFTPLUS: - dt = softplus(dt) - - A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) - dA_cumsum = tl.cumsum(A * dt, axis=0) - - # Cross-step continuity for old_dA_cumsum: rectangle precompute runs - # only on the nowrite path (write_buf == buf_active, write_offset == PNAT). - # Add the running tail from buf_active[head_idx, PNAT-1] so the buffer - # holds one continuous cumsum across back-to-back nowrites. PNAT is - # scalar/uniform, use scalar if to short-circuit the load at PNAT=0. - if prev_num_accepted_tokens == 0: - prev_total = 0.0 - else: - last_cumsum_ptr = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_active * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T - ) - prev_total = tl.load(last_cumsum_ptr).to(tl.float32) - - # Store dt and dA_cumsum to write_buf at [write_offset, write_offset+T) - # for next step's replay/rectangle use. - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - tl.store( - old_dt_base + (write_offset + offs_t) * stride_old_dt_T, - dt, - mask=t_mask, - ) - - old_dA_cumsum_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - tl.store( - old_dA_cumsum_base + (write_offset + offs_t) * stride_old_dA_cumsum_T, - dA_cumsum + prev_total, - mask=t_mask, - ) - - # ---- Hoisted: cache-only loads independent of conv1d ---- - # old_B (group-level, BLOCK_K × BLOCK_DSTATE = ~8KB tile) and the - # decay_vec_full per-head pre-compute (which writes to DRAM and doesn't - # need cross-gdc_wait variables) are issued BEFORE gdc_wait so their - # HBM latency overlaps with conv1d. Per-head factor_dt/exp_diff stay - # below gdc_wait — they need cross-iteration spans, which Triton can't - # express without a DRAM round-trip; the per-head LOADS in the post- - # wait loop are small and cheap, so leave them. - group_idx = first_head // nheads_ngroups_ratio - - # Group-level: old B from active buffer at [0, PNAT) of the K-axis. - old_B_read_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + buf_active * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - old_B_load = tl.load( - old_B_read_base - + safe_old_k[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - mask=is_old_k[:, None] & n_mask[None, :], - other=0.0, - ) - - # Pre-wait: vectorized across HEADS_PER_BLOCK heads. Compute decay_vec_full - # (H, T) and combo = factor_dt * exp_diff (H, T, K). Store decay_vec_full; - # combo_block stays in registers across gdc_wait — used directly post-wait - # to compute rect_CB_scaled without a global memory roundtrip. - offs_h = tl.arange(0, HEADS_PER_BLOCK) - heads_block = first_head + offs_h # (H,) - - # Per-head bases (H,) — broadcast with offs_k or offs_t for 2D loads. - old_dt_read_h = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + buf_active * stride_old_dt_dbuf - + heads_block * stride_old_dt_head - ) - old_dA_cumsum_read_h = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_active * stride_old_dA_cumsum_dbuf - + heads_block * stride_old_dA_cumsum_head - ) - old_dt_write_h = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + heads_block * stride_old_dt_head - ) - old_dA_cumsum_write_h = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + heads_block * stride_old_dA_cumsum_head - ) - - # (H, K) loads at [0, PNAT) — old data from previous step. - hk_mask = is_old_k[None, :] # (1, K) - old_dt_all = tl.load( - old_dt_read_h[:, None] + safe_old_k[None, :] * stride_old_dt_T, - mask=hk_mask, other=0.0, - ).to(tl.float32) - old_dA_cumsum_all = tl.load( - old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, - mask=hk_mask, other=0.0, - ).to(tl.float32) - # (H, T) loads at [PNAT, PNAT+T) — this step's dA_cumsum_new from loop 1. - # With the cross-step continuity fix in loop 1, the values stored at - # [PNAT, PNAT+T) are the continuous cumsum (prefix + per-step new - # cumsum) — i.e., continuous_cumsum[PNAT..PNAT+T-1] in global indexing. - ht_mask = t_mask[None, :] # (1, T) - dA_cumsum_new = tl.load( - old_dA_cumsum_write_h[:, None] - + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T, - mask=ht_mask, other=0.0, - ).to(tl.float32) - # (H, K) loads at K_NEW_SHIFT-shifted positions for new tokens. - hkn_mask = is_new_k[None, :] - dt_at_kn = tl.load( - old_dt_write_h[:, None] - + (write_offset + safe_k_new)[None, :] * stride_old_dt_T, - mask=hkn_mask, other=0.0, - ).to(tl.float32) - dA_cumsum_at_kn = tl.load( - old_dA_cumsum_write_h[:, None] - + (write_offset + safe_k_new)[None, :] * stride_old_dA_cumsum_T, - mask=hkn_mask, other=0.0, - ).to(tl.float32) - - # decay_vec_full[t] = exp(continuous_cumsum[PNAT+t]) — directly the - # continuous value now stored at buffer position write_offset+t. Was - # decomposed as total_decay * exp(per_step_new[t]) when the buffer held - # per-step (non-continuous) cumsum; with the continuity fix the value - # IS continuous_cumsum[PNAT+t] so no decomposition is needed. - decay_vec_full_block = tl.exp(dA_cumsum_new) # (H, T) - decay_vec_addrs = ( - decay_vec_ptr - + pid_b * stride_dv_batch - + heads_block[:, None] * stride_dv_head - + offs_t[None, :] * stride_dv_t - ) # (H, T) - tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) - - # combo_block = factor_dt * exp_diff — (H, T, K). Stays in registers - # across gdc_wait. With continuous cumsum in the buffer, s_k for any k - # (old or new) is simply -continuous_cumsum[k]; exp_diff[t, k] then - # equals exp(continuous_cumsum[PNAT+t] - continuous_cumsum[k]) — the - # decay weight for token k's contribution to the output at position - # PNAT+t. No need to subtract any "total" — the dA_cumsum_new[t] term - # already carries the full prefix. - # - # Numerical note: pre-fix this kernel computed `total - old_dA[k]` - # (small-minus-small) then summed `+ dA_cumsum_new[t]` (also small, - # per-step). Post-fix `s_k = -old_dA[k]` is large-magnitude positive - # and `dA_cumsum_new[t]` is large-magnitude negative; their sum - # cancels back to the same small value. Cancellation error is bounded - # by ulp(max_magnitude) ≈ 2^-23 · |continuous_cumsum| — negligible - # for max_window ≤ ~1024. Still one exp on the sum (not two muls of - # exps), so no overflow regression vs the original formulation. - factor_dt = tl.where(is_old_k[None, :], old_dt_all, dt_at_kn) # (H, K) - s_k = tl.where( - is_old_k[None, :], - -old_dA_cumsum_all, - -dA_cumsum_at_kn, - ) # (H, K) - # exp_diff (H, T, K) = exp(s_k (H, 1, K) + dA_cumsum_new (H, T, 1)). - exp_diff = tl.exp(s_k[:, None, :] + dA_cumsum_new[:, :, None]) - combo_block = factor_dt[:, None, :] * exp_diff # (H, T, K) - - # ---- gdc_wait: from here on we depend on conv1d's outputs ---- - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - # Conv1d outputs: B and C - C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group - B_new_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group - - C_all = tl.load( - C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - B_new_orig = tl.load( - B_new_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - B_new_shifted = tl.load( - B_new_base + safe_k_new[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, - mask=is_new_k[:, None] & n_mask[None, :], - other=0.0, - ) - # Disjoint masks: old at [0, PNAT), new at [K_NEW_SHIFT, K_NEW_SHIFT+T). - B_combined = old_B_load + B_new_shifted - raw_rect_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_combined).to(tl.bfloat16)) - - # Append new B to cache at [PNAT, PNAT+T) of write_buf (once per group). - if first_head % nheads_ngroups_ratio == 0: - old_B_write_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + write_buf * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - tl.store( - old_B_write_base - + (write_offset + offs_t)[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - B_new_orig, - mask=t_mask[:, None] & n_mask[None, :], - ) - - # Causal mask (BLOCK_SIZE_T × BLOCK_SIZE_K, shared across heads). - # Approach C: new tokens at runtime [PNAT, PNAT+T) instead of K_NEW_SHIFT. - t_idx_2d = offs_t[:, None] - k_idx_2d = offs_k[None, :] - is_old_k_2d = k_idx_2d < prev_num_accepted_tokens - k_new_idx_2d = k_idx_2d - prev_num_accepted_tokens - is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) - causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] - - # Post-wait vectorized: combo_block (H, T, K) is still live in registers. - # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store as - # one (H, T, K) tile. - rect_CB_scaled_block = tl.where( - causal_combined[None, :, :], - raw_rect_CB[None, :, :] * combo_block, - 0.0, - ) # (H, T, K) - cb_scaled_addrs = ( - cb_scaled_ptr - + pid_b * stride_cb_batch - + heads_block[:, None, None] * stride_cb_head - + offs_t[None, :, None] * stride_cb_t - + offs_k[None, None, :] * stride_cb_j - ) # (H, T, K) - cb_store_mask_3d = ( - (offs_t[None, :, None] < BLOCK_SIZE_T) - & (offs_k[None, None, :] < BLOCK_SIZE_K) - ) # (1, T, K) → broadcasts to (H, T, K) - tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) - - -# Rectangle precompute kernel. Thin wrapper around _rectangle_precompute_impl -# that carries the @triton.heuristics for constexpr derivation; called from -# the Python wrapper on the rectangle nowrite path. -@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.jit() -def _dynamic_precompute_kernel( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, - decay_vec_ptr, - # Cache pointers - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, - # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. - RECTANGLE: tl.constexpr, -): - # Hoisted PDL signal: fire as the first thing every program does. - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - - pid_b = tl.program_id(axis=0) - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH - # write_checkpoint is now runtime in replay precompute, so a single - # call site handles both write and nowrite for the replay branch. - # Take rectangle only when RECTANGLE is True AND this slot doesn't - # need write; everything else funnels into replay. - if needs_write_runtime or not RECTANGLE: - _replay_precompute_impl( - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - cb_scaled_ptr, - decay_vec_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - T, - dstate, - nheads_ngroups_ratio, - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - DT_SOFTPLUS, - HAS_DT_BIAS, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - LAUNCH_WITH_PDL, - HEADS_PER_BLOCK, - needs_write_runtime, - ) - else: - _rectangle_precompute_impl( - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - cb_scaled_ptr, - decay_vec_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - T, - MAX_REPLAY_BUFFER_LENGTH, - dstate, - nheads_ngroups_ratio, - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - DT_SOFTPLUS, - HAS_DT_BIAS, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - BLOCK_SIZE_K, - LAUNCH_WITH_PDL, - HEADS_PER_BLOCK, - ) - - -# Main kernel: tl.dot replay + precomputed CB output. -# Grid: (cdiv(dim, M), batch, nheads). - - -@triton.jit() -def _persistent_main_impl( - # Per-work-unit indices (computed by the persistent wrapper). - # `pid_b` is the post-perm slot index (caller has already applied any - # slot permutation and slot_offset). - pid_m, - pid_b, - pid_h, - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or - # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor - # USE_TMA_STORE is enabled (kernel ignores it via constexpr). - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - rand_seed_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_WINDOW: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, - QUANT_MAX: tl.constexpr, - WRITE_CHECKPOINT: tl.constexpr, - # IS_DYNAMIC: kept in the signature for caller-side bookkeeping (the - # outer _persistent_main_kernel still inspects it to decide the slot- - # IS_DYNAMIC: when True (persistent_dynamic), is_write is per-slot from - # PNAT. When False (persistent_main), is_write is constexpr from - # WRITE_CHECKPOINT. See also WC_IS_CONSTEXPR below. - IS_DYNAMIC: tl.constexpr, - # WC_IS_CONSTEXPR: when True, force is_write = WRITE_CHECKPOINT (constexpr) - # regardless of IS_DYNAMIC. Callers in RECT=1 use this in the is_w=True - # arm of _persistent_main_kernel (we know all slots that reach this call - # need is_write=True because is_w was the PNAT-derived runtime check, and - # this arm only fires when is_w is True). Passing WRITE_CHECKPOINT=True - # as a literal at the call site + WC_IS_CONSTEXPR=True here lets the inner - # body DCE the nowrite path under IS_DYNAMIC=True too — same codegen - # quality as persistent_main mode (-3.7% measured at b=1024 dyn-shape). - # When False (RECT=0 callers, where both write and nowrite slots are - # dispatched to ONE call), use the original runtime is_write under - # IS_DYNAMIC=True; avoids the binary-doubling regression that two - # specialized calls would cause. - WC_IS_CONSTEXPR: tl.constexpr = False, - # TMA flags — picked inside body based on is_write. When is_write is - # constexpr (either IS_DYNAMIC=False or WC_IS_CONSTEXPR=True), the - # use_tma_load = USE_TMA_LOAD_WRITE if is_write else USE_TMA_LOAD_NOWRITE - # ternary constexpr-folds and only one TMA load form survives. - USE_TMA_LOAD_WRITE: tl.constexpr = False, - USE_TMA_LOAD_NOWRITE: tl.constexpr = False, - USE_TMA_STORE: tl.constexpr = False, -): - # IS_DYNAMIC: kernel-mode label, used by the OUTER _persistent_main_kernel - # to decide slot-range derivation and outer is_w dispatch strategy - # (constexpr WC for persistent_main; runtime is_w split -> 2 specialized - # impl calls for persistent_dynamic). Inside this impl, IS_DYNAMIC is - # NOT consulted at runtime -- WRITE_CHECKPOINT is the only constexpr that - # gates the write/nowrite codegen, in BOTH modes. - - # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized - # state dtype (int8 / int16 / float8e4nv) and only those. - tl.static_assert( - (QUANT_MAX > 0.0) - == ( - (state_ptr.dtype.element_ty == tl.int8) - or (state_ptr.dtype.element_ty == tl.int16) - or (state_ptr.dtype.element_ty == tl.float8e4nv) - ), - "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", - ) - - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - # Resolve is_write: see WC_IS_CONSTEXPR / IS_DYNAMIC docs in the param - # list above. Three cases: - # - WC_IS_CONSTEXPR=True (RECT=1 is_w=True arm callers): use WC - # constexpr. Caller knows the slot needs write; inner DCEs nowrite - # paths. Avoids the binary-doubling overhead that calling the impl - # twice would cause, while still constexpr-DCEing the nowrite half. - # - IS_DYNAMIC=True (RECT=0 caller, persistent_dynamic): runtime - # branch on PNAT. Both write and nowrite codegen live in one body - # (no bloat) — same as the pre-refactor behavior. - # - IS_DYNAMIC=False (persistent_main): WC constexpr from caller. - if WC_IS_CONSTEXPR: - is_write: tl.constexpr = WRITE_CHECKPOINT - elif IS_DYNAMIC: - is_write = (prev_num_accepted_tokens + T) > MAX_REPLAY_BUFFER_LENGTH - else: - is_write = WRITE_CHECKPOINT - if is_write: - write_buf = 1 - active_buf # noqa: F841 - write_offset = 0 - else: - write_buf = active_buf # noqa: F841 - write_offset = prev_num_accepted_tokens - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) - m_mask = offs_m < dim - n_mask = offs_n < dstate - t_mask = offs_t < T - - # Load state. state_tma_descriptor is a host-built tensor_descriptor - # over a flat (cache*nheads*dim, dstate) view of state when any TMA - # path is enabled; raw `state_ptr` is the underlying tensor and is - # always passed. state_ptrs / state_ptr_raw are the raw-pointer view - # used for !TMA load and store paths. offs_y is the flat row index - # for TMA load/store; computed unconditionally (cheap int math; DCE'd - # when no TMA path is reachable). - state_mask = m_mask[:, None] & n_mask[None, :] - offs_y = ( - cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) - + pid_h * dim - + pid_m * BLOCK_SIZE_M - ) - state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - ) - # Load state. Branch on is_write (constexpr = WRITE_CHECKPOINT in BOTH - # modes after the outer-dispatch refactor), then constexpr-pick TMA-vs- - # tl.load per side. Outer `if` DCE's, only the matching side's - # constexpr-gated load survives -- same compile-time picking for both - # persistent_main and persistent_dynamic (the latter dispatches at the - # outer kernel level so each impl instance sees a constexpr WC). - if is_write: - if USE_TMA_LOAD_WRITE: - state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) - else: - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) - else: - if USE_TMA_LOAD_NOWRITE: - state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) - else: - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) - if QUANT_MAX > 0.0: - state_scales_base = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - ) - decode_scale = tl.load( - state_scales_base + offs_m * stride_state_scales_dim, - mask=m_mask, - other=1.0, - ).to(tl.float32) - state = state * decode_scale[:, None] - - # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) - group_idx = pid_h // nheads_ngroups_ratio - - old_window_mask = offs_window < prev_num_accepted_tokens - - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + active_buf * stride_old_dt_dbuf - + pid_h * stride_old_dt_head - ) - old_dt_all = tl.load( - old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 - ).to(tl.float32) - - old_dA_cumsum_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + active_buf * stride_old_dA_cumsum_dbuf - + pid_h * stride_old_dA_cumsum_head - ) - old_dA_cumsum_all = tl.load( - old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, - mask=old_window_mask, other=0.0, - ).to(tl.float32) - - prev_k_idx = tl.minimum( - tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 - ) - total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( - tl.float32 - ) - - coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all - - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head - old_x_all = tl.load( - old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - mask=old_window_mask[:, None] & m_mask[None, :], - other=0.0, - ) - - old_B_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + active_buf * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - old_B_all = tl.load( - old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - mask=old_window_mask[:, None] & n_mask[None, :], - other=0.0, - ).to(tl.float32) - - dB_scaled = coeff[:, None] * old_B_all - - total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) - state *= total_decay - - state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) - - if is_write: - if USE_RS_ROUNDING: - # Generate random tensor for stochastic rounding. The amount of - # randomness needed depends on the SR codegen path: - # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) - # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) - # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs - # int16 SR (24b + bitrev32): 1 b32 per 2 outputs - # The PTX cvt.rs.* instructions consume a single 32-bit random - # and split the bits internally for 2 or 4 conversions. Generate - # only what's actually consumed and broadcast to fill the unused - # slots — saves Philox rounds proportionally. - if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: - RAND_DIVISOR: tl.constexpr = 4 # fp8 SR - elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: - RAND_DIVISOR: tl.constexpr = 4 # int8 SR - elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: - RAND_DIVISOR: tl.constexpr = 2 # int16 SR - elif QUANT_MAX == 0.0: - RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) - else: - RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized - - rand_seed = tl.load(rand_seed_ptr) - base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head - # Number of unique randoms per row = dstate / RAND_DIVISOR. - # randint4x emits 4 randoms per offset, so use that / 4 offsets. - offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) - rand_offsets_q = ( - base_rand - + offs_m[:, None] * stride_state_dim - + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) - ) # (M, dstate / (4*RAND_DIVISOR)) - if PHILOX_ROUNDS > 0: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) - else: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) - r01 = tl.join(r0, r1) - r23 = tl.join(r2, r3) - r0123 = tl.join(r01, r23) - rand_compact = tl.reshape( - r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) - ) - # Broadcast each unique rand to RAND_DIVISOR adjacent positions. - # Pack-group (pack=2 fp16 / pack=4 fp8) consumes adjacent positions; - # the unique rand lands at the asm's read slot; duplicates feed - # the dead slots. Triton's broadcast_to is stride-0 in IR. - if RAND_DIVISOR > 1: - rand_3d = rand_compact[:, :, None] - rand_3d = tl.broadcast_to( - rand_3d, - (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), - ) - rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) - else: - rand = rand_compact - - if QUANT_MAX > 0.0: - amax = tl.max(tl.abs(state), axis=1) - encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) - decode_scale = 1.0 / encode_scale - state_scales_ptrs = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - + offs_m * stride_state_scales_dim - ) - tl.store(state_scales_ptrs, decode_scale, mask=m_mask) - state_q = state * encode_scale[:, None] - if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): - _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_q_fp8sr) - else: - tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) - else: - if USE_RS_ROUNDING: - tl.static_assert( - (state_ptrs.dtype.element_ty == tl.int8) - or (state_ptrs.dtype.element_ty == tl.int16), - "Quantized SR fall-through expects int8 or int16; " - "fp8 SR is handled by the prior branch.", - ) - if state_ptrs.dtype.element_ty == tl.int8: - state_q = _stochastic_round_int8_packed( - state_q, rand, offs_n[None, :] - ) - else: - state_q = _stochastic_round_int16_packed( - state_q, rand, offs_n[None, :] - ) - elif state_ptrs.dtype.element_ty != tl.float8e4nv: - tl.static_assert( - (state_ptrs.dtype.element_ty == tl.int8) - or (state_ptrs.dtype.element_ty == tl.int16), - "Quantized RN with explicit round() expects int8 or int16.", - ) - state_q = tl.extra.cuda.libdevice.round(state_q) - state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) - _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_q_cast) - else: - tl.store(state_ptrs, _state_q_cast, mask=state_mask) - elif USE_RS_ROUNDING: - tl.static_assert( - state_ptrs.dtype.element_ty == tl.float16, - "Non-quantized SR only supports fp16 state.", - ) - _state_sr = _stochastic_round_fp16x2(state, rand) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_sr) - else: - tl.store(state_ptrs, _state_sr, mask=state_mask) - else: - _state_cast = state.to(state_ptrs.dtype.element_ty) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_cast) - else: - tl.store(state_ptrs, _state_cast, mask=state_mask) - - # Phase 2: Output using precomputed CB_scaled and decay_vec - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head - C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group - if HAS_Z: - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - - if HAS_D: - D = tl.load( - D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - C_all = tl.load( - C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - - x_all = tl.load( - x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=t_mask[:, None] & m_mask[None, :], - other=0.0, - ) - tl.store( - old_x_base - + (write_offset + offs_t)[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - x_all, - mask=t_mask[:, None] & m_mask[None, :], - ) - x_all = x_all.to(tl.float32) - - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head - CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), - other=0.0, - ).to(tl.float32) - - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( - tl.float32 - ) - - init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] - cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) - out_all = init_out + cb_out - - if HAS_D: - out_all = out_all + x_all * D[None, :] - - if HAS_Z: - z_all = tl.load( - z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, - mask=t_mask[:, None] & m_mask[None, :], other=0.0, - ).to(tl.float32) - out_all_z = out_all * z_all * tl.sigmoid(z_all) - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) - else: - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) - - -# `_persistent_rectangle_impl`: rectangle nowrite path for the persistent -# kernel. Body is a copy of `_rectangle_main_impl` with `pid_m`/`pid_b`/`pid_h` -# lifted to args (same pattern as `_persistent_main_impl` vs `_replay_main_impl`). -# Called only for nowrite slots when the kernel runs with RECTANGLE=True. -# Dropped from the rect impl: LAUNCH_DEPENDENT_KERNELS / USE_PERM / REVERSE_PERM -# (kernel-level, signalled once at top); slot_perm_ptr (kernel resolves perm). -@triton.jit() -def _persistent_rectangle_impl( - # Per-work-unit indices (computed by the persistent wrapper). - pid_m, - pid_b, - pid_h, - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as - # replay path). Used when USE_TMA_LOAD; ignored otherwise. - state_tma_descriptor, - state_scales_ptr, # only consulted when QUANT_MAX > 0 - old_x_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides (rectangle (batch, nheads, T, K)) - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - QUANT_MAX: tl.constexpr, - USE_TMA_LOAD: tl.constexpr = False, -): - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - write_offset = prev_num_accepted_tokens - - # Static rectangle K-axis layout (matches precompute). - K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_k = tl.arange(0, BLOCK_SIZE_K) - m_mask = offs_m < dim - n_mask = offs_n < dstate - t_mask = offs_t < T - - # K-axis masks (approach C: PNAT-runtime offset, matches precompute). - is_old_k = offs_k < prev_num_accepted_tokens - safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - prev_num_accepted_tokens - is_new_k = (k_new_idx >= 0) & (k_new_idx < T) - safe_k_new = tl.where(is_new_k, k_new_idx, 0) - - # Load state. Quant scale hoist: defer `* decode_scale` post-matmul. - if USE_TMA_LOAD: - offs_y = ( - cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) - + pid_h * dim - + pid_m * BLOCK_SIZE_M - ) - state = state_tma_descriptor.load([offs_y, 0]) - else: - state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - ) - state_mask = m_mask[:, None] & n_mask[None, :] - state = tl.load(state_ptrs, mask=state_mask, other=0.0) - if QUANT_MAX > 0.0: - state_scales_base = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - ) - decode_scale = tl.load( - state_scales_base + offs_m * stride_state_scales_dim, - mask=m_mask, other=1.0, - ).to(tl.float32) - else: - state = state.to(tl.float32) - - # Group / pointer offset setup - group_idx = pid_h // nheads_ngroups_ratio - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head - C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group - if HAS_Z: - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head - - if HAS_D: - D = tl.load( - D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - - # Hoist: old_x doesn't depend on conv1d/precompute; load before gdc_wait. - old_x_load = tl.load( - old_x_base - + safe_old_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - mask=is_old_k[:, None] & m_mask[None, :], - other=0.0, - ).to(tl.float32) - - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - C_all = tl.load( - C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - x_K = tl.load( - x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=is_new_k[:, None] & m_mask[None, :], - other=0.0, - ) - tl.store( - old_x_base - + offs_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - x_K, - mask=is_new_k[:, None] & m_mask[None, :], - ) - - x_K_f32 = x_K.to(tl.float32) - x_combined = old_x_load + x_K_f32 - - if HAS_D or HAS_Z: - sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) - x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) - else: - x_all = x_K_f32 # placeholder; unused - - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head - CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), - other=0.0, - ).to(tl.float32) - - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec_full = tl.load( - decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 - ).to(tl.float32) - - state_out = ( - tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) - * decay_vec_full[:, None] - ) - if QUANT_MAX > 0.0: - state_out = state_out * decode_scale[None, :] - - token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) - - out_all = state_out + token_out - - if HAS_D: - out_all = out_all + x_all * D[None, :] - - if HAS_Z: - z_all = tl.load( - z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, - mask=t_mask[:, None] & m_mask[None, :], other=0.0, - ).to(tl.float32) - out_all_z = out_all * z_all * tl.sigmoid(z_all) - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) - else: - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) - - -# Persistent main kernel: 1D grid, persistent CTA loop. -# Heuristics mirror those of `_checkpointing_main_kernel`. -@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) -@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics( - {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} -) -@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_WINDOW": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.heuristics( - {"NUM_PID_M_BLOCKS": lambda args: triton.cdiv(args["dim"], args["BLOCK_SIZE_M"])} -) -@triton.jit() -def _persistent_main_kernel( - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view. - # Shared across BOTH the replay path (consumed by _persistent_main_impl - # when USE_TMA_LOAD_*/STORE) AND the rectangle path (consumed by - # _persistent_rectangle_impl when USE_TMA_LOAD) — same descriptor, same - # block_shape, just gated by separate constexprs per impl. Wrapper sets - # this to a TensorDescriptor when ANY of the three TMA flags is on, else - # to `state_ptr` (raw); each impl ignores it via its own constexpr when - # not consuming it. - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - rand_seed_ptr, - pad_slot_id, - # Persistent-loop work-distribution scalars. Caller pre-sorts the batch - # write-first; the kernel uses (n_writes, batch_total, WRITE_CHECKPOINT) - # to derive its own slot range. Write half processes [0, n_writes), - # nowrite half processes [n_writes, batch_total). - # - # n_writes_ptr is a device pointer to a (1,) int32 tensor. Reading - # from device memory (rather than taking a Python int kernel arg) is - # required so mix-mode benchmarking can vary n_writes per iter inside - # a captured CUDA graph — the source tensor's contents change, the - # pointer doesn't. Cost: one int load per kernel launch (~negligible). - # When IS_DYNAMIC=True the value is unused (Triton DCEs the load). - n_writes_ptr, # int32 *: device-side count of write-mode slots - batch_total, # int32: total slot count - nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_WINDOW: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, - QUANT_MAX: tl.constexpr, - WRITE_CHECKPOINT: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - USE_PERM: tl.constexpr, - # NUM_PERSISTENT: runtime int (not constexpr). Used ONLY as the loop - # stride in `tl.range(pid, total_work, NUM_PERSISTENT, ...)`. Making it - # runtime collapses the cta_per_sm tuning dim from the kernel's compile - # signature: 8 CPS values used to mean 8x recompiles; now they share one - # compiled kernel. Work decomposition (pid_m, pid_b_local, pid_h) does - # NOT depend on NUM_PERSISTENT — it uses constexpr NUM_PID_M_BLOCKS and - # runtime n_slots_local — so loop unrolling and flatten=/num_stages=/ - # warp_specialize= optimizations on `tl.range` operate independently of - # the stride value. - NUM_PERSISTENT, - NUM_LOOP_STAGES: tl.constexpr, - NUM_PID_M_BLOCKS: tl.constexpr, - FLATTEN: tl.constexpr, - WARP_SPECIALIZE: tl.constexpr, - IS_DYNAMIC: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr = 16, # rectangle K-axis (heuristic-derived) - RECTANGLE: tl.constexpr = False, # when True, dispatch nowrite slots to _persistent_rectangle_impl - # 3 TMA toggles per the 3 live paths per-compilation: - # USE_TMA_LOAD_WRITE — replay-style state load when is_write - # USE_TMA_LOAD_NOWRITE — nowrite-path state load (rect when RECTANGLE, - # else replay-nowrite) - # USE_TMA_STORE — replay-style state store (only fires on write - # path; no-op when not is_write) - # Wrapper picks USE_TMA_LOAD_NOWRITE = _use_tma_rect_load (if rectangle) - # or _use_tma_replay_nowrite_load (if not). - USE_TMA_LOAD_WRITE: tl.constexpr = False, - USE_TMA_LOAD_NOWRITE: tl.constexpr = False, - USE_TMA_STORE: tl.constexpr = False, -): - # PDL signal: fire once at kernel entry (not per work unit). - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - - # Load runtime n_writes from device memory. Read once at kernel entry; - # used only by the !IS_DYNAMIC slot-range derivation below. Triton - # DCEs the load when IS_DYNAMIC=True (n_writes is dead there). - n_writes = tl.load(n_writes_ptr) - - # Derive this kernel's slot range. Two modes: - # IS_DYNAMIC=False (persistent_main): caller pre-sorts and splits halves; - # slot range is [0, n_writes) when WRITE_CHECKPOINT else [n_writes, batch_total) - # IS_DYNAMIC=True (persistent_dynamic): single launch covers full batch; - # each work-item dispatches via runtime PNAT check inside the impl. - if IS_DYNAMIC: - slot_lo = 0 - slot_hi = batch_total - else: - if WRITE_CHECKPOINT: - slot_lo = 0 - slot_hi = n_writes - else: - slot_lo = n_writes - slot_hi = batch_total - n_slots_local = slot_hi - slot_lo - - pid = tl.program_id(axis=0) - total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads - - # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) - # with pid_m varying fastest (M-tile cache locality on state load), then - # slot, then head — mirrors the existing 3D grid's axis ordering - # (axis=0 fastest = pid_m). - for tile_id in tl.range( - pid, total_work, NUM_PERSISTENT, - flatten=FLATTEN, num_stages=NUM_LOOP_STAGES, warp_specialize=WARP_SPECIALIZE, - ): - pid_m = tile_id % NUM_PID_M_BLOCKS - pid_b_local = (tile_id // NUM_PID_M_BLOCKS) % n_slots_local - pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) - # Translate local slot index → global slot index. When USE_PERM is - # set, the caller-provided slot_perm gives the original slot index - # for the post-sort position. - pid_b_grid = pid_b_local + slot_lo - if USE_PERM: - pid_b = tl.load(slot_perm_ptr + pid_b_grid) - else: - pid_b = pid_b_grid - - # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle - # impl. Pre-load cache_batch_idx + PNAT here only for the RECTANGLE - # path's branch decision. Both impls re-load and handle pad_slot_id - # internally (Triton's L1 cache makes the duplicate loads ~free). - if RECTANGLE: - if HAS_CACHE_BATCH_INDICES: - cbi_pre = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - is_pad = cbi_pre == pad_slot_id - else: - cbi_pre = pid_b.to(tl.int64) - is_pad = False - if not is_pad: - pnat_pre = tl.load(prev_num_accepted_tokens_ptr + cbi_pre) - if IS_DYNAMIC: - is_w = (pnat_pre + T) > MAX_REPLAY_BUFFER_LENGTH - else: - is_w = WRITE_CHECKPOINT - if is_w: - # Pass WRITE_CHECKPOINT=True constexpr to specialize this - # impl call for the write path. Under IS_DYNAMIC=True, the - # kernel-level WRITE_CHECKPOINT is False (launcher default), - # but the OUTER is_w branch we are inside narrows the - # runtime path to writes-only, so we override to True here - # so the impl's constexpr-gated `if is_write:` blocks DCE - # to the write-only codegen. Under IS_DYNAMIC=False - # (persistent_main), the kernel-level WRITE_CHECKPOINT is - # itself True for this half (write half launches with - # WC=True), and the outer is_w = WRITE_CHECKPOINT = True - # constexpr-folds; passing literal True here is consistent - # and constexpr-equivalent. - _persistent_main_impl( - pid_m, pid_b, pid_h, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, - stride_old_B_group, stride_old_B_dstate, - stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, - stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) - True, # WC_IS_CONSTEXPR — force inner to use WC constexpr - # 3 TMA flags: write-load fires here (we're in the - # is_write branch), nowrite-load is dead (no slot - # reaches it), store fires (write path). - USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, - ) - else: - # Rectangle nowrite: pass state_ptr (raw, always) + - # state_tma_descriptor (the single unified descriptor — - # same memory replay paths use). Rect impl gates use - # of the descriptor via its USE_TMA_LOAD constexpr. - _persistent_rectangle_impl( - pid_m, pid_b, pid_h, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - state_batch_indices_ptr, pad_slot_id, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, - LAUNCH_WITH_PDL, QUANT_MAX, - USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle - ) - # else: pad slot — skip both impls (both would early-return anyway) - else: - # No rectangle path — single _persistent_main_impl call covers - # both write and nowrite slots via WC constexpr (non-dynamic) or - # runtime is_write (IS_DYNAMIC=True). Pass all 3 TMA flags; - # impl picks USE_TMA_LOAD_WRITE vs USE_TMA_LOAD_NOWRITE based on - # its computed is_write — constexpr-folds when is_write is - # constexpr (non-dyn), runtime branch when IS_DYNAMIC=True. - # (Reverted from outer two-call dispatch: that doubled the - # compiled body size under IS_DYNAMIC=True and regressed RECT=0 - # perf by ~+24%.) - _persistent_main_impl( - pid_m, pid_b, pid_h, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, - stride_old_B_group, stride_old_B_dstate, - stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, - stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - WRITE_CHECKPOINT, IS_DYNAMIC, - False, # WC_IS_CONSTEXPR=False — RECT=0 has both write/nowrite slots in one call - USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, - ) - - -# ============================================================================ -# Python wrapper -# ============================================================================ - - -_QUANT_MAX_BY_DTYPE = { - torch.int8: 127.0, - torch.int16: 32767.0, - torch.float8_e4m3fn: 448.0, -} - - -# --------------------------------------------------------------------------- -# Default tunings — looked up by (effective_batch, dtype, sr) when the caller -# leaves mode/knobs as None. -# -# Effective batch = raw_batch × nheads_per_rank. Our sweep was at TP=8 with -# the standard Mamba2 nheads; at call time we compute it from the input -# tensor shape so callers at other TP / nheads pick up the right cell. -# -# Schema: dict[(dtype_str, sr_str)] → list[(eff_batch_threshold, mode, knobs)] -# sorted by threshold ascending. Lookup finds the first threshold ≥ eff_b -# (so missing intermediate batches fall up to the next tuned cell). If -# eff_b exceeds the largest threshold, use the largest entry. -# -# Each `knobs` dict only contains keys for the chosen mode; the wrapper -# unpacks them with the same name as the matching kwargs. Caller-provided -# kwargs always win over table values. -# -# This table is intentionally NOT parameterized by T or max_window. Our -# sweep was T=6, max_window=16. Callers outside that regime silently get -# the same numbers — they may be suboptimal but they're correct. -# -# Source: audit_v2.py --emit-tuning. Auto-generated from per-cell search -# winners (best of pd / pm by bucket_expected_renorm). Sweep was TP=8 with -# NHEADS=128 → nheads_per_rank=16; thresholds are in effective_batch units. -# Missing dtype/SR combos (fp16/RN, int8/RN, fp8/*) fall back via the -# _resolve_tuning chain — RN→SR for same dtype, then fp8→int8/SR. -_DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { - ("fp32", "RN"): [ - ( 16, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 8, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.22us - ( 32, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.17us - ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.08us - ( 128, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=8, score=9.00us - ( 256, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=16, score=10.92us - ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=13.53us - ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=19.50us - ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=30.28us - ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=50.32us - ( 8192, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=90.99us - (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=171.69us - ], - ("fp16", "SR"): [ - ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.16us - ( 32, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.01us - ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=7.95us - ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 4, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=8.87us - ( 256, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 2, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.28us - ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=12.90us - ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=16.71us - ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=25.71us - ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=39.80us - ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=71.34us - (16384, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=133.51us - ], - ("int8", "SR"): [ - ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.34us - ( 32, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.36us - ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 2, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.40us - ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 10, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=9.37us - ( 256, "persistent_dynamic", {'_block_size_m': 16, '_cta_per_sm': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages': 1, '_num_stages': 4, '_num_warps': 1, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.02us - ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=32, score=13.15us - ( 1024, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=17.82us - ( 2048, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=27.01us - ( 4096, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=43.23us - ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=77.01us - (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=140.43us - ], -} - - -# Knob names that map between the modes' single-value (pd) and split-value -# (pm) namespaces. Used by `_bridge_tuning_knobs` when caller forces a mode -# different from the table's recommendation. -_PD_TO_PM_SPLIT_MAP = { # pd unsplit knob → (pm_write_knob, pm_nowrite_knob) - "_block_size_m": ("_block_size_m_write", "_block_size_m_nowrite"), - "_num_warps": ("_num_warps_write", "_num_warps_nowrite"), - "_num_stages": ("_num_stages_write", "_num_stages_nowrite"), - # CPS / LS are persistent-loop knobs; pd uses _cta_per_sm + _num_loop_stages - # as unsplit, pm uses _cta_per_sm_write/_nowrite + _num_loop_stages_write/_nowrite. - "_cta_per_sm": ("_cta_per_sm_write", "_cta_per_sm_nowrite"), - "_num_loop_stages": ("_num_loop_stages_write", "_num_loop_stages_nowrite"), -} - - -def _bridge_tuning_knobs(knobs: dict, from_mode: str, to_mode: str) -> dict: - """Convert a tuning dict between pd ↔ pm knob namespaces. - - pd → pm: copy each unsplit value to both write/nowrite split knobs; drop - the unsplit form (pm doesn't read it). - pm → pd: take the nowrite split value as the unsplit knob; drop the - write/nowrite split forms (pd doesn't read them). - Shape knobs that exist in both modes (_heads_per_block, _flatten, - _warp_specialize, TMA flags, rectangle_for_nowrite) carry over unchanged. - """ - out = dict(knobs) - if from_mode == "persistent_dynamic" and to_mode == "persistent_main": - for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): - if unsplit in out: - out.setdefault(pm_w, out[unsplit]) - out.setdefault(pm_nw, out[unsplit]) - del out[unsplit] - elif from_mode == "persistent_main" and to_mode == "persistent_dynamic": - for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): - if pm_nw in out: - out.setdefault(unsplit, out[pm_nw]) - out.pop(pm_w, None) - out.pop(pm_nw, None) - return out - - -def _resolve_tuning( - batch: int, nheads_per_rank: int, dt_str: str, sr_str: str, -) -> tuple[str, dict] | None: - """Look up the default mode + knobs for this (eff_batch, dt, sr) cell. - - Returns (mode, knobs_dict) or None if the table has no entry covering - this dtype/sr (including the fp8→int8/SR and dtype/RN→dtype/SR fallbacks). - Returning None lets the wrapper fall back to caller-provided kwargs or - kernel-side defaults. - """ - eff_b = batch * max(1, nheads_per_rank) - # Lookup chain. Order: - # 1. Exact (dt, sr). - # 2. (dt, SR) if RN missing for that dtype. - # 3. Cross-dtype fallback for dtypes we haven't tuned: - # bf16 / int16 → fp16/SR - # fp8 → int8/SR - # Unknown dtype → raise. - valid_dtypes = {"fp32", "fp16", "bf16", "int8", "int16", "fp8"} - if dt_str not in valid_dtypes: - raise ValueError( - f"checkpointing_state_update: unsupported state dtype {dt_str!r}; " - f"expected one of {sorted(valid_dtypes)}" - ) - keys_to_try = [(dt_str, sr_str)] - if sr_str == "RN": - keys_to_try.append((dt_str, "SR")) - if dt_str in ("bf16", "int16"): - keys_to_try.append(("fp16", "SR")) - elif dt_str == "fp8": - keys_to_try.append(("int8", "SR")) - entries = None - for k in keys_to_try: - if k in _DEFAULT_TUNING: - entries = _DEFAULT_TUNING[k] - break - if entries is None: - return None - # Find first threshold ≥ eff_b; if none, use largest entry. - for thresh, mode, knobs in entries: - if eff_b <= thresh: - return mode, dict(knobs) - thresh, mode, knobs = entries[-1] - return mode, dict(knobs) - - -def checkpointing_state_update( - state: torch.Tensor, - old_x: torch.Tensor, - old_B: torch.Tensor, - old_dt: torch.Tensor, - old_dA_cumsum: torch.Tensor, - cache_buf_idx: torch.Tensor, - prev_num_accepted_tokens: torch.Tensor, - x: torch.Tensor, - dt: torch.Tensor, - A: torch.Tensor, - B: torch.Tensor, - C: torch.Tensor, - out: torch.Tensor, - # Required persistent-mode plumbing (REQUIRED for both pd and pm; pd - # ignores both internally but the wrapper still demands them): - # n_writes : (1,) int32 device tensor with the count of write-mode - # slots in the batch. pm uses it to size the two halves; - # pd ignores it (per-slot runtime PNAT check). - # slot_perm : (batch,) int32 device tensor remapping grid pid → slot. - # pm uses it to cluster writes first (kernel grid step is - # write_half then nowrite_half); pd ignores it. Callers - # that don't care about ordering should pass arange(batch). - n_writes: torch.Tensor, - slot_perm: torch.Tensor, - D: torch.Tensor | None = None, - z: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - dt_softplus: bool = False, - state_batch_indices: torch.Tensor | None = None, - pad_slot_id: int = PAD_SLOT_ID, - rand_seed: torch.Tensor | None = None, - philox_rounds: int = 10, - state_scales: torch.Tensor | None = None, - launch_with_pdl=False, - use_internal_pdl=True, - write_checkpoint: bool = True, - rectangle_for_nowrite: bool | None = None, - mode: str | None = None, - _block_size_m: int | None = None, - _num_warps: int | None = None, - _num_stages: int | None = None, - _precompute_num_warps: int | None = None, - _precompute_num_stages: int | None = None, - _heads_per_block: int | None = None, - _maxnreg: int | None = None, - _num_ctas: int | None = None, - # Per-main knobs (override shared values for one half of the dl-family / - # persistent_main launches). Default None = tied to the shared value - # (backward compat). The two main kernels (write vs nowrite) have - # different per-slot work — write does a state shift + store, nowrite - # just appends — so the optimum (M, W, S, H) can differ. Precompute - # knobs are intentionally NOT split: shared precompute wins (cheaper - # launch, hotter precompute outputs in L2). Persistent CPS / LS knobs - # are also split per-main since the two persistent_main launches have - # different grid sizes. - _block_size_m_write: int | None = None, - _block_size_m_nowrite: int | None = None, - _num_warps_write: int | None = None, - _num_warps_nowrite: int | None = None, - _num_stages_write: int | None = None, - _num_stages_nowrite: int | None = None, - # Note: heads_per_block / precompute_num_warps are NOT split — they only - # affect the precompute kernel, which is shared across write/nowrite. - # TMA state-tensor toggles — 4 independent paths (see CHECKPOINTING_DESIGN.md - # item #17 for measured perf profiles). Each is False=raw load/store, True= - # use a host-built TMA tensor_descriptor for that path. - _use_tma_rect_load: bool | None = None, # rect kernel's state load (nowrite-only) - _use_tma_replay_write_load: bool | None = None, # replay-style state load when WC=True - _use_tma_replay_write_store: bool | None = None, # replay-style state store when WC=True - _use_tma_replay_nowrite_load: bool | None = None, # replay-style state load when WC=False - # Persistent-mode tuning kwargs (consulted for both pd and pm; pd uses - # _cta_per_sm / _num_loop_stages, pm uses the _write/_nowrite splits): - # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally - # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. - # _num_loop_stages : int — `num_stages` arg on the inner `tl.range(...)` - # persistent loop. Note: this is loop-level, NOT the kernel-arg - # `num_stages` (which only pipelines dot-feeding loads). - # _flatten : bool — `flatten` arg on `tl.range(...)`. - # _warp_specialize : bool — `warp_specialize` arg on `tl.range(...)`. - _cta_per_sm: int | None = None, - _num_loop_stages: int | None = None, - _flatten: bool | None = None, - _warp_specialize: bool | None = None, - # Per-main persistent-specific knobs. Same rationale as the BLOCK_SIZE_M - # split above: the two persistent_main launches (write half vs nowrite - # half) have different grid sizes and per-work-item costs, so they may - # want different cta_per_sm / num_loop_stages. - _cta_per_sm_write: int | None = None, - _cta_per_sm_nowrite: int | None = None, - _num_loop_stages_write: int | None = None, - _num_loop_stages_nowrite: int | None = None, -): - """ - Replay SSM state update with precomputed CB and tl.dot fast-forward. - - Two-kernel architecture: - 1. Precompute kernel: computes CB_scaled and decay_vec from B, C, dt, A. - Writes processed dt/dA_cumsum/B to double-buffered cache for next step. - 2. Main kernel: replays old tokens via tl.dot fast-forward on cached data, - then computes output using precomputed CB_scaled and new x/C inputs. - - PDL (Programmatic Dependent Launch) chain: - conv1d → (external PDL) → precompute → (internal PDL) → main - External PDL: precompute starts while conv1d is running; gdc_wait() - in precompute blocks until conv1d completes before loading B/C. - Internal PDL: main starts while precompute is running; main's replay - phase uses only cached data from the previous step. gdc_wait() in - main blocks until precompute completes before loading conv1d outputs - (x, C) and precompute outputs (CB_scaled, decay_vec). - - Uses double-buffered cache tensors. cache_buf_idx[slot] indicates which - buffer (0 or 1) to READ from for replay. The WRITE buffer is 1 - read. - Caller must flip cache_buf_idx[slot] after each call. - - Arguments: - state: (cache, nheads, dim, dstate) in-place. After the call, contains - the state after replaying prev_num_accepted_tokens old tokens. - old_x: (cache, T, nheads, dim) bf16 — old x cache (single-buffered). - old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. - old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. - old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. - cache_buf_idx: (cache,) int32 — which buffer to read (0 or 1). - prev_num_accepted_tokens: (cache,) int32. - x: (batch, T, nheads, dim) new token inputs. - dt: (batch, T, nheads, dim) with stride(-1)==0 (tie_hdim). - A: (nheads, dim, dstate) with stride(-1)==0, stride(-2)==0 (tie_hdim). - B: (batch, T, ngroups, dstate). - C: (batch, T, ngroups, dstate). - out: (batch, T, nheads, dim) preallocated output. - D: (nheads, dim) optional feed-through parameter. - z: (batch, T, nheads, dim) optional silu gate. - dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). - state_batch_indices: (batch,) optional cache slot mapping. - rand_seed: optional single-element int64 CUDA tensor for Philox PRNG seed. - When provided, state is stochastically rounded on store. Supported - for state.dtype in (fp16, int8, int16, fp8_e4m3fn); other dtypes - silently use deterministic rounding. fp16+SR and fp8+SR both - require sm_100a (Blackwell B200+) — wrapper asserts this loudly. - philox_rounds: number of Philox PRNG rounds (default 10). - state_scales: required when state.dtype in (int8, int16, fp8_e4m3fn). - Shape (cache_size, nheads, dim), fp32. Per-(head, dim) channel - decode scale (= 1 / encode_scale). The kernel writes scales on - checkpoint steps and reads them on load (broadcast over dstate). - Ignored for non-quantized state dtypes. - launch_with_pdl: enable external PDL (conv1d → precompute chain). - Defaults False; caller opts in when the upstream chain is PDL-safe. - Ignored on hardware that doesn't support PDL (sm < 90). - use_internal_pdl: enable internal PDL (precompute → main overlap). - Defaults True; override for testing only. - Ignored on hardware that doesn't support PDL (sm < 90). - - _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, - _precompute_num_warps, _precompute_num_stages, _heads_per_block, - _maxnreg, _num_ctas) are benchmark-only overrides; production callers - should leave them None to use the heuristic-tuned defaults. - """ - # PDL needs sm >= 90. - if get_sm_version() < 90: - launch_with_pdl = False - use_internal_pdl = False - - # Mode selection: - # mode=None (default): look up the table-tuned mode + knobs for this - # (effective_batch, dtype, sr) cell. See `_resolve_tuning` above. - # mode="persistent_dynamic": single persistent-CTA kernel covering the - # full batch. Each work-item dispatches via runtime PNAT check - # (is_write = (pnat + T) > MAX). No write/nowrite split. - # slot_perm is honored but optional. write_checkpoint is ignored. - # mode="persistent_main": persistent-CTA kernel with two launches - # (write half + nowrite half). Caller MUST pre-sort slot_perm - # write-first; the n_writes tensor partitions the persistent loop - # into the two halves with the right WRITE_CHECKPOINT constexpr - # each time. RECTANGLE constexpr (= rectangle_for_nowrite) picks - # rect vs replay for the nowrite half. write_checkpoint is ignored. - # Note: mode-and-knob resolution from the default-tuning table happens - # below, after we have `batch` and `nheads`. - - # --- Hardware support gates --- - # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX - # instructions (Ada Lovelace introduced them; Hopper/Blackwell carry them). - if state.dtype == torch.float8_e4m3fn: - assert get_sm_version() >= 89, ( - "fp8_e4m3fn state requires SM 89+ (Ada Lovelace / Hopper / Blackwell) " - f"for fp32↔fp8 cvt PTX instructions; current SM is {get_sm_version()}." - ) - - # PTX cvt.rs.* (stochastic rounding) family lands on Blackwell only. - # Wrapper fails loud; framework decides fall-back (e.g. drop SR, use RN). - # int8 / int16 SR uses pure-Triton libdevice.floor + uniform noise — no - # PTX SR instruction needed, runs anywhere. - if rand_seed is not None: - if state.dtype == torch.float16: - assert get_sm_version() >= 100, ( - "fp16 stochastic rounding (PTX cvt.rs.f16x2.f32) requires " - f"sm_100a (Blackwell B200+); current SM is {get_sm_version()}." - ) - elif state.dtype == torch.float8_e4m3fn: - assert get_sm_version() >= 100, ( - "fp8 stochastic rounding (PTX cvt.rs.satfinite.e4m3x4.f32) " - f"requires sm_100a (Blackwell B200+); current SM is {get_sm_version()}." - ) - - # --- Unsqueeze inputs to canonical shapes --- - if state.dim() == 3: - state = state.unsqueeze(1) - if x.dim() == 2: - x = x.unsqueeze(1) - if x.dim() == 3: - x = x.unsqueeze(1) - if dt.dim() == 2: - dt = dt.unsqueeze(1) - if dt.dim() == 3: - dt = dt.unsqueeze(1) - if A.dim() == 2: - A = A.unsqueeze(0) - if B.dim() == 2: - B = B.unsqueeze(1) - if B.dim() == 3: - B = B.unsqueeze(1) - if C.dim() == 2: - C = C.unsqueeze(1) - if C.dim() == 3: - C = C.unsqueeze(1) - if D is not None and D.dim() == 1: - D = D.unsqueeze(0) - if z is not None: - if z.dim() == 2: - z = z.unsqueeze(1) - if z.dim() == 3: - z = z.unsqueeze(1) - if dt_bias is not None and dt_bias.dim() == 1: - dt_bias = dt_bias.unsqueeze(0) - if out.dim() == 2: - out = out.unsqueeze(1) - if out.dim() == 3: - out = out.unsqueeze(1) - - cache_size, nheads, dim, dstate = state.shape - batch, T, _, _ = x.shape - ngroups = B.shape[2] - assert nheads % ngroups == 0 - - # --- Quantization plumbing (needed for SR/RN classification below) --- - # QUANT_MAX > 0 ⇔ state is int8 / int16 / fp8_e4m3fn. Kernel-entry - # static_assert on the Triton side mirrors this invariant. - quant_max = _QUANT_MAX_BY_DTYPE.get(state.dtype, 0.0) - is_quantized = quant_max > 0.0 - - # --- Default-tuning lookup --- - # Resolve (mode, knobs) from the table when caller leaves them None. - # Caller-provided kwargs always win. If the caller forces a mode that - # differs from the table's recommendation for this cell, we BRIDGE the - # table's knobs into the forced mode's knob namespace rather than fall - # back to (likely-terrible) kernel defaults: - # table pd → forced pm: copy each unsplit pd knob (M, W, S, CPS, LS) - # to both write and nowrite split knobs. - # table pm → forced pd: take the nowrite split values (Mnw, Wnw, Snw, - # CPSnw, LSnw) as the unsplit knobs. - # Empty table → no-op (caller passes whatever, mode falls back to pd). - _dt_str = { - torch.float32: "fp32", - torch.float16: "fp16", - torch.bfloat16: "bf16", - torch.int8: "int8", - torch.int16: "int16", - torch.float8_e4m3fn: "fp8", - }.get(state.dtype, str(state.dtype)) - _sr_str = "SR" if (rand_seed is not None and is_quantized) else "RN" - _table_entry = _resolve_tuning(batch, nheads, _dt_str, _sr_str) - if _table_entry is not None: - _table_mode, _table_knobs = _table_entry - if mode is None: - mode = _table_mode - if mode != _table_mode: - # Bridge across modes — see header comment above. - _table_knobs = _bridge_tuning_knobs(_table_knobs, _table_mode, mode) - # Fill None-valued kwargs from table. We can't reliably mutate - # locals() for re-read, so re-bind each kwarg explicitly. - if rectangle_for_nowrite is None and "rectangle_for_nowrite" in _table_knobs: - rectangle_for_nowrite = bool(_table_knobs["rectangle_for_nowrite"]) - _block_size_m = _block_size_m if _block_size_m is not None else _table_knobs.get("_block_size_m") - _num_warps = _num_warps if _num_warps is not None else _table_knobs.get("_num_warps") - _num_stages = _num_stages if _num_stages is not None else _table_knobs.get("_num_stages") - _heads_per_block = _heads_per_block if _heads_per_block is not None else _table_knobs.get("_heads_per_block") - _precompute_num_warps = _precompute_num_warps if _precompute_num_warps is not None else _table_knobs.get("_precompute_num_warps") - _precompute_num_stages = _precompute_num_stages if _precompute_num_stages is not None else _table_knobs.get("_precompute_num_stages") - _block_size_m_write = _block_size_m_write if _block_size_m_write is not None else _table_knobs.get("_block_size_m_write") - _block_size_m_nowrite = _block_size_m_nowrite if _block_size_m_nowrite is not None else _table_knobs.get("_block_size_m_nowrite") - _num_warps_write = _num_warps_write if _num_warps_write is not None else _table_knobs.get("_num_warps_write") - _num_warps_nowrite = _num_warps_nowrite if _num_warps_nowrite is not None else _table_knobs.get("_num_warps_nowrite") - _num_stages_write = _num_stages_write if _num_stages_write is not None else _table_knobs.get("_num_stages_write") - _num_stages_nowrite = _num_stages_nowrite if _num_stages_nowrite is not None else _table_knobs.get("_num_stages_nowrite") - _cta_per_sm = _cta_per_sm if _cta_per_sm is not None else _table_knobs.get("_cta_per_sm") - _num_loop_stages = _num_loop_stages if _num_loop_stages is not None else _table_knobs.get("_num_loop_stages") - # Split-form resolution for pm's per-half knobs. Without these the - # table's _num_loop_stages_{write,nowrite} and _cta_per_sm_{write, - # nowrite} values are dead — pm reads the split forms but the - # wrapper would leave them None, falling through to Triton defaults - # (or our hardcoded `or 1` / `or 2` per-mode fallbacks). - _num_loop_stages_write = _num_loop_stages_write if _num_loop_stages_write is not None else _table_knobs.get("_num_loop_stages_write") - _num_loop_stages_nowrite = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _table_knobs.get("_num_loop_stages_nowrite") - _cta_per_sm_write = _cta_per_sm_write if _cta_per_sm_write is not None else _table_knobs.get("_cta_per_sm_write") - _cta_per_sm_nowrite = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _table_knobs.get("_cta_per_sm_nowrite") - _flatten = _flatten if _flatten is not None else _table_knobs.get("_flatten") - _warp_specialize = _warp_specialize if _warp_specialize is not None else _table_knobs.get("_warp_specialize") - _use_tma_rect_load = _use_tma_rect_load or bool(_table_knobs.get("_use_tma_rect_load", False)) - _use_tma_replay_write_load = _use_tma_replay_write_load or bool(_table_knobs.get("_use_tma_replay_write_load", False)) - _use_tma_replay_write_store = _use_tma_replay_write_store or bool(_table_knobs.get("_use_tma_replay_write_store", False)) - _use_tma_replay_nowrite_load = _use_tma_replay_nowrite_load or bool(_table_knobs.get("_use_tma_replay_nowrite_load", False)) - # Final defaults if neither caller nor table set them (empty table case). - if mode is None: - mode = "persistent_dynamic" - if rectangle_for_nowrite is None: - rectangle_for_nowrite = False - assert mode in ("persistent_dynamic", "persistent_main"), ( - f"unknown mode {mode!r}; expected 'persistent_dynamic' or 'persistent_main'" - ) - if is_quantized: - assert state_scales is not None, ( - f"state.dtype={state.dtype} requires state_scales tensor " - "(shape (cache_size, nheads, dim), fp32)." - ) - assert state_scales.shape == (cache_size, nheads, dim), ( - f"state_scales shape mismatch: expected {(cache_size, nheads, dim)}, " - f"got {state_scales.shape}." - ) - assert state_scales.dtype == torch.float32, ( - f"state_scales must be fp32, got {state_scales.dtype}." - ) - assert state_scales.device == state.device - - # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the - # placeholder degenerate case max_window = T (every step is a checkpoint - # step). For real replay-style checkpointing, max_window > T and - # `prev_num_accepted_tokens` can be 0..max_window. Window-axis kernel - # tiles (BLOCK_SIZE_WINDOW, BLOCK_SIZE_K) are derived independently from - # MAX_REPLAY_BUFFER_LENGTH so max_window can exceed BLOCK_SIZE_T freely. - max_window = old_x.shape[1] - assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" - - assert x.shape == (batch, T, nheads, dim) - assert dt.shape == x.shape - assert A.shape == (nheads, dim, dstate) - assert B.shape == (batch, T, ngroups, dstate) - assert C.shape == B.shape - assert old_x.shape == (cache_size, max_window, nheads, dim) - assert old_B.shape == (cache_size, 2, max_window, ngroups, dstate) - assert old_dt.shape == (cache_size, 2, nheads, max_window) - assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_window) - assert cache_buf_idx.shape == (cache_size,) - assert prev_num_accepted_tokens.shape == (cache_size,) - - tie_hdim = ( - A.stride(-1) == 0 - and A.stride(-2) == 0 - and dt.stride(-1) == 0 - and (dt_bias is None or dt_bias.stride(-1) == 0) - ) - assert tie_hdim - - device = x.device - BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) - # Rectangle K-axis bound = window (max_window). Computed unconditionally - # so the launch sites can refer to it; only used on the rectangle path. - BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) - - # Allocate precomputed intermediates (per-call, not cached). Always - # allocate (T, K) — the largest layout that any path uses. Replay-style - # paths only touch the first T columns; rectangle/dynamic use the full K. - # The few extra unused columns per row are negligible (~6KB per layer at - # production sizes) and let the dispatch helpers share one buffer. - cb_scaled = torch.empty( - batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 - ) - decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) - - z_strides = ( - (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) if z is not None else (0, 0, 0, 0) - ) - - # Kernel tuning: BLOCK_SIZE_M, num_warps, HEADS_PER_BLOCK, precompute_num_warps. - # Dtype-aware heuristic from B200 sweeps (batch 1-512, T=6/32, TP=8, conv1d + - # chained PDL). Keyed on total_heads, BLOCK_SIZE_T, and state dtype; 16-bit - # states prefer different tiles from fp32 due to lower bandwidth. Philox - # gets its own branch — stochastic rounding shifts compute toward CUDA cores, - # so small-batch configs want more warps to hide the extra work. - total_heads = batch * nheads - heads_per_group = nheads // ngroups - state_is_16bit = state.dtype in (torch.float16, torch.bfloat16) - use_philox = rand_seed is not None - if BLOCK_SIZE_T <= 16: - if use_philox and state_is_16bit: - # Philox: more warps at small batch to hide CUDA core work. - # At large batch, converges to non-Philox fp16 config. - if total_heads <= 16: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 4, 4, 4, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - elif state_is_16bit: - if total_heads <= 16: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 4, 1 - elif total_heads <= 64: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 2, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 1, - min(2, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - else: # fp32 state (no Philox — fp32 doesn't need stochastic rounding) - if total_heads <= 32: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 4, 1 - elif total_heads <= 64: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - elif total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 2, 2, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 1, 2, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 2, - min(2, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - else: # T > 16 - if state_is_16bit: - if total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 16, - 1, - 4, - min(2, heads_per_group), - ) - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 1, - min(4, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 4, - min(2, heads_per_group), - ) - else: # fp32 state - if total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 2, - 4, - min(2, heads_per_group), - ) - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 2, - min(4, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 4, - min(2, heads_per_group), - ) - if _block_size_m is not None: - BLOCK_SIZE_M = _block_size_m - if _num_warps is not None: - num_warps = _num_warps - if _heads_per_block is not None: - # Cap at heads_per_group: HEADS_PER_BLOCK divides the kernel's head - # axis, so a table value larger than the model's heads-per-group - # would overshoot. Protects callers running smaller models than - # the one we tuned against. - heads_per_block = min(_heads_per_block, heads_per_group) - if _precompute_num_warps is not None: - precompute_num_warps = _precompute_num_warps - - # Per-main knob resolution: each _*_{write,nowrite} arg, if not None, - # overrides the corresponding shared value for ONE main launch only. - # Default (None) = tied to shared value (current behavior). - BLOCK_SIZE_M_WRITE = _block_size_m_write if _block_size_m_write is not None else BLOCK_SIZE_M - BLOCK_SIZE_M_NOWRITE = _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M - NUM_WARPS_WRITE = _num_warps_write if _num_warps_write is not None else num_warps - NUM_WARPS_NOWRITE = _num_warps_nowrite if _num_warps_nowrite is not None else num_warps - NUM_STAGES_WRITE = _num_stages_write if _num_stages_write is not None else _num_stages - NUM_STAGES_NOWRITE = _num_stages_nowrite if _num_stages_nowrite is not None else _num_stages - # Persistent-only per-main: - CTA_PER_SM_WRITE = _cta_per_sm_write if _cta_per_sm_write is not None else _cta_per_sm - CTA_PER_SM_NOWRITE = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _cta_per_sm - NUM_LOOP_STAGES_WRITE = _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages - NUM_LOOP_STAGES_NOWRITE = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages - - HAS_CACHE_BATCH_INDICES = state_batch_indices is not None - - assert nheads % heads_per_block == 0, ( - f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" - ) - assert heads_per_block <= heads_per_group, ( - f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" - ) - - # state_scales pointer + strides: real tensor when quantized, otherwise - # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). - if is_quantized: - state_scales_arg = state_scales - state_scales_strides = ( - state_scales.stride(0), - state_scales.stride(1), - state_scales.stride(2), - ) - else: - state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 - state_scales_strides = (0, 0, 0) - - # Per-path TMA descriptors for state — write-side and nowrite-side. Each - # kernel launch consumes the descriptor whose block_shape[0] matches its - # BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need - # distinct descriptors; otherwise the descriptor's block_shape[0] would - # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot / arithmetic - # on the loaded tile fails shape inference at compile time - # ("Cannot make_shape_compatible: incompatible dimensions"). When Mw == - # Mnw (tied, the common case) the two descriptors are the same object. - # Same memory (state's flat 2D view, shape (cache*nheads*dim, dstate)) - # and same dstate block_shape — only block_shape[0] differs. - # When no TMA flag is on, both variables hold the raw `state` tensor as a - # dummy; kernels never reference it because their constexprs are all - # False (Triton DCEs the dead branches). - # `triton.set_allocator()` must run before any descriptor-using launch. - if (_use_tma_rect_load or _use_tma_replay_write_load - or _use_tma_replay_write_store or _use_tma_replay_nowrite_load): - from triton.tools.tensor_descriptor import TensorDescriptor - _ensure_tma_allocator() - assert state.is_contiguous(), "TMA state requires contiguous state" - assert state.stride(-1) == 1, "TMA state requires inner stride 1" - _state_flat = state.view(-1, state.shape[-1]) - _dstate_pow2 = triton.next_power_of_2(dstate) - state_tma_descriptor_write = TensorDescriptor.from_tensor( - _state_flat, block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], - ) - if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: - state_tma_descriptor_nowrite = state_tma_descriptor_write - else: - state_tma_descriptor_nowrite = TensorDescriptor.from_tensor( - _state_flat, block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], - ) - else: - state_tma_descriptor_write = state # dummy; all consuming constexprs False - state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False - - # Slot permutation — pointer + USE_PERM gate. Always required: the - # persistent_main kernel reads pid_b through slot_perm to walk the - # write-first sorted batch. persistent_dynamic forces USE_PERM=False - # at the call site (see launch_persistent_dynamic_main below) so the - # perm value doesn't matter for pd, but the tensor must still be valid. - assert isinstance(slot_perm, torch.Tensor), ( - f"slot_perm must be a torch.Tensor, got {type(slot_perm).__name__}" - ) - assert slot_perm.device == device, ( - f"slot_perm must be on device {device}, got {slot_perm.device}" - ) - assert slot_perm.dtype in (torch.int32, torch.int64), ( - f"slot_perm must be int32/int64, got {slot_perm.dtype}" - ) - assert slot_perm.shape == (batch,), ( - f"slot_perm must have shape (batch={batch},), got {tuple(slot_perm.shape)}" - ) - assert isinstance(n_writes, torch.Tensor), ( - f"n_writes must be a torch.Tensor, got {type(n_writes).__name__}" - ) - assert n_writes.device == device, ( - f"n_writes must be on device {device}, got {n_writes.device}" - ) - assert n_writes.dtype == torch.int32, ( - f"n_writes must be int32, got {n_writes.dtype}" - ) - assert n_writes.shape == (1,), ( - f"n_writes must have shape (1,), got {tuple(n_writes.shape)}" - ) - slot_perm_arg = slot_perm - use_perm = True - - precomp_grid = (batch, nheads // heads_per_block) - d_strides = (D.stride(0), D.stride(1)) if D is not None else (0, 0) - - # ---- Launch helpers (close over locals) ------------------------------- - # Each helper is a thin closure that calls one Triton kernel with the - # full positional + kwarg argument list. Mode-dependent constexprs - # (write_checkpoint, early_out, rectangle) are passed in. - - def launch_dynamic_precompute(rectangle: bool): - _dynamic_precompute_kernel[precomp_grid]( - dt, dt_bias, A, B, C, - cb_scaled, decay_vec, - old_B, old_dt, old_dA_cumsum, - cache_buf_idx, prev_num_accepted_tokens, - state_batch_indices, pad_slot_id, - T, max_window, dstate, nheads // ngroups, - dt.stride(0), dt.stride(1), dt.stride(2), - dt_bias.stride(0) if dt_bias is not None else 0, - A.stride(0), - B.stride(0), B.stride(1), B.stride(2), B.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - dt_softplus, - HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, - LAUNCH_WITH_PDL=launch_with_pdl, - LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, - HEADS_PER_BLOCK=heads_per_block, - RECTANGLE=rectangle, - num_warps=precompute_num_warps, - **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), - launch_pdl=launch_with_pdl, - ) - - # ---- launch_persistent_main ------------------------------------------ - # Persistent-CTA main kernel. Single launch covers `n_slots` slots - # starting at `slot_offset`. Caller invokes twice: once for the write - # half (slot_offset=0, n_slots=n_writes, write_checkpoint=True) and - # once for the nowrite half (slot_offset=n_writes, - # n_slots=batch-n_writes, write_checkpoint=False). Hard-sort - # contract: caller has pre-sorted slots so [0, n_writes) are writes - # and [n_writes, batch) are nowrites. - - # Resolve persistent-mode bench knobs. Defaults: cta_per_sm = 1 - # (one CTA per SM, matches upstream `_p_matmul_ogs.py`); num_loop_stages - # = 2 (matches in-tree `swiglu` precedent for non-dot persistent loops); - # flatten = True (canonical Triton 3.6 idiom); warp_specialize = False. - _num_sms = torch.cuda.get_device_properties(device).multi_processor_count - cta_per_sm_arg = _cta_per_sm if _cta_per_sm else 1 - num_persistent_arg = cta_per_sm_arg * _num_sms - num_loop_stages_arg = _num_loop_stages if _num_loop_stages else 2 - flatten_arg = True if _flatten is None else bool(_flatten) - warp_specialize_arg = False if _warp_specialize is None else bool(_warp_specialize) - # Per-launch work-item count. At small batch, total_work may be < the - # full persistent grid; capping `grid` at `min(NUM_PERSISTENT, total_work)` - # avoids launching empty CTAs that pay setup cost for no work. Correctness: - # the kernel's `tl.range(pid, total_work, NUM_PERSISTENT)` ensures each - # tile_id is covered exactly once across all live pids in [0, grid) when - # grid <= NUM_PERSISTENT (each CTA does 1 tile; loop step >= total_work - # exits immediately) AND when grid == NUM_PERSISTENT (each CTA loops over - # multiple tiles). NUM_PERSISTENT is now a runtime int (see kernel def - # docstring at _persistent_main_kernel) so changing cta_per_sm does NOT - # trigger a new Triton compile — same kernel binary, different loop step. - # (Named UPPERCASE for historical Triton-style consistency only; not - # constexpr.) - _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M - - def launch_persistent_main(write_checkpoint: bool, - *, - launch_dependent_kernels: bool = False, - rectangle: bool = False): - # `n_writes` (wrapper-level) is the (1,) int32 device tensor with the - # write count. Both halves always launch; the kernel's runtime PNAT - # check iterates only the slots that belong to its half. - _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE - _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE - _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE - _cps = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE - _cps = _cps if _cps else 1 - _nls = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE - _nls = _nls if _nls else 2 - _num_persistent = _cps * _num_sms - _num_pid_m_local = (dim + _bsm - 1) // _bsm - # Grid sizing: cap at min(full persistent grid, upper-bound total work). - # We use `batch` as the upper bound on slots-per-half — overcounting - # by a few CTAs is fine since the kernel's runtime check only - # iterates the slots that actually belong to its half. - _total_work_launch = max(1, batch * _num_pid_m_local * nheads) - grid = (min(_num_persistent, _total_work_launch),) - # Per-path TMA descriptor — block_shape[0] must match _bsm. - _desc = (state_tma_descriptor_write if write_checkpoint - else state_tma_descriptor_nowrite) - _persistent_main_kernel[grid]( - state, _desc, state_scales_arg, old_x, - old_B, old_dt, old_dA_cumsum, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, - n_writes, batch, nheads, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - _bsm, - LAUNCH_WITH_PDL=use_internal_pdl, - PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, - QUANT_MAX=quant_max, - WRITE_CHECKPOINT=write_checkpoint, - LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, - USE_PERM=use_perm, - NUM_PERSISTENT=_num_persistent, - NUM_LOOP_STAGES=_nls, - FLATTEN=flatten_arg, - WARP_SPECIALIZE=warp_specialize_arg, - IS_DYNAMIC=False, - RECTANGLE=rectangle, - # 3 TMA flags. IS_DYNAMIC=False: WC fixed per launch; impl - # constexpr-folds the LOAD pick. When WC=True (write half), - # NOWRITE_LOAD is dummy False; when WC=False, WRITE_LOAD/STORE - # dummy False. NOWRITE_LOAD picks rect-load (RECTANGLE) or - # replay-nowrite-load. - USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), - USE_TMA_LOAD_NOWRITE=bool( - (_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load) - and not write_checkpoint - ), - USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), - num_warps=_nw, - **({"num_stages": _ns} if _ns else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), - launch_pdl=use_internal_pdl, - ) - - def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, - launch_dependent_kernels: bool = False, - rectangle: bool = False): - # Single-launch persistent kernel covering the whole batch with - # runtime per-slot WRITE_CHECKPOINT branch. No half-split, no - # n_writes needed (the kernel ignores n_writes_dev when - # IS_DYNAMIC=True; Triton DCEs the load). is_write is computed - # at runtime per work-item from the loaded PNAT. - # We still pass `n_writes_dev` (the same tensor the persistent_main - # path uses) so the kernel signature is uniform; the value is - # immaterial. - # Grid sizing: cap at total_work (= batch * num_pid_m * nheads) for - # the dynamic case (full-batch coverage); see launch_persistent_main - # comment for correctness rationale. - _total_work_launch = max(1, batch * _num_pid_m * nheads) - grid = (min(num_persistent_arg, _total_work_launch),) - # Persistent-dynamic kernel uses a single BLOCK_SIZE_M (same as the - # wrapper's BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE tied convention), so - # the write-side descriptor matches. Both write and nowrite slots - # in this kernel share that BSM. - _persistent_main_kernel[grid]( - state, state_tma_descriptor_write, state_scales_arg, old_x, - old_B, old_dt, old_dA_cumsum, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, - n_writes, batch, nheads, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - BLOCK_SIZE_M, - LAUNCH_WITH_PDL=use_internal_pdl, - PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, - QUANT_MAX=quant_max, - WRITE_CHECKPOINT=False, - LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, - # persistent_dynamic forces USE_PERM=False regardless of caller- - # provided slot_perm — our pd tuning runs all happened with - # SORT=0 (no slot_perm passed), so honoring slot_perm here would - # silently shift pd to an untimed code path. Revisit if/when - # we benchmark pd with slot_perm. - USE_PERM=False, - NUM_PERSISTENT=num_persistent_arg, - NUM_LOOP_STAGES=num_loop_stages_arg, - FLATTEN=flatten_arg, - WARP_SPECIALIZE=warp_specialize_arg, - IS_DYNAMIC=True, - RECTANGLE=rectangle, - # 3 TMA flags. IS_DYNAMIC=True: is_write is runtime per slot; - # impl's load TMA picks per-slot (constexpr ternary becomes a - # runtime branch — both load forms emitted, ~negligible cost). - # NOWRITE_LOAD picks rect-load when RECTANGLE, else - # replay-nowrite-load. STORE only fires on runtime is_write. - USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load), - USE_TMA_LOAD_NOWRITE=bool( - _use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load - ), - USE_TMA_STORE=bool(_use_tma_replay_write_store), - num_warps=num_warps, - **({"num_stages": _num_stages} if _num_stages else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), - launch_pdl=use_internal_pdl, - ) - - # ---- Mode dispatch ---------------------------------------------------- - with torch.cuda.device(device.index): - if mode == "persistent_dynamic": - # Single-launch persistent kernel covering the full batch. Each - # work-item dispatches via runtime PNAT check. Kernel ignores - # n_writes (Triton DCEs the load) when IS_DYNAMIC=True; we still - # pass the wrapper-provided tensor as required by the signature. - launch_dynamic_precompute(rectangle=rectangle_for_nowrite) - launch_persistent_dynamic_main( - n_writes, - launch_dependent_kernels=False, - rectangle=rectangle_for_nowrite, - ) - elif mode == "persistent_main": - # Persistent-CTA main kernel. One shared dynamic_precompute - # (per-slot dispatch via PNAT) feeds two persistent_main - # launches (write half + nowrite half). Both halves ALWAYS - # launch; the kernel's runtime check iterates only the slots - # belonging to its half (write: [0, n_writes), nowrite: - # [n_writes, batch)). - # - # Caller-provided contract: `n_writes` is a (1,) int32 device - # tensor (the kernel reads it at runtime, after the precompute); - # `slot_perm` is a (batch,) int32 device tensor pre-sorted - # write-first. - launch_dynamic_precompute(rectangle=rectangle_for_nowrite) - launch_persistent_main( - write_checkpoint=True, - launch_dependent_kernels=True, - rectangle=False, # write always replay-style - ) - launch_persistent_main( - write_checkpoint=False, - launch_dependent_kernels=False, - rectangle=rectangle_for_nowrite, - ) - else: - raise ValueError( - f"mode={mode!r} is not supported. Supported modes: " - f"'persistent_dynamic', 'persistent_main'." - ) diff --git a/tensorrt_llm/_torch/modules/mamba/_v3_baseline_checkpointing_state_update_slim.py b/tensorrt_llm/_torch/modules/mamba/_v3_baseline_checkpointing_state_update_slim.py deleted file mode 100644 index 961dbf69bd7e..000000000000 --- a/tensorrt_llm/_torch/modules/mamba/_v3_baseline_checkpointing_state_update_slim.py +++ /dev/null @@ -1,3173 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. -# -# Adapted from: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py -# SPDX-FileCopyrightText: Copyright contributors to the sglang project -# -# Copyright (c) 2024, Tri Dao, Albert Gu. -# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py - -import torch -import triton -import triton.language as tl - -from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID -from tensorrt_llm._utils import get_sm_version - -from .softplus import softplus - - -# Lazy global allocator for Triton TMA tensor descriptors. Required by any -# host- or device-built tensor_descriptor; without it Triton raises at first -# launch. See TMA backlog item #17 / scratch experiment notes. -_TMA_ALLOCATOR_SET = False - - -def _ensure_tma_allocator() -> None: - global _TMA_ALLOCATOR_SET - if _TMA_ALLOCATOR_SET: - return - - def _alloc_fn(size, alignment, stream): - # Triton expects an int8 buffer of `size` bytes; alignment is enforced - # by the allocator returning a buffer satisfying it (PyTorch's - # cudaMalloc-backed tensors are 256B-aligned, so we're fine). - return torch.empty(size, device="cuda", dtype=torch.int8) - - triton.set_allocator(_alloc_fn) - _TMA_ALLOCATOR_SET = True - - -@triton.jit -def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: - """Stochastic rounding: fp32 pair → fp16x2 using Philox random bits. - - Uses PTX cvt.rs.f16x2.f32 which rounds each fp32 value to fp16 using - the random bits to break ties, avoiding systematic rounding bias that - accumulates over many decode steps with fp16 state. - - Adapted from flashinfer (Apache-2.0, vLLM/mamba lineage). - """ - return tl.inline_asm_elementwise( - asm="""{ - cvt.rs.f16x2.f32 $0, $2, $1, $3; - }""", - constraints=("=r,r,r,r,r"), - args=(x, rand), - dtype=tl.float16, - is_pure=True, - pack=2, - ) - - -@triton.jit -def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: - """Stochastic rounding: fp32 quad → fp8 e4m3 using Philox random bits. - - Uses PTX cvt.rs.satfinite.e4m3x4.f32 which combines stochastic rounding - and saturating cast in a single op (output is final fp8, no separate - clamp needed). The reversed source-register order {$4,$3,$2,$1} is - load-bearing — PTX packs leftmost source into the high byte but Triton's - pack=4 is little-endian, so the natural {$1,$2,$3,$4} order would - silently shuffle every group of 4 contiguous outputs. - - Requires SM_100a+ (Blackwell B200). Caller must gate at the wrapper - level — this kernel does not check. - - Adapted from vLLM PR #40012 (Apache-2.0). - """ - return tl.inline_asm_elementwise( - asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", - constraints="=r,r,r,r,r,r,r,r,r", - args=(x, rand), - dtype=tl.float8e4nv, - is_pure=True, - pack=4, - ) - - -@triton.jit -def _bitrev32(x: tl.tensor) -> tl.tensor: - return tl.inline_asm_elementwise( - asm="brev.b32 $0, $1;", - constraints="=r,r", - args=(x,), - dtype=tl.uint32, - is_pure=True, - pack=1, - ) - - -@triton.jit -def _stochastic_round_int8_packed( - x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor -) -> tl.tensor: - """Stochastic rounding for int8 using one random uint32 per 4 values.""" - low = rand & 0x0000FFFF - high = (rand >> 16) & 0x0000FFFF - low_rev = _bitrev32(low) >> 16 - high_rev = _bitrev32(high) >> 16 - rand_pos = offs_n & 3 - rand16 = tl.where( - rand_pos == 0, - low, - tl.where(rand_pos == 1, low_rev, tl.where(rand_pos == 2, high, high_rev)), - ) - rand01 = rand16.to(tl.float32) * (1.0 / float(1 << 16)) - return tl.extra.cuda.libdevice.floor(x + rand01) - - -@triton.jit -def _stochastic_round_int16_packed( - x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor -) -> tl.tensor: - """Stochastic rounding for int16 using one random uint32 per 2 values.""" - rand_bits = tl.where((offs_n & 1) == 0, rand, _bitrev32(rand)) - rand01 = (rand_bits & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) - return tl.extra.cuda.libdevice.floor(x + rand01) - - -# Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, -# old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. -# Grid: (batch, nheads // HEADS_PER_BLOCK). - - -@triton.jit() -def _replay_precompute_impl( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, - decay_vec_ptr, - # Cache pointers (both buffers reachable via stride_*_dbuf). This - # kernel writes to either the active (= cache_buf_idx) or inactive - # (= 1 - cache_buf_idx) buffer depending on WRITE_CHECKPOINT — see - # comment block at top of kernel body. - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - # Double-buffer index (per cache slot) — selects this step's "active" - # buffer (= where the historical inputs for this step live). - cache_buf_idx_ptr, - # Per-request accepted-tokens count (already-cached old tokens at - # [0, PNAT) of the active buffer; new tokens this step go after them - # on no-checkpoint steps). - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides: (cache, 2, T, ngroups, dstate) - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, - # Checkpointing flag — selects target buffer + offset for new-token - # cache writes. See "Cache write semantics" block below. - # Runtime (not constexpr): the only WRITE_CHECKPOINT-dependent code in - # this body is the write_buf/write_offset selection, which is plain - # arithmetic — no constexpr-shaped tile or whole-block gate. Letting - # it be runtime lets the dynamic dispatch kernel call us once with the - # per-slot needs_write flag instead of inlining two specializations. - write_checkpoint, -): - pid_b = tl.program_id(axis=0) - pid_hg = tl.program_id(axis=1) # head-group index - first_head = pid_hg * HEADS_PER_BLOCK - - # Resolve cache index for writes - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - - # --- Cache write semantics --- - # cache_buf_idx names this step's "active" buffer — the one with the - # historical inputs at [0, PNAT). The other buffer is "staging". - # - # Where do we write new tokens this step? - # WRITE_CHECKPOINT=False (no overflow): append to ACTIVE buffer at - # offset [PNAT : PNAT+T). Caller does NOT flip cache_buf_idx - # afterward; PNAT_next = PNAT + accepted. [0, PNAT) preserved. - # WRITE_CHECKPOINT=True (would overflow): write to STAGING buffer at - # [0, T). Caller flips cache_buf_idx afterward; next step's - # active = the one we just wrote. PNAT_next = accepted. Old - # data in the previous active buffer is folded into state via - # the replay update and discarded. This matches today's replay - # kernel behavior exactly. - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - if write_checkpoint: - write_buf = 1 - buf_active - write_offset = 0 - else: - write_buf = buf_active - write_offset = prev_num_accepted_tokens - - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - t_mask = offs_t < T - n_mask = offs_n < dstate - - # Causal mask is shared across all heads (depends only on offs_t) - causal_mask = offs_t[:, None] >= offs_t[None, :] - valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] - - # --- Vectorized pre-wait phase across HEADS_PER_BLOCK heads --- - # Compute dt, dA_cumsum, decay_vec as (H, T) tiles. Pre-compute - # scale_combo = decay_matrix * dt[:, None, :] as an (H, T, T) tile that - # stays in registers across gdc_wait — eliminates the post-wait reload - # of dt + dA_cumsum and the per-head loop. - offs_h = tl.arange(0, HEADS_PER_BLOCK) - heads_block = first_head + offs_h # (H,) - - # Load dt (H, T) - dt_addrs = ( - dt_ptr + pid_b * stride_dt_batch - + heads_block[:, None] * stride_dt_head - + offs_t[None, :] * stride_dt_T - ) - dt = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) - if HAS_DT_BIAS: - dt_bias = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) - dt = dt + dt_bias[:, None] - if DT_SOFTPLUS: - dt = softplus(dt) - - A = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) # (H,) - dA_cumsum = tl.cumsum(A[:, None] * dt, axis=1) # (H, T) - decay_vec = tl.exp(dA_cumsum) # (H, T) - - # Cross-step continuity for old_dA_cumsum: when appending to active_buf at - # offset PNAT > 0, the previous step left a running cumsum at [0, PNAT) - # whose tail value lives at active_buf[head, PNAT-1]. Add that tail to - # this step's per-step-restarted cumsum before storing so the buffer - # holds one continuous cumsum across N back-to-back nowrites. Write path - # (write_buf = 1 - buf_active, write_offset = 0) starts fresh, no prefix. - # Both branches are on scalar runtime values (write_checkpoint and PNAT), - # uniform across the block — use scalar if to short-circuit the load. - if write_checkpoint or prev_num_accepted_tokens == 0: - prev_total = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) - else: - last_cumsum_ptrs = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_active * stride_old_dA_cumsum_dbuf - + heads_block * stride_old_dA_cumsum_head - + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T - ) - prev_total = tl.load(last_cumsum_ptrs).to(tl.float32) - - # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of write_buf. - old_dt_addrs = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + heads_block[:, None] * stride_old_dt_head - + (write_offset + offs_t)[None, :] * stride_old_dt_T - ) - tl.store(old_dt_addrs, dt, mask=t_mask[None, :]) - - old_dA_cumsum_addrs = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + heads_block[:, None] * stride_old_dA_cumsum_head - + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T - ) - tl.store(old_dA_cumsum_addrs, dA_cumsum + prev_total[:, None], mask=t_mask[None, :]) - - # decay_vec scratch — always at offs_t. - decay_vec_addrs = ( - decay_vec_ptr + pid_b * stride_dv_batch - + heads_block[:, None] * stride_dv_head - + offs_t[None, :] * stride_dv_t - ) - tl.store(decay_vec_addrs, decay_vec, mask=t_mask[None, :]) - - # scale_combo (H, T, T) = exp(dA_cumsum[h, t1] - dA_cumsum[h, t2]) * dt[h, t2] - # Stays live across gdc_wait — used post-wait to compute CB_scaled. - decay_matrix = tl.exp(dA_cumsum[:, :, None] - dA_cumsum[:, None, :]) # (H, T, T) - scale_combo = decay_matrix * dt[:, None, :] # (H, T, T) - - # --- Wait for upstream kernel (external PDL) before loading B and C --- - # All dt processing above is independent of conv1d outputs. - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- - group_idx = first_head // nheads_ngroups_ratio - C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group - B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group - - C_all = tl.load( - C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - B_all = tl.load( - B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - - # Compute raw CB once — shared across all heads in this block - raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) - - # Store B to cache at [write_offset : write_offset+T) of write_buf. - if first_head % nheads_ngroups_ratio == 0: - old_B_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + write_buf * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - tl.store( - old_B_base - + (write_offset + offs_t)[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - B_all, - mask=t_mask[:, None] & n_mask[None, :], - ) - - # --- Vectorized post-wait phase: scale_combo (H, T, T) is still live in - # registers from pre-wait; multiply by raw_CB (T, T), apply causal mask, - # store as one (H, T, T) tile. --- - CB_scaled_block = tl.where( - valid_mask[None, :, :], - raw_CB[None, :, :] * scale_combo, - 0.0, - ) # (H, T, T) - cb_scaled_addrs = ( - cb_scaled_ptr + pid_b * stride_cb_batch - + heads_block[:, None, None] * stride_cb_head - + offs_t[None, :, None] * stride_cb_t - + offs_t[None, None, :] * stride_cb_j - ) # (H, T, T) - cb_store_mask = ( - (offs_t[None, :, None] < BLOCK_SIZE_T) - & (offs_t[None, None, :] < BLOCK_SIZE_T) - ) - tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) - - -# Replay-style precompute kernel. Thin wrapper around _replay_precompute_impl -# that carries the @triton.heuristics for constexpr derivation; called from -# the Python wrapper on the replay-style path (write or replay-nowrite). -@triton.jit() -def _rectangle_precompute_impl( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle - decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) - # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite - # path: read from buf_active at [0, PNAT), write new tokens at - # [PNAT, PNAT+T) of buf_active (same buffer). - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # rectangle K-axis bound - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides (rectangle: (batch, nheads, T, K)) - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides: (cache, 2, T_max, ngroups, dstate) - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides: (cache, 2, nheads, T_max) - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides: (cache, 2, nheads, T_max) - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, -): - pid_b = tl.program_id(axis=0) - pid_hg = tl.program_id(axis=1) - first_head = pid_hg * HEADS_PER_BLOCK - - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - write_buf = buf_active - write_offset = prev_num_accepted_tokens - - # Static rectangle K-axis layout: old at [0, PNAT) (mask is_old_k); - # new at [MAX-T, MAX) at compile-time shift K_NEW_SHIFT = MAX - T. - # PNAT + T <= MAX is guaranteed on the nowrite path → no overlap. - K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T - - offs_t = tl.arange(0, BLOCK_SIZE_T) # T-axis (output rows) - offs_k = tl.arange(0, BLOCK_SIZE_K) # K-axis (rectangle input cols) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - t_mask = offs_t < T - n_mask = offs_n < dstate - - # K-axis masks (approach C: runtime PNAT-offset instead of K_NEW_SHIFT) - # Old at [0, PNAT), new at [PNAT, PNAT+T). Cache and matmul share rows. - is_old_k = offs_k < prev_num_accepted_tokens - safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - prev_num_accepted_tokens - is_new_k = (k_new_idx >= 0) & (k_new_idx < T) - safe_k_new = tl.where(is_new_k, k_new_idx, 0) - - # V2: Loop 1 vectorized across HEADS_PER_BLOCK — single (H, T) tile, - # no per-head loop. Mirrors the replay precompute's pre-wait layout. - offs_h_lp1 = tl.arange(0, HEADS_PER_BLOCK) - heads_block_lp1 = first_head + offs_h_lp1 # (H,) - - dt_addrs_v = ( - dt_ptr + pid_b * stride_dt_batch - + heads_block_lp1[:, None] * stride_dt_head - + offs_t[None, :] * stride_dt_T - ) - dt_v = tl.load(dt_addrs_v, mask=t_mask[None, :], other=0.0).to(tl.float32) - if HAS_DT_BIAS: - dt_bias_v = tl.load(dt_bias_ptr + heads_block_lp1 * stride_dt_bias_head).to(tl.float32) - dt_v = dt_v + dt_bias_v[:, None] - if DT_SOFTPLUS: - dt_v = softplus(dt_v) - - A_v = tl.load(A_ptr + heads_block_lp1 * stride_A_head).to(tl.float32) # (H,) - dA_cumsum_v = tl.cumsum(A_v[:, None] * dt_v, axis=1) # (H, T) - - # Cross-step continuity: hoisted (H,) prefix load. - if prev_num_accepted_tokens == 0: - prev_total_v = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) - else: - prev_total_ptrs_v = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_active * stride_old_dA_cumsum_dbuf - + heads_block_lp1 * stride_old_dA_cumsum_head - + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T - ) - prev_total_v = tl.load(prev_total_ptrs_v).to(tl.float32) - - # Coalesced (H, T) stores — replaces HPB per-head scalar stores. - old_dt_addrs_v = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + heads_block_lp1[:, None] * stride_old_dt_head - + (write_offset + offs_t)[None, :] * stride_old_dt_T - ) - tl.store(old_dt_addrs_v, dt_v, mask=t_mask[None, :]) - - old_dA_cumsum_addrs_v = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + heads_block_lp1[:, None] * stride_old_dA_cumsum_head - + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T - ) - tl.store( - old_dA_cumsum_addrs_v, - dA_cumsum_v + prev_total_v[:, None], - mask=t_mask[None, :], - ) - - # ---- Hoisted: cache-only loads independent of conv1d ---- - # old_B (group-level, BLOCK_K × BLOCK_DSTATE = ~8KB tile) and the - # decay_vec_full per-head pre-compute (which writes to DRAM and doesn't - # need cross-gdc_wait variables) are issued BEFORE gdc_wait so their - # HBM latency overlaps with conv1d. Per-head factor_dt/exp_diff stay - # below gdc_wait — they need cross-iteration spans, which Triton can't - # express without a DRAM round-trip; the per-head LOADS in the post- - # wait loop are small and cheap, so leave them. - group_idx = first_head // nheads_ngroups_ratio - - # Group-level: old B from active buffer at [0, PNAT) of the K-axis. - old_B_read_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + buf_active * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - old_B_load = tl.load( - old_B_read_base - + safe_old_k[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - mask=is_old_k[:, None] & n_mask[None, :], - other=0.0, - ) - - # Pre-wait: vectorized across HEADS_PER_BLOCK heads. Compute decay_vec_full - # (H, T) and combo = factor_dt * exp_diff (H, T, K). Store decay_vec_full; - # combo_block stays in registers across gdc_wait — used directly post-wait - # to compute rect_CB_scaled without a global memory roundtrip. - offs_h = tl.arange(0, HEADS_PER_BLOCK) - heads_block = first_head + offs_h # (H,) - - # Per-head bases (H,) — broadcast with offs_k or offs_t for 2D loads. - old_dt_read_h = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + buf_active * stride_old_dt_dbuf - + heads_block * stride_old_dt_head - ) - old_dA_cumsum_read_h = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_active * stride_old_dA_cumsum_dbuf - + heads_block * stride_old_dA_cumsum_head - ) - old_dt_write_h = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + write_buf * stride_old_dt_dbuf - + heads_block * stride_old_dt_head - ) - old_dA_cumsum_write_h = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + write_buf * stride_old_dA_cumsum_dbuf - + heads_block * stride_old_dA_cumsum_head - ) - - # (H, K) loads at [0, PNAT) — old data from previous step. - hk_mask = is_old_k[None, :] # (1, K) - old_dt_all = tl.load( - old_dt_read_h[:, None] + safe_old_k[None, :] * stride_old_dt_T, - mask=hk_mask, other=0.0, - ).to(tl.float32) - old_dA_cumsum_all = tl.load( - old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, - mask=hk_mask, other=0.0, - ).to(tl.float32) - # V3: use loop-1 registers directly instead of reloading dA_cumsum_new - # from buffer. dA_cumsum_v + prev_total_v[:, None] IS what the buffer - # holds at positions [PNAT, PNAT+T). Saves an (H, T) DRAM round-trip - # per kernel call. - ht_mask = t_mask[None, :] # (1, T) - dA_cumsum_new = dA_cumsum_v + prev_total_v[:, None] # (H, T) - # (H, K) loads at K_NEW_SHIFT-shifted positions for new tokens. - hkn_mask = is_new_k[None, :] - dt_at_kn = tl.load( - old_dt_write_h[:, None] - + (write_offset + safe_k_new)[None, :] * stride_old_dt_T, - mask=hkn_mask, other=0.0, - ).to(tl.float32) - dA_cumsum_at_kn = tl.load( - old_dA_cumsum_write_h[:, None] - + (write_offset + safe_k_new)[None, :] * stride_old_dA_cumsum_T, - mask=hkn_mask, other=0.0, - ).to(tl.float32) - - # decay_vec_full[t] = exp(continuous_cumsum[PNAT+t]) — directly the - # continuous value now stored at buffer position write_offset+t. Was - # decomposed as total_decay * exp(per_step_new[t]) when the buffer held - # per-step (non-continuous) cumsum; with the continuity fix the value - # IS continuous_cumsum[PNAT+t] so no decomposition is needed. - decay_vec_full_block = tl.exp(dA_cumsum_new) # (H, T) - decay_vec_addrs = ( - decay_vec_ptr - + pid_b * stride_dv_batch - + heads_block[:, None] * stride_dv_head - + offs_t[None, :] * stride_dv_t - ) # (H, T) - tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) - - # combo_block = factor_dt * exp_diff — (H, T, K). Stays in registers - # across gdc_wait. With continuous cumsum in the buffer, s_k for any k - # (old or new) is simply -continuous_cumsum[k]; exp_diff[t, k] then - # equals exp(continuous_cumsum[PNAT+t] - continuous_cumsum[k]) — the - # decay weight for token k's contribution to the output at position - # PNAT+t. No need to subtract any "total" — the dA_cumsum_new[t] term - # already carries the full prefix. - # - # Numerical note: pre-fix this kernel computed `total - old_dA[k]` - # (small-minus-small) then summed `+ dA_cumsum_new[t]` (also small, - # per-step). Post-fix `s_k = -old_dA[k]` is large-magnitude positive - # and `dA_cumsum_new[t]` is large-magnitude negative; their sum - # cancels back to the same small value. Cancellation error is bounded - # by ulp(max_magnitude) ≈ 2^-23 · |continuous_cumsum| — negligible - # for max_window ≤ ~1024. Still one exp on the sum (not two muls of - # exps), so no overflow regression vs the original formulation. - factor_dt = tl.where(is_old_k[None, :], old_dt_all, dt_at_kn) # (H, K) - s_k = tl.where( - is_old_k[None, :], - -old_dA_cumsum_all, - -dA_cumsum_at_kn, - ) # (H, K) - # exp_diff (H, T, K) = exp(s_k (H, 1, K) + dA_cumsum_new (H, T, 1)). - exp_diff = tl.exp(s_k[:, None, :] + dA_cumsum_new[:, :, None]) - combo_block = factor_dt[:, None, :] * exp_diff # (H, T, K) - - # ---- gdc_wait: from here on we depend on conv1d's outputs ---- - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - # Conv1d outputs: B and C - C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group - B_new_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group - - C_all = tl.load( - C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - B_new_orig = tl.load( - B_new_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - B_new_shifted = tl.load( - B_new_base + safe_k_new[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, - mask=is_new_k[:, None] & n_mask[None, :], - other=0.0, - ) - # Disjoint masks: old at [0, PNAT), new at [K_NEW_SHIFT, K_NEW_SHIFT+T). - B_combined = old_B_load + B_new_shifted - raw_rect_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_combined).to(tl.bfloat16)) - - # Append new B to cache at [PNAT, PNAT+T) of write_buf (once per group). - if first_head % nheads_ngroups_ratio == 0: - old_B_write_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + write_buf * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - tl.store( - old_B_write_base - + (write_offset + offs_t)[:, None] * stride_old_B_T - + offs_n[None, :] * stride_old_B_dstate, - B_new_orig, - mask=t_mask[:, None] & n_mask[None, :], - ) - - # Causal mask (BLOCK_SIZE_T × BLOCK_SIZE_K, shared across heads). - # Approach C: new tokens at runtime [PNAT, PNAT+T) instead of K_NEW_SHIFT. - t_idx_2d = offs_t[:, None] - k_idx_2d = offs_k[None, :] - is_old_k_2d = k_idx_2d < prev_num_accepted_tokens - k_new_idx_2d = k_idx_2d - prev_num_accepted_tokens - is_new_causal_2d = (k_new_idx_2d >= 0) & (k_new_idx_2d < T) & (k_new_idx_2d <= t_idx_2d) - causal_combined = (is_old_k_2d | is_new_causal_2d) & t_mask[:, None] - - # Post-wait vectorized: combo_block (H, T, K) is still live in registers. - # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store as - # one (H, T, K) tile. - rect_CB_scaled_block = tl.where( - causal_combined[None, :, :], - raw_rect_CB[None, :, :] * combo_block, - 0.0, - ) # (H, T, K) - cb_scaled_addrs = ( - cb_scaled_ptr - + pid_b * stride_cb_batch - + heads_block[:, None, None] * stride_cb_head - + offs_t[None, :, None] * stride_cb_t - + offs_k[None, None, :] * stride_cb_j - ) # (H, T, K) - cb_store_mask_3d = ( - (offs_t[None, :, None] < BLOCK_SIZE_T) - & (offs_k[None, None, :] < BLOCK_SIZE_K) - ) # (1, T, K) → broadcasts to (H, T, K) - tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) - - -# Rectangle precompute kernel. Thin wrapper around _rectangle_precompute_impl -# that carries the @triton.heuristics for constexpr derivation; called from -# the Python wrapper on the rectangle nowrite path. -@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.jit() -def _dynamic_precompute_kernel( - # Input pointers - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - # Output pointers - cb_scaled_ptr, - decay_vec_ptr, - # Cache pointers - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # dt strides - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - # B strides - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # cb_scaled strides — wrapper allocates (T, K), so stride_cb_t = K - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # Meta-parameters - DT_SOFTPLUS: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - HEADS_PER_BLOCK: tl.constexpr, - # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. - RECTANGLE: tl.constexpr, -): - # Hoisted PDL signal: fire as the first thing every program does. - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - - pid_b = tl.program_id(axis=0) - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH - # write_checkpoint is now runtime in replay precompute, so a single - # call site handles both write and nowrite for the replay branch. - # Take rectangle only when RECTANGLE is True AND this slot doesn't - # need write; everything else funnels into replay. - if needs_write_runtime or not RECTANGLE: - _replay_precompute_impl( - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - cb_scaled_ptr, - decay_vec_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - T, - dstate, - nheads_ngroups_ratio, - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - DT_SOFTPLUS, - HAS_DT_BIAS, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - LAUNCH_WITH_PDL, - HEADS_PER_BLOCK, - needs_write_runtime, - ) - else: - _rectangle_precompute_impl( - dt_ptr, - dt_bias_ptr, - A_ptr, - B_ptr, - C_ptr, - cb_scaled_ptr, - decay_vec_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - cache_buf_idx_ptr, - prev_num_accepted_tokens_ptr, - state_batch_indices_ptr, - pad_slot_id, - T, - MAX_REPLAY_BUFFER_LENGTH, - dstate, - nheads_ngroups_ratio, - stride_dt_batch, - stride_dt_T, - stride_dt_head, - stride_dt_bias_head, - stride_A_head, - stride_B_batch, - stride_B_T, - stride_B_group, - stride_B_dstate, - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - stride_dv_batch, - stride_dv_head, - stride_dv_t, - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - DT_SOFTPLUS, - HAS_DT_BIAS, - HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, - BLOCK_SIZE_T, - BLOCK_SIZE_K, - LAUNCH_WITH_PDL, - HEADS_PER_BLOCK, - ) - - -# Main kernel: tl.dot replay + precomputed CB output. -# Grid: (cdiv(dim, M), batch, nheads). - - -@triton.jit() -def _persistent_main_impl( - # Per-work-unit indices (computed by the persistent wrapper). - # `pid_b` is the post-perm slot index (caller has already applied any - # slot permutation and slot_offset). - pid_m, - pid_b, - pid_h, - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or - # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor - # USE_TMA_STORE is enabled (kernel ignores it via constexpr). - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - rand_seed_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_WINDOW: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, - QUANT_MAX: tl.constexpr, - WRITE_CHECKPOINT: tl.constexpr, - # IS_DYNAMIC: kept in the signature for caller-side bookkeeping (the - # outer _persistent_main_kernel still inspects it to decide the slot- - # IS_DYNAMIC: when True (persistent_dynamic), is_write is per-slot from - # PNAT. When False (persistent_main), is_write is constexpr from - # WRITE_CHECKPOINT. See also WC_IS_CONSTEXPR below. - IS_DYNAMIC: tl.constexpr, - # WC_IS_CONSTEXPR: when True, force is_write = WRITE_CHECKPOINT (constexpr) - # regardless of IS_DYNAMIC. Callers in RECT=1 use this in the is_w=True - # arm of _persistent_main_kernel (we know all slots that reach this call - # need is_write=True because is_w was the PNAT-derived runtime check, and - # this arm only fires when is_w is True). Passing WRITE_CHECKPOINT=True - # as a literal at the call site + WC_IS_CONSTEXPR=True here lets the inner - # body DCE the nowrite path under IS_DYNAMIC=True too — same codegen - # quality as persistent_main mode (-3.7% measured at b=1024 dyn-shape). - # When False (RECT=0 callers, where both write and nowrite slots are - # dispatched to ONE call), use the original runtime is_write under - # IS_DYNAMIC=True; avoids the binary-doubling regression that two - # specialized calls would cause. - WC_IS_CONSTEXPR: tl.constexpr = False, - # TMA flags — picked inside body based on is_write. When is_write is - # constexpr (either IS_DYNAMIC=False or WC_IS_CONSTEXPR=True), the - # use_tma_load = USE_TMA_LOAD_WRITE if is_write else USE_TMA_LOAD_NOWRITE - # ternary constexpr-folds and only one TMA load form survives. - USE_TMA_LOAD_WRITE: tl.constexpr = False, - USE_TMA_LOAD_NOWRITE: tl.constexpr = False, - USE_TMA_STORE: tl.constexpr = False, -): - # IS_DYNAMIC: kernel-mode label, used by the OUTER _persistent_main_kernel - # to decide slot-range derivation and outer is_w dispatch strategy - # (constexpr WC for persistent_main; runtime is_w split -> 2 specialized - # impl calls for persistent_dynamic). Inside this impl, IS_DYNAMIC is - # NOT consulted at runtime -- WRITE_CHECKPOINT is the only constexpr that - # gates the write/nowrite codegen, in BOTH modes. - - # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized - # state dtype (int8 / int16 / float8e4nv) and only those. - tl.static_assert( - (QUANT_MAX > 0.0) - == ( - (state_ptr.dtype.element_ty == tl.int8) - or (state_ptr.dtype.element_ty == tl.int16) - or (state_ptr.dtype.element_ty == tl.float8e4nv) - ), - "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", - ) - - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - # Resolve is_write: see WC_IS_CONSTEXPR / IS_DYNAMIC docs in the param - # list above. Three cases: - # - WC_IS_CONSTEXPR=True (RECT=1 is_w=True arm callers): use WC - # constexpr. Caller knows the slot needs write; inner DCEs nowrite - # paths. Avoids the binary-doubling overhead that calling the impl - # twice would cause, while still constexpr-DCEing the nowrite half. - # - IS_DYNAMIC=True (RECT=0 caller, persistent_dynamic): runtime - # branch on PNAT. Both write and nowrite codegen live in one body - # (no bloat) — same as the pre-refactor behavior. - # - IS_DYNAMIC=False (persistent_main): WC constexpr from caller. - if WC_IS_CONSTEXPR: - is_write: tl.constexpr = WRITE_CHECKPOINT - elif IS_DYNAMIC: - is_write = (prev_num_accepted_tokens + T) > MAX_REPLAY_BUFFER_LENGTH - else: - is_write = WRITE_CHECKPOINT - if is_write: - write_buf = 1 - active_buf # noqa: F841 - write_offset = 0 - else: - write_buf = active_buf # noqa: F841 - write_offset = prev_num_accepted_tokens - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) - m_mask = offs_m < dim - n_mask = offs_n < dstate - t_mask = offs_t < T - - # Load state. state_tma_descriptor is a host-built tensor_descriptor - # over a flat (cache*nheads*dim, dstate) view of state when any TMA - # path is enabled; raw `state_ptr` is the underlying tensor and is - # always passed. state_ptrs / state_ptr_raw are the raw-pointer view - # used for !TMA load and store paths. offs_y is the flat row index - # for TMA load/store; computed unconditionally (cheap int math; DCE'd - # when no TMA path is reachable). - state_mask = m_mask[:, None] & n_mask[None, :] - offs_y = ( - cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) - + pid_h * dim - + pid_m * BLOCK_SIZE_M - ) - state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - ) - # Load state. Branch on is_write (constexpr = WRITE_CHECKPOINT in BOTH - # modes after the outer-dispatch refactor), then constexpr-pick TMA-vs- - # tl.load per side. Outer `if` DCE's, only the matching side's - # constexpr-gated load survives -- same compile-time picking for both - # persistent_main and persistent_dynamic (the latter dispatches at the - # outer kernel level so each impl instance sees a constexpr WC). - if is_write: - if USE_TMA_LOAD_WRITE: - state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) - else: - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) - else: - if USE_TMA_LOAD_NOWRITE: - state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) - else: - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) - if QUANT_MAX > 0.0: - state_scales_base = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - ) - decode_scale = tl.load( - state_scales_base + offs_m * stride_state_scales_dim, - mask=m_mask, - other=1.0, - ).to(tl.float32) - state = state * decode_scale[:, None] - - # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) - group_idx = pid_h // nheads_ngroups_ratio - - old_window_mask = offs_window < prev_num_accepted_tokens - - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + active_buf * stride_old_dt_dbuf - + pid_h * stride_old_dt_head - ) - old_dt_all = tl.load( - old_dt_base + offs_window * stride_old_dt_T, mask=old_window_mask, other=0.0 - ).to(tl.float32) - - old_dA_cumsum_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + active_buf * stride_old_dA_cumsum_dbuf - + pid_h * stride_old_dA_cumsum_head - ) - old_dA_cumsum_all = tl.load( - old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, - mask=old_window_mask, other=0.0, - ).to(tl.float32) - - prev_k_idx = tl.minimum( - tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 - ) - total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( - tl.float32 - ) - - coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all - - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head - old_x_all = tl.load( - old_x_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - mask=old_window_mask[:, None] & m_mask[None, :], - other=0.0, - ) - - old_B_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + active_buf * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - old_B_all = tl.load( - old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - mask=old_window_mask[:, None] & n_mask[None, :], - other=0.0, - ).to(tl.float32) - - dB_scaled = coeff[:, None] * old_B_all - - total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) - state *= total_decay - - state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) - - if is_write: - if USE_RS_ROUNDING: - # Generate random tensor for stochastic rounding. The amount of - # randomness needed depends on the SR codegen path: - # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) - # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) - # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs - # int16 SR (24b + bitrev32): 1 b32 per 2 outputs - # The PTX cvt.rs.* instructions consume a single 32-bit random - # and split the bits internally for 2 or 4 conversions. Generate - # only what's actually consumed and broadcast to fill the unused - # slots — saves Philox rounds proportionally. - if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: - RAND_DIVISOR: tl.constexpr = 4 # fp8 SR - elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: - RAND_DIVISOR: tl.constexpr = 4 # int8 SR - elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: - RAND_DIVISOR: tl.constexpr = 2 # int16 SR - elif QUANT_MAX == 0.0: - RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) - else: - RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized - - rand_seed = tl.load(rand_seed_ptr) - base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head - # Number of unique randoms per row = dstate / RAND_DIVISOR. - # randint4x emits 4 randoms per offset, so use that / 4 offsets. - offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) - rand_offsets_q = ( - base_rand - + offs_m[:, None] * stride_state_dim - + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) - ) # (M, dstate / (4*RAND_DIVISOR)) - if PHILOX_ROUNDS > 0: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) - else: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) - r01 = tl.join(r0, r1) - r23 = tl.join(r2, r3) - r0123 = tl.join(r01, r23) - rand_compact = tl.reshape( - r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) - ) - # Broadcast each unique rand to RAND_DIVISOR adjacent positions. - # Pack-group (pack=2 fp16 / pack=4 fp8) consumes adjacent positions; - # the unique rand lands at the asm's read slot; duplicates feed - # the dead slots. Triton's broadcast_to is stride-0 in IR. - if RAND_DIVISOR > 1: - rand_3d = rand_compact[:, :, None] - rand_3d = tl.broadcast_to( - rand_3d, - (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), - ) - rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) - else: - rand = rand_compact - - if QUANT_MAX > 0.0: - amax = tl.max(tl.abs(state), axis=1) - encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) - decode_scale = 1.0 / encode_scale - state_scales_ptrs = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - + offs_m * stride_state_scales_dim - ) - tl.store(state_scales_ptrs, decode_scale, mask=m_mask) - state_q = state * encode_scale[:, None] - if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): - _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_q_fp8sr) - else: - tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) - else: - if USE_RS_ROUNDING: - tl.static_assert( - (state_ptrs.dtype.element_ty == tl.int8) - or (state_ptrs.dtype.element_ty == tl.int16), - "Quantized SR fall-through expects int8 or int16; " - "fp8 SR is handled by the prior branch.", - ) - if state_ptrs.dtype.element_ty == tl.int8: - state_q = _stochastic_round_int8_packed( - state_q, rand, offs_n[None, :] - ) - else: - state_q = _stochastic_round_int16_packed( - state_q, rand, offs_n[None, :] - ) - elif state_ptrs.dtype.element_ty != tl.float8e4nv: - tl.static_assert( - (state_ptrs.dtype.element_ty == tl.int8) - or (state_ptrs.dtype.element_ty == tl.int16), - "Quantized RN with explicit round() expects int8 or int16.", - ) - state_q = tl.extra.cuda.libdevice.round(state_q) - state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) - _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_q_cast) - else: - tl.store(state_ptrs, _state_q_cast, mask=state_mask) - elif USE_RS_ROUNDING: - tl.static_assert( - state_ptrs.dtype.element_ty == tl.float16, - "Non-quantized SR only supports fp16 state.", - ) - _state_sr = _stochastic_round_fp16x2(state, rand) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_sr) - else: - tl.store(state_ptrs, _state_sr, mask=state_mask) - else: - _state_cast = state.to(state_ptrs.dtype.element_ty) - if USE_TMA_STORE: - state_tma_descriptor.store([offs_y, 0], _state_cast) - else: - tl.store(state_ptrs, _state_cast, mask=state_mask) - - # Phase 2: Output using precomputed CB_scaled and decay_vec - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head - C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group - if HAS_Z: - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - - if HAS_D: - D = tl.load( - D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - C_all = tl.load( - C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - - x_all = tl.load( - x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=t_mask[:, None] & m_mask[None, :], - other=0.0, - ) - tl.store( - old_x_base - + (write_offset + offs_t)[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - x_all, - mask=t_mask[:, None] & m_mask[None, :], - ) - x_all = x_all.to(tl.float32) - - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head - CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), - other=0.0, - ).to(tl.float32) - - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( - tl.float32 - ) - - init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] - cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) - out_all = init_out + cb_out - - if HAS_D: - out_all = out_all + x_all * D[None, :] - - if HAS_Z: - z_all = tl.load( - z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, - mask=t_mask[:, None] & m_mask[None, :], other=0.0, - ).to(tl.float32) - out_all_z = out_all * z_all * tl.sigmoid(z_all) - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) - else: - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) - - -# `_persistent_rectangle_impl`: rectangle nowrite path for the persistent -# kernel. Body is a copy of `_rectangle_main_impl` with `pid_m`/`pid_b`/`pid_h` -# lifted to args (same pattern as `_persistent_main_impl` vs `_replay_main_impl`). -# Called only for nowrite slots when the kernel runs with RECTANGLE=True. -# Dropped from the rect impl: LAUNCH_DEPENDENT_KERNELS / USE_PERM / REVERSE_PERM -# (kernel-level, signalled once at top); slot_perm_ptr (kernel resolves perm). -@triton.jit() -def _persistent_rectangle_impl( - # Per-work-unit indices (computed by the persistent wrapper). - pid_m, - pid_b, - pid_h, - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as - # replay path). Used when USE_TMA_LOAD; ignored otherwise. - state_tma_descriptor, - state_scales_ptr, # only consulted when QUANT_MAX > 0 - old_x_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - pad_slot_id, - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides (rectangle (batch, nheads, T, K)) - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - QUANT_MAX: tl.constexpr, - USE_TMA_LOAD: tl.constexpr = False, -): - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - # Nowrite-only: write_offset = PNAT (new tokens append at [PNAT, PNAT+T)). - buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - write_offset = prev_num_accepted_tokens - - # Static rectangle K-axis layout (matches precompute). - K_NEW_SHIFT: tl.constexpr = MAX_REPLAY_BUFFER_LENGTH - T - - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) - offs_k = tl.arange(0, BLOCK_SIZE_K) - m_mask = offs_m < dim - n_mask = offs_n < dstate - t_mask = offs_t < T - - # K-axis masks (approach C: PNAT-runtime offset, matches precompute). - is_old_k = offs_k < prev_num_accepted_tokens - safe_old_k = tl.where(is_old_k, offs_k, 0) - k_new_idx = offs_k - prev_num_accepted_tokens - is_new_k = (k_new_idx >= 0) & (k_new_idx < T) - safe_k_new = tl.where(is_new_k, k_new_idx, 0) - - # Load state. Quant scale hoist: defer `* decode_scale` post-matmul. - if USE_TMA_LOAD: - offs_y = ( - cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) - + pid_h * dim - + pid_m * BLOCK_SIZE_M - ) - state = state_tma_descriptor.load([offs_y, 0]) - else: - state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate - ) - state_mask = m_mask[:, None] & n_mask[None, :] - state = tl.load(state_ptrs, mask=state_mask, other=0.0) - if QUANT_MAX > 0.0: - state_scales_base = ( - state_scales_ptr - + cache_batch_idx * stride_state_scales_cache - + pid_h * stride_state_scales_head - ) - decode_scale = tl.load( - state_scales_base + offs_m * stride_state_scales_dim, - mask=m_mask, other=1.0, - ).to(tl.float32) - else: - state = state.to(tl.float32) - - # Group / pointer offset setup - group_idx = pid_h // nheads_ngroups_ratio - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head - C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group - if HAS_Z: - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head - - if HAS_D: - D = tl.load( - D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - - # Hoist: old_x doesn't depend on conv1d/precompute; load before gdc_wait. - old_x_load = tl.load( - old_x_base - + safe_old_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - mask=is_old_k[:, None] & m_mask[None, :], - other=0.0, - ).to(tl.float32) - - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - - C_all = tl.load( - C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) - x_K = tl.load( - x_ptr + safe_k_new[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=is_new_k[:, None] & m_mask[None, :], - other=0.0, - ) - tl.store( - old_x_base - + offs_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, - x_K, - mask=is_new_k[:, None] & m_mask[None, :], - ) - - x_K_f32 = x_K.to(tl.float32) - x_combined = old_x_load + x_K_f32 - - if HAS_D or HAS_Z: - sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) - x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) - else: - x_all = x_K_f32 # placeholder; unused - - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head - CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_k[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_k[None, :] < BLOCK_SIZE_K), - other=0.0, - ).to(tl.float32) - - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec_full = tl.load( - decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 - ).to(tl.float32) - - state_out = ( - tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) - * decay_vec_full[:, None] - ) - if QUANT_MAX > 0.0: - state_out = state_out * decode_scale[None, :] - - token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_combined.to(tl.bfloat16)) - - out_all = state_out + token_out - - if HAS_D: - out_all = out_all + x_all * D[None, :] - - if HAS_Z: - z_all = tl.load( - z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, - mask=t_mask[:, None] & m_mask[None, :], other=0.0, - ).to(tl.float32) - out_all_z = out_all * z_all * tl.sigmoid(z_all) - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all_z, mask=t_mask[:, None] & m_mask[None, :]) - else: - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) - - -# Persistent main kernel: 1D grid, persistent CTA loop. -# Heuristics mirror those of `_checkpointing_main_kernel`. -@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) -@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics( - {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} -) -@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) -@triton.heuristics( - {"BLOCK_SIZE_WINDOW": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} -) -@triton.heuristics( - {"NUM_PID_M_BLOCKS": lambda args: triton.cdiv(args["dim"], args["BLOCK_SIZE_M"])} -) -@triton.jit() -def _persistent_main_kernel( - # Pointers - state_ptr, - # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view. - # Shared across BOTH the replay path (consumed by _persistent_main_impl - # when USE_TMA_LOAD_*/STORE) AND the rectangle path (consumed by - # _persistent_rectangle_impl when USE_TMA_LOAD) — same descriptor, same - # block_shape, just gated by separate constexprs per impl. Wrapper sets - # this to a TensorDescriptor when ANY of the three TMA flags is on, else - # to `state_ptr` (raw); each impl ignores it via its own constexpr when - # not consuming it. - state_tma_descriptor, - state_scales_ptr, - old_x_ptr, - old_B_ptr, - old_dt_ptr, - old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, - cache_buf_idx_ptr, - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - cb_scaled_ptr, - decay_vec_ptr, - state_batch_indices_ptr, - slot_perm_ptr, - rand_seed_ptr, - pad_slot_id, - # Persistent-loop work-distribution scalars. Caller pre-sorts the batch - # write-first; the kernel uses (n_writes, batch_total, WRITE_CHECKPOINT) - # to derive its own slot range. Write half processes [0, n_writes), - # nowrite half processes [n_writes, batch_total). - # - # n_writes_ptr is a device pointer to a (1,) int32 tensor. Reading - # from device memory (rather than taking a Python int kernel arg) is - # required so mix-mode benchmarking can vary n_writes per iter inside - # a captured CUDA graph — the source tensor's contents change, the - # pointer doesn't. Cost: one int load per kernel launch (~negligible). - # When IS_DYNAMIC=True the value is unused (Triton DCEs the load). - n_writes_ptr, # int32 *: device-side count of write-mode slots - batch_total, # int32: total slot count - nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) - # Dimensions - T: tl.constexpr, - MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, - dim: tl.constexpr, - dstate: tl.constexpr, - nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # state_scales strides - stride_state_scales_cache, - stride_state_scales_head, - stride_state_scales_dim, - # old_x strides - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, - # C strides - stride_C_batch, - stride_C_T, - stride_C_group, - stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides - stride_cb_batch, - stride_cb_head, - stride_cb_t, - stride_cb_j, - # decay_vec strides - stride_dv_batch, - stride_dv_head, - stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, - BLOCK_SIZE_DSTATE: tl.constexpr, - BLOCK_SIZE_T: tl.constexpr, - BLOCK_SIZE_WINDOW: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, - QUANT_MAX: tl.constexpr, - WRITE_CHECKPOINT: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, - USE_PERM: tl.constexpr, - # NUM_PERSISTENT: runtime int (not constexpr). Used ONLY as the loop - # stride in `tl.range(pid, total_work, NUM_PERSISTENT, ...)`. Making it - # runtime collapses the cta_per_sm tuning dim from the kernel's compile - # signature: 8 CPS values used to mean 8x recompiles; now they share one - # compiled kernel. Work decomposition (pid_m, pid_b_local, pid_h) does - # NOT depend on NUM_PERSISTENT — it uses constexpr NUM_PID_M_BLOCKS and - # runtime n_slots_local — so loop unrolling and flatten=/num_stages=/ - # warp_specialize= optimizations on `tl.range` operate independently of - # the stride value. - NUM_PERSISTENT, - NUM_LOOP_STAGES: tl.constexpr, - NUM_PID_M_BLOCKS: tl.constexpr, - FLATTEN: tl.constexpr, - WARP_SPECIALIZE: tl.constexpr, - IS_DYNAMIC: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr = 16, # rectangle K-axis (heuristic-derived) - RECTANGLE: tl.constexpr = False, # when True, dispatch nowrite slots to _persistent_rectangle_impl - # 3 TMA toggles per the 3 live paths per-compilation: - # USE_TMA_LOAD_WRITE — replay-style state load when is_write - # USE_TMA_LOAD_NOWRITE — nowrite-path state load (rect when RECTANGLE, - # else replay-nowrite) - # USE_TMA_STORE — replay-style state store (only fires on write - # path; no-op when not is_write) - # Wrapper picks USE_TMA_LOAD_NOWRITE = _use_tma_rect_load (if rectangle) - # or _use_tma_replay_nowrite_load (if not). - USE_TMA_LOAD_WRITE: tl.constexpr = False, - USE_TMA_LOAD_NOWRITE: tl.constexpr = False, - USE_TMA_STORE: tl.constexpr = False, -): - # PDL signal: fire once at kernel entry (not per work unit). - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() - - # Load runtime n_writes from device memory. Read once at kernel entry; - # used only by the !IS_DYNAMIC slot-range derivation below. Triton - # DCEs the load when IS_DYNAMIC=True (n_writes is dead there). - n_writes = tl.load(n_writes_ptr) - - # Derive this kernel's slot range. Two modes: - # IS_DYNAMIC=False (persistent_main): caller pre-sorts and splits halves; - # slot range is [0, n_writes) when WRITE_CHECKPOINT else [n_writes, batch_total) - # IS_DYNAMIC=True (persistent_dynamic): single launch covers full batch; - # each work-item dispatches via runtime PNAT check inside the impl. - if IS_DYNAMIC: - slot_lo = 0 - slot_hi = batch_total - else: - if WRITE_CHECKPOINT: - slot_lo = 0 - slot_hi = n_writes - else: - slot_lo = n_writes - slot_hi = batch_total - n_slots_local = slot_hi - slot_lo - - pid = tl.program_id(axis=0) - total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads - - # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) - # with pid_m varying fastest (M-tile cache locality on state load), then - # slot, then head — mirrors the existing 3D grid's axis ordering - # (axis=0 fastest = pid_m). - for tile_id in tl.range( - pid, total_work, NUM_PERSISTENT, - flatten=FLATTEN, num_stages=NUM_LOOP_STAGES, warp_specialize=WARP_SPECIALIZE, - ): - pid_m = tile_id % NUM_PID_M_BLOCKS - pid_b_local = (tile_id // NUM_PID_M_BLOCKS) % n_slots_local - pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) - # Translate local slot index → global slot index. When USE_PERM is - # set, the caller-provided slot_perm gives the original slot index - # for the post-sort position. - pid_b_grid = pid_b_local + slot_lo - if USE_PERM: - pid_b = tl.load(slot_perm_ptr + pid_b_grid) - else: - pid_b = pid_b_grid - - # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle - # impl. Pre-load cache_batch_idx + PNAT here only for the RECTANGLE - # path's branch decision. Both impls re-load and handle pad_slot_id - # internally (Triton's L1 cache makes the duplicate loads ~free). - if RECTANGLE: - if HAS_CACHE_BATCH_INDICES: - cbi_pre = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - is_pad = cbi_pre == pad_slot_id - else: - cbi_pre = pid_b.to(tl.int64) - is_pad = False - if not is_pad: - pnat_pre = tl.load(prev_num_accepted_tokens_ptr + cbi_pre) - if IS_DYNAMIC: - is_w = (pnat_pre + T) > MAX_REPLAY_BUFFER_LENGTH - else: - is_w = WRITE_CHECKPOINT - if is_w: - # Pass WRITE_CHECKPOINT=True constexpr to specialize this - # impl call for the write path. Under IS_DYNAMIC=True, the - # kernel-level WRITE_CHECKPOINT is False (launcher default), - # but the OUTER is_w branch we are inside narrows the - # runtime path to writes-only, so we override to True here - # so the impl's constexpr-gated `if is_write:` blocks DCE - # to the write-only codegen. Under IS_DYNAMIC=False - # (persistent_main), the kernel-level WRITE_CHECKPOINT is - # itself True for this half (write half launches with - # WC=True), and the outer is_w = WRITE_CHECKPOINT = True - # constexpr-folds; passing literal True here is consistent - # and constexpr-equivalent. - _persistent_main_impl( - pid_m, pid_b, pid_h, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, - stride_old_B_group, stride_old_B_dstate, - stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, - stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) - True, # WC_IS_CONSTEXPR — force inner to use WC constexpr - # 3 TMA flags: write-load fires here (we're in the - # is_write branch), nowrite-load is dead (no slot - # reaches it), store fires (write path). - USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, - ) - else: - # Rectangle nowrite: pass state_ptr (raw, always) + - # state_tma_descriptor (the single unified descriptor — - # same memory replay paths use). Rect impl gates use - # of the descriptor via its USE_TMA_LOAD constexpr. - _persistent_rectangle_impl( - pid_m, pid_b, pid_h, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - state_batch_indices_ptr, pad_slot_id, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, - LAUNCH_WITH_PDL, QUANT_MAX, - USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle - ) - # else: pad slot — skip both impls (both would early-return anyway) - else: - # No rectangle path — single _persistent_main_impl call covers - # both write and nowrite slots via WC constexpr (non-dynamic) or - # runtime is_write (IS_DYNAMIC=True). Pass all 3 TMA flags; - # impl picks USE_TMA_LOAD_WRITE vs USE_TMA_LOAD_NOWRITE based on - # its computed is_write — constexpr-folds when is_write is - # constexpr (non-dyn), runtime branch when IS_DYNAMIC=True. - # (Reverted from outer two-call dispatch: that doubled the - # compiled body size under IS_DYNAMIC=True and regressed RECT=0 - # perf by ~+24%.) - _persistent_main_impl( - pid_m, pid_b, pid_h, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - state_batch_indices_ptr, rand_seed_ptr, pad_slot_id, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, - stride_old_B_group, stride_old_B_dstate, - stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, - stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - WRITE_CHECKPOINT, IS_DYNAMIC, - False, # WC_IS_CONSTEXPR=False — RECT=0 has both write/nowrite slots in one call - USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, - ) - - -# ============================================================================ -# Python wrapper -# ============================================================================ - - -_QUANT_MAX_BY_DTYPE = { - torch.int8: 127.0, - torch.int16: 32767.0, - torch.float8_e4m3fn: 448.0, -} - - -# --------------------------------------------------------------------------- -# Default tunings — looked up by (effective_batch, dtype, sr) when the caller -# leaves mode/knobs as None. -# -# Effective batch = raw_batch × nheads_per_rank. Our sweep was at TP=8 with -# the standard Mamba2 nheads; at call time we compute it from the input -# tensor shape so callers at other TP / nheads pick up the right cell. -# -# Schema: dict[(dtype_str, sr_str)] → list[(eff_batch_threshold, mode, knobs)] -# sorted by threshold ascending. Lookup finds the first threshold ≥ eff_b -# (so missing intermediate batches fall up to the next tuned cell). If -# eff_b exceeds the largest threshold, use the largest entry. -# -# Each `knobs` dict only contains keys for the chosen mode; the wrapper -# unpacks them with the same name as the matching kwargs. Caller-provided -# kwargs always win over table values. -# -# This table is intentionally NOT parameterized by T or max_window. Our -# sweep was T=6, max_window=16. Callers outside that regime silently get -# the same numbers — they may be suboptimal but they're correct. -# -# Source: audit_v2.py --emit-tuning. Auto-generated from per-cell search -# winners (best of pd / pm by bucket_expected_renorm). Sweep was TP=8 with -# NHEADS=128 → nheads_per_rank=16; thresholds are in effective_batch units. -# Missing dtype/SR combos (fp16/RN, int8/RN, fp8/*) fall back via the -# _resolve_tuning chain — RN→SR for same dtype, then fp8→int8/SR. -_DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { - ("fp32", "RN"): [ - ( 16, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 8, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.22us - ( 32, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.17us - ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.08us - ( 128, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=8, score=9.00us - ( 256, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=16, score=10.92us - ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=13.53us - ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=19.50us - ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=30.28us - ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=50.32us - ( 8192, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=90.99us - (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=171.69us - ], - ("fp16", "SR"): [ - ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.16us - ( 32, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.01us - ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=7.95us - ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 4, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=8.87us - ( 256, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 2, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.28us - ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=12.90us - ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=16.71us - ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=25.71us - ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=39.80us - ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=71.34us - (16384, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=133.51us - ], - ("int8", "SR"): [ - ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.34us - ( 32, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.36us - ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 2, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.40us - ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 10, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=9.37us - ( 256, "persistent_dynamic", {'_block_size_m': 16, '_cta_per_sm': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages': 1, '_num_stages': 4, '_num_warps': 1, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.02us - ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=32, score=13.15us - ( 1024, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=17.82us - ( 2048, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=27.01us - ( 4096, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=43.23us - ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=77.01us - (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=140.43us - ], -} - - -# Knob names that map between the modes' single-value (pd) and split-value -# (pm) namespaces. Used by `_bridge_tuning_knobs` when caller forces a mode -# different from the table's recommendation. -_PD_TO_PM_SPLIT_MAP = { # pd unsplit knob → (pm_write_knob, pm_nowrite_knob) - "_block_size_m": ("_block_size_m_write", "_block_size_m_nowrite"), - "_num_warps": ("_num_warps_write", "_num_warps_nowrite"), - "_num_stages": ("_num_stages_write", "_num_stages_nowrite"), - # CPS / LS are persistent-loop knobs; pd uses _cta_per_sm + _num_loop_stages - # as unsplit, pm uses _cta_per_sm_write/_nowrite + _num_loop_stages_write/_nowrite. - "_cta_per_sm": ("_cta_per_sm_write", "_cta_per_sm_nowrite"), - "_num_loop_stages": ("_num_loop_stages_write", "_num_loop_stages_nowrite"), -} - - -def _bridge_tuning_knobs(knobs: dict, from_mode: str, to_mode: str) -> dict: - """Convert a tuning dict between pd ↔ pm knob namespaces. - - pd → pm: copy each unsplit value to both write/nowrite split knobs; drop - the unsplit form (pm doesn't read it). - pm → pd: take the nowrite split value as the unsplit knob; drop the - write/nowrite split forms (pd doesn't read them). - Shape knobs that exist in both modes (_heads_per_block, _flatten, - _warp_specialize, TMA flags, rectangle_for_nowrite) carry over unchanged. - """ - out = dict(knobs) - if from_mode == "persistent_dynamic" and to_mode == "persistent_main": - for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): - if unsplit in out: - out.setdefault(pm_w, out[unsplit]) - out.setdefault(pm_nw, out[unsplit]) - del out[unsplit] - elif from_mode == "persistent_main" and to_mode == "persistent_dynamic": - for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): - if pm_nw in out: - out.setdefault(unsplit, out[pm_nw]) - out.pop(pm_w, None) - out.pop(pm_nw, None) - return out - - -def _resolve_tuning( - batch: int, nheads_per_rank: int, dt_str: str, sr_str: str, -) -> tuple[str, dict] | None: - """Look up the default mode + knobs for this (eff_batch, dt, sr) cell. - - Returns (mode, knobs_dict) or None if the table has no entry covering - this dtype/sr (including the fp8→int8/SR and dtype/RN→dtype/SR fallbacks). - Returning None lets the wrapper fall back to caller-provided kwargs or - kernel-side defaults. - """ - eff_b = batch * max(1, nheads_per_rank) - # Lookup chain. Order: - # 1. Exact (dt, sr). - # 2. (dt, SR) if RN missing for that dtype. - # 3. Cross-dtype fallback for dtypes we haven't tuned: - # bf16 / int16 → fp16/SR - # fp8 → int8/SR - # Unknown dtype → raise. - valid_dtypes = {"fp32", "fp16", "bf16", "int8", "int16", "fp8"} - if dt_str not in valid_dtypes: - raise ValueError( - f"checkpointing_state_update: unsupported state dtype {dt_str!r}; " - f"expected one of {sorted(valid_dtypes)}" - ) - keys_to_try = [(dt_str, sr_str)] - if sr_str == "RN": - keys_to_try.append((dt_str, "SR")) - if dt_str in ("bf16", "int16"): - keys_to_try.append(("fp16", "SR")) - elif dt_str == "fp8": - keys_to_try.append(("int8", "SR")) - entries = None - for k in keys_to_try: - if k in _DEFAULT_TUNING: - entries = _DEFAULT_TUNING[k] - break - if entries is None: - return None - # Find first threshold ≥ eff_b; if none, use largest entry. - for thresh, mode, knobs in entries: - if eff_b <= thresh: - return mode, dict(knobs) - thresh, mode, knobs = entries[-1] - return mode, dict(knobs) - - -def checkpointing_state_update( - state: torch.Tensor, - old_x: torch.Tensor, - old_B: torch.Tensor, - old_dt: torch.Tensor, - old_dA_cumsum: torch.Tensor, - cache_buf_idx: torch.Tensor, - prev_num_accepted_tokens: torch.Tensor, - x: torch.Tensor, - dt: torch.Tensor, - A: torch.Tensor, - B: torch.Tensor, - C: torch.Tensor, - out: torch.Tensor, - # Required persistent-mode plumbing (REQUIRED for both pd and pm; pd - # ignores both internally but the wrapper still demands them): - # n_writes : (1,) int32 device tensor with the count of write-mode - # slots in the batch. pm uses it to size the two halves; - # pd ignores it (per-slot runtime PNAT check). - # slot_perm : (batch,) int32 device tensor remapping grid pid → slot. - # pm uses it to cluster writes first (kernel grid step is - # write_half then nowrite_half); pd ignores it. Callers - # that don't care about ordering should pass arange(batch). - n_writes: torch.Tensor, - slot_perm: torch.Tensor, - D: torch.Tensor | None = None, - z: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - dt_softplus: bool = False, - state_batch_indices: torch.Tensor | None = None, - pad_slot_id: int = PAD_SLOT_ID, - rand_seed: torch.Tensor | None = None, - philox_rounds: int = 10, - state_scales: torch.Tensor | None = None, - launch_with_pdl=False, - use_internal_pdl=True, - write_checkpoint: bool = True, - rectangle_for_nowrite: bool | None = None, - mode: str | None = None, - _block_size_m: int | None = None, - _num_warps: int | None = None, - _num_stages: int | None = None, - _precompute_num_warps: int | None = None, - _precompute_num_stages: int | None = None, - _heads_per_block: int | None = None, - _maxnreg: int | None = None, - _num_ctas: int | None = None, - # Per-main knobs (override shared values for one half of the dl-family / - # persistent_main launches). Default None = tied to the shared value - # (backward compat). The two main kernels (write vs nowrite) have - # different per-slot work — write does a state shift + store, nowrite - # just appends — so the optimum (M, W, S, H) can differ. Precompute - # knobs are intentionally NOT split: shared precompute wins (cheaper - # launch, hotter precompute outputs in L2). Persistent CPS / LS knobs - # are also split per-main since the two persistent_main launches have - # different grid sizes. - _block_size_m_write: int | None = None, - _block_size_m_nowrite: int | None = None, - _num_warps_write: int | None = None, - _num_warps_nowrite: int | None = None, - _num_stages_write: int | None = None, - _num_stages_nowrite: int | None = None, - # Note: heads_per_block / precompute_num_warps are NOT split — they only - # affect the precompute kernel, which is shared across write/nowrite. - # TMA state-tensor toggles — 4 independent paths (see CHECKPOINTING_DESIGN.md - # item #17 for measured perf profiles). Each is False=raw load/store, True= - # use a host-built TMA tensor_descriptor for that path. - _use_tma_rect_load: bool | None = None, # rect kernel's state load (nowrite-only) - _use_tma_replay_write_load: bool | None = None, # replay-style state load when WC=True - _use_tma_replay_write_store: bool | None = None, # replay-style state store when WC=True - _use_tma_replay_nowrite_load: bool | None = None, # replay-style state load when WC=False - # Persistent-mode tuning kwargs (consulted for both pd and pm; pd uses - # _cta_per_sm / _num_loop_stages, pm uses the _write/_nowrite splits): - # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally - # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. - # _num_loop_stages : int — `num_stages` arg on the inner `tl.range(...)` - # persistent loop. Note: this is loop-level, NOT the kernel-arg - # `num_stages` (which only pipelines dot-feeding loads). - # _flatten : bool — `flatten` arg on `tl.range(...)`. - # _warp_specialize : bool — `warp_specialize` arg on `tl.range(...)`. - _cta_per_sm: int | None = None, - _num_loop_stages: int | None = None, - _flatten: bool | None = None, - _warp_specialize: bool | None = None, - # Per-main persistent-specific knobs. Same rationale as the BLOCK_SIZE_M - # split above: the two persistent_main launches (write half vs nowrite - # half) have different grid sizes and per-work-item costs, so they may - # want different cta_per_sm / num_loop_stages. - _cta_per_sm_write: int | None = None, - _cta_per_sm_nowrite: int | None = None, - _num_loop_stages_write: int | None = None, - _num_loop_stages_nowrite: int | None = None, -): - """ - Replay SSM state update with precomputed CB and tl.dot fast-forward. - - Two-kernel architecture: - 1. Precompute kernel: computes CB_scaled and decay_vec from B, C, dt, A. - Writes processed dt/dA_cumsum/B to double-buffered cache for next step. - 2. Main kernel: replays old tokens via tl.dot fast-forward on cached data, - then computes output using precomputed CB_scaled and new x/C inputs. - - PDL (Programmatic Dependent Launch) chain: - conv1d → (external PDL) → precompute → (internal PDL) → main - External PDL: precompute starts while conv1d is running; gdc_wait() - in precompute blocks until conv1d completes before loading B/C. - Internal PDL: main starts while precompute is running; main's replay - phase uses only cached data from the previous step. gdc_wait() in - main blocks until precompute completes before loading conv1d outputs - (x, C) and precompute outputs (CB_scaled, decay_vec). - - Uses double-buffered cache tensors. cache_buf_idx[slot] indicates which - buffer (0 or 1) to READ from for replay. The WRITE buffer is 1 - read. - Caller must flip cache_buf_idx[slot] after each call. - - Arguments: - state: (cache, nheads, dim, dstate) in-place. After the call, contains - the state after replaying prev_num_accepted_tokens old tokens. - old_x: (cache, T, nheads, dim) bf16 — old x cache (single-buffered). - old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. - old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. - old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. - cache_buf_idx: (cache,) int32 — which buffer to read (0 or 1). - prev_num_accepted_tokens: (cache,) int32. - x: (batch, T, nheads, dim) new token inputs. - dt: (batch, T, nheads, dim) with stride(-1)==0 (tie_hdim). - A: (nheads, dim, dstate) with stride(-1)==0, stride(-2)==0 (tie_hdim). - B: (batch, T, ngroups, dstate). - C: (batch, T, ngroups, dstate). - out: (batch, T, nheads, dim) preallocated output. - D: (nheads, dim) optional feed-through parameter. - z: (batch, T, nheads, dim) optional silu gate. - dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). - state_batch_indices: (batch,) optional cache slot mapping. - rand_seed: optional single-element int64 CUDA tensor for Philox PRNG seed. - When provided, state is stochastically rounded on store. Supported - for state.dtype in (fp16, int8, int16, fp8_e4m3fn); other dtypes - silently use deterministic rounding. fp16+SR and fp8+SR both - require sm_100a (Blackwell B200+) — wrapper asserts this loudly. - philox_rounds: number of Philox PRNG rounds (default 10). - state_scales: required when state.dtype in (int8, int16, fp8_e4m3fn). - Shape (cache_size, nheads, dim), fp32. Per-(head, dim) channel - decode scale (= 1 / encode_scale). The kernel writes scales on - checkpoint steps and reads them on load (broadcast over dstate). - Ignored for non-quantized state dtypes. - launch_with_pdl: enable external PDL (conv1d → precompute chain). - Defaults False; caller opts in when the upstream chain is PDL-safe. - Ignored on hardware that doesn't support PDL (sm < 90). - use_internal_pdl: enable internal PDL (precompute → main overlap). - Defaults True; override for testing only. - Ignored on hardware that doesn't support PDL (sm < 90). - - _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, - _precompute_num_warps, _precompute_num_stages, _heads_per_block, - _maxnreg, _num_ctas) are benchmark-only overrides; production callers - should leave them None to use the heuristic-tuned defaults. - """ - # PDL needs sm >= 90. - if get_sm_version() < 90: - launch_with_pdl = False - use_internal_pdl = False - - # Mode selection: - # mode=None (default): look up the table-tuned mode + knobs for this - # (effective_batch, dtype, sr) cell. See `_resolve_tuning` above. - # mode="persistent_dynamic": single persistent-CTA kernel covering the - # full batch. Each work-item dispatches via runtime PNAT check - # (is_write = (pnat + T) > MAX). No write/nowrite split. - # slot_perm is honored but optional. write_checkpoint is ignored. - # mode="persistent_main": persistent-CTA kernel with two launches - # (write half + nowrite half). Caller MUST pre-sort slot_perm - # write-first; the n_writes tensor partitions the persistent loop - # into the two halves with the right WRITE_CHECKPOINT constexpr - # each time. RECTANGLE constexpr (= rectangle_for_nowrite) picks - # rect vs replay for the nowrite half. write_checkpoint is ignored. - # Note: mode-and-knob resolution from the default-tuning table happens - # below, after we have `batch` and `nheads`. - - # --- Hardware support gates --- - # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX - # instructions (Ada Lovelace introduced them; Hopper/Blackwell carry them). - if state.dtype == torch.float8_e4m3fn: - assert get_sm_version() >= 89, ( - "fp8_e4m3fn state requires SM 89+ (Ada Lovelace / Hopper / Blackwell) " - f"for fp32↔fp8 cvt PTX instructions; current SM is {get_sm_version()}." - ) - - # PTX cvt.rs.* (stochastic rounding) family lands on Blackwell only. - # Wrapper fails loud; framework decides fall-back (e.g. drop SR, use RN). - # int8 / int16 SR uses pure-Triton libdevice.floor + uniform noise — no - # PTX SR instruction needed, runs anywhere. - if rand_seed is not None: - if state.dtype == torch.float16: - assert get_sm_version() >= 100, ( - "fp16 stochastic rounding (PTX cvt.rs.f16x2.f32) requires " - f"sm_100a (Blackwell B200+); current SM is {get_sm_version()}." - ) - elif state.dtype == torch.float8_e4m3fn: - assert get_sm_version() >= 100, ( - "fp8 stochastic rounding (PTX cvt.rs.satfinite.e4m3x4.f32) " - f"requires sm_100a (Blackwell B200+); current SM is {get_sm_version()}." - ) - - # --- Unsqueeze inputs to canonical shapes --- - if state.dim() == 3: - state = state.unsqueeze(1) - if x.dim() == 2: - x = x.unsqueeze(1) - if x.dim() == 3: - x = x.unsqueeze(1) - if dt.dim() == 2: - dt = dt.unsqueeze(1) - if dt.dim() == 3: - dt = dt.unsqueeze(1) - if A.dim() == 2: - A = A.unsqueeze(0) - if B.dim() == 2: - B = B.unsqueeze(1) - if B.dim() == 3: - B = B.unsqueeze(1) - if C.dim() == 2: - C = C.unsqueeze(1) - if C.dim() == 3: - C = C.unsqueeze(1) - if D is not None and D.dim() == 1: - D = D.unsqueeze(0) - if z is not None: - if z.dim() == 2: - z = z.unsqueeze(1) - if z.dim() == 3: - z = z.unsqueeze(1) - if dt_bias is not None and dt_bias.dim() == 1: - dt_bias = dt_bias.unsqueeze(0) - if out.dim() == 2: - out = out.unsqueeze(1) - if out.dim() == 3: - out = out.unsqueeze(1) - - cache_size, nheads, dim, dstate = state.shape - batch, T, _, _ = x.shape - ngroups = B.shape[2] - assert nheads % ngroups == 0 - - # --- Quantization plumbing (needed for SR/RN classification below) --- - # QUANT_MAX > 0 ⇔ state is int8 / int16 / fp8_e4m3fn. Kernel-entry - # static_assert on the Triton side mirrors this invariant. - quant_max = _QUANT_MAX_BY_DTYPE.get(state.dtype, 0.0) - is_quantized = quant_max > 0.0 - - # --- Default-tuning lookup --- - # Resolve (mode, knobs) from the table when caller leaves them None. - # Caller-provided kwargs always win. If the caller forces a mode that - # differs from the table's recommendation for this cell, we BRIDGE the - # table's knobs into the forced mode's knob namespace rather than fall - # back to (likely-terrible) kernel defaults: - # table pd → forced pm: copy each unsplit pd knob (M, W, S, CPS, LS) - # to both write and nowrite split knobs. - # table pm → forced pd: take the nowrite split values (Mnw, Wnw, Snw, - # CPSnw, LSnw) as the unsplit knobs. - # Empty table → no-op (caller passes whatever, mode falls back to pd). - _dt_str = { - torch.float32: "fp32", - torch.float16: "fp16", - torch.bfloat16: "bf16", - torch.int8: "int8", - torch.int16: "int16", - torch.float8_e4m3fn: "fp8", - }.get(state.dtype, str(state.dtype)) - _sr_str = "SR" if (rand_seed is not None and is_quantized) else "RN" - _table_entry = _resolve_tuning(batch, nheads, _dt_str, _sr_str) - if _table_entry is not None: - _table_mode, _table_knobs = _table_entry - if mode is None: - mode = _table_mode - if mode != _table_mode: - # Bridge across modes — see header comment above. - _table_knobs = _bridge_tuning_knobs(_table_knobs, _table_mode, mode) - # Fill None-valued kwargs from table. We can't reliably mutate - # locals() for re-read, so re-bind each kwarg explicitly. - if rectangle_for_nowrite is None and "rectangle_for_nowrite" in _table_knobs: - rectangle_for_nowrite = bool(_table_knobs["rectangle_for_nowrite"]) - _block_size_m = _block_size_m if _block_size_m is not None else _table_knobs.get("_block_size_m") - _num_warps = _num_warps if _num_warps is not None else _table_knobs.get("_num_warps") - _num_stages = _num_stages if _num_stages is not None else _table_knobs.get("_num_stages") - _heads_per_block = _heads_per_block if _heads_per_block is not None else _table_knobs.get("_heads_per_block") - _precompute_num_warps = _precompute_num_warps if _precompute_num_warps is not None else _table_knobs.get("_precompute_num_warps") - _precompute_num_stages = _precompute_num_stages if _precompute_num_stages is not None else _table_knobs.get("_precompute_num_stages") - _block_size_m_write = _block_size_m_write if _block_size_m_write is not None else _table_knobs.get("_block_size_m_write") - _block_size_m_nowrite = _block_size_m_nowrite if _block_size_m_nowrite is not None else _table_knobs.get("_block_size_m_nowrite") - _num_warps_write = _num_warps_write if _num_warps_write is not None else _table_knobs.get("_num_warps_write") - _num_warps_nowrite = _num_warps_nowrite if _num_warps_nowrite is not None else _table_knobs.get("_num_warps_nowrite") - _num_stages_write = _num_stages_write if _num_stages_write is not None else _table_knobs.get("_num_stages_write") - _num_stages_nowrite = _num_stages_nowrite if _num_stages_nowrite is not None else _table_knobs.get("_num_stages_nowrite") - _cta_per_sm = _cta_per_sm if _cta_per_sm is not None else _table_knobs.get("_cta_per_sm") - _num_loop_stages = _num_loop_stages if _num_loop_stages is not None else _table_knobs.get("_num_loop_stages") - # Split-form resolution for pm's per-half knobs. Without these the - # table's _num_loop_stages_{write,nowrite} and _cta_per_sm_{write, - # nowrite} values are dead — pm reads the split forms but the - # wrapper would leave them None, falling through to Triton defaults - # (or our hardcoded `or 1` / `or 2` per-mode fallbacks). - _num_loop_stages_write = _num_loop_stages_write if _num_loop_stages_write is not None else _table_knobs.get("_num_loop_stages_write") - _num_loop_stages_nowrite = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _table_knobs.get("_num_loop_stages_nowrite") - _cta_per_sm_write = _cta_per_sm_write if _cta_per_sm_write is not None else _table_knobs.get("_cta_per_sm_write") - _cta_per_sm_nowrite = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _table_knobs.get("_cta_per_sm_nowrite") - _flatten = _flatten if _flatten is not None else _table_knobs.get("_flatten") - _warp_specialize = _warp_specialize if _warp_specialize is not None else _table_knobs.get("_warp_specialize") - _use_tma_rect_load = _use_tma_rect_load or bool(_table_knobs.get("_use_tma_rect_load", False)) - _use_tma_replay_write_load = _use_tma_replay_write_load or bool(_table_knobs.get("_use_tma_replay_write_load", False)) - _use_tma_replay_write_store = _use_tma_replay_write_store or bool(_table_knobs.get("_use_tma_replay_write_store", False)) - _use_tma_replay_nowrite_load = _use_tma_replay_nowrite_load or bool(_table_knobs.get("_use_tma_replay_nowrite_load", False)) - # Final defaults if neither caller nor table set them (empty table case). - if mode is None: - mode = "persistent_dynamic" - if rectangle_for_nowrite is None: - rectangle_for_nowrite = False - assert mode in ("persistent_dynamic", "persistent_main"), ( - f"unknown mode {mode!r}; expected 'persistent_dynamic' or 'persistent_main'" - ) - if is_quantized: - assert state_scales is not None, ( - f"state.dtype={state.dtype} requires state_scales tensor " - "(shape (cache_size, nheads, dim), fp32)." - ) - assert state_scales.shape == (cache_size, nheads, dim), ( - f"state_scales shape mismatch: expected {(cache_size, nheads, dim)}, " - f"got {state_scales.shape}." - ) - assert state_scales.dtype == torch.float32, ( - f"state_scales must be fp32, got {state_scales.dtype}." - ) - assert state_scales.device == state.device - - # Cache T-axis = MAX_WINDOW (the replay buffer capacity). For the - # placeholder degenerate case max_window = T (every step is a checkpoint - # step). For real replay-style checkpointing, max_window > T and - # `prev_num_accepted_tokens` can be 0..max_window. Window-axis kernel - # tiles (BLOCK_SIZE_WINDOW, BLOCK_SIZE_K) are derived independently from - # MAX_REPLAY_BUFFER_LENGTH so max_window can exceed BLOCK_SIZE_T freely. - max_window = old_x.shape[1] - assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" - - assert x.shape == (batch, T, nheads, dim) - assert dt.shape == x.shape - assert A.shape == (nheads, dim, dstate) - assert B.shape == (batch, T, ngroups, dstate) - assert C.shape == B.shape - assert old_x.shape == (cache_size, max_window, nheads, dim) - assert old_B.shape == (cache_size, 2, max_window, ngroups, dstate) - assert old_dt.shape == (cache_size, 2, nheads, max_window) - assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_window) - assert cache_buf_idx.shape == (cache_size,) - assert prev_num_accepted_tokens.shape == (cache_size,) - - tie_hdim = ( - A.stride(-1) == 0 - and A.stride(-2) == 0 - and dt.stride(-1) == 0 - and (dt_bias is None or dt_bias.stride(-1) == 0) - ) - assert tie_hdim - - device = x.device - BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) - # Rectangle K-axis bound = window (max_window). Computed unconditionally - # so the launch sites can refer to it; only used on the rectangle path. - BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), 16) - - # Allocate precomputed intermediates (per-call, not cached). Always - # allocate (T, K) — the largest layout that any path uses. Replay-style - # paths only touch the first T columns; rectangle/dynamic use the full K. - # The few extra unused columns per row are negligible (~6KB per layer at - # production sizes) and let the dispatch helpers share one buffer. - cb_scaled = torch.empty( - batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 - ) - decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) - - z_strides = ( - (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) if z is not None else (0, 0, 0, 0) - ) - - # Kernel tuning: BLOCK_SIZE_M, num_warps, HEADS_PER_BLOCK, precompute_num_warps. - # Dtype-aware heuristic from B200 sweeps (batch 1-512, T=6/32, TP=8, conv1d + - # chained PDL). Keyed on total_heads, BLOCK_SIZE_T, and state dtype; 16-bit - # states prefer different tiles from fp32 due to lower bandwidth. Philox - # gets its own branch — stochastic rounding shifts compute toward CUDA cores, - # so small-batch configs want more warps to hide the extra work. - total_heads = batch * nheads - heads_per_group = nheads // ngroups - state_is_16bit = state.dtype in (torch.float16, torch.bfloat16) - use_philox = rand_seed is not None - if BLOCK_SIZE_T <= 16: - if use_philox and state_is_16bit: - # Philox: more warps at small batch to hide CUDA core work. - # At large batch, converges to non-Philox fp16 config. - if total_heads <= 16: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 4, 4, 4, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - elif state_is_16bit: - if total_heads <= 16: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 4, 1 - elif total_heads <= 64: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 2, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 1, - min(2, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - else: # fp32 state (no Philox — fp32 doesn't need stochastic rounding) - if total_heads <= 32: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 4, 1 - elif total_heads <= 64: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - elif total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 2, 2, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 1, 2, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 2, - min(2, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - else: # T > 16 - if state_is_16bit: - if total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 16, - 1, - 4, - min(2, heads_per_group), - ) - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 1, - min(4, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 4, - min(2, heads_per_group), - ) - else: # fp32 state - if total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 2, - 4, - min(2, heads_per_group), - ) - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 2, - min(4, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 4, - min(2, heads_per_group), - ) - if _block_size_m is not None: - BLOCK_SIZE_M = _block_size_m - if _num_warps is not None: - num_warps = _num_warps - if _heads_per_block is not None: - # Cap at heads_per_group: HEADS_PER_BLOCK divides the kernel's head - # axis, so a table value larger than the model's heads-per-group - # would overshoot. Protects callers running smaller models than - # the one we tuned against. - heads_per_block = min(_heads_per_block, heads_per_group) - if _precompute_num_warps is not None: - precompute_num_warps = _precompute_num_warps - - # Per-main knob resolution: each _*_{write,nowrite} arg, if not None, - # overrides the corresponding shared value for ONE main launch only. - # Default (None) = tied to shared value (current behavior). - BLOCK_SIZE_M_WRITE = _block_size_m_write if _block_size_m_write is not None else BLOCK_SIZE_M - BLOCK_SIZE_M_NOWRITE = _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M - NUM_WARPS_WRITE = _num_warps_write if _num_warps_write is not None else num_warps - NUM_WARPS_NOWRITE = _num_warps_nowrite if _num_warps_nowrite is not None else num_warps - NUM_STAGES_WRITE = _num_stages_write if _num_stages_write is not None else _num_stages - NUM_STAGES_NOWRITE = _num_stages_nowrite if _num_stages_nowrite is not None else _num_stages - # Persistent-only per-main: - CTA_PER_SM_WRITE = _cta_per_sm_write if _cta_per_sm_write is not None else _cta_per_sm - CTA_PER_SM_NOWRITE = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _cta_per_sm - NUM_LOOP_STAGES_WRITE = _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages - NUM_LOOP_STAGES_NOWRITE = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages - - HAS_CACHE_BATCH_INDICES = state_batch_indices is not None - - assert nheads % heads_per_block == 0, ( - f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" - ) - assert heads_per_block <= heads_per_group, ( - f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" - ) - - # state_scales pointer + strides: real tensor when quantized, otherwise - # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). - if is_quantized: - state_scales_arg = state_scales - state_scales_strides = ( - state_scales.stride(0), - state_scales.stride(1), - state_scales.stride(2), - ) - else: - state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 - state_scales_strides = (0, 0, 0) - - # Per-path TMA descriptors for state — write-side and nowrite-side. Each - # kernel launch consumes the descriptor whose block_shape[0] matches its - # BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need - # distinct descriptors; otherwise the descriptor's block_shape[0] would - # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot / arithmetic - # on the loaded tile fails shape inference at compile time - # ("Cannot make_shape_compatible: incompatible dimensions"). When Mw == - # Mnw (tied, the common case) the two descriptors are the same object. - # Same memory (state's flat 2D view, shape (cache*nheads*dim, dstate)) - # and same dstate block_shape — only block_shape[0] differs. - # When no TMA flag is on, both variables hold the raw `state` tensor as a - # dummy; kernels never reference it because their constexprs are all - # False (Triton DCEs the dead branches). - # `triton.set_allocator()` must run before any descriptor-using launch. - if (_use_tma_rect_load or _use_tma_replay_write_load - or _use_tma_replay_write_store or _use_tma_replay_nowrite_load): - from triton.tools.tensor_descriptor import TensorDescriptor - _ensure_tma_allocator() - assert state.is_contiguous(), "TMA state requires contiguous state" - assert state.stride(-1) == 1, "TMA state requires inner stride 1" - _state_flat = state.view(-1, state.shape[-1]) - _dstate_pow2 = triton.next_power_of_2(dstate) - state_tma_descriptor_write = TensorDescriptor.from_tensor( - _state_flat, block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], - ) - if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: - state_tma_descriptor_nowrite = state_tma_descriptor_write - else: - state_tma_descriptor_nowrite = TensorDescriptor.from_tensor( - _state_flat, block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], - ) - else: - state_tma_descriptor_write = state # dummy; all consuming constexprs False - state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False - - # Slot permutation — pointer + USE_PERM gate. Always required: the - # persistent_main kernel reads pid_b through slot_perm to walk the - # write-first sorted batch. persistent_dynamic forces USE_PERM=False - # at the call site (see launch_persistent_dynamic_main below) so the - # perm value doesn't matter for pd, but the tensor must still be valid. - assert isinstance(slot_perm, torch.Tensor), ( - f"slot_perm must be a torch.Tensor, got {type(slot_perm).__name__}" - ) - assert slot_perm.device == device, ( - f"slot_perm must be on device {device}, got {slot_perm.device}" - ) - assert slot_perm.dtype in (torch.int32, torch.int64), ( - f"slot_perm must be int32/int64, got {slot_perm.dtype}" - ) - assert slot_perm.shape == (batch,), ( - f"slot_perm must have shape (batch={batch},), got {tuple(slot_perm.shape)}" - ) - assert isinstance(n_writes, torch.Tensor), ( - f"n_writes must be a torch.Tensor, got {type(n_writes).__name__}" - ) - assert n_writes.device == device, ( - f"n_writes must be on device {device}, got {n_writes.device}" - ) - assert n_writes.dtype == torch.int32, ( - f"n_writes must be int32, got {n_writes.dtype}" - ) - assert n_writes.shape == (1,), ( - f"n_writes must have shape (1,), got {tuple(n_writes.shape)}" - ) - slot_perm_arg = slot_perm - use_perm = True - - precomp_grid = (batch, nheads // heads_per_block) - d_strides = (D.stride(0), D.stride(1)) if D is not None else (0, 0) - - # ---- Launch helpers (close over locals) ------------------------------- - # Each helper is a thin closure that calls one Triton kernel with the - # full positional + kwarg argument list. Mode-dependent constexprs - # (write_checkpoint, early_out, rectangle) are passed in. - - def launch_dynamic_precompute(rectangle: bool): - _dynamic_precompute_kernel[precomp_grid]( - dt, dt_bias, A, B, C, - cb_scaled, decay_vec, - old_B, old_dt, old_dA_cumsum, - cache_buf_idx, prev_num_accepted_tokens, - state_batch_indices, pad_slot_id, - T, max_window, dstate, nheads // ngroups, - dt.stride(0), dt.stride(1), dt.stride(2), - dt_bias.stride(0) if dt_bias is not None else 0, - A.stride(0), - B.stride(0), B.stride(1), B.stride(2), B.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - dt_softplus, - HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, - LAUNCH_WITH_PDL=launch_with_pdl, - LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, - HEADS_PER_BLOCK=heads_per_block, - RECTANGLE=rectangle, - num_warps=precompute_num_warps, - **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), - launch_pdl=launch_with_pdl, - ) - - # ---- launch_persistent_main ------------------------------------------ - # Persistent-CTA main kernel. Single launch covers `n_slots` slots - # starting at `slot_offset`. Caller invokes twice: once for the write - # half (slot_offset=0, n_slots=n_writes, write_checkpoint=True) and - # once for the nowrite half (slot_offset=n_writes, - # n_slots=batch-n_writes, write_checkpoint=False). Hard-sort - # contract: caller has pre-sorted slots so [0, n_writes) are writes - # and [n_writes, batch) are nowrites. - - # Resolve persistent-mode bench knobs. Defaults: cta_per_sm = 1 - # (one CTA per SM, matches upstream `_p_matmul_ogs.py`); num_loop_stages - # = 2 (matches in-tree `swiglu` precedent for non-dot persistent loops); - # flatten = True (canonical Triton 3.6 idiom); warp_specialize = False. - _num_sms = torch.cuda.get_device_properties(device).multi_processor_count - cta_per_sm_arg = _cta_per_sm if _cta_per_sm else 1 - num_persistent_arg = cta_per_sm_arg * _num_sms - num_loop_stages_arg = _num_loop_stages if _num_loop_stages else 2 - flatten_arg = True if _flatten is None else bool(_flatten) - warp_specialize_arg = False if _warp_specialize is None else bool(_warp_specialize) - # Per-launch work-item count. At small batch, total_work may be < the - # full persistent grid; capping `grid` at `min(NUM_PERSISTENT, total_work)` - # avoids launching empty CTAs that pay setup cost for no work. Correctness: - # the kernel's `tl.range(pid, total_work, NUM_PERSISTENT)` ensures each - # tile_id is covered exactly once across all live pids in [0, grid) when - # grid <= NUM_PERSISTENT (each CTA does 1 tile; loop step >= total_work - # exits immediately) AND when grid == NUM_PERSISTENT (each CTA loops over - # multiple tiles). NUM_PERSISTENT is now a runtime int (see kernel def - # docstring at _persistent_main_kernel) so changing cta_per_sm does NOT - # trigger a new Triton compile — same kernel binary, different loop step. - # (Named UPPERCASE for historical Triton-style consistency only; not - # constexpr.) - _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M - - def launch_persistent_main(write_checkpoint: bool, - *, - launch_dependent_kernels: bool = False, - rectangle: bool = False): - # `n_writes` (wrapper-level) is the (1,) int32 device tensor with the - # write count. Both halves always launch; the kernel's runtime PNAT - # check iterates only the slots that belong to its half. - _bsm = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE - _nw = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE - _ns = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE - _cps = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE - _cps = _cps if _cps else 1 - _nls = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE - _nls = _nls if _nls else 2 - _num_persistent = _cps * _num_sms - _num_pid_m_local = (dim + _bsm - 1) // _bsm - # Grid sizing: cap at min(full persistent grid, upper-bound total work). - # We use `batch` as the upper bound on slots-per-half — overcounting - # by a few CTAs is fine since the kernel's runtime check only - # iterates the slots that actually belong to its half. - _total_work_launch = max(1, batch * _num_pid_m_local * nheads) - grid = (min(_num_persistent, _total_work_launch),) - # Per-path TMA descriptor — block_shape[0] must match _bsm. - _desc = (state_tma_descriptor_write if write_checkpoint - else state_tma_descriptor_nowrite) - _persistent_main_kernel[grid]( - state, _desc, state_scales_arg, old_x, - old_B, old_dt, old_dA_cumsum, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, - n_writes, batch, nheads, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - _bsm, - LAUNCH_WITH_PDL=use_internal_pdl, - PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, - QUANT_MAX=quant_max, - WRITE_CHECKPOINT=write_checkpoint, - LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, - USE_PERM=use_perm, - NUM_PERSISTENT=_num_persistent, - NUM_LOOP_STAGES=_nls, - FLATTEN=flatten_arg, - WARP_SPECIALIZE=warp_specialize_arg, - IS_DYNAMIC=False, - RECTANGLE=rectangle, - # 3 TMA flags. IS_DYNAMIC=False: WC fixed per launch; impl - # constexpr-folds the LOAD pick. When WC=True (write half), - # NOWRITE_LOAD is dummy False; when WC=False, WRITE_LOAD/STORE - # dummy False. NOWRITE_LOAD picks rect-load (RECTANGLE) or - # replay-nowrite-load. - USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), - USE_TMA_LOAD_NOWRITE=bool( - (_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load) - and not write_checkpoint - ), - USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), - num_warps=_nw, - **({"num_stages": _ns} if _ns else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), - launch_pdl=use_internal_pdl, - ) - - def launch_persistent_dynamic_main(n_writes_dev: torch.Tensor, - launch_dependent_kernels: bool = False, - rectangle: bool = False): - # Single-launch persistent kernel covering the whole batch with - # runtime per-slot WRITE_CHECKPOINT branch. No half-split, no - # n_writes needed (the kernel ignores n_writes_dev when - # IS_DYNAMIC=True; Triton DCEs the load). is_write is computed - # at runtime per work-item from the loaded PNAT. - # We still pass `n_writes_dev` (the same tensor the persistent_main - # path uses) so the kernel signature is uniform; the value is - # immaterial. - # Grid sizing: cap at total_work (= batch * num_pid_m * nheads) for - # the dynamic case (full-batch coverage); see launch_persistent_main - # comment for correctness rationale. - _total_work_launch = max(1, batch * _num_pid_m * nheads) - grid = (min(num_persistent_arg, _total_work_launch),) - # Persistent-dynamic kernel uses a single BLOCK_SIZE_M (same as the - # wrapper's BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE tied convention), so - # the write-side descriptor matches. Both write and nowrite slots - # in this kernel share that BSM. - _persistent_main_kernel[grid]( - state, state_tma_descriptor_write, state_scales_arg, old_x, - old_B, old_dt, old_dA_cumsum, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, slot_perm_arg, rand_seed, pad_slot_id, - n_writes, batch, nheads, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - BLOCK_SIZE_M, - LAUNCH_WITH_PDL=use_internal_pdl, - PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, - QUANT_MAX=quant_max, - WRITE_CHECKPOINT=False, - LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, - # persistent_dynamic forces USE_PERM=False regardless of caller- - # provided slot_perm — our pd tuning runs all happened with - # SORT=0 (no slot_perm passed), so honoring slot_perm here would - # silently shift pd to an untimed code path. Revisit if/when - # we benchmark pd with slot_perm. - USE_PERM=False, - NUM_PERSISTENT=num_persistent_arg, - NUM_LOOP_STAGES=num_loop_stages_arg, - FLATTEN=flatten_arg, - WARP_SPECIALIZE=warp_specialize_arg, - IS_DYNAMIC=True, - RECTANGLE=rectangle, - # 3 TMA flags. IS_DYNAMIC=True: is_write is runtime per slot; - # impl's load TMA picks per-slot (constexpr ternary becomes a - # runtime branch — both load forms emitted, ~negligible cost). - # NOWRITE_LOAD picks rect-load when RECTANGLE, else - # replay-nowrite-load. STORE only fires on runtime is_write. - USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load), - USE_TMA_LOAD_NOWRITE=bool( - _use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load - ), - USE_TMA_STORE=bool(_use_tma_replay_write_store), - num_warps=num_warps, - **({"num_stages": _num_stages} if _num_stages else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), - launch_pdl=use_internal_pdl, - ) - - # ---- Mode dispatch ---------------------------------------------------- - with torch.cuda.device(device.index): - if mode == "persistent_dynamic": - # Single-launch persistent kernel covering the full batch. Each - # work-item dispatches via runtime PNAT check. Kernel ignores - # n_writes (Triton DCEs the load) when IS_DYNAMIC=True; we still - # pass the wrapper-provided tensor as required by the signature. - launch_dynamic_precompute(rectangle=rectangle_for_nowrite) - launch_persistent_dynamic_main( - n_writes, - launch_dependent_kernels=False, - rectangle=rectangle_for_nowrite, - ) - elif mode == "persistent_main": - # Persistent-CTA main kernel. One shared dynamic_precompute - # (per-slot dispatch via PNAT) feeds two persistent_main - # launches (write half + nowrite half). Both halves ALWAYS - # launch; the kernel's runtime check iterates only the slots - # belonging to its half (write: [0, n_writes), nowrite: - # [n_writes, batch)). - # - # Caller-provided contract: `n_writes` is a (1,) int32 device - # tensor (the kernel reads it at runtime, after the precompute); - # `slot_perm` is a (batch,) int32 device tensor pre-sorted - # write-first. - launch_dynamic_precompute(rectangle=rectangle_for_nowrite) - launch_persistent_main( - write_checkpoint=True, - launch_dependent_kernels=True, - rectangle=False, # write always replay-style - ) - launch_persistent_main( - write_checkpoint=False, - launch_dependent_kernels=False, - rectangle=rectangle_for_nowrite, - ) - else: - raise ValueError( - f"mode={mode!r} is not supported. Supported modes: " - f"'persistent_dynamic', 'persistent_main'." - ) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 0e387c26efd0..c1e48888b829 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -220,14 +220,11 @@ def __init__(self, max_batch_size: int, chunk_size: int): dtype=torch.int32, device="cuda") - self.replay_work_items = torch.zeros( - max_batch_size, - REPLAY_WORK_ITEM_WIDTH, - dtype=torch.int32, - device="cuda") - self.replay_n_writes = torch.zeros(1, - dtype=torch.int32, - device="cuda") + self.replay_work_items = torch.zeros(max_batch_size, + REPLAY_WORK_ITEM_WIDTH, + dtype=torch.int32, + device="cuda") + self.replay_n_writes = torch.zeros(1, dtype=torch.int32, device="cuda") self.replay_num_decodes = 0 # Pre-allocated buffers. @@ -261,9 +258,8 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, if num_decodes == 0: return - position_in_decode_batch = torch.arange(num_decodes, - dtype=torch.int32, - device=self.state_indices.device) + position_in_decode_batch = torch.arange( + num_decodes, dtype=torch.int32, device=self.state_indices.device) cache_slot = self.state_indices[num_contexts:batch_size] cache_slot_idx = cache_slot.to(torch.long) pnat = prev_num_accepted_tokens[cache_slot_idx].to(torch.int32) @@ -272,8 +268,7 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, writes = (pnat + replay_step_width > replay_history_size) writes_i32 = writes.to(torch.int32) write_offsets = torch.cumsum(writes_i32, dim=0) - writes_i32 - n_writes = torch.sum(writes_i32, dim=0, - keepdim=True).to(torch.int32) + n_writes = torch.sum(writes_i32, dim=0, keepdim=True).to(torch.int32) no_write_offsets = position_in_decode_batch - write_offsets output_offsets = torch.where(writes, write_offsets, n_writes + no_write_offsets) @@ -285,8 +280,9 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, work_items[:, REPLAY_WORK_CACHE_SLOT].scatter_(0, output_offsets, cache_slot) work_items[:, REPLAY_WORK_PNAT].scatter_(0, output_offsets, pnat) - work_items[:, REPLAY_WORK_CACHE_BUF_IDX].scatter_( - 0, output_offsets, active_cache_buf_idx) + work_items[:, + REPLAY_WORK_CACHE_BUF_IDX].scatter_(0, output_offsets, + active_cache_buf_idx) self.replay_n_writes.copy_(n_writes) def prepare(self, attn_metadata: AttentionMetadata): diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index f7872fffd111..b42bedcfa5bc 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -35,8 +35,7 @@ ) from .softplus import softplus -_REPLAY_WORK_POSITION_IN_DECODE_BATCH = tl.constexpr( - REPLAY_WORK_POSITION_IN_DECODE_BATCH) +_REPLAY_WORK_POSITION_IN_DECODE_BATCH = tl.constexpr(REPLAY_WORK_POSITION_IN_DECODE_BATCH) _REPLAY_WORK_CACHE_SLOT = tl.constexpr(REPLAY_WORK_CACHE_SLOT) _REPLAY_WORK_PNAT = tl.constexpr(REPLAY_WORK_PNAT) _REPLAY_WORK_CACHE_BUF_IDX = tl.constexpr(REPLAY_WORK_CACHE_BUF_IDX) @@ -125,9 +124,7 @@ def _bitrev32(x: tl.tensor) -> tl.tensor: @triton.jit -def _stochastic_round_int8_packed( - x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor -) -> tl.tensor: +def _stochastic_round_int8_packed(x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor) -> tl.tensor: """Stochastic rounding for int8 using one random uint32 per 4 values.""" low = rand & 0x0000FFFF high = (rand >> 16) & 0x0000FFFF @@ -144,9 +141,7 @@ def _stochastic_round_int8_packed( @triton.jit -def _stochastic_round_int16_packed( - x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor -) -> tl.tensor: +def _stochastic_round_int16_packed(x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor) -> tl.tensor: """Stochastic rounding for int16 using one random uint32 per 2 values.""" rand_bits = tl.where((offs_n & 1) == 0, rand, _bitrev32(rand)) rand01 = (rand_bits & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) @@ -256,7 +251,6 @@ def _replay_precompute_impl( else: cache_batch_idx = pid_b.to(tl.int64) - # --- Cache write semantics --- # cache_buf_idx names this step's "active" buffer — the one with the # historical inputs at [0, PNAT). The other buffer is "staging". @@ -298,7 +292,8 @@ def _replay_precompute_impl( # Load dt (H, T) dt_addrs = ( - dt_ptr + pid_b * stride_dt_batch + dt_ptr + + pid_b * stride_dt_batch + heads_block[:, None] * stride_dt_head + offs_t[None, :] * stride_dt_T ) @@ -354,7 +349,8 @@ def _replay_precompute_impl( # decay_vec scratch — always at offs_t. decay_vec_addrs = ( - decay_vec_ptr + pid_b * stride_dv_batch + decay_vec_ptr + + pid_b * stride_dv_batch + heads_block[:, None] * stride_dv_head + offs_t[None, :] * stride_dv_t ) @@ -414,15 +410,13 @@ def _replay_precompute_impl( 0.0, ) # (H, T, T) cb_scaled_addrs = ( - cb_scaled_ptr + pid_b * stride_cb_batch + cb_scaled_ptr + + pid_b * stride_cb_batch + heads_block[:, None, None] * stride_cb_head + offs_t[None, :, None] * stride_cb_t + offs_t[None, None, :] * stride_cb_j ) # (H, T, T) - cb_store_mask = ( - (offs_t[None, :, None] < BLOCK_SIZE_T) - & (offs_t[None, None, :] < BLOCK_SIZE_T) - ) + cb_store_mask = (offs_t[None, :, None] < BLOCK_SIZE_T) & (offs_t[None, None, :] < BLOCK_SIZE_T) tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) @@ -438,8 +432,8 @@ def _rectangle_precompute_impl( B_ptr, C_ptr, # Output pointers - cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle - decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) + cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle + decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite # path: read from buf_active at [0, PNAT), write new tokens at # [PNAT, PNAT+T) of buf_active (same buffer). @@ -635,11 +629,13 @@ def _rectangle_precompute_impl( hk_mask = is_old_k[None, :] # (1, K) old_dt_all = tl.load( old_dt_read_h[:, None] + safe_old_k[None, :] * stride_old_dt_T, - mask=hk_mask, other=0.0, + mask=hk_mask, + other=0.0, ).to(tl.float32) old_dA_cumsum_all = tl.load( old_dA_cumsum_read_h[:, None] + safe_old_k[None, :] * stride_old_dA_cumsum_T, - mask=hk_mask, other=0.0, + mask=hk_mask, + other=0.0, ).to(tl.float32) # Use loop-1 registers for this step's newly appended tokens. These are # exactly the values stored above at [PNAT, PNAT+T); reloading them here @@ -654,9 +650,7 @@ def _rectangle_precompute_impl( # wrapper passes this as an explicit constexpr so fast-path compilations do # not carry the fallback branch. if USE_GATHER_FOR_NEW_TOKENS: - new_token_gather_idx = tl.broadcast_to( - safe_k_new[None, :], (HEADS_PER_BLOCK, BLOCK_SIZE_K) - ) + new_token_gather_idx = tl.broadcast_to(safe_k_new[None, :], (HEADS_PER_BLOCK, BLOCK_SIZE_K)) dt_at_kn = tl.where( is_new_k[None, :], tl.gather(dt_new, new_token_gather_idx, axis=1), @@ -669,9 +663,8 @@ def _rectangle_precompute_impl( ) # (H, K) else: new_token_selector = ( - (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) - & is_new_k[None, :] - ) + offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens) + ) & is_new_k[None, :] dt_at_kn = tl.sum( tl.where(new_token_selector[None, :, :], dt_new[:, :, None], 0.0), axis=1, @@ -772,9 +765,8 @@ def _rectangle_precompute_impl( + offs_t[None, :, None] * stride_cb_t + offs_k[None, None, :] * stride_cb_j ) # (H, T, K) - cb_store_mask_3d = ( - (offs_t[None, :, None] < BLOCK_SIZE_T) - & (offs_k[None, None, :] < BLOCK_SIZE_K) + cb_store_mask_3d = (offs_t[None, :, None] < BLOCK_SIZE_T) & ( + offs_k[None, None, :] < BLOCK_SIZE_K ) # (1, T, K) → broadcasts to (H, T, K) tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) @@ -786,8 +778,7 @@ def _rectangle_precompute_impl( @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) @triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) @triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} + {"BLOCK_SIZE_K": lambda args: max(triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} ) @triton.jit() def _dynamic_precompute_kernel( @@ -1256,7 +1247,8 @@ def _persistent_main_impl( ) old_dA_cumsum_all = tl.load( old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, - mask=old_window_mask, other=0.0, + mask=old_window_mask, + other=0.0, ).to(tl.float32) prev_k_idx = tl.minimum( @@ -1283,7 +1275,9 @@ def _persistent_main_impl( + pid_h * stride_old_x_head ) old_x_all = tl.load( - old_x_read_base + offs_window[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, + old_x_read_base + + offs_window[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, mask=old_window_mask[:, None] & m_mask[None, :], other=0.0, ) @@ -1347,9 +1341,7 @@ def _persistent_main_impl( r01 = tl.join(r0, r1) r23 = tl.join(r2, r3) r0123 = tl.join(r01, r23) - rand_compact = tl.reshape( - r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR) - ) + rand_compact = tl.reshape(r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR)) # Broadcast each unique rand to RAND_DIVISOR adjacent positions. # Pack-group (pack=2 fp16 / pack=4 fp8) consumes adjacent positions; # the unique rand lands at the asm's read slot; duplicates feed @@ -1391,13 +1383,9 @@ def _persistent_main_impl( "fp8 SR is handled by the prior branch.", ) if state_ptrs.dtype.element_ty == tl.int8: - state_q = _stochastic_round_int8_packed( - state_q, rand, offs_n[None, :] - ) + state_q = _stochastic_round_int8_packed(state_q, rand, offs_n[None, :]) else: - state_q = _stochastic_round_int16_packed( - state_q, rand, offs_n[None, :] - ) + state_q = _stochastic_round_int16_packed(state_q, rand, offs_n[None, :]) elif state_ptrs.dtype.element_ty != tl.float8e4nv: tl.static_assert( (state_ptrs.dtype.element_ty == tl.int8) @@ -1485,7 +1473,8 @@ def _persistent_main_impl( if HAS_Z: z_all = tl.load( z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, - mask=t_mask[:, None] & m_mask[None, :], other=0.0, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, ).to(tl.float32) out_all_z = out_all * z_all * tl.sigmoid(z_all) out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim @@ -1515,7 +1504,7 @@ def _persistent_rectangle_impl( # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as # replay path). Used when USE_TMA_LOAD; ignored otherwise. state_tma_descriptor, - state_scales_ptr, # only consulted when QUANT_MAX > 0 + state_scales_ptr, # only consulted when QUANT_MAX > 0 old_x_ptr, x_ptr, C_ptr, @@ -1617,9 +1606,13 @@ def _persistent_rectangle_impl( ) state = state_tma_descriptor.load([offs_y, 0]) else: - state_ptr_local = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptr_local = ( + state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + ) state_ptrs = ( - state_ptr_local + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + state_ptr_local + + offs_m[:, None] * stride_state_dim + + offs_n[None, :] * stride_state_dstate ) state_mask = m_mask[:, None] & n_mask[None, :] state = tl.load(state_ptrs, mask=state_mask, other=0.0) @@ -1631,7 +1624,8 @@ def _persistent_rectangle_impl( ) decode_scale = tl.load( state_scales_base + offs_m * stride_state_scales_dim, - mask=m_mask, other=1.0, + mask=m_mask, + other=1.0, ).to(tl.float32) else: state = state.to(tl.float32) @@ -1662,9 +1656,7 @@ def _persistent_rectangle_impl( # Hoist: old_x doesn't depend on conv1d/precompute; load before gdc_wait. old_x_load = tl.load( - old_x_read_base - + safe_old_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, + old_x_read_base + safe_old_k[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, mask=is_old_k[:, None] & m_mask[None, :], other=0.0, ).to(tl.float32) @@ -1683,9 +1675,7 @@ def _persistent_rectangle_impl( other=0.0, ) tl.store( - old_x_write_base - + offs_k[:, None] * stride_old_x_T - + offs_m[None, :] * stride_old_x_dim, + old_x_write_base + offs_k[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, x_K, mask=is_new_k[:, None] & m_mask[None, :], ) @@ -1694,7 +1684,7 @@ def _persistent_rectangle_impl( x_combined = old_x_load + x_K_f32 if HAS_D or HAS_Z: - sel_tk = (offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens)) + sel_tk = offs_t[:, None] == (offs_k[None, :] - prev_num_accepted_tokens) x_all = tl.dot(sel_tk.to(tl.bfloat16), x_K.to(tl.bfloat16)) else: x_all = x_K_f32 # placeholder; unused @@ -1707,13 +1697,12 @@ def _persistent_rectangle_impl( ).to(tl.float32) decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec_full = tl.load( - decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0 - ).to(tl.float32) + decay_vec_full = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + tl.float32 + ) state_out = ( - tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) - * decay_vec_full[:, None] + tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec_full[:, None] ) if QUANT_MAX > 0.0: state_out = state_out * decode_scale[None, :] @@ -1728,7 +1717,8 @@ def _persistent_rectangle_impl( if HAS_Z: z_all = tl.load( z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, - mask=t_mask[:, None] & m_mask[None, :], other=0.0, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, ).to(tl.float32) out_all_z = out_all * z_all * tl.sigmoid(z_all) out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim @@ -1749,12 +1739,14 @@ def _persistent_rectangle_impl( @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) @triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) @triton.heuristics( - {"BLOCK_SIZE_WINDOW": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} + { + "BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16 + ) + } ) @triton.heuristics( - {"BLOCK_SIZE_K": lambda args: max( - triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} + {"BLOCK_SIZE_K": lambda args: max(triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), 16)} ) @triton.heuristics( {"NUM_PID_M_BLOCKS": lambda args: triton.cdiv(args["dim"], args["BLOCK_SIZE_M"])} @@ -1798,9 +1790,9 @@ def _persistent_main_kernel( # device memory keeps the pointer stable across CUDA graph replay while # allowing the value to change between iterations. # When IS_DYNAMIC=True the value is unused (Triton DCEs the load). - n_writes_ptr, # int32 *: device-side count of write-mode slots - batch_total, # int32: total slot count - nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) + n_writes_ptr, # int32 *: device-side count of write-mode slots + batch_total, # int32: total slot count + nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) # Dimensions T: tl.constexpr, MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, @@ -1942,8 +1934,12 @@ def _persistent_main_kernel( # slot, then head — mirrors the existing 3D grid's axis ordering # (axis=0 fastest = pid_m). for tile_id in tl.range( - pid, total_work, NUM_PERSISTENT, - flatten=FLATTEN, num_stages=NUM_LOOP_STAGES, warp_specialize=WARP_SPECIALIZE, + pid, + total_work, + NUM_PERSISTENT, + flatten=FLATTEN, + num_stages=NUM_LOOP_STAGES, + warp_specialize=WARP_SPECIALIZE, ): pid_m = tile_id % NUM_PID_M_BLOCKS pid_b_local = (tile_id // NUM_PID_M_BLOCKS) % n_slots_local @@ -1961,25 +1957,15 @@ def _persistent_main_kernel( else: work_item_base = replay_work_items_ptr + work_item_idx * _REPLAY_WORK_ITEM_WIDTH if USE_REPLAY_CACHE_SLOT: - pid_b = tl.load( - work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH - ) - cache_batch_idx = tl.load( - work_item_base + _REPLAY_WORK_CACHE_SLOT - ).to(tl.int64) + pid_b = tl.load(work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH) + cache_batch_idx = tl.load(work_item_base + _REPLAY_WORK_CACHE_SLOT).to(tl.int64) pnat = tl.load(work_item_base + _REPLAY_WORK_PNAT) - active_buf = tl.load( - work_item_base + _REPLAY_WORK_CACHE_BUF_IDX - ).to(tl.int32) + active_buf = tl.load(work_item_base + _REPLAY_WORK_CACHE_BUF_IDX).to(tl.int32) else: - pid_b = tl.load( - work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH - ) + pid_b = tl.load(work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH) cache_batch_idx = work_item_idx.to(tl.int64) pnat = tl.load(work_item_base + _REPLAY_WORK_PNAT) - active_buf = tl.load( - work_item_base + _REPLAY_WORK_CACHE_BUF_IDX - ).to(tl.int32) + active_buf = tl.load(work_item_base + _REPLAY_WORK_CACHE_BUF_IDX).to(tl.int32) # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle # impl. `replay_work_items` carries the cache slot, PNAT and active # buffer for persistent_main; persistent_dynamic resolves those once @@ -2003,91 +1989,267 @@ def _persistent_main_kernel( # constexpr-folds; passing literal True here is consistent # and constexpr-equivalent. _persistent_main_impl( - pid_m, pid_b, pid_h, - cache_batch_idx, active_buf, pnat, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, + pid_m, + pid_b, + pid_h, + cache_batch_idx, + active_buf, + pnat, + state_ptr, + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, rand_seed_ptr, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_dbuf, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, - stride_old_B_group, stride_old_B_dstate, - stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, - stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - True, IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_dbuf, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + True, + IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) True, # WRITE_CHECKPOINT_IS_CONSTEXPR - USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + USE_TMA_LOAD_WRITE, + USE_TMA_LOAD_NOWRITE, + USE_TMA_STORE, ) else: _persistent_rectangle_impl( - pid_m, pid_b, pid_h, - cache_batch_idx, active_buf, pnat, - state_ptr, state_tma_descriptor, state_scales_ptr, + pid_m, + pid_b, + pid_h, + cache_batch_idx, + active_buf, + pnat, + state_ptr, + state_tma_descriptor, + state_scales_ptr, old_x_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_dbuf, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, - LAUNCH_WITH_PDL, QUANT_MAX, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_dbuf, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + QUANT_MAX, USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle ) else: _persistent_main_impl( - pid_m, pid_b, pid_h, - cache_batch_idx, active_buf, pnat, - state_ptr, state_tma_descriptor, state_scales_ptr, - old_x_ptr, old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - x_ptr, C_ptr, D_ptr, z_ptr, out_ptr, - cb_scaled_ptr, decay_vec_ptr, + pid_m, + pid_b, + pid_h, + cache_batch_idx, + active_buf, + pnat, + state_ptr, + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, rand_seed_ptr, - T, MAX_REPLAY_BUFFER_LENGTH, dim, dstate, nheads_ngroups_ratio, - stride_state_batch, stride_state_head, stride_state_dim, stride_state_dstate, - stride_state_scales_cache, stride_state_scales_head, stride_state_scales_dim, - stride_old_x_cache, stride_old_x_dbuf, stride_old_x_T, stride_old_x_head, stride_old_x_dim, - stride_old_B_cache, stride_old_B_dbuf, stride_old_B_T, - stride_old_B_group, stride_old_B_dstate, - stride_old_dt_cache, stride_old_dt_dbuf, stride_old_dt_head, stride_old_dt_T, - stride_old_dA_cumsum_cache, stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, stride_old_dA_cumsum_T, - stride_x_batch, stride_x_T, stride_x_head, stride_x_dim, - stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - stride_D_head, stride_D_dim, - stride_z_batch, stride_z_T, stride_z_head, stride_z_dim, - stride_out_batch, stride_out_T, stride_out_head, stride_out_dim, - stride_cb_batch, stride_cb_head, stride_cb_t, stride_cb_j, - stride_dv_batch, stride_dv_head, stride_dv_t, - BLOCK_SIZE_M, HAS_D, HAS_Z, HAS_CACHE_BATCH_INDICES, - BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - WRITE_CHECKPOINT, IS_DYNAMIC, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_dbuf, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + HAS_CACHE_BATCH_INDICES, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL, + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + WRITE_CHECKPOINT, + IS_DYNAMIC, False, # WRITE_CHECKPOINT_IS_CONSTEXPR - USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, USE_TMA_STORE, + USE_TMA_LOAD_WRITE, + USE_TMA_LOAD_NOWRITE, + USE_TMA_STORE, ) @@ -2131,43 +2293,830 @@ def _persistent_main_kernel( # _resolve_tuning chain — RN→SR for same dtype, then fp8→int8/SR. _DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { ("fp32", "RN"): [ - ( 16, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 8, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.22us - ( 32, "persistent_main", {'_block_size_m_nowrite': 16, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.17us - ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.08us - ( 128, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=8, score=9.00us - ( 256, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 1, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=16, score=10.92us - ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=13.53us - ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=19.50us - ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=30.28us - ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=50.32us - ( 8192, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=90.99us - (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=171.69us + ( + 16, + "persistent_main", + { + "_block_size_m_nowrite": 16, + "_block_size_m_write": 8, + "_cta_per_sm_nowrite": 4, + "_cta_per_sm_write": 1, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 1, + "_num_stages_write": 2, + "_num_warps_nowrite": 2, + "_num_warps_write": 2, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=1, score=6.22us + ( + 32, + "persistent_main", + { + "_block_size_m_nowrite": 16, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 9, + "_cta_per_sm_write": 7, + "_flatten": False, + "_heads_per_block": 1, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 4, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, + "_num_warps_write": 2, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=2, score=7.17us + ( + 64, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 9, + "_cta_per_sm_write": 4, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 3, + "_num_stages_nowrite": 1, + "_num_stages_write": 2, + "_num_warps_nowrite": 4, + "_num_warps_write": 4, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=4, score=8.08us + ( + 128, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 6, + "_cta_per_sm_write": 9, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 2, + "_num_stages_write": 1, + "_num_warps_nowrite": 1, + "_num_warps_write": 2, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=8, score=9.00us + ( + 256, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 9, + "_cta_per_sm_write": 1, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 3, + "_num_stages_write": 4, + "_num_warps_nowrite": 4, + "_num_warps_write": 2, + "_precompute_num_warps": 16, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=16, score=10.92us + ( + 512, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 10, + "_cta_per_sm_write": 4, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 1, + "_num_warps_nowrite": 1, + "_num_warps_write": 2, + "_precompute_num_warps": 16, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=32, score=13.53us + ( + 1024, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 4, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 4, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=64, score=19.50us + ( + 2048, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 2, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=128, score=30.28us + ( + 4096, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 4, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 4, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 2, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=256, score=50.32us + ( + 8192, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 10, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 2, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=512, score=90.99us + ( + 16384, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 9, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 1, + "_num_stages_write": 3, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 2, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=1024, score=171.69us ], ("fp16", "SR"): [ - ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 3, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.16us - ( 32, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 9, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.01us - ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=7.95us - ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 4, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=8.87us - ( 256, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 2, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 1, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.28us - ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 2, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=32, score=12.90us - ( 1024, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=16.71us - ( 2048, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 2, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=25.71us - ( 4096, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=39.80us - ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 2, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=71.34us - (16384, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=133.51us + ( + 16, + "persistent_main", + { + "_block_size_m_nowrite": 8, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 5, + "_cta_per_sm_write": 7, + "_flatten": False, + "_heads_per_block": 1, + "_num_loop_stages_nowrite": 4, + "_num_loop_stages_write": 3, + "_num_stages_nowrite": 1, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, + "_num_warps_write": 4, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=1, score=6.16us + ( + 32, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 9, + "_cta_per_sm_write": 4, + "_flatten": False, + "_heads_per_block": 1, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 2, + "_num_stages_write": 4, + "_num_warps_nowrite": 4, + "_num_warps_write": 2, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=2, score=7.01us + ( + 64, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 5, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 1, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, + "_num_warps_write": 2, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=4, score=7.95us + ( + 128, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 4, + "_flatten": False, + "_heads_per_block": 1, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 3, + "_num_stages_write": 1, + "_num_warps_nowrite": 4, + "_num_warps_write": 4, + "_precompute_num_warps": 4, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=8, score=8.87us + ( + 256, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 10, + "_cta_per_sm_write": 2, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 1, + "_num_stages_write": 1, + "_num_warps_nowrite": 1, + "_num_warps_write": 4, + "_precompute_num_warps": 16, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=16, score=10.28us + ( + 512, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 7, + "_cta_per_sm_write": 6, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 2, + "_num_warps_nowrite": 1, + "_num_warps_write": 2, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=32, score=12.90us + ( + 1024, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 7, + "_cta_per_sm_write": 7, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 2, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=64, score=16.71us + ( + 2048, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 7, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 1, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 2, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=128, score=25.71us + ( + 4096, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 7, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 4, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=256, score=39.80us + ( + 8192, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 7, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 2, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=512, score=71.34us + ( + 16384, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 7, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 1, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=1024, score=133.51us ], ("int8", "SR"): [ - ( 16, "persistent_main", {'_block_size_m_nowrite': 8, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 5, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 1, '_num_loop_stages_nowrite': 4, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1, score=6.34us - ( 32, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 8, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 4, '_num_stages_write': 2, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=2, score=7.36us - ( 64, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 2, '_cta_per_sm_write': 4, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 8, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=4, score=8.40us - ( 128, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 10, '_flatten': False, '_heads_per_block': 2, '_num_loop_stages_nowrite': 2, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 4, '_num_warps_write': 4, '_precompute_num_warps': 16, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=8, score=9.37us - ( 256, "persistent_dynamic", {'_block_size_m': 16, '_cta_per_sm': 8, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages': 1, '_num_stages': 4, '_num_warps': 1, '_precompute_num_warps': 16, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': True, '_use_tma_replay_write_load': False, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': False}), # raw_batch=16, score=10.02us - ( 512, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 16, '_cta_per_sm_nowrite': 7, '_cta_per_sm_write': 9, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 2, '_num_stages_write': 3, '_num_warps_nowrite': 2, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=32, score=13.15us - ( 1024, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 32, '_cta_per_sm_nowrite': 10, '_cta_per_sm_write': 7, '_flatten': False, '_heads_per_block': 16, '_num_loop_stages_nowrite': 1, '_num_loop_stages_write': 1, '_num_stages_nowrite': 4, '_num_stages_write': 3, '_num_warps_nowrite': 1, '_num_warps_write': 1, '_precompute_num_warps': 8, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': False, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=64, score=17.82us - ( 2048, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=128, score=27.01us - ( 4096, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 4, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 2, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': False, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=256, score=43.23us - ( 8192, "persistent_main", {'_block_size_m_nowrite': 32, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 8, '_cta_per_sm_write': 6, '_flatten': False, '_heads_per_block': 4, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 3, '_num_stages_write': 1, '_num_warps_nowrite': 1, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=512, score=77.01us - (16384, "persistent_main", {'_block_size_m_nowrite': 64, '_block_size_m_write': 64, '_cta_per_sm_nowrite': 6, '_cta_per_sm_write': 3, '_flatten': False, '_heads_per_block': 8, '_num_loop_stages_nowrite': 3, '_num_loop_stages_write': 2, '_num_stages_nowrite': 1, '_num_stages_write': 4, '_num_warps_nowrite': 2, '_num_warps_write': 4, '_precompute_num_warps': 1, '_use_tma_rect_load': True, '_use_tma_replay_nowrite_load': False, '_use_tma_replay_write_load': True, '_use_tma_replay_write_store': True, '_warp_specialize': False, 'rectangle_for_nowrite': True}), # raw_batch=1024, score=140.43us + ( + 16, + "persistent_main", + { + "_block_size_m_nowrite": 8, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 5, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 1, + "_num_loop_stages_nowrite": 4, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 2, + "_num_stages_write": 3, + "_num_warps_nowrite": 2, + "_num_warps_write": 4, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=1, score=6.34us + ( + 32, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 10, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 4, + "_num_stages_write": 2, + "_num_warps_nowrite": 4, + "_num_warps_write": 4, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=2, score=7.36us + ( + 64, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 2, + "_cta_per_sm_write": 4, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 2, + "_num_stages_write": 3, + "_num_warps_nowrite": 4, + "_num_warps_write": 4, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=4, score=8.40us + ( + 128, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 10, + "_cta_per_sm_write": 10, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 3, + "_num_stages_write": 4, + "_num_warps_nowrite": 4, + "_num_warps_write": 4, + "_precompute_num_warps": 16, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=8, score=9.37us + ( + 256, + "persistent_dynamic", + { + "_block_size_m": 16, + "_cta_per_sm": 8, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 16, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=16, score=10.02us + ( + 512, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 7, + "_cta_per_sm_write": 9, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 3, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=32, score=13.15us + ( + 1024, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 10, + "_cta_per_sm_write": 7, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=64, score=17.82us + ( + 2048, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 4, + "_cta_per_sm_write": 3, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 3, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, + "_num_warps_write": 4, + "_precompute_num_warps": 1, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=128, score=27.01us + ( + 4096, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 4, + "_cta_per_sm_write": 3, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 3, + "_num_stages_write": 2, + "_num_warps_nowrite": 2, + "_num_warps_write": 4, + "_precompute_num_warps": 1, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=256, score=43.23us + ( + 8192, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 6, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 3, + "_num_stages_write": 1, + "_num_warps_nowrite": 1, + "_num_warps_write": 4, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=512, score=77.01us + ( + 16384, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 6, + "_cta_per_sm_write": 3, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 1, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, + "_num_warps_write": 4, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=1024, score=140.43us ], } @@ -2177,11 +3126,11 @@ def _persistent_main_kernel( # different from the table's recommendation. _PD_TO_PM_SPLIT_MAP = { # pd unsplit knob → (pm_write_knob, pm_nowrite_knob) "_block_size_m": ("_block_size_m_write", "_block_size_m_nowrite"), - "_num_warps": ("_num_warps_write", "_num_warps_nowrite"), - "_num_stages": ("_num_stages_write", "_num_stages_nowrite"), + "_num_warps": ("_num_warps_write", "_num_warps_nowrite"), + "_num_stages": ("_num_stages_write", "_num_stages_nowrite"), # CPS / LS are persistent-loop knobs; pd uses _cta_per_sm + _num_loop_stages # as unsplit, pm uses _cta_per_sm_write/_nowrite + _num_loop_stages_write/_nowrite. - "_cta_per_sm": ("_cta_per_sm_write", "_cta_per_sm_nowrite"), + "_cta_per_sm": ("_cta_per_sm_write", "_cta_per_sm_nowrite"), "_num_loop_stages": ("_num_loop_stages_write", "_num_loop_stages_nowrite"), } @@ -2213,7 +3162,10 @@ def _bridge_tuning_knobs(knobs: dict, from_mode: str, to_mode: str) -> dict: def _resolve_tuning( - batch: int, nheads_per_rank: int, dt_str: str, sr_str: str, + batch: int, + nheads_per_rank: int, + dt_str: str, + sr_str: str, ) -> tuple[str, dict] | None: """Look up the default mode + knobs for this (eff_batch, dt, sr) cell. @@ -2321,10 +3273,10 @@ def replay_selective_state_update( # TMA state-tensor toggles — 4 independent paths (see replay design notes # item #17 for measured perf profiles). Each is False=raw load/store, True= # use a host-built TMA tensor_descriptor for that path. - _use_tma_rect_load: bool | None = None, # rect kernel's state load (nowrite-only) - _use_tma_replay_write_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=True + _use_tma_rect_load: bool | None = None, # rect kernel's state load (nowrite-only) + _use_tma_replay_write_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=True _use_tma_replay_write_store: bool | None = None, # SSM state store when WRITE_CHECKPOINT=True - _use_tma_replay_nowrite_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=False + _use_tma_replay_nowrite_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=False _use_replay_cache_slot: bool = True, # Persistent-mode tuning kwargs (consulted for both pd and pm; pd uses # _cta_per_sm / _num_loop_stages, pm uses the _write/_nowrite splits): @@ -2543,41 +3495,62 @@ def replay_selective_state_update( # locals() for re-read, so re-bind each kwarg explicitly. if rectangle_for_nowrite is None and "rectangle_for_nowrite" in _table_knobs: rectangle_for_nowrite = bool(_table_knobs["rectangle_for_nowrite"]) - _block_size_m = _block_size_m if _block_size_m is not None else _table_knobs.get("_block_size_m") + _block_size_m = ( + _block_size_m if _block_size_m is not None else _table_knobs.get("_block_size_m") + ) _num_warps = _num_warps if _num_warps is not None else _table_knobs.get("_num_warps") _num_stages = _num_stages if _num_stages is not None else _table_knobs.get("_num_stages") - _heads_per_block = _heads_per_block if _heads_per_block is not None else _table_knobs.get("_heads_per_block") + _heads_per_block = ( + _heads_per_block + if _heads_per_block is not None + else _table_knobs.get("_heads_per_block") + ) _precompute_num_warps = ( _precompute_num_warps if _precompute_num_warps is not None - else _table_knobs.get("_precompute_num_warps")) + else _table_knobs.get("_precompute_num_warps") + ) _precompute_num_stages = ( _precompute_num_stages if _precompute_num_stages is not None - else _table_knobs.get("_precompute_num_stages")) + else _table_knobs.get("_precompute_num_stages") + ) _block_size_m_write = ( _block_size_m_write if _block_size_m_write is not None - else _table_knobs.get("_block_size_m_write")) + else _table_knobs.get("_block_size_m_write") + ) _block_size_m_nowrite = ( _block_size_m_nowrite if _block_size_m_nowrite is not None - else _table_knobs.get("_block_size_m_nowrite")) - _num_warps_write = _num_warps_write if _num_warps_write is not None else _table_knobs.get("_num_warps_write") + else _table_knobs.get("_block_size_m_nowrite") + ) + _num_warps_write = ( + _num_warps_write + if _num_warps_write is not None + else _table_knobs.get("_num_warps_write") + ) _num_warps_nowrite = ( _num_warps_nowrite if _num_warps_nowrite is not None - else _table_knobs.get("_num_warps_nowrite")) + else _table_knobs.get("_num_warps_nowrite") + ) _num_stages_write = ( _num_stages_write if _num_stages_write is not None - else _table_knobs.get("_num_stages_write")) + else _table_knobs.get("_num_stages_write") + ) _num_stages_nowrite = ( _num_stages_nowrite if _num_stages_nowrite is not None - else _table_knobs.get("_num_stages_nowrite")) + else _table_knobs.get("_num_stages_nowrite") + ) _cta_per_sm = _cta_per_sm if _cta_per_sm is not None else _table_knobs.get("_cta_per_sm") - _num_loop_stages = _num_loop_stages if _num_loop_stages is not None else _table_knobs.get("_num_loop_stages") + _num_loop_stages = ( + _num_loop_stages + if _num_loop_stages is not None + else _table_knobs.get("_num_loop_stages") + ) # persistent_main uses split write/nowrite tuning knobs. _num_loop_stages_write = ( _num_loop_stages_write @@ -2600,15 +3573,23 @@ def replay_selective_state_update( else _table_knobs.get("_cta_per_sm_nowrite") ) _flatten = _flatten if _flatten is not None else _table_knobs.get("_flatten") - _warp_specialize = _warp_specialize if _warp_specialize is not None else _table_knobs.get("_warp_specialize") + _warp_specialize = ( + _warp_specialize + if _warp_specialize is not None + else _table_knobs.get("_warp_specialize") + ) if _use_tma_rect_load is None: _use_tma_rect_load = bool(_table_knobs.get("_use_tma_rect_load", False)) if _use_tma_replay_write_load is None: _use_tma_replay_write_load = bool(_table_knobs.get("_use_tma_replay_write_load", False)) if _use_tma_replay_write_store is None: - _use_tma_replay_write_store = bool(_table_knobs.get("_use_tma_replay_write_store", False)) + _use_tma_replay_write_store = bool( + _table_knobs.get("_use_tma_replay_write_store", False) + ) if _use_tma_replay_nowrite_load is None: - _use_tma_replay_nowrite_load = bool(_table_knobs.get("_use_tma_replay_nowrite_load", False)) + _use_tma_replay_nowrite_load = bool( + _table_knobs.get("_use_tma_replay_nowrite_load", False) + ) # Final defaults if neither caller nor table set them (empty table case). if mode is None: mode = "persistent_dynamic" @@ -2802,8 +3783,7 @@ def replay_selective_state_update( heads_per_block = int(_heads_per_block) assert heads_per_block > 0, "heads_per_block must be positive" heads_per_block = min(heads_per_block, heads_per_group) - while (heads_per_group % heads_per_block != 0 - or heads_per_block & (heads_per_block - 1) != 0): + while heads_per_group % heads_per_block != 0 or heads_per_block & (heads_per_block - 1) != 0: heads_per_block -= 1 if _precompute_num_warps is not None: precompute_num_warps = _precompute_num_warps @@ -2812,7 +3792,9 @@ def replay_selective_state_update( # overrides the corresponding shared value for ONE main launch only. # Default (None) = tied to shared value (current behavior). BLOCK_SIZE_M_WRITE = _block_size_m_write if _block_size_m_write is not None else BLOCK_SIZE_M - BLOCK_SIZE_M_NOWRITE = _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M + BLOCK_SIZE_M_NOWRITE = ( + _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M + ) NUM_WARPS_WRITE = _num_warps_write if _num_warps_write is not None else num_warps NUM_WARPS_NOWRITE = _num_warps_nowrite if _num_warps_nowrite is not None else num_warps NUM_STAGES_WRITE = _num_stages_write if _num_stages_write is not None else _num_stages @@ -2820,8 +3802,12 @@ def replay_selective_state_update( # Persistent-only per-main: CTA_PER_SM_WRITE = _cta_per_sm_write if _cta_per_sm_write is not None else _cta_per_sm CTA_PER_SM_NOWRITE = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _cta_per_sm - NUM_LOOP_STAGES_WRITE = _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages - NUM_LOOP_STAGES_NOWRITE = _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages + NUM_LOOP_STAGES_WRITE = ( + _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages + ) + NUM_LOOP_STAGES_NOWRITE = ( + _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages + ) HAS_CACHE_BATCH_INDICES = state_batch_indices is not None @@ -2862,36 +3848,41 @@ def replay_selective_state_update( # dummy; kernels never reference it because their constexprs are all # False (Triton DCEs the dead branches). # `triton.set_allocator()` must run before any descriptor-using launch. - if (_use_tma_rect_load or _use_tma_replay_write_load - or _use_tma_replay_write_store or _use_tma_replay_nowrite_load): + if ( + _use_tma_rect_load + or _use_tma_replay_write_load + or _use_tma_replay_write_store + or _use_tma_replay_nowrite_load + ): from triton.tools.tensor_descriptor import TensorDescriptor + _ensure_tma_allocator() assert state.is_contiguous(), "TMA state requires contiguous state" assert state.stride(-1) == 1, "TMA state requires inner stride 1" _state_flat = state.view(-1, state.shape[-1]) _dstate_pow2 = triton.next_power_of_2(dstate) state_tma_descriptor_write = TensorDescriptor.from_tensor( - _state_flat, block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], + _state_flat, + block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], ) if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: state_tma_descriptor_nowrite = state_tma_descriptor_write else: state_tma_descriptor_nowrite = TensorDescriptor.from_tensor( - _state_flat, block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], + _state_flat, + block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], ) else: - state_tma_descriptor_write = state # dummy; all consuming constexprs False + state_tma_descriptor_write = state # dummy; all consuming constexprs False state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False # Work items are sorted write-first for persistent_main. Each row carries # decode-batch position, cache slot, PNAT, and active cache buffer index. assert isinstance(replay_work_items, torch.Tensor), ( - "replay_work_items must be a torch.Tensor, got " - f"{type(replay_work_items).__name__}" + f"replay_work_items must be a torch.Tensor, got {type(replay_work_items).__name__}" ) assert replay_work_items.device == device, ( - f"replay_work_items must be on device {device}, got " - f"{replay_work_items.device}" + f"replay_work_items must be on device {device}, got {replay_work_items.device}" ) assert replay_work_items.dtype == torch.int32, ( f"replay_work_items must be int32, got {replay_work_items.dtype}" @@ -2905,15 +3896,9 @@ def replay_selective_state_update( assert isinstance(n_writes, torch.Tensor), ( f"n_writes must be a torch.Tensor, got {type(n_writes).__name__}" ) - assert n_writes.device == device, ( - f"n_writes must be on device {device}, got {n_writes.device}" - ) - assert n_writes.dtype == torch.int32, ( - f"n_writes must be int32, got {n_writes.dtype}" - ) - assert n_writes.shape == (1,), ( - f"n_writes must have shape (1,), got {tuple(n_writes.shape)}" - ) + assert n_writes.device == device, f"n_writes must be on device {device}, got {n_writes.device}" + assert n_writes.dtype == torch.int32, f"n_writes must be int32, got {n_writes.dtype}" + assert n_writes.shape == (1,), f"n_writes must have shape (1,), got {tuple(n_writes.shape)}" replay_work_items_arg = replay_work_items precomp_grid = (batch, nheads // heads_per_block) @@ -2926,26 +3911,56 @@ def replay_selective_state_update( def launch_dynamic_precompute(rectangle: bool): _dynamic_precompute_kernel[precomp_grid]( - dt, dt_bias, A, B, C, - cb_scaled, decay_vec, - old_B, old_dt, old_dA_cumsum, - cache_buf_idx, prev_num_accepted_tokens, + dt, + dt_bias, + A, + B, + C, + cb_scaled, + decay_vec, + old_B, + old_dt, + old_dA_cumsum, + cache_buf_idx, + prev_num_accepted_tokens, state_batch_indices, - T, max_window, dstate, nheads // ngroups, - dt.stride(0), dt.stride(1), dt.stride(2), + T, + max_window, + dstate, + nheads // ngroups, + dt.stride(0), + dt.stride(1), + dt.stride(2), dt_bias.stride(0) if dt_bias is not None else 0, A.stride(0), - B.stride(0), B.stride(1), B.stride(2), B.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), + B.stride(0), + B.stride(1), + B.stride(2), + B.stride(3), + C.stride(0), + C.stride(1), + C.stride(2), + C.stride(3), + cb_scaled.stride(0), + cb_scaled.stride(1), + cb_scaled.stride(2), + cb_scaled.stride(3), + decay_vec.stride(0), + decay_vec.stride(1), + decay_vec.stride(2), + old_B.stride(0), + old_B.stride(1), + old_B.stride(2), + old_B.stride(3), + old_B.stride(4), + old_dt.stride(0), + old_dt.stride(1), + old_dt.stride(2), + old_dt.stride(3), + old_dA_cumsum.stride(0), + old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), + old_dA_cumsum.stride(3), dt_softplus, HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, LAUNCH_WITH_PDL=launch_with_pdl, @@ -2985,10 +4000,9 @@ def launch_dynamic_precompute(rectangle: bool): # constexpr.) _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M - def launch_persistent_main(write_checkpoint: bool, - *, - launch_dependent_kernels: bool = False, - rectangle: bool = False): + def launch_persistent_main( + write_checkpoint: bool, *, launch_dependent_kernels: bool = False, rectangle: bool = False + ): # `n_writes` is the (1,) int32 device tensor with the write count. # Both halves always launch; an empty half has a zero-length slot # range and the persistent loop does no work. @@ -3009,34 +4023,86 @@ def launch_persistent_main(write_checkpoint: bool, grid = (min(num_persistent, total_work_launch),) # Per-path TMA descriptor — block_shape[0] must match block_size_m. selected_state_tma_descriptor = ( - state_tma_descriptor_write if write_checkpoint - else state_tma_descriptor_nowrite) + state_tma_descriptor_write if write_checkpoint else state_tma_descriptor_nowrite + ) _persistent_main_kernel[grid]( - state, selected_state_tma_descriptor, state_scales_arg, old_x, - old_B, old_dt, old_dA_cumsum, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, replay_work_items_arg, rand_seed, - n_writes, batch, nheads, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), old_x.stride(4), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + state, + selected_state_tma_descriptor, + state_scales_arg, + old_x, + old_B, + old_dt, + old_dA_cumsum, + prev_num_accepted_tokens, + cache_buf_idx, + x, + C, + D, + z, + out, + cb_scaled, + decay_vec, + state_batch_indices, + replay_work_items_arg, + rand_seed, + n_writes, + batch, + nheads, + T, + max_window, + dim, + dstate, + nheads // ngroups, + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + state_scales_strides[0], + state_scales_strides[1], + state_scales_strides[2], + old_x.stride(0), + old_x.stride(1), + old_x.stride(2), + old_x.stride(3), + old_x.stride(4), + old_B.stride(0), + old_B.stride(1), + old_B.stride(2), + old_B.stride(3), + old_B.stride(4), + old_dt.stride(0), + old_dt.stride(1), + old_dt.stride(2), + old_dt.stride(3), + old_dA_cumsum.stride(0), + old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), + old_dA_cumsum.stride(3), + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + C.stride(0), + C.stride(1), + C.stride(2), + C.stride(3), + d_strides[0], + d_strides[1], + z_strides[0], + z_strides[1], + z_strides[2], + z_strides[3], + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + cb_scaled.stride(0), + cb_scaled.stride(1), + cb_scaled.stride(2), + cb_scaled.stride(3), + decay_vec.stride(0), + decay_vec.stride(1), + decay_vec.stride(2), block_size_m, LAUNCH_WITH_PDL=use_internal_pdl, PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, @@ -3068,9 +4134,11 @@ def launch_persistent_main(write_checkpoint: bool, launch_pdl=use_internal_pdl, ) - def launch_persistent_dynamic_main(n_writes_tensor: torch.Tensor, - launch_dependent_kernels: bool = False, - rectangle: bool = False): + def launch_persistent_dynamic_main( + n_writes_tensor: torch.Tensor, + launch_dependent_kernels: bool = False, + rectangle: bool = False, + ): # Single-launch persistent kernel covering the whole batch with # runtime per-slot WRITE_CHECKPOINT branch. No half-split, no # n_writes needed (the kernel ignores n_writes_tensor when @@ -3088,31 +4156,83 @@ def launch_persistent_dynamic_main(n_writes_tensor: torch.Tensor, # the write-side descriptor matches. Both write and nowrite slots # in this kernel share that BSM. _persistent_main_kernel[grid]( - state, state_tma_descriptor_write, state_scales_arg, old_x, - old_B, old_dt, old_dA_cumsum, - prev_num_accepted_tokens, cache_buf_idx, - x, C, D, z, out, - cb_scaled, decay_vec, - state_batch_indices, replay_work_items_arg, rand_seed, - n_writes_tensor, batch, nheads, - T, max_window, dim, dstate, nheads // ngroups, - state.stride(0), state.stride(1), state.stride(2), state.stride(3), - state_scales_strides[0], state_scales_strides[1], state_scales_strides[2], - old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), old_x.stride(4), - old_B.stride(0), old_B.stride(1), old_B.stride(2), - old_B.stride(3), old_B.stride(4), - old_dt.stride(0), old_dt.stride(1), - old_dt.stride(2), old_dt.stride(3), - old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), - old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - x.stride(0), x.stride(1), x.stride(2), x.stride(3), - C.stride(0), C.stride(1), C.stride(2), C.stride(3), - d_strides[0], d_strides[1], - z_strides[0], z_strides[1], z_strides[2], z_strides[3], - out.stride(0), out.stride(1), out.stride(2), out.stride(3), - cb_scaled.stride(0), cb_scaled.stride(1), - cb_scaled.stride(2), cb_scaled.stride(3), - decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), + state, + state_tma_descriptor_write, + state_scales_arg, + old_x, + old_B, + old_dt, + old_dA_cumsum, + prev_num_accepted_tokens, + cache_buf_idx, + x, + C, + D, + z, + out, + cb_scaled, + decay_vec, + state_batch_indices, + replay_work_items_arg, + rand_seed, + n_writes_tensor, + batch, + nheads, + T, + max_window, + dim, + dstate, + nheads // ngroups, + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + state_scales_strides[0], + state_scales_strides[1], + state_scales_strides[2], + old_x.stride(0), + old_x.stride(1), + old_x.stride(2), + old_x.stride(3), + old_x.stride(4), + old_B.stride(0), + old_B.stride(1), + old_B.stride(2), + old_B.stride(3), + old_B.stride(4), + old_dt.stride(0), + old_dt.stride(1), + old_dt.stride(2), + old_dt.stride(3), + old_dA_cumsum.stride(0), + old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), + old_dA_cumsum.stride(3), + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + C.stride(0), + C.stride(1), + C.stride(2), + C.stride(3), + d_strides[0], + d_strides[1], + z_strides[0], + z_strides[1], + z_strides[2], + z_strides[3], + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + cb_scaled.stride(0), + cb_scaled.stride(1), + cb_scaled.stride(2), + cb_scaled.stride(3), + decay_vec.stride(0), + decay_vec.stride(1), + decay_vec.stride(2), BLOCK_SIZE_M, LAUNCH_WITH_PDL=use_internal_pdl, PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index c184d24da97d..4b9c77405e6b 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -83,8 +83,7 @@ def get_state_indices(self, *args, **kwargs) -> torch.Tensor: ... def get_replay_state_update_metadata( - self - ) -> Optional[ReplayStateUpdateMetadata]: + self) -> Optional[ReplayStateUpdateMetadata]: """Return replay metadata tensors and fixed replay sizes.""" return None @@ -616,8 +615,7 @@ def use_replay_state_update(self) -> bool: return self._use_replay_state_update def get_replay_state_update_metadata( - self - ) -> Optional[ReplayStateUpdateMetadata]: + self) -> Optional[ReplayStateUpdateMetadata]: if (not self._use_replay_state_update or not isinstance(self.mamba_cache, self.SpeculativeState) or self.mamba_cache.prev_num_accepted_tokens is None @@ -824,8 +822,7 @@ def use_replay_state_update(self) -> bool: return getattr(self._impl, 'use_replay_state_update', False) def get_replay_state_update_metadata( - self - ) -> Optional[ReplayStateUpdateMetadata]: + self) -> Optional[ReplayStateUpdateMetadata]: get_metadata = getattr(self._impl, 'get_replay_state_update_metadata', None) if get_metadata is None: @@ -1709,12 +1706,10 @@ def use_replay_state_update(self) -> bool: return self._use_replay_state_update def get_replay_state_update_metadata( - self - ) -> Optional[ReplayStateUpdateMetadata]: + self) -> Optional[ReplayStateUpdateMetadata]: if (not self._use_replay_state_update or self.prev_num_accepted_tokens is None - or self.cache_buf_idx is None - or self.replay_step_width is None + or self.cache_buf_idx is None or self.replay_step_width is None or self.replay_history_size is None): return None return ReplayStateUpdateMetadata( diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index d07198cad7e4..8bc3f507ad37 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -227,9 +227,7 @@ def _load(mod_name: str, file_name: str): # from mamba2_metadata at import time. Stub those constants instead of # importing the full metadata module and its scheduler/attention deps. metadata_stub = types.ModuleType(f"{mamba_pkg}.mamba2_metadata") - metadata_stub.REPLAY_WORK_POSITION_IN_DECODE_BATCH = ( - REPLAY_WORK_POSITION_IN_DECODE_BATCH - ) + metadata_stub.REPLAY_WORK_POSITION_IN_DECODE_BATCH = REPLAY_WORK_POSITION_IN_DECODE_BATCH metadata_stub.REPLAY_WORK_CACHE_SLOT = REPLAY_WORK_CACHE_SLOT metadata_stub.REPLAY_WORK_PNAT = REPLAY_WORK_PNAT metadata_stub.REPLAY_WORK_CACHE_BUF_IDX = REPLAY_WORK_CACHE_BUF_IDX @@ -328,9 +326,7 @@ def _resolve_prev_ks(args, mtp_len: int) -> list[int]: upper = getattr(args, "max_window", 0) or mtp_len if getattr(args, "prev_tokens_int", None): return sorted(set(max(0, min(upper, int(v))) for v in args.prev_tokens_int)) - return sorted( - set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs) - ) + return sorted(set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs)) def _al_counts_to_distribution( @@ -343,8 +339,7 @@ def _al_counts_to_distribution( raise SystemExit(f"no numeric AL rows found in {label}") if max(al_to_count) > T: raise SystemExit( - f"AL histogram {label} contains accepted length {max(al_to_count)} " - f"> T={T}" + f"AL histogram {label} contains accepted length {max(al_to_count)} > T={T}" ) if min(al_to_count) < 0: raise SystemExit(f"AL histogram {label} contains negative accepted lengths") @@ -379,9 +374,7 @@ def _load_al_distribution(path: Path, T: int, column: int = 1) -> np.ndarray: count = float(row[column]) except (IndexError, ValueError): continue - al_to_count[accepted_length] = ( - al_to_count.get(accepted_length, 0.0) + count - ) + al_to_count[accepted_length] = al_to_count.get(accepted_length, 0.0) + count return _al_counts_to_distribution(al_to_count, T, str(path)) @@ -420,8 +413,7 @@ def _markov_stationary(al_dist: np.ndarray, T: int, window: int) -> np.ndarray: idx = int(np.argmin(np.abs(eigvals - 1.0))) if abs(eigvals[idx] - 1.0) > 1e-6: raise SystemExit( - "stationary distribution eigensolve failed: closest eigenvalue " - f"to 1 is {eigvals[idx]}" + f"stationary distribution eigensolve failed: closest eigenvalue to 1 is {eigvals[idx]}" ) pi = np.real(eigvecs[:, idx]) @@ -471,27 +463,21 @@ def _build_replay_work_items_cpu( n_samples, batch = samples.shape write_mask = samples + T > window order = np.argsort(-write_mask.astype(np.int8), kind="stable", axis=1) - positions = np.broadcast_to( - np.arange(batch, dtype=np.int32), (n_samples, batch) - ) + positions = np.broadcast_to(np.arange(batch, dtype=np.int32), (n_samples, batch)) if cache_buf_idx_samples is None: cache_buf_idx_samples = np.zeros_like(samples, dtype=np.int32) else: cache_buf_idx_samples = np.asarray(cache_buf_idx_samples, dtype=np.int32) if cache_buf_idx_samples.ndim == 1: - cache_buf_idx_samples = np.broadcast_to( - cache_buf_idx_samples[None, :], samples.shape - ) + cache_buf_idx_samples = np.broadcast_to(cache_buf_idx_samples[None, :], samples.shape) if cache_buf_idx_samples.shape != samples.shape: raise ValueError( "cache_buf_idx_samples shape must match pnat_samples, got " f"{cache_buf_idx_samples.shape} and {samples.shape}" ) - work_items = np.empty( - (n_samples, batch, REPLAY_WORK_ITEM_WIDTH), dtype=np.int32 - ) + work_items = np.empty((n_samples, batch, REPLAY_WORK_ITEM_WIDTH), dtype=np.int32) work_items[:, :, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = np.take_along_axis( positions, order, axis=1 ) @@ -558,8 +544,7 @@ def _build_tensors( device = "cuda" # Cache lookup — grow batch in place if needed; else return views. - cache_key = (state_dtype, act_dtype, max_window, mtp_len, - nheads, head_dim, d_state, ngroups) + cache_key = (state_dtype, act_dtype, max_window, mtp_len, nheads, head_dim, d_state, ngroups) cached = _TENSOR_CACHE.get(cache_key) if cached is not None and cached["max_batch"] >= batch: # Hit — return slices for current batch. @@ -637,9 +622,7 @@ def _build_tensors( else: state0 = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) else: - state0 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=state_dtype - ) + state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=state_dtype) state_scales0 = None # --- Cache tensors for replay kernel --- @@ -669,12 +652,8 @@ def _build_tensors( # prev_tokens placeholder — overwritten per-run prev_tokens = torch.zeros(batch, device=device, dtype=torch.int32) state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) - replay_work_items = torch.empty( - batch, REPLAY_WORK_ITEM_WIDTH, device=device, dtype=torch.int32 - ) - replay_work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = ( - state_batch_indices - ) + replay_work_items = torch.empty(batch, REPLAY_WORK_ITEM_WIDTH, device=device, dtype=torch.int32) + replay_work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = state_batch_indices replay_work_items[:, REPLAY_WORK_CACHE_SLOT] = state_batch_indices replay_work_items[:, REPLAY_WORK_PNAT] = 0 replay_work_items[:, REPLAY_WORK_CACHE_BUF_IDX] = 0 @@ -800,8 +779,7 @@ def _build_tensors( ) -def _sr_modes_for_dtype(state_dtype: torch.dtype, - requested_modes: list[str]) -> list[str]: +def _sr_modes_for_dtype(state_dtype: torch.dtype, requested_modes: list[str]) -> list[str]: """Return the rounding modes that should run for one state dtype.""" if state_dtype == torch.float32: return ["RN"] @@ -830,9 +808,7 @@ def _kernels_per_iter_incremental( elif mode == "persistent_main": k = 3 # 1 dynamic_precomp + write-main + nowrite-main else: - raise ValueError( - f"mode must be resolved before CUPTI kernel counting, got {mode!r}" - ) + raise ValueError(f"mode must be resolved before CUPTI kernel counting, got {mode!r}") if with_conv1d: k += 1 return k @@ -1003,23 +979,27 @@ def _parse_cupti_buffer_ptr(libcupti, buffer_ptr: int, valid_size: int, *, inclu zero_ts_names[name] = zero_ts_names.get(name, 0) + 1 continue if include_names: - records.append(( - name, - int(kernel.start), - int(kernel.end), - int(kernel.correlation_id), - 0, - int(kernel.graph_node_id), - int(kernel.stream_id), - )) + records.append( + ( + name, + int(kernel.start), + int(kernel.end), + int(kernel.correlation_id), + 0, + int(kernel.graph_node_id), + int(kernel.stream_id), + ) + ) else: - records.append(( - int(kernel.start), - int(kernel.end), - int(kernel.correlation_id), - int(kernel.graph_node_id), - int(kernel.stream_id), - )) + records.append( + ( + int(kernel.start), + int(kernel.end), + int(kernel.correlation_id), + int(kernel.graph_node_id), + int(kernel.stream_id), + ) + ) elif result == _CUPTI_ERROR_MAX_LIMIT_REACHED: break elif result == _CUPTI_ERROR_INVALID_KIND: @@ -1082,13 +1062,17 @@ def _cupti_parser_worker(input_queue, output_queue, ready_event) -> None: include_names=False, ) records_by_generation.setdefault(generation, []).extend(records) - zero_ts_by_generation[generation] = zero_ts_by_generation.get(generation, 0) + zero_ts_count + zero_ts_by_generation[generation] = ( + zero_ts_by_generation.get(generation, 0) + zero_ts_count + ) ctypes.memset(parser_ptr, 0, len(shm.buf)) except Exception as exc: # pragma: no cover - diagnostic worker path output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) finally: del shared_char - output_queue.put({"kind": "buffer_done", "generation": generation, "buffer_id": buffer_id}) + output_queue.put( + {"kind": "buffer_done", "generation": generation, "buffer_id": buffer_id} + ) elif kind == "finish": if len(item) == 4: _, generation, filter_plan, stats_request = item @@ -1116,21 +1100,25 @@ def _cupti_parser_worker(input_queue, output_queue, ready_event) -> None: ) parser_stats_ms = 1000.0 * (time.perf_counter() - stats_start_s) filtered_records = [] - output_queue.put({ - "kind": "finish_done", - "generation": generation, - "records": filtered_records, - "zero_ts_count": zero_ts_count, - "zero_ts_names": {}, - "raw_record_count": len(raw_records), - "stats": stats, - "stats_ready": stats_ready, - "parser_stats_ms": parser_stats_ms, - }) + output_queue.put( + { + "kind": "finish_done", + "generation": generation, + "records": filtered_records, + "zero_ts_count": zero_ts_count, + "zero_ts_names": {}, + "raw_record_count": len(raw_records), + "stats": stats, + "stats_ready": stats_ready, + "parser_stats_ms": parser_stats_ms, + } + ) except Exception as exc: # pragma: no cover - diagnostic worker path output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) else: - output_queue.put({"kind": "error", "generation": -1, "error": f"unknown parser message {kind!r}"}) + output_queue.put( + {"kind": "error", "generation": -1, "error": f"unknown parser message {kind!r}"} + ) for shm in shared_blocks.values(): shm.close() @@ -1219,9 +1207,11 @@ def __init__(self) -> None: if spawn_ok: last_err = None break - last_err = (f"attempt {_spawn_attempt + 1}: " - f"alive={self._parse_process.is_alive()}, " - f"exitcode={self._parse_process.exitcode}") + last_err = ( + f"attempt {_spawn_attempt + 1}: " + f"alive={self._parse_process.is_alive()}, " + f"exitcode={self._parse_process.exitcode}" + ) try: if self._parse_process.is_alive(): self._parse_process.terminate() @@ -1240,10 +1230,12 @@ def __init__(self) -> None: self._request_callback = self._request_callback_type(self._request_buffer) self._complete_callback = self._complete_callback_type(self._complete_buffer) - self._check(self._libcupti.cuptiActivityRegisterCallbacks( - self._request_callback, - self._complete_callback, - )) + self._check( + self._libcupti.cuptiActivityRegisterCallbacks( + self._request_callback, + self._complete_callback, + ) + ) self._check(self._libcupti.cuptiActivityEnable(_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL)) atexit.register(self.close) @@ -1397,9 +1389,7 @@ def _begin( phase_start_s = time.perf_counter() if collect_timing else 0.0 self._set_flush_period_ms(flush_period_ms) if collect_timing: - start_timing["period_enable_ms"] = 1000.0 * ( - time.perf_counter() - phase_start_s - ) + start_timing["period_enable_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) self._last_start_timing = start_timing return generation @@ -1452,9 +1442,7 @@ def stop_async( phase_start_s = time.perf_counter() if collect_timing else 0.0 self._set_flush_period_ms(0) if collect_timing: - stop_timing["period_disable_ms"] = 1000.0 * ( - time.perf_counter() - phase_start_s - ) + stop_timing["period_disable_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) phase_start_s = time.perf_counter() if collect_timing else 0.0 self._flush(0) if collect_timing: @@ -1480,9 +1468,7 @@ def wait_for_generation_result( result = self._finish_results.pop(generation, None) if result is not None: if collect_timing: - stop_timing["parser_wait_ms"] = 1000.0 * ( - time.perf_counter() - phase_start_s - ) + stop_timing["parser_wait_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) stop_timing["total_ms"] = ( stop_timing.get("period_disable_ms", 0.0) + stop_timing.get("flush_ms", 0.0) @@ -1557,10 +1543,7 @@ def _stats_from_spans(spans_us: list[float]) -> dict: def _binomial_pmf(n: int, p: float) -> tuple[float, ...]: - return tuple( - math.comb(n, k) * (p**k) * ((1.0 - p)**(n - k)) - for k in range(n + 1) - ) + return tuple(math.comb(n, k) * (p**k) * ((1.0 - p) ** (n - k)) for k in range(n + 1)) def _kmix_bucket_score( @@ -1595,10 +1578,16 @@ def _kmix_bucket_score( return numerator / denominator if denominator > 0.0 else None -def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, - zero_ts_count: int = 0, - zero_ts_names: dict | None = None, - include_details: bool = True): +def _stats_from_cupti_records( + records, + warmup, + iters, + tag, + expected_K, + zero_ts_count: int = 0, + zero_ts_names: dict | None = None, + include_details: bool = True, +): """Bin a flat CUPTI kernel record stream into per-iter spans + per-kernel relative timestamps. Used by both graph and eager CUPTI paths. @@ -1611,7 +1600,8 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, them out). """ records = [ - r for r in records + r + for r in records if r[0] is not None and any(s in r[0] for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS) ] records.sort(key=lambda r: r[1]) # by start_ns @@ -1621,6 +1611,7 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, expected_total = expected_K * expected_iters if total != expected_total: from collections import Counter + name_counts = dict(Counter(r[0] for r in records)) # Non-fatal: skip this cell instead of killing the whole sweep. # Mismatch may be a CUPTI dropped-records issue (rare configs), @@ -1659,16 +1650,17 @@ def _stats_from_cupti_records(records, warmup, iters, tag, expected_K, flush=True, ) if len(records) > 30: - print(f" ... ({len(records) - 30} more records elided)", - file=sys.stderr, flush=True) + print( + f" ... ({len(records) - 30} more records elided)", file=sys.stderr, flush=True + ) return None K = expected_K - timed = records[warmup * K:] + timed = records[warmup * K :] spans_us: list[float] = [] per_kernel: dict[str, dict[str, list[float]]] = {} for i in range(iters): - chunk = timed[i * K:(i + 1) * K] + chunk = timed[i * K : (i + 1) * K] iter_start_ns = min(r[1] for r in chunk) iter_end_ns = max(r[2] for r in chunk) spans_us.append((iter_end_ns - iter_start_ns) / 1000.0) @@ -1719,7 +1711,6 @@ def attach(self, stats: dict | None) -> None: class _PendingCuptiStats: - def __init__( self, timer: CuptiKernelTimer, @@ -1830,8 +1821,9 @@ def _graph_group_iters(args, total_iters: int, pre_iter_fn, pre_iter_group_facto return max(1, int(requested)) -def _get_cupti_filter_plan(timer: CuptiKernelTimer, graph, cache_key: tuple | None, - group_iters: int) -> tuple[int, tuple[str | None, ...]]: +def _get_cupti_filter_plan( + timer: CuptiKernelTimer, graph, cache_key: tuple | None, group_iters: int +) -> tuple[int, tuple[str | None, ...]]: full_cache_key = None if cache_key is None else (cache_key, group_iters) if full_cache_key is not None: cached = _CUPTI_FILTER_PLAN_CACHE.get(full_cache_key) @@ -1945,9 +1937,10 @@ def _time_kernel_cuda_graph( host_timing.stop("graph_preload_ms") plan_cache_key = None if cupti_plan_key is None else (cupti_plan_key, group_iters) - host_timing.add("cupti_plan_cached", ( - plan_cache_key is not None and plan_cache_key in _CUPTI_FILTER_PLAN_CACHE - )) + host_timing.add( + "cupti_plan_cached", + (plan_cache_key is not None and plan_cache_key in _CUPTI_FILTER_PLAN_CACHE), + ) host_timing.start() records_per_replay, ordinal_names = _get_cupti_filter_plan( timer, @@ -2140,8 +2133,11 @@ def _run_kernel_untimed(args, run_fn, reset_fn, tag: str) -> dict: if args.cuda_graph: # Eager warmup before capture (Triton autotune) - reset_fn(); run_fn(); torch.cuda.synchronize() - reset_fn(); torch.cuda.synchronize() + reset_fn() + run_fn() + torch.cuda.synchronize() + reset_fn() + torch.cuda.synchronize() g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): for _ in range(warmup + iters): @@ -2172,7 +2168,10 @@ def _run_kernel_untimed(args, run_fn, reset_fn, tag: str) -> dict: def _time_kernel( - args, run_fn, reset_fn, tag: str, + args, + run_fn, + reset_fn, + tag: str, *, expected_K: int, pre_iter_fn=None, @@ -2200,7 +2199,10 @@ def _time_kernel( return _run_kernel_untimed(args, run_fn, reset_fn, tag) if args.cuda_graph: return _time_kernel_cuda_graph( - args, run_fn, reset_fn, tag, + args, + run_fn, + reset_fn, + tag, expected_K=expected_K, pre_iter_fn=pre_iter_fn, pre_iter_group_factory=pre_iter_group_factory, @@ -2208,7 +2210,10 @@ def _time_kernel( cupti_plan_key=cupti_plan_key, ) return _time_kernel_eager( - args, run_fn, reset_fn, tag, + args, + run_fn, + reset_fn, + tag, expected_K=expected_K, pre_iter_fn=pre_iter_fn, iters_override=iters_override, @@ -2242,12 +2247,16 @@ def _warm_one_config(args, cfg, baseline_fn) -> None: function references across processes. """ outer_cfg, inner_overrides_or_list = cfg - overrides_list = (inner_overrides_or_list - if isinstance(inner_overrides_or_list, list) - else [inner_overrides_or_list]) - (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, - rect, mode, hardcode_sort) = outer_cfg + overrides_list = ( + inner_overrides_or_list + if isinstance(inner_overrides_or_list, list) + else [inner_overrides_or_list] + ) + (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, rect, mode, hardcode_sort) = ( + outer_cfg + ) import argparse as _ap + for inner_overrides in overrides_list: # Fresh clone per entry: prevents knob-value leakage between # consecutive cells in a CPS-grouped task (entries may set @@ -2256,19 +2265,30 @@ def _warm_one_config(args, cfg, baseline_fn) -> None: for k, v in inner_overrides.items(): setattr(args_copy, k, v) _bench_config( - args_copy, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn, - sr_mode=sr_mode, rectangle_for_nowrite=rect, mode=mode, + args_copy, + batch, + mtp_len, + prev_ks, + state_dtype, + act_dtype, + baseline_fn, + sr_mode=sr_mode, + rectangle_for_nowrite=rect, + mode=mode, hardcode_sort=hardcode_sort, warmup_only=True, ) -def _compile_warmup_phase(args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, - baseline_fn, max_workers: int) -> None: +def _compile_warmup_phase( + args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, baseline_fn, max_workers: int +) -> None: _cw_t0 = time.perf_counter() + def _cw(label: str) -> None: dt = time.perf_counter() - _cw_t0 print(f"[compile-warmup] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) + _cw("entered _compile_warmup_phase") """Parallel compile-warmup using a ProcessPoolExecutor with `spawn` start method. @@ -2292,8 +2312,8 @@ def _cw(label: str) -> None: avoid pickling complications; the parent compiles the baseline kernel itself before launching the pool when applicable. """ - from concurrent.futures import ProcessPoolExecutor import multiprocessing + from concurrent.futures import ProcessPoolExecutor sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) @@ -2321,19 +2341,22 @@ def _cw(label: str) -> None: for sr_mode in _sr_modes_for_dtype(state_dtype, sr_modes_list): for mode in modes_list: for rect in rect_list: - can_sort = ( - args.mix_csv is not None - or getattr(args, "pmix", False) - ) - effective_hsort_list = ( - hsort_list if can_sort else [False] - ) + can_sort = args.mix_csv is not None or getattr(args, "pmix", False) + effective_hsort_list = hsort_list if can_sort else [False] for hardcode_sort in effective_hsort_list: - configs.append(( - batch, mtp_len, prev_ks, - state_dtype, act_dtype, sr_mode, - rect, mode, hardcode_sort, - )) + configs.append( + ( + batch, + mtp_len, + prev_ks, + state_dtype, + act_dtype, + sr_mode, + rect, + mode, + hardcode_sort, + ) + ) # Enumerate inner-knob signatures. Two paths: # (1) --cell-list mode (preferred when set): pull exactly the cells @@ -2402,14 +2425,15 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): # tasks then failed the cell-list filter inside the worker and # wasted dispatch overhead. This path is O(|unique cell groups|). from collections import defaultdict as _dd + cell_groups: dict = _dd(list) for tup in cell_set: d = dict(zip(cell_keys, tup)) cell_outer = ( - "SR" if d.get("SR", 0) else "RN", # sr_mode - bool(d.get("RECT", 0)), # rect - d.get("MODE", "persistent_dynamic"), # mode - bool(d.get("HSORT", 0)), # hardcode_sort + "SR" if d.get("SR", 0) else "RN", # sr_mode + bool(d.get("RECT", 0)), # rect + d.get("MODE", "persistent_dynamic"), # mode + bool(d.get("HSORT", 0)), # hardcode_sort ) inner = {} for k, v in d.items(): @@ -2440,31 +2464,48 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): # ============== SWEEP-ARGS PATH ============== # Build inner_dicts via cartesian over knob axes, then cross with # the `configs` outer cartesian. Existing behavior. - m_pairs = _split_or_pair("block_size_m", "block_size_m_write", "block_size_m_nowrite") - w_pairs = _split_or_pair("num_warps", "num_warps_write", "num_warps_nowrite") - ns_pairs = _split_or_pair("num_stages", "num_stages_write", "num_stages_nowrite") - cps_pairs = _split_or_pair("cta_per_sm", "cta_per_sm_write", "cta_per_sm_nowrite") - ls_pairs = _split_or_pair("num_loop_stages", "num_loop_stages_write", "num_loop_stages_nowrite") - pw_vals = _ps(args.precompute_num_warps) - ps_vals = _ps(args.precompute_num_stages) - h_vals = _ps(args.heads_per_block) - mr_vals = _ps(args.maxnreg) - ct_vals = _ps(args.num_ctas) - fl_vals = _ps(args.flatten) - wsp_vals = _ps(args.warp_specialize) - trl_vals = _ps(args.use_tma_rect_load) - twl_vals = _ps(args.use_tma_replay_write_load) - tnl_vals = _ps(args.use_tma_replay_nowrite_load) - tws_vals = _ps(args.use_tma_replay_write_store) + m_pairs = _split_or_pair("block_size_m", "block_size_m_write", "block_size_m_nowrite") + w_pairs = _split_or_pair("num_warps", "num_warps_write", "num_warps_nowrite") + ns_pairs = _split_or_pair("num_stages", "num_stages_write", "num_stages_nowrite") + cps_pairs = _split_or_pair("cta_per_sm", "cta_per_sm_write", "cta_per_sm_nowrite") + ls_pairs = _split_or_pair( + "num_loop_stages", "num_loop_stages_write", "num_loop_stages_nowrite" + ) + pw_vals = _ps(args.precompute_num_warps) + ps_vals = _ps(args.precompute_num_stages) + h_vals = _ps(args.heads_per_block) + mr_vals = _ps(args.maxnreg) + ct_vals = _ps(args.num_ctas) + fl_vals = _ps(args.flatten) + wsp_vals = _ps(args.warp_specialize) + trl_vals = _ps(args.use_tma_rect_load) + twl_vals = _ps(args.use_tma_replay_write_load) + tnl_vals = _ps(args.use_tma_replay_nowrite_load) + tws_vals = _ps(args.use_tma_replay_write_store) import itertools as _it + inner_dicts = [] - for ((mw, mnw), (ww, wnw), (sw, snw), (cw, cnw), (lw, lnw), - pw, ps_, h, mr, ct, fl, wsp, - trl, twl, tnl, tws) in _it.product( - m_pairs, w_pairs, ns_pairs, cps_pairs, ls_pairs, - pw_vals, ps_vals, h_vals, mr_vals, ct_vals, - fl_vals, wsp_vals, - trl_vals, twl_vals, tnl_vals, tws_vals): + for (mw, mnw), (ww, wnw), (sw, snw), (cw, cnw), ( + lw, + lnw, + ), pw, ps_, h, mr, ct, fl, wsp, trl, twl, tnl, tws in _it.product( + m_pairs, + w_pairs, + ns_pairs, + cps_pairs, + ls_pairs, + pw_vals, + ps_vals, + h_vals, + mr_vals, + ct_vals, + fl_vals, + wsp_vals, + trl_vals, + twl_vals, + tnl_vals, + tws_vals, + ): d = {} for k, v in ( ("block_size_m_write", mw), @@ -2508,14 +2549,17 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): # Shuffle ACROSS tasks (preserve within-group CPS sequence for in-process # cache adjacency — within-group order is intentional, not shuffled). import random as _r + _r.shuffle(tasks) _cw(f"built {len(tasks)} tasks covering {n_total_cells} cells in {n_groups} groups") - print(f"[compile-warmup] {len(tasks)} compile tasks " - f"({n_outer_used} outer × {n_groups} cell-groups " - f"covering {n_total_cells} cells, CPS-grouped" - + (", per-cell outer" if cell_set else "") - + f") across {max_workers} processes (ProcessPoolExecutor, {_MP_START_METHOD} start)") + print( + f"[compile-warmup] {len(tasks)} compile tasks " + f"({n_outer_used} outer × {n_groups} cell-groups " + f"covering {n_total_cells} cells, CPS-grouped" + + (", per-cell outer" if cell_set else "") + + f") across {max_workers} processes (ProcessPoolExecutor, {_MP_START_METHOD} start)" + ) t0 = time.perf_counter() ctx = multiprocessing.get_context(_MP_START_METHOD) @@ -2526,10 +2570,7 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): # baseline_fn=None: workers compile only the replay kernel. # Baseline kernels (if any) get compiled lazily in the parent during # the timing phase — usually just one extra compile, negligible. - futures = { - ex.submit(_warm_one_config, args, task, None): task - for task in tasks - } + futures = {ex.submit(_warm_one_config, args, task, None): task for task in tasks} _cw(f"submitted {len(futures)} tasks, waiting for results") _n_done = 0 for fut in futures: @@ -2539,17 +2580,18 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): errors.append((futures[fut], e)) _n_done += 1 # Progress beacons at 10/25/50/75/100% to gauge effective parallelism. - if _n_done in (max(1, len(futures)//10), - max(1, len(futures)//4), - max(1, len(futures)//2), - max(1, (3*len(futures))//4), - len(futures)): + if _n_done in ( + max(1, len(futures) // 10), + max(1, len(futures) // 4), + max(1, len(futures) // 2), + max(1, (3 * len(futures)) // 4), + len(futures), + ): _cw(f"{_n_done}/{len(futures)} tasks complete") if errors: for cfg, e in errors: - print(f"[compile-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", - file=sys.stderr) + print(f"[compile-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", file=sys.stderr) raise errors[0][1] print(f"[compile-warmup] done in {time.perf_counter() - t0:.1f}s") @@ -2632,7 +2674,7 @@ def _bench_config( head_dim = args.head_dim d_state = args.d_state with_conv1d = getattr(args, "with_conv1d", False) - use_philox = (sr_mode == "SR") + use_philox = sr_mode == "SR" # SR rounding: allow fp16 and the quantized dtypes (int8/int16/fp8). # bf16/fp32 SR is not supported (no PTX path for bf16; fp32 doesn't need @@ -2643,9 +2685,7 @@ def _bench_config( if state_dtype not in _SR_SUPPORTED_DTYPES: return rand_seed = torch.randint(0, 2**62, (1,), device="cuda", dtype=torch.int64) - mode = _resolve_effective_replay_mode( - args, batch, state_dtype, use_philox, mode - ) + mode = _resolve_effective_replay_mode(args, batch, state_dtype, use_philox, mode) if mode != "persistent_main" and nowrite_first: return @@ -2703,7 +2743,10 @@ def _baseline_supports() -> bool: return False assert is_pr3324_baseline, args.baseline if state_dtype not in ( - torch.float32, torch.float16, torch.int8, torch.float8_e4m3fn, + torch.float32, + torch.float16, + torch.int8, + torch.float8_e4m3fn, ): return False return not (use_philox and state_dtype == torch.float32) @@ -2764,11 +2807,13 @@ def _parse_sweep(val): num_loop_stages_values = _parse_sweep(args.num_loop_stages) flatten_values = _parse_sweep(args.flatten) warp_specialize_values = _parse_sweep(args.warp_specialize) + # Per-main split-knob sweeps. Default = same as the shared sweep (so each # combo is tied). When set independently, the inner loop sweeps the # cross-product (write × nowrite); --skip-diagonal drops the tied subset. def _split_or_share(split_csv, shared_values): return _parse_sweep(split_csv) if split_csv else shared_values + block_size_m_write_values = _split_or_share(args.block_size_m_write, block_size_m_values) block_size_m_nowrite_values = _split_or_share(args.block_size_m_nowrite, block_size_m_values) num_warps_write_values = _split_or_share(args.num_warps_write, num_warps_values) @@ -2777,19 +2822,31 @@ def _split_or_share(split_csv, shared_values): num_stages_nowrite_values = _split_or_share(args.num_stages_nowrite, num_stages_values) cta_per_sm_write_values = _split_or_share(args.cta_per_sm_write, cta_per_sm_values) cta_per_sm_nowrite_values = _split_or_share(args.cta_per_sm_nowrite, cta_per_sm_values) - num_loop_stages_write_values = _split_or_share(args.num_loop_stages_write, num_loop_stages_values) - num_loop_stages_nowrite_values = _split_or_share(args.num_loop_stages_nowrite, num_loop_stages_values) + num_loop_stages_write_values = _split_or_share( + args.num_loop_stages_write, num_loop_stages_values + ) + num_loop_stages_nowrite_values = _split_or_share( + args.num_loop_stages_nowrite, num_loop_stages_values + ) # Whether any *_write / *_nowrite knob was independently set — used by # --skip-diagonal to know if the cross-product is non-trivial. Without # any split, the per-main values == shared values and skip-diagonal is # a no-op (which is correct). - _any_split = any(getattr(args, name) for name in ( - "block_size_m_write", "block_size_m_nowrite", - "num_warps_write", "num_warps_nowrite", - "num_stages_write", "num_stages_nowrite", - "cta_per_sm_write", "cta_per_sm_nowrite", - "num_loop_stages_write", "num_loop_stages_nowrite", - )) + _any_split = any( + getattr(args, name) + for name in ( + "block_size_m_write", + "block_size_m_nowrite", + "num_warps_write", + "num_warps_nowrite", + "num_stages_write", + "num_stages_nowrite", + "cta_per_sm_write", + "cta_per_sm_nowrite", + "num_loop_stages_write", + "num_loop_stages_nowrite", + ) + ) # TMA toggles — independent 0/1 sweep per path. The skip-dupe at the # top of the inner loop body collapses cells where a flag's path is # unreachable for the current rectangle_for_nowrite setting. @@ -2812,13 +2869,15 @@ def _split_or_share(split_csv, shared_values): for prev_k in prev_ks: # Persistent modes dispatch per-slot from PNAT, so any prev_k # <= max_window is valid. - scenarios.append({ - "label": f"k{prev_k}", - "print_label": prev_k, - "fill": prev_k, - "pre_iter": None, - "iters": None, # use args.iters - }) + scenarios.append( + { + "label": f"k{prev_k}", + "print_label": prev_k, + "fill": prev_k, + "pre_iter": None, + "iters": None, # use args.iters + } + ) # Mix scenario: bench pre-bakes per-iter PNAT samples, n_writes, and # replay_work_items. Grouped graph capture copies window rows into the # persistent kernel-input tensors before each in-graph L2 flush, so timed @@ -2835,15 +2894,15 @@ def _split_or_share(split_csv, shared_values): ) samples_gpu = torch.from_numpy(src).to(device=device, dtype=torch.int32) - n_writes_per_iter_all, replay_work_items_samples_cpu = ( - _build_replay_work_items_cpu(src, mtp_len, max_window) + n_writes_per_iter_all, replay_work_items_samples_cpu = _build_replay_work_items_cpu( + src, mtp_len, max_window ) n_writes_samples_gpu = torch.from_numpy(n_writes_per_iter_all).to( device=device, dtype=torch.int32 ) - replay_work_items_samples_gpu = torch.from_numpy( - replay_work_items_samples_cpu - ).to(device=device, dtype=torch.int32) + replay_work_items_samples_gpu = torch.from_numpy(replay_work_items_samples_cpu).to( + device=device, dtype=torch.int32 + ) n_writes_mix = torch.zeros(1, dtype=torch.int32, device=device) def _mix_pre_iter( @@ -2856,7 +2915,7 @@ def _mix_pre_iter( _rwi=replay_work_items_buf, ): _pt.copy_(_s[i]) - _nw.copy_(_ns[i:i + 1]) + _nw.copy_(_ns[i : i + 1]) _rwi.copy_(_wi[i]) def _mix_pre_iter_group_factory( @@ -2869,11 +2928,11 @@ def _mix_pre_iter_group_factory( _rwi=replay_work_items_buf, ): sample_window = torch.empty( - (group_iters, _s.shape[1]), device=_s.device, dtype=_s.dtype, - ) - n_writes_window = torch.empty( - (group_iters,), device=_ns.device, dtype=_ns.dtype + (group_iters, _s.shape[1]), + device=_s.device, + dtype=_s.dtype, ) + n_writes_window = torch.empty((group_iters,), device=_ns.device, dtype=_ns.dtype) work_items_window = torch.empty( (group_iters, _wi.shape[1], _wi.shape[2]), device=_wi.device, @@ -2889,26 +2948,28 @@ def _pre_replay(replay_idx): def _graph_pre_iter(j): _pt.copy_(sample_window[j]) - _nw.copy_(n_writes_window[j:j + 1]) + _nw.copy_(n_writes_window[j : j + 1]) _rwi.copy_(work_items_window[j]) return _pre_replay, _graph_pre_iter # Mix iters override: if --mix-iters set, use it; else use args.iters. mix_iters = getattr(args, "mix_iters", None) - scenarios.append({ - "label": f"mix{mix_label}", - "print_label": "mix", - "fill": None, - "pre_iter": _mix_pre_iter, - "pre_iter_group_factory": _mix_pre_iter_group_factory, - "iters": mix_iters, # None => use args.iters - "n_writes": n_writes_mix, - # Full per-iter n_writes array (size = warmup + iters). Used by - # the JSON-detailed output to pair each iter's span with its - # mix composition for post-hoc bucketing analysis. - "n_writes_per_iter": n_writes_per_iter_all, - }) + scenarios.append( + { + "label": f"mix{mix_label}", + "print_label": "mix", + "fill": None, + "pre_iter": _mix_pre_iter, + "pre_iter_group_factory": _mix_pre_iter_group_factory, + "iters": mix_iters, # None => use args.iters + "n_writes": n_writes_mix, + # Full per-iter n_writes array (size = warmup + iters). Used by + # the JSON-detailed output to pair each iter's span with its + # mix composition for post-hoc bucketing analysis. + "n_writes_per_iter": n_writes_per_iter_all, + } + ) # Pure scenarios use one constant n_writes/work-items row. Mix scenarios # update both tensors per iter. persistent_main always launches both @@ -2926,9 +2987,7 @@ def _graph_pre_iter(j): device=state_work.device, dtype=torch.int32 ) replay_work_items_buf.copy_( - torch.from_numpy(work_items_cpu).to( - device=state_work.device, dtype=torch.int32 - ) + torch.from_numpy(work_items_cpu).to(device=state_work.device, dtype=torch.int32) ) prev_k_for_print = scn["print_label"] scenario_pre_iter = scn["pre_iter"] @@ -2946,18 +3005,13 @@ def _graph_pre_iter(j): else: nw_full = scn.get("n_writes_per_iter") if nw_full is not None: - scenario_per_iter_nw = ( - nw_full[args.warmup:args.warmup + eff_iters].tolist() - ) + scenario_per_iter_nw = nw_full[args.warmup : args.warmup + eff_iters].tolist() if baseline_fn is not None and is_pr3324_baseline: baseline_suffix_parts = [f"SR={int(use_philox)}"] hsort_list_for_tags = getattr(args, "hardcode_sort_list", [False]) hsort_in_cell_list = "HSORT" in getattr(args, "_cell_list_keys", ()) - emit_hsort_tag = ( - hardcode_sort or len(hsort_list_for_tags) > 1 - or hsort_in_cell_list - ) + emit_hsort_tag = hardcode_sort or len(hsort_list_for_tags) > 1 or hsort_in_cell_list if scn["fill"] is None and emit_hsort_tag: baseline_suffix_parts.append(f"HSORT={1 if hardcode_sort else 0}") baseline_sweep_suffix = ",".join(baseline_suffix_parts) @@ -3017,9 +3071,7 @@ def _run_pr3324_baseline(): D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=( - state_batch_indices if args.use_cache_slot else None - ), + state_batch_indices=(state_batch_indices if args.use_cache_slot else None), state_scale=state_scales_work, rand_seed=rand_seed, philox_rounds=args.philox_rounds, @@ -3070,8 +3122,7 @@ def _run_pr3324_baseline(): act_dtype_name=act_dtype_name, sweep_suffix=baseline_sweep_suffix, per_iter_nw=scenario_per_iter_nw, - kmix_bucket_write_frac=mix_write_frac - if scn["fill"] is None else None, + kmix_bucket_write_frac=mix_write_frac if scn["fill"] is None else None, skipped_tag=base_tag, ) @@ -3084,16 +3135,23 @@ def _run_pr3324_baseline(): # pair this with --skip-diagonal to drop the tied subset). if _any_split: _iter_axes = ( - block_size_m_write_values, block_size_m_nowrite_values, - num_warps_write_values, num_warps_nowrite_values, - num_stages_write_values, num_stages_nowrite_values, + block_size_m_write_values, + block_size_m_nowrite_values, + num_warps_write_values, + num_warps_nowrite_values, + num_stages_write_values, + num_stages_nowrite_values, precompute_num_warps_values, precompute_num_stages_values, heads_per_block_values, - maxnreg_values, num_ctas_values, - cta_per_sm_write_values, cta_per_sm_nowrite_values, - num_loop_stages_write_values, num_loop_stages_nowrite_values, - flatten_values, warp_specialize_values, + maxnreg_values, + num_ctas_values, + cta_per_sm_write_values, + cta_per_sm_nowrite_values, + num_loop_stages_write_values, + num_loop_stages_nowrite_values, + flatten_values, + warp_specialize_values, use_tma_rect_load_values, use_tma_replay_write_load_values, use_tma_replay_nowrite_load_values, @@ -3103,16 +3161,23 @@ def _run_pr3324_baseline(): # Tied: one value per shared knob. Wrap in single-element list for # uniform iteration; the body sets w/nw both to the shared value. _iter_axes = ( - block_size_m_values, [None], - num_warps_values, [None], - num_stages_values, [None], + block_size_m_values, + [None], + num_warps_values, + [None], + num_stages_values, + [None], precompute_num_warps_values, precompute_num_stages_values, heads_per_block_values, - maxnreg_values, num_ctas_values, - cta_per_sm_values, [None], - num_loop_stages_values, [None], - flatten_values, warp_specialize_values, + maxnreg_values, + num_ctas_values, + cta_per_sm_values, + [None], + num_loop_stages_values, + [None], + flatten_values, + warp_specialize_values, use_tma_rect_load_values, use_tma_replay_write_load_values, use_tma_replay_nowrite_load_values, @@ -3136,23 +3201,35 @@ def _run_pr3324_baseline(): # its assigned one — turning compile-warmup into 28-way duplication. # (Observed: 256 tasks in 233s under that bug vs ~18s correct.) if getattr(args, "_cell_list_set", None) and not warmup_only: + def _gen_from_cell_list(): keys = args._cell_list_keys for tup in args._cell_list_set: d = dict(zip(keys, tup)) yield ( - d.get("Mw"), d.get("Mnw"), - d.get("Ww"), d.get("Wnw"), - d.get("Sw"), d.get("Snw"), - d.get("pW"), d.get("pS"), + d.get("Mw"), + d.get("Mnw"), + d.get("Ww"), + d.get("Wnw"), + d.get("Sw"), + d.get("Snw"), + d.get("pW"), + d.get("pS"), d.get("H"), - d.get("R"), d.get("CT"), - d.get("CPSw"), d.get("CPSnw"), - d.get("LSw"), d.get("LSnw"), - d.get("FL"), d.get("WS"), - d.get("TMARL"), d.get("TMAWL"), - d.get("TMANL"), d.get("TMAWS"), + d.get("R"), + d.get("CT"), + d.get("CPSw"), + d.get("CPSnw"), + d.get("LSw"), + d.get("LSnw"), + d.get("FL"), + d.get("WS"), + d.get("TMARL"), + d.get("TMAWL"), + d.get("TMANL"), + d.get("TMAWS"), ) + _iter_source = _gen_from_cell_list() else: _iter_source = itertools.product(*_iter_axes) @@ -3190,12 +3267,16 @@ def _gen_from_cell_list(): num_loop_stages_nw = num_loop_stages_w # Skip-diagonal: when split is on, drop the tied subset (same as a # prior shared-knob sweep would cover). - if _any_split and args.skip_diagonal and ( - block_size_m_w == block_size_m_nw and - num_warps_w == num_warps_nw and - num_stages_w == num_stages_nw and - cta_per_sm_w == cta_per_sm_nw and - num_loop_stages_w == num_loop_stages_nw + if ( + _any_split + and args.skip_diagonal + and ( + block_size_m_w == block_size_m_nw + and num_warps_w == num_warps_nw + and num_stages_w == num_stages_nw + and cta_per_sm_w == cta_per_sm_nw + and num_loop_stages_w == num_loop_stages_nw + ) ): continue # Backward-compat aliases used by the existing body below. When @@ -3217,12 +3298,20 @@ def _gen_from_cell_list(): _write_path = True # both halves exist for persistent modes _rect_path = rectangle_for_nowrite _replay_nowrite_path = not rectangle_for_nowrite + def _set(v): # flag set to a non-zero sweep value return v is not None and v != 0 - if (_set(use_tma_rect_load) and not _rect_path - or _set(use_tma_replay_write_load) and not _write_path - or _set(use_tma_replay_nowrite_load) and not _replay_nowrite_path - or _set(use_tma_replay_write_store) and not _write_path): + + if ( + _set(use_tma_rect_load) + and not _rect_path + or _set(use_tma_replay_write_load) + and not _write_path + or _set(use_tma_replay_nowrite_load) + and not _replay_nowrite_path + or _set(use_tma_replay_write_store) + and not _write_path + ): continue assert scenario_n_writes is not None @@ -3323,6 +3412,7 @@ def _run_incr( ) parts = [] + # Tuned wrapper knobs emit "auto" when unset, meaning the wrapper # resolves them from _DEFAULT_TUNING per cell. pS/R/CT are not # tuning-table knobs today, so leave them out unless explicitly @@ -3343,6 +3433,7 @@ def _emit_split(name_w, name_nw, val_w, val_nw): else: parts.append(f"{name_w}={_val(val_w)}") parts.append(f"{name_nw}={_val(val_nw)}") + _emit_split("Mw", "Mnw", block_size_m_w, block_size_m_nw) _emit_split("Ww", "Wnw", num_warps_w, num_warps_nw) _emit_split("Sw", "Snw", num_stages_w, num_stages_nw) @@ -3374,7 +3465,9 @@ def _emit_split(name_w, name_nw, val_w, val_nw): parts.append(f"TMANL={_val(use_tma_replay_nowrite_load)}") parts.append(f"TMAWS={_val(use_tma_replay_write_store)}") parts.append(f"SR={1 if use_philox else 0}") - parts.append(f"RECT={'auto' if rectangle_for_nowrite is None else (1 if rectangle_for_nowrite else 0)}") + parts.append( + f"RECT={'auto' if rectangle_for_nowrite is None else (1 if rectangle_for_nowrite else 0)}" + ) nowrite_first_list_for_tags = getattr(args, "nowrite_first_list", [False]) nowrite_first_in_cell_list = "NWF" in getattr(args, "_cell_list_keys", ()) if nowrite_first or len(nowrite_first_list_for_tags) > 1 or nowrite_first_in_cell_list: @@ -3408,8 +3501,13 @@ def _emit_split(name_w, name_nw, val_w, val_nw): # re-evaluating the inner scenario loop; conservatively skip # only when the prev_k_for_print's specific key is done. _resume_key = _build_json_key( - "replay", batch, mtp_len, prev_k_for_print, - state_dtype_name, sweep_suffix, args.tp_size, + "replay", + batch, + mtp_len, + prev_k_for_print, + state_dtype_name, + sweep_suffix, + args.tp_size, ) if _resume_key in done_keys: continue @@ -3433,7 +3531,8 @@ def _emit_split(name_w, name_nw, val_w, val_nw): retry_budget = 0 if defer_results else max(0, getattr(args, "cupti_retry", 1)) stats = None expected_K = _kernels_per_iter_incremental( - mode, with_conv1d=with_conv1d, + mode, + with_conv1d=with_conv1d, ) plan_key = ( "incremental", @@ -3457,7 +3556,10 @@ def _emit_split(name_w, name_nw, val_w, val_nw): ) for attempt in range(retry_budget + 1): stats = _time_kernel( - args, _run_incr, reset_fn, sweep_tag, + args, + _run_incr, + reset_fn, + sweep_tag, expected_K=expected_K, pre_iter_fn=scenario_pre_iter, pre_iter_group_factory=scenario_pre_iter_group_factory, @@ -3491,22 +3593,23 @@ def _emit_split(name_w, name_nw, val_w, val_nw): act_dtype_name=act_dtype_name, sweep_suffix=sweep_suffix, per_iter_nw=per_iter_nw, - kmix_bucket_write_frac=mix_write_frac - if scn["fill"] is None else None, + kmix_bucket_write_frac=mix_write_frac if scn["fill"] is None else None, skipped_tag=sweep_tag, ) # Map full torch dtype name → short tag used in JSON keys (matches collect.py). _DTYPE_SHORT = { - "float32": "fp32", "bfloat16": "bf16", "float16": "fp16", - "int8": "int8", "int16": "int16", "float8_e4m3fn": "fp8", + "float32": "fp32", + "bfloat16": "bf16", + "float16": "fp16", + "int8": "int8", + "int16": "int16", + "float8_e4m3fn": "fp8", } -def _build_json_key( - kernel_name, batch, mtp_len, prev_k, state_dtype_name, sweep_suffix, tp_size -): +def _build_json_key(kernel_name, batch, mtp_len, prev_k, state_dtype_name, sweep_suffix, tp_size): """Build a key matching collect.py's kernel_data.json convention: incremental/{batch}/{mtp}/{sd}/k{k}/{sweep_parts}/tp{tp} @@ -3528,9 +3631,7 @@ def _build_json_key( # sweep_suffix format: " M=4,W=1,S=1,SR=0,RECT=0" # collect.py format: "M4_W1_S1_SR0_RECT0" # Strip leading/trailing whitespace, drop '=', commas → underscores. - parts.append( - sweep_suffix.strip().replace("=", "").replace(",", "_") - ) + parts.append(sweep_suffix.strip().replace("=", "").replace(",", "_")) parts.append(f"tp{tp_size}") return "/".join(parts) @@ -3578,8 +3679,13 @@ def _print_row( ) if jsonl_path is not None: key = _build_json_key( - kernel_name, batch, mtp_len, prev_k, state_dtype_name, - sweep_suffix, tp_size, + kernel_name, + batch, + mtp_len, + prev_k, + state_dtype_name, + sweep_suffix, + tp_size, ) if json_detailed: row_stats = stats @@ -3587,8 +3693,13 @@ def _print_row( row_stats = { k: stats[k] for k in ( - "median", "p95", "p99", "n", "iters_us", - "n_writes_per_iter", "kmix_bucket_score", + "median", + "p95", + "p99", + "n", + "iters_us", + "n_writes_per_iter", + "kmix_bucket_score", ) if k in stats } @@ -3607,6 +3718,7 @@ def _print_row( # segments (cells/sec, downtime between bench invocations) without # needing to instrument the bench's outer loops separately. import time as _time + rec = {"key": key, "stats": row_stats, "t": _time.time()} if jsonl_host is not None: rec["host"] = jsonl_host @@ -3722,29 +3834,29 @@ def _submit_result_job( # For split (write/nowrite) knobs, we use Xw / Xnw keys. Tied forms (M, W, # S, CPS, LS) accepted on load and expanded to their w/nw variants. _CELL_LIST_KEY_TO_ARG = { - "Mw": "block_size_m_write", - "Mnw": "block_size_m_nowrite", - "Ww": "num_warps_write", - "Wnw": "num_warps_nowrite", - "Sw": "num_stages_write", - "Snw": "num_stages_nowrite", - "CPSw": "cta_per_sm_write", + "Mw": "block_size_m_write", + "Mnw": "block_size_m_nowrite", + "Ww": "num_warps_write", + "Wnw": "num_warps_nowrite", + "Sw": "num_stages_write", + "Snw": "num_stages_nowrite", + "CPSw": "cta_per_sm_write", "CPSnw": "cta_per_sm_nowrite", - "LSw": "num_loop_stages_write", - "LSnw": "num_loop_stages_nowrite", - "pW": "precompute_num_warps", - "pS": "precompute_num_stages", - "H": "heads_per_block", - "R": "maxnreg", - "CT": "num_ctas", - "FL": "flatten", - "WS": "warp_specialize", + "LSw": "num_loop_stages_write", + "LSnw": "num_loop_stages_nowrite", + "pW": "precompute_num_warps", + "pS": "precompute_num_stages", + "H": "heads_per_block", + "R": "maxnreg", + "CT": "num_ctas", + "FL": "flatten", + "WS": "warp_specialize", "TMARL": "use_tma_rect_load", "TMAWL": "use_tma_replay_write_load", "TMANL": "use_tma_replay_nowrite_load", "TMAWS": "use_tma_replay_write_store", - "RECT": "rectangle_for_nowrite", - "NWF": "nowrite_first", + "RECT": "rectangle_for_nowrite", + "NWF": "nowrite_first", "HSORT": "hardcode_sort", # MODE and SR get special handling (string values): # MODE → args.modes (single mode name) @@ -3753,13 +3865,14 @@ def _submit_result_job( # Split-knob tied form: "M" expands to both "Mw" and "Mnw". _CELL_LIST_TIED_EXPANSIONS = { - "M": ("Mw", "Mnw"), - "W": ("Ww", "Wnw"), - "S": ("Sw", "Snw"), + "M": ("Mw", "Mnw"), + "W": ("Ww", "Wnw"), + "S": ("Sw", "Snw"), "CPS": ("CPSw", "CPSnw"), - "LS": ("LSw", "LSnw"), + "LS": ("LSw", "LSnw"), } + def _normalize_cell(cell: dict) -> dict: """Expand tied-form keys (M, W, S, CPS, LS) to their w/nw variants. Returns a new dict with only canonical split-or-plain keys. @@ -3787,10 +3900,7 @@ def _load_cell_list_into_args(args) -> None: print("[cell-list] empty list — nothing to time", file=sys.stderr) return allowed_keys = set(_CELL_LIST_KEY_TO_ARG) | {"MODE", "SR"} - unknown_keys = sorted({ - key for cell in cells for key in cell - if key not in allowed_keys - }) + unknown_keys = sorted({key for cell in cells for key in cell if key not in allowed_keys}) if unknown_keys: sys.exit(f"--cell-list: unknown keys {unknown_keys}") # All cells must share the same key set (uniform schema) @@ -3818,19 +3928,22 @@ def _load_cell_list_into_args(args) -> None: elif key == "SR": args.sr_modes = ",".join(sorted({"SR" if v else "RN" for v in vals})) else: - print(f"[cell-list] WARNING: unknown key {key!r} in cells; " - f"will not override any args.* attribute (the value will " - f"still be matched in the filter if a matching local var " - f"is in scope)", file=sys.stderr) + print( + f"[cell-list] WARNING: unknown key {key!r} in cells; " + f"will not override any args.* attribute (the value will " + f"still be matched in the filter if a matching local var " + f"is in scope)", + file=sys.stderr, + ) # Canonical key order (sorted) for tuple matching in the inner loop args._cell_list_keys = tuple(sorted(keys0)) - args._cell_list_set = { - tuple(c[k] for k in args._cell_list_keys) for c in cells - } - print(f"[cell-list] loaded {len(cells)} cells with keys " - f"{list(args._cell_list_keys)}; overrode args.* to auto-cover", - file=sys.stderr) + args._cell_list_set = {tuple(c[k] for k in args._cell_list_keys) for c in cells} + print( + f"[cell-list] loaded {len(cells)} cells with keys " + f"{list(args._cell_list_keys)}; overrode args.* to auto-cover", + file=sys.stderr, + ) # Maps cell-list key → name of the local variable in _bench_config's inner @@ -3838,32 +3951,32 @@ def _load_cell_list_into_args(args) -> None: # Keep in sync with the loop-variable names; the filter is lenient about # missing names (it picks them up from the inner scope at runtime). _CELL_LIST_KEY_TO_LOCAL = { - "Mw": "block_size_m_w", - "Mnw": "block_size_m_nw", - "Ww": "num_warps_w", - "Wnw": "num_warps_nw", - "Sw": "num_stages_w", - "Snw": "num_stages_nw", - "CPSw": "cta_per_sm_w", + "Mw": "block_size_m_w", + "Mnw": "block_size_m_nw", + "Ww": "num_warps_w", + "Wnw": "num_warps_nw", + "Sw": "num_stages_w", + "Snw": "num_stages_nw", + "CPSw": "cta_per_sm_w", "CPSnw": "cta_per_sm_nw", - "LSw": "num_loop_stages_w", - "LSnw": "num_loop_stages_nw", - "pW": "precompute_num_warps", - "pS": "precompute_num_stages", - "H": "heads_per_block", - "R": "maxnreg", - "CT": "num_ctas", - "FL": "flatten", - "WS": "warp_specialize", + "LSw": "num_loop_stages_w", + "LSnw": "num_loop_stages_nw", + "pW": "precompute_num_warps", + "pS": "precompute_num_stages", + "H": "heads_per_block", + "R": "maxnreg", + "CT": "num_ctas", + "FL": "flatten", + "WS": "warp_specialize", "TMARL": "use_tma_rect_load", "TMAWL": "use_tma_replay_write_load", "TMANL": "use_tma_replay_nowrite_load", "TMAWS": "use_tma_replay_write_store", - "RECT": "rectangle_for_nowrite", - "NWF": "nowrite_first", - "MODE": "mode", + "RECT": "rectangle_for_nowrite", + "NWF": "nowrite_first", + "MODE": "mode", "HSORT": "hardcode_sort", - "SR": "use_philox", + "SR": "use_philox", } @@ -3896,9 +4009,11 @@ def _run_benchmark(args) -> None: # run can later attribute wall time to setup vs compile-warmup vs prewarm # vs timing. Single-line format makes log-grepping trivial. _phase_t0 = time.perf_counter() + def _phase(label: str) -> None: dt = time.perf_counter() - _phase_t0 print(f"[phase] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) + _phase("enter _run_benchmark") # Pending-results FIFO for srxl's deferred CUPTI parsing pipeline. Each @@ -3928,9 +4043,10 @@ def _phase(label: str) -> None: args._done_keys: set[str] = set() args._baseline_seen_keys: set[str] = set() args._jsonl_host = None # hostname stamp for the current run - args._jsonl_gpu = None # GPU device id stamp (current process visibility) + args._jsonl_gpu = None # GPU device id stamp (current process visibility) if getattr(args, "json_output", None): import socket + args._jsonl_host = socket.gethostname() # Capture GPU id once at startup. Used by the oracle-cache layer in # search_driver to attribute timings to a specific (host, gpu) pair @@ -3964,9 +4080,11 @@ def _phase(label: str) -> None: if rec_host: host_counts[rec_host] = host_counts.get(rec_host, 0) + 1 if n_loaded: - host_summary = ", ".join( - f"{h}={n}" for h, n in sorted(host_counts.items()) - ) if host_counts else "(no host stamps)" + host_summary = ( + ", ".join(f"{h}={n}" for h, n in sorted(host_counts.items())) + if host_counts + else "(no host stamps)" + ) print( f"[resume] {args._jsonl_path}: loaded {n_loaded} prior " f"cell results across hosts [{host_summary}]; sweep will " @@ -4093,8 +4211,13 @@ def _phase(label: str) -> None: _phase("about to enter compile-warmup") if args.compile_threads > 0: _compile_warmup_phase( - args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, - baseline_fn, max_workers=args.compile_threads, + args, + batch_sizes, + mtp_lengths, + state_dtypes, + act_dtypes, + baseline_fn, + max_workers=args.compile_threads, ) _phase("returned from compile-warmup") @@ -4109,8 +4232,14 @@ def _phase(label: str) -> None: for act_dtype in act_dtypes: for mtp_len in mtp_lengths: _build_tensors( - _max_batch, mtp_len, state_dtype, act_dtype, - args.tp_nheads, args.head_dim, args.d_state, args.tp_ngroups, + _max_batch, + mtp_len, + state_dtype, + act_dtype, + args.tp_nheads, + args.head_dim, + args.d_state, + args.tp_ngroups, max_window=getattr(args, "max_window", None) or None, ) _phase("done tensor prewarm — entering timing") @@ -4164,9 +4293,7 @@ def _phase(label: str) -> None: mix_label = mix_csv.stem # T (= mtp_len) varies per cell; load once with the LARGEST mtp so # we have enough columns; the loader normalizes the dist anyway. - mix_al = _load_al_distribution( - mix_csv, T=max(mtp_lengths), column=args.mix_csv_column - ) + mix_al = _load_al_distribution(mix_csv, T=max(mtp_lengths), column=args.mix_csv_column) for batch in batch_sizes: for mtp_len in mtp_lengths: @@ -4184,23 +4311,23 @@ def _phase(label: str) -> None: if mix_al is not None: _max_window = getattr(args, "max_window", 0) or mtp_len mix_pi = _markov_stationary(mix_al, mtp_len, _max_window) - mix_write_frac = float( - mix_pi[_max_window - mtp_len + 1:].sum() - ) + mix_write_frac = float(mix_pi[_max_window - mtp_len + 1 :].sum()) _max_iters = max(args.iters, getattr(args, "mix_iters", None) or args.iters) mix_samples_cpu = _sample_steady_state_pnat( - mix_al, T=mtp_len, window=_max_window, batch=batch, - K=args.warmup + _max_iters, seed=args.mix_seed, + mix_al, + T=mtp_len, + window=_max_window, + batch=batch, + K=args.warmup + _max_iters, + seed=args.mix_seed, ) if any(hsort_list): # write-first stable argsort: kind='stable' preserves # original-slot order within each mode group. - write_mask = ( - mix_samples_cpu + mtp_len > _max_window - ).astype(np.int8) # 1 = write, 0 = nowrite - perm_idx = np.argsort( - -write_mask, kind="stable", axis=-1 - ).astype(np.int32) + write_mask = (mix_samples_cpu + mtp_len > _max_window).astype( + np.int8 + ) # 1 = write, 0 = nowrite + perm_idx = np.argsort(-write_mask, kind="stable", axis=-1).astype(np.int32) # Apply the perm to the prev_tokens samples themselves. # Result row i = mix_samples_cpu[i] reordered such # that write-mode entries come first. @@ -4215,17 +4342,19 @@ def _phase(label: str) -> None: for rect in rect_list: for nowrite_first in nowrite_first_list: can_sort = mix_samples_cpu is not None - cell_list_active = bool( - getattr(args, "_cell_list_keys", ()) - ) + cell_list_active = bool(getattr(args, "_cell_list_keys", ())) effective_hsort_list = ( - hsort_list if (can_sort or cell_list_active) - else [False] + hsort_list if (can_sort or cell_list_active) else [False] ) for hardcode_sort in effective_hsort_list: _bench_config( - args, batch, mtp_len, prev_ks, - state_dtype, act_dtype, baseline_fn, + args, + batch, + mtp_len, + prev_ks, + state_dtype, + act_dtype, + baseline_fn, sr_mode=sr_mode, rectangle_for_nowrite=rect, mode=mode, @@ -4246,8 +4375,7 @@ def _phase(label: str) -> None: # host stamps). No clean-exit `.json` write — use `jsonl_to_json.py` to # materialize a snapshot when an analyzer wants one. if args.json_output and args._jsonl_path is not None: - print(f"\nJSONL results: {args._jsonl_path} " - f"(meta sidecar: {args.json_output}.meta.json)") + print(f"\nJSONL results: {args._jsonl_path} (meta sidecar: {args.json_output}.meta.json)") # Write the skipped-cells sidecar. Caller can convert this list to a # --cell-list JSON (one dict per skipped cell) to drive a retry pass in @@ -4269,12 +4397,18 @@ def _phase(label: str) -> None: with open(tmp, "w") as f: json.dump(payload, f, indent=2) os.replace(tmp, skipped_path) - print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " - f"tags written to: {skipped_path}", file=sys.stderr) + print( + f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " + f"tags written to: {skipped_path}", + file=sys.stderr, + ) elif args._skipped_cells: # No output path but there are skipped cells — emit a stderr summary. - print(f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " - f"first 5: {args._skipped_cells[:5]}", file=sys.stderr) + print( + f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " + f"first 5: {args._skipped_cells[:5]}", + file=sys.stderr, + ) # CLI @@ -4328,12 +4462,16 @@ def _parse_args() -> argparse.Namespace: default="bf16", help="Comma-separated activation dtypes for x/B/C/dt: fp32,bf16", ) - parser.add_argument("--warmup", type=int, default=4, - help="Number of warmup iterations. Default aligns with " - "the graph group-iters (default 4 for mix scenarios) so " - "warmup + iters / mix-iters lands on a clean multiple " - "without per-args rounding overhead. Earlier default of " - "20 was overkill for steady-state warming.") + parser.add_argument( + "--warmup", + type=int, + default=4, + help="Number of warmup iterations. Default aligns with " + "the graph group-iters (default 4 for mix scenarios) so " + "warmup + iters / mix-iters lands on a clean multiple " + "without per-args rounding overhead. Earlier default of " + "20 was overkill for steady-state warming.", + ) parser.add_argument("--iters", type=int, default=100, help="Number of timed iterations") parser.add_argument( "--compile-threads", @@ -4521,48 +4659,70 @@ def _parse_args() -> argparse.Namespace: help="Override num_stages for the main kernel (comma-separated sweep).", ) parser.add_argument( - "--block-size-m-write", type=str, default=None, + "--block-size-m-write", + type=str, + default=None, help="Sweep BLOCK_SIZE_M for the WRITE main only (overrides --block-size-m " "for the write half). Tied to --block-size-m if unset.", ) parser.add_argument( - "--block-size-m-nowrite", type=str, default=None, + "--block-size-m-nowrite", + type=str, + default=None, help="Sweep BLOCK_SIZE_M for the NOWRITE main only. Tied to --block-size-m if unset.", ) parser.add_argument( - "--num-warps-write", type=str, default=None, + "--num-warps-write", + type=str, + default=None, help="Sweep num_warps for the WRITE main only. Tied to --num-warps if unset.", ) parser.add_argument( - "--num-warps-nowrite", type=str, default=None, + "--num-warps-nowrite", + type=str, + default=None, help="Sweep num_warps for the NOWRITE main only. Tied to --num-warps if unset.", ) parser.add_argument( - "--num-stages-write", type=str, default=None, + "--num-stages-write", + type=str, + default=None, help="Sweep num_stages for the WRITE main only. Tied to --num-stages if unset.", ) parser.add_argument( - "--num-stages-nowrite", type=str, default=None, + "--num-stages-nowrite", + type=str, + default=None, help="Sweep num_stages for the NOWRITE main only. Tied to --num-stages if unset.", ) parser.add_argument( - "--cta-per-sm-write", type=str, default=None, + "--cta-per-sm-write", + type=str, + default=None, help="Sweep cta_per_sm for the WRITE persistent_main only. Tied to --cta-per-sm if unset.", ) parser.add_argument( - "--cta-per-sm-nowrite", type=str, default=None, + "--cta-per-sm-nowrite", + type=str, + default=None, help="Sweep cta_per_sm for the NOWRITE persistent_main only. Tied to --cta-per-sm if unset.", ) parser.add_argument( - "--num-loop-stages-write", type=str, default=None, + "--num-loop-stages-write", + type=str, + default=None, help="Sweep num_loop_stages for the WRITE persistent_main only. Tied to --num-loop-stages if unset.", ) parser.add_argument( - "--num-loop-stages-nowrite", type=str, default=None, + "--num-loop-stages-nowrite", + type=str, + default=None, help="Sweep num_loop_stages for the NOWRITE persistent_main only. Tied to --num-loop-stages if unset.", ) parser.add_argument( - "--skip-diagonal", action=argparse.BooleanOptionalAction, default=False, + "--skip-diagonal", + action=argparse.BooleanOptionalAction, + default=False, help="When sweeping any per-main *_write / *_nowrite knobs, skip cells " "where ALL splittable knobs satisfy write_value == nowrite_value (i.e. " "the 'diagonal' that's already covered by a prior shared-knob sweep). " @@ -4846,9 +5006,7 @@ def _parse_args() -> argparse.Namespace: f"built-in T{DEFAULT_PMIX_T} distribution or --mix-csv for a custom histogram." ) if args.pmix: - mtp_lengths_for_pmix = [ - int(x) for x in args.mtp_lengths.split(",") if x.strip() - ] + mtp_lengths_for_pmix = [int(x) for x in args.mtp_lengths.split(",") if x.strip()] if any(t != DEFAULT_PMIX_T for t in mtp_lengths_for_pmix): parser.error( f"--pmix uses the built-in T{DEFAULT_PMIX_T} histogram; " @@ -4870,16 +5028,20 @@ def _parse_args() -> argparse.Namespace: _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE, getattr(args, "cuda_graph_group_iters", None) or 0, ) + def _round_iters_to_group(name, val): total = args.warmup + val if total % _group_for_rounding == 0: return val new_total = ((total + _group_for_rounding - 1) // _group_for_rounding) * _group_for_rounding new_val = new_total - args.warmup - print(f"[bench] rounding --{name} {val} → {new_val} so warmup+{name} " - f"({new_total}) is a multiple of graph group_iters={_group_for_rounding}", - file=sys.stderr) + print( + f"[bench] rounding --{name} {val} → {new_val} so warmup+{name} " + f"({new_total}) is a multiple of graph group_iters={_group_for_rounding}", + file=sys.stderr, + ) return new_val + args.iters = _round_iters_to_group("iters", args.iters) if getattr(args, "mix_iters", None): args.mix_iters = _round_iters_to_group("mix-iters", args.mix_iters) @@ -4895,7 +5057,7 @@ def _round_iters_to_group(name, val): # Backward-compat: --philox-rounding implies --sr-modes SR if --sr-modes # was left at the default. If both are set explicitly, error. - sr_modes_default = (args.sr_modes == "RN") + sr_modes_default = args.sr_modes == "RN" if args.philox_rounding: if not sr_modes_default and args.sr_modes != "SR": parser.error( @@ -4943,9 +5105,7 @@ def _round_iters_to_group(name, val): hsort_modes = [ v.strip() - for v in ( - args.hardcode_sort if args.hardcode_sort is not None else "0" - ).split(",") + for v in (args.hardcode_sort if args.hardcode_sort is not None else "0").split(",") if v.strip() ] hsort_list = [] @@ -4956,10 +5116,7 @@ def _round_iters_to_group(name, val): if not hsort_list: hsort_list = [False] if not args.use_cache_slot and not all(hsort_list): - parser.error( - "--no-use-cache-slot is a diagnostic shortcut and requires " - "--hardcode-sort 1" - ) + parser.error("--no-use-cache-slot is a diagnostic shortcut and requires --hardcode-sort 1") args.hardcode_sort_list = hsort_list # mode=None means "let the wrapper resolve from _DEFAULT_TUNING". Same @@ -4969,13 +5126,12 @@ def _round_iters_to_group(name, val): else: modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] valid_modes = { - "persistent_main", "persistent_dynamic", + "persistent_main", + "persistent_dynamic", } for m in modes_raw: if m not in valid_modes: - parser.error( - f"--modes value must be one of {sorted(valid_modes)}, got {m!r}" - ) + parser.error(f"--modes value must be one of {sorted(valid_modes)}, got {m!r}") args.modes_list = modes_raw if modes_raw else [None] return args @@ -5014,13 +5170,17 @@ def close(self): sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) mp.set_start_method("forkserver", force=True) try: - mp.set_forkserver_preload([ - "benchmark_replay_selective_state_update", - ]) + mp.set_forkserver_preload( + [ + "benchmark_replay_selective_state_update", + ] + ) except Exception as _e: - print(f"[warn] set_forkserver_preload failed: {_e!r}; " - f"forks will still work but pay full import cost", - file=sys.stderr) + print( + f"[warn] set_forkserver_preload failed: {_e!r}; " + f"forks will still work but pay full import cost", + file=sys.stderr, + ) _MP_START_METHOD = _args.mp_start_method _out_path = None diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py index 86f52d35cfda..e03eb6151eea 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py @@ -28,7 +28,6 @@ import jinja2 import torch - from flashinfer.compilation_context import CompilationContext from flashinfer.jit import env as jit_env from flashinfer.jit.core import JitSpec, gen_jit_spec @@ -138,9 +137,7 @@ def _gen_module( with open(_CSRC_DIR / "checkpointing_ssu_customize_config.jinja") as file: config_template = jinja2.Template(file.read()) - state_scale_type = ( - _DTYPE_MAP[state_scale_dtype] if state_scale_dtype is not None else "void" - ) + state_scale_type = _DTYPE_MAP[state_scale_dtype] if state_scale_dtype is not None else "void" config = config_template.render( state_dtype=_DTYPE_MAP[state_dtype], input_dtype=_DTYPE_MAP[input_dtype], @@ -247,9 +244,7 @@ def checkpointing_ssu( if state_scale is None: raise ValueError(f"state dtype {state.dtype} requires state_scale") elif state_scale is not None: - raise ValueError( - f"state_scale must be None for non-quantized state dtype {state.dtype}" - ) + raise ValueError(f"state_scale must be None for non-quantized state dtype {state.dtype}") if cu_seqlens is not None: npredicted = max_seqlen if max_seqlen is not None else old_x.size(1) else: diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu index df0e7e189274..e9ff3f98c188 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu @@ -25,522 +25,530 @@ using namespace flashinfer; using tvm::ffi::Optional; -namespace flashinfer::mamba::checkpointing { - -void checkpointing_ssu( - TensorView state, // (state_cache_size, nheads, dim, dstate) - TensorView x, // (batch, NPREDICTED, nheads, dim) / (1, total_tokens, nheads, dim) under varlen - TensorView dt, // (batch, NPREDICTED, nheads, dim) tie_hdim / (1, total_tokens, nheads, dim) - TensorView A, // (nheads, dim, dstate) tie_hdim - TensorView B, // (batch, NPREDICTED, ngroups, dstate) / (1, total_tokens, ngroups, dstate) - TensorView C, // same as B - TensorView output, // same layout as x +namespace flashinfer::mamba::checkpointing +{ + +void checkpointing_ssu(TensorView state, // (state_cache_size, nheads, dim, dstate) + TensorView x, // (batch, NPREDICTED, nheads, dim) / (1, total_tokens, nheads, dim) under varlen + TensorView dt, // (batch, NPREDICTED, nheads, dim) tie_hdim / (1, total_tokens, nheads, dim) + TensorView A, // (nheads, dim, dstate) tie_hdim + TensorView B, // (batch, NPREDICTED, ngroups, dstate) / (1, total_tokens, ngroups, dstate) + TensorView C, // same as B + TensorView output, // same layout as x // Cache tensors - TensorView old_x, // (state_cache_size, MAX_WINDOW, nheads, dim) - TensorView old_B, // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) - TensorView old_dt, // (state_cache_size, 2, nheads, MAX_WINDOW) f32 - TensorView old_cumAdt, // (state_cache_size, 2, nheads, MAX_WINDOW) f32 - TensorView cache_buf_idx, // (state_cache_size,) int32 - TensorView prev_num_accepted, // (state_cache_size,) int32 + TensorView old_x, // (state_cache_size, MAX_WINDOW, nheads, dim) + TensorView old_B, // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) + TensorView old_dt, // (state_cache_size, 2, nheads, MAX_WINDOW) f32 + TensorView old_cumAdt, // (state_cache_size, 2, nheads, MAX_WINDOW) f32 + TensorView cache_buf_idx, // (state_cache_size,) int32 + TensorView prev_num_accepted, // (state_cache_size,) int32 // Optional tensors - Optional D, // (nheads, dim) - Optional z, // same layout as x - Optional dt_bias, // (nheads, dim) tie_hdim + Optional D, // (nheads, dim) + Optional z, // same layout as x + Optional dt_bias, // (nheads, dim) tie_hdim bool dt_softplus, - Optional state_batch_indices, // (batch,) int32 + Optional state_batch_indices, // (batch,) int32 int64_t pad_slot_id, - Optional state_scale, // (state_cache_size, nheads, dim) f32 - Optional rand_seed, // single int64 - int64_t d_split, // v12 §59: per-head DIM split factor (1, 2, or 4) - Optional cu_seqlens) { // (batch+1,) int32, varlen mode - - bool const is_varlen = cu_seqlens.has_value(); - - // ── Extract dimensions ── - auto const state_cache_size = state.size(0); - auto const nheads = state.size(1); - auto const dim = state.size(2); - auto const dstate = state.size(3); - auto const max_window = old_x.size(1); - auto const ngroups = B.size(2); - - // In non-varlen mode, batch = x.size(0) and npredicted = x.size(1) (the - // 4D batched layout). In varlen, the JIT compile-time NPREDICTED is the - // max seq_len the caller commits to; the wrapper stamped it into the JIT - // URI from `max_seqlen` and we read it back as the constexpr `NPREDICTED` - // (validated against runtime cu_seqlens on the host side below). `batch` - // = number of sequences = `cu_seqlens.size(0) - 1`. - int64_t batch; - int64_t npredicted; - if (is_varlen) { - auto const& cs = cu_seqlens.value(); - CHECK_CUDA(cs); - CHECK_DIM(1, cs); - FLASHINFER_CHECK( - cs.size(0) >= 2, - "cu_seqlens must have shape (batch+1,) with batch >= 1, got size(0)=", cs.size(0)); - FLASHINFER_CHECK(cs.dtype().code == kDLInt && cs.dtype().bits == 32, - "cu_seqlens must be int32"); - CHECK_CONTIGUOUS(cs); - batch = cs.size(0) - 1; - npredicted = NPREDICTED; // JIT-stamped — wrapper ensures max(seq_lens) <= NPREDICTED. - } else { - batch = x.size(0); - npredicted = x.size(1); - } - - // ── JIT compile-time / runtime cross-check ── - // NPREDICTED and MAX_WINDOW are JIT compile-time constants stamped by the - // wrapper. In non-varlen NPREDICTED = x.shape[1]; in varlen NPREDICTED = - // user-supplied `max_seqlen` (upper bound on every cu_seqlens diff). - FLASHINFER_CHECK(npredicted == NPREDICTED, is_varlen ? "max_seqlen=" : "x.size(1)=", npredicted, - " must equal JIT NPREDICTED=", NPREDICTED); - FLASHINFER_CHECK(max_window == MAX_WINDOW, "old_x.size(1)=", max_window, - " must equal JIT MAX_WINDOW=", MAX_WINDOW); - FLASHINFER_CHECK(npredicted <= max_window, "npredicted=", npredicted, - " must be <= max_window=", max_window); - - // ── Validate state ── - CHECK_CUDA(state); - CHECK_DIM(4, state); - { - auto s = state.strides(); - auto sz = state.sizes(); - FLASHINFER_CHECK(s[3] == 1, "state dim 3 (dstate) must have stride 1, got ", s[3]); - FLASHINFER_CHECK(s[2] == sz[3], "state dim 2 (dim) must be contiguous with dim 3, got stride ", - s[2], " expected ", sz[3]); - FLASHINFER_CHECK(s[1] == sz[2] * sz[3], - "state dim 1 (nheads) must be contiguous with dim 2, got stride ", s[1], - " expected ", sz[2] * sz[3]); - } - - // ── Validate x ── - // Non-varlen: shape (batch, NPREDICTED, nheads, dim). - // Varlen : shape (1, total_tokens, nheads, dim) — batch axis collapsed, - // token axis is the outer iteration. The kernel reads x via - // `bos * x_stride_token + …` so x_stride_token is the per-token - // stride in either layout (= nheads*dim for contig). - CHECK_CUDA(x); - CHECK_DIM(4, x); - if (is_varlen) { - FLASHINFER_CHECK(x.size(0) == 1, "varlen: x.size(0)=", x.size(0), " must be 1"); - } else { - FLASHINFER_CHECK(x.size(0) == batch, "x.size(0)=", x.size(0), " must equal batch=", batch); - FLASHINFER_CHECK(x.size(1) == npredicted, "x.size(1)=", x.size(1), - " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(x.size(2) == nheads, "x.size(2)=", x.size(2), " must equal nheads=", nheads); - FLASHINFER_CHECK(x.size(3) == dim, "x.size(3)=", x.size(3), " must equal dim=", dim); - CHECK_LAST_DIM_CONTIGUOUS(x); - FLASHINFER_CHECK(x.stride(2) == dim, "x.stride(2)=", x.stride(2), " must equal dim=", dim, - " ((nheads, dim) must be contiguous)"); - - // In varlen, all per-token tensors share the flattened token axis — use - // x.size(1) as the canonical total_tokens and cross-check the others below. - int64_t const total_tokens = is_varlen ? x.size(1) : 0; - - // ── Validate dt ── - CHECK_CUDA(dt); - CHECK_DIM(4, dt); - if (is_varlen) { - FLASHINFER_CHECK(dt.size(0) == 1, "varlen: dt.size(0)=", dt.size(0), " must be 1"); - FLASHINFER_CHECK(dt.size(1) == total_tokens, "varlen: dt.size(1)=", dt.size(1), - " must equal x.size(1)=", total_tokens); - } else { - FLASHINFER_CHECK(dt.size(0) == batch, "dt.size(0)=", dt.size(0), " must equal batch=", batch); - FLASHINFER_CHECK(dt.size(1) == npredicted, "dt.size(1)=", dt.size(1), - " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(dt.size(2) == nheads, "dt.size(2)=", dt.size(2), " must equal nheads=", nheads); - FLASHINFER_CHECK(dt.size(3) == dim, "dt.size(3)=", dt.size(3), " must equal dim=", dim); - FLASHINFER_CHECK(dt.stride(2) == 1, "dt.stride(2) must be 1 (tie_hdim), got ", dt.stride(2)); - FLASHINFER_CHECK(dt.stride(3) == 0, "dt.stride(3) must be 0 (tie_hdim), got ", dt.stride(3)); - - // ── Validate A: (nheads, dim, dstate) tie_hdim ── - CHECK_CUDA(A); - CHECK_DIM(3, A); - FLASHINFER_CHECK(A.size(0) == nheads, "A.size(0)=", A.size(0), " must equal nheads=", nheads); - FLASHINFER_CHECK(A.size(1) == dim, "A.size(1)=", A.size(1), " must equal dim=", dim); - FLASHINFER_CHECK(A.size(2) == dstate, "A.size(2)=", A.size(2), " must equal dstate=", dstate); - FLASHINFER_CHECK(A.stride(0) == 1, "A.stride(0) must be 1, got ", A.stride(0)); - FLASHINFER_CHECK(A.stride(1) == 0, "A.stride(1) must be 0 (tie_hdim), got ", A.stride(1)); - FLASHINFER_CHECK(A.stride(2) == 0, "A.stride(2) must be 0 (tie_hdim), got ", A.stride(2)); - - // ── Validate B ── - CHECK_CUDA(B); - CHECK_DIM(4, B); - if (is_varlen) { - FLASHINFER_CHECK(B.size(0) == 1, "varlen: B.size(0)=", B.size(0), " must be 1"); - FLASHINFER_CHECK(B.size(1) == total_tokens, "varlen: B.size(1)=", B.size(1), - " must equal x.size(1)=", total_tokens); - } else { - FLASHINFER_CHECK(B.size(0) == batch, "B.size(0)=", B.size(0), " must equal batch=", batch); - FLASHINFER_CHECK(B.size(1) == npredicted, "B.size(1)=", B.size(1), - " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(B.size(3) == dstate, "B.size(3)=", B.size(3), " must equal dstate=", dstate); - CHECK_LAST_DIM_CONTIGUOUS(B); - FLASHINFER_CHECK(B.stride(2) == dstate, "B.stride(2)=", B.stride(2), - " must equal dstate=", dstate, " ((ngroups, dstate) must be contiguous)"); - FLASHINFER_CHECK(nheads % ngroups == 0, "nheads=", nheads, - " must be divisible by ngroups=", ngroups); - - // ── Validate C ── - CHECK_CUDA(C); - CHECK_DIM(4, C); - if (is_varlen) { - FLASHINFER_CHECK(C.size(0) == 1, "varlen: C.size(0)=", C.size(0), " must be 1"); - FLASHINFER_CHECK(C.size(1) == total_tokens, "varlen: C.size(1)=", C.size(1), - " must equal x.size(1)=", total_tokens); - } else { - FLASHINFER_CHECK(C.size(0) == batch, "C.size(0)=", C.size(0), " must equal batch=", batch); - FLASHINFER_CHECK(C.size(1) == npredicted, "C.size(1)=", C.size(1), - " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(C.size(2) == ngroups, "C.size(2)=", C.size(2), " must equal ngroups=", ngroups); - FLASHINFER_CHECK(C.size(3) == dstate, "C.size(3)=", C.size(3), " must equal dstate=", dstate); - CHECK_LAST_DIM_CONTIGUOUS(C); - FLASHINFER_CHECK(C.stride(2) == dstate, "C.stride(2)=", C.stride(2), - " must equal dstate=", dstate, " ((ngroups, dstate) must be contiguous)"); - - // ── Validate output ── - CHECK_CUDA(output); - CHECK_DIM(4, output); - if (is_varlen) { - FLASHINFER_CHECK(output.size(0) == 1, "varlen: output.size(0)=", output.size(0), " must be 1"); - FLASHINFER_CHECK(output.size(1) == total_tokens, "varlen: output.size(1)=", output.size(1), - " must equal x.size(1)=", total_tokens); - } else { - FLASHINFER_CHECK(output.size(0) == batch, "output.size(0)=", output.size(0), - " must equal batch=", batch); - FLASHINFER_CHECK(output.size(1) == npredicted, "output.size(1)=", output.size(1), - " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(output.size(2) == nheads, "output.size(2)=", output.size(2), - " must equal nheads=", nheads); - FLASHINFER_CHECK(output.size(3) == dim, "output.size(3)=", output.size(3), - " must equal dim=", dim); - CHECK_LAST_DIM_CONTIGUOUS(output); - FLASHINFER_CHECK(output.stride(2) == dim, "output.stride(2)=", output.stride(2), - " must equal dim=", dim, " ((nheads, dim) must be contiguous)"); - - // ── Validate cache tensors ── - // old_x: kernel uses `head * DIM + d_tile_off` → (nheads, dim) contig. - CHECK_CUDA(old_x); - CHECK_DIM(4, old_x); // (state_cache_size, MAX_WINDOW, nheads, dim) - FLASHINFER_CHECK(old_x.size(0) == state_cache_size, "old_x.size(0)=", old_x.size(0), - " must equal state_cache_size=", state_cache_size); - FLASHINFER_CHECK(old_x.size(1) == max_window, "old_x.size(1)=", old_x.size(1), - " must equal max_window=", max_window); - FLASHINFER_CHECK(old_x.size(2) == nheads, "old_x.size(2)=", old_x.size(2), - " must equal nheads=", nheads); - FLASHINFER_CHECK(old_x.size(3) == dim, "old_x.size(3)=", old_x.size(3), " must equal dim=", dim); - CHECK_LAST_DIM_CONTIGUOUS(old_x); - FLASHINFER_CHECK(old_x.stride(2) == dim, "old_x.stride(2)=", old_x.stride(2), - " must equal dim=", dim, " ((nheads, dim) must be contiguous)"); - - // old_B: kernel uses `group_idx * DSTATE` → (ngroups, dstate) contig. - CHECK_CUDA(old_B); - CHECK_DIM(5, old_B); // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) - FLASHINFER_CHECK(old_B.size(0) == state_cache_size, "old_B.size(0)=", old_B.size(0), - " must equal state_cache_size=", state_cache_size); - FLASHINFER_CHECK(old_B.size(1) == 2, "old_B.size(1) must be 2 (double-buffered), got ", - old_B.size(1)); - FLASHINFER_CHECK(old_B.size(2) == max_window, "old_B.size(2)=", old_B.size(2), - " must equal max_window=", max_window); - FLASHINFER_CHECK(old_B.size(3) == ngroups, "old_B.size(3)=", old_B.size(3), - " must equal ngroups=", ngroups); - FLASHINFER_CHECK(old_B.size(4) == dstate, "old_B.size(4)=", old_B.size(4), - " must equal dstate=", dstate); - CHECK_LAST_DIM_CONTIGUOUS(old_B); - FLASHINFER_CHECK(old_B.stride(3) == dstate, "old_B.stride(3)=", old_B.stride(3), - " must equal dstate=", dstate, " ((ngroups, dstate) must be contiguous)"); - - // old_dt: kernel only assumes last-dim contig (head row). - CHECK_CUDA(old_dt); - CHECK_DIM(4, old_dt); // (state_cache_size, 2, nheads, MAX_WINDOW) - FLASHINFER_CHECK(old_dt.size(0) == state_cache_size, "old_dt.size(0)=", old_dt.size(0), - " must equal state_cache_size=", state_cache_size); - FLASHINFER_CHECK(old_dt.size(1) == 2, "old_dt.size(1) must be 2, got ", old_dt.size(1)); - FLASHINFER_CHECK(old_dt.size(2) == nheads, "old_dt.size(2)=", old_dt.size(2), - " must equal nheads=", nheads); - FLASHINFER_CHECK(old_dt.size(3) == max_window, "old_dt.size(3)=", old_dt.size(3), - " must equal max_window=", max_window); - CHECK_LAST_DIM_CONTIGUOUS(old_dt); - - // old_cumAdt: same as old_dt. - CHECK_CUDA(old_cumAdt); - CHECK_DIM(4, old_cumAdt); // (state_cache_size, 2, nheads, MAX_WINDOW) - FLASHINFER_CHECK(old_cumAdt.size(0) == state_cache_size, - "old_cumAdt.size(0)=", old_cumAdt.size(0), - " must equal state_cache_size=", state_cache_size); - FLASHINFER_CHECK(old_cumAdt.size(1) == 2, "old_cumAdt.size(1) must be 2, got ", - old_cumAdt.size(1)); - FLASHINFER_CHECK(old_cumAdt.size(2) == nheads, "old_cumAdt.size(2)=", old_cumAdt.size(2), - " must equal nheads=", nheads); - FLASHINFER_CHECK(old_cumAdt.size(3) == max_window, "old_cumAdt.size(3)=", old_cumAdt.size(3), - " must equal max_window=", max_window); - CHECK_LAST_DIM_CONTIGUOUS(old_cumAdt); - - CHECK_CUDA(cache_buf_idx); - CHECK_DIM(1, cache_buf_idx); - FLASHINFER_CHECK(cache_buf_idx.size(0) == state_cache_size, - "cache_buf_idx.size(0)=", cache_buf_idx.size(0), - " must equal state_cache_size=", state_cache_size); - CHECK_CONTIGUOUS(cache_buf_idx); - - CHECK_CUDA(prev_num_accepted); - CHECK_DIM(1, prev_num_accepted); - FLASHINFER_CHECK(prev_num_accepted.size(0) == state_cache_size, - "prev_num_accepted.size(0)=", prev_num_accepted.size(0), - " must equal state_cache_size=", state_cache_size); - CHECK_CONTIGUOUS(prev_num_accepted); - - // ── Validate optional D ── - if (D.has_value()) { - auto& Dv = D.value(); - CHECK_CUDA(Dv); - CHECK_DIM(2, Dv); - FLASHINFER_CHECK(Dv.size(0) == nheads, "D.size(0)=", Dv.size(0), " must equal nheads=", nheads); - FLASHINFER_CHECK(Dv.size(1) == dim, "D.size(1)=", Dv.size(1), " must equal dim=", dim); - FLASHINFER_CHECK(Dv.stride(0) == 1, "D.stride(0) must be 1 (tie_hdim), got ", Dv.stride(0)); - FLASHINFER_CHECK(Dv.stride(1) == 0, "D.stride(1) must be 0 (tie_hdim), got ", Dv.stride(1)); - } - - // ── Validate optional dt_bias ── - if (dt_bias.has_value()) { - auto& db = dt_bias.value(); - CHECK_CUDA(db); - CHECK_DIM(2, db); - FLASHINFER_CHECK(db.size(0) == nheads, "dt_bias.size(0)=", db.size(0), - " must equal nheads=", nheads); - FLASHINFER_CHECK(db.size(1) == dim, "dt_bias.size(1)=", db.size(1), " must equal dim=", dim); - FLASHINFER_CHECK(db.stride(0) == 1, "dt_bias.stride(0) must be 1 (tie_hdim), got ", - db.stride(0)); - FLASHINFER_CHECK(db.stride(1) == 0, "dt_bias.stride(1) must be 0 (tie_hdim), got ", - db.stride(1)); - } - - // ── Validate optional z: same layout/contig rules as x ── - if (z.has_value()) { - auto& zv = z.value(); - CHECK_CUDA(zv); - CHECK_DIM(4, zv); - if (is_varlen) { - FLASHINFER_CHECK(zv.size(0) == 1, "varlen: z.size(0)=", zv.size(0), " must be 1"); - FLASHINFER_CHECK(zv.size(1) == total_tokens, "varlen: z.size(1)=", zv.size(1), - " must equal x.size(1)=", total_tokens); - } else { - FLASHINFER_CHECK(zv.size(0) == batch, "z.size(0)=", zv.size(0), " must equal batch=", batch); - FLASHINFER_CHECK(zv.size(1) == npredicted, "z.size(1)=", zv.size(1), - " must equal npredicted=", npredicted); + Optional state_scale, // (state_cache_size, nheads, dim) f32 + Optional rand_seed, // single int64 + int64_t d_split, // v12 §59: per-head DIM split factor (1, 2, or 4) + Optional cu_seqlens) +{ // (batch+1,) int32, varlen mode + + bool const is_varlen = cu_seqlens.has_value(); + + // ── Extract dimensions ── + auto const state_cache_size = state.size(0); + auto const nheads = state.size(1); + auto const dim = state.size(2); + auto const dstate = state.size(3); + auto const max_window = old_x.size(1); + auto const ngroups = B.size(2); + + // In non-varlen mode, batch = x.size(0) and npredicted = x.size(1) (the + // 4D batched layout). In varlen, the JIT compile-time NPREDICTED is the + // max seq_len the caller commits to; the wrapper stamped it into the JIT + // URI from `max_seqlen` and we read it back as the constexpr `NPREDICTED` + // (validated against runtime cu_seqlens on the host side below). `batch` + // = number of sequences = `cu_seqlens.size(0) - 1`. + int64_t batch; + int64_t npredicted; + if (is_varlen) + { + auto const& cs = cu_seqlens.value(); + CHECK_CUDA(cs); + CHECK_DIM(1, cs); + FLASHINFER_CHECK( + cs.size(0) >= 2, "cu_seqlens must have shape (batch+1,) with batch >= 1, got size(0)=", cs.size(0)); + FLASHINFER_CHECK(cs.dtype().code == kDLInt && cs.dtype().bits == 32, "cu_seqlens must be int32"); + CHECK_CONTIGUOUS(cs); + batch = cs.size(0) - 1; + npredicted = NPREDICTED; // JIT-stamped — wrapper ensures max(seq_lens) <= NPREDICTED. + } + else + { + batch = x.size(0); + npredicted = x.size(1); + } + + // ── JIT compile-time / runtime cross-check ── + // NPREDICTED and MAX_WINDOW are JIT compile-time constants stamped by the + // wrapper. In non-varlen NPREDICTED = x.shape[1]; in varlen NPREDICTED = + // user-supplied `max_seqlen` (upper bound on every cu_seqlens diff). + FLASHINFER_CHECK(npredicted == NPREDICTED, is_varlen ? "max_seqlen=" : "x.size(1)=", npredicted, + " must equal JIT NPREDICTED=", NPREDICTED); + FLASHINFER_CHECK(max_window == MAX_WINDOW, "old_x.size(1)=", max_window, " must equal JIT MAX_WINDOW=", MAX_WINDOW); + FLASHINFER_CHECK(npredicted <= max_window, "npredicted=", npredicted, " must be <= max_window=", max_window); + + // ── Validate state ── + CHECK_CUDA(state); + CHECK_DIM(4, state); + { + auto s = state.strides(); + auto sz = state.sizes(); + FLASHINFER_CHECK(s[3] == 1, "state dim 3 (dstate) must have stride 1, got ", s[3]); + FLASHINFER_CHECK( + s[2] == sz[3], "state dim 2 (dim) must be contiguous with dim 3, got stride ", s[2], " expected ", sz[3]); + FLASHINFER_CHECK(s[1] == sz[2] * sz[3], "state dim 1 (nheads) must be contiguous with dim 2, got stride ", s[1], + " expected ", sz[2] * sz[3]); + } + + // ── Validate x ── + // Non-varlen: shape (batch, NPREDICTED, nheads, dim). + // Varlen : shape (1, total_tokens, nheads, dim) — batch axis collapsed, + // token axis is the outer iteration. The kernel reads x via + // `bos * x_stride_token + …` so x_stride_token is the per-token + // stride in either layout (= nheads*dim for contig). + CHECK_CUDA(x); + CHECK_DIM(4, x); + if (is_varlen) + { + FLASHINFER_CHECK(x.size(0) == 1, "varlen: x.size(0)=", x.size(0), " must be 1"); } - FLASHINFER_CHECK(zv.size(2) == nheads, "z.size(2)=", zv.size(2), " must equal nheads=", nheads); - FLASHINFER_CHECK(zv.size(3) == dim, "z.size(3)=", zv.size(3), " must equal dim=", dim); - CHECK_LAST_DIM_CONTIGUOUS(zv); - FLASHINFER_CHECK(zv.stride(2) == dim, "z.stride(2)=", zv.stride(2), " must equal dim=", dim, - " ((nheads, dim) must be contiguous)"); - } - - // ── Validate optional state_batch_indices ── - if (state_batch_indices.has_value()) { - auto& sbi = state_batch_indices.value(); - CHECK_CUDA(sbi); - CHECK_DIM(1, sbi); - FLASHINFER_CHECK(sbi.size(0) == batch, "state_batch_indices.size(0)=", sbi.size(0), - " must equal batch=", batch); - CHECK_CONTIGUOUS(sbi); - } - - // ── Validate optional state_scale: (state_cache_size, nheads, dim) ── - // Inner two dims (nheads, dim) must be contiguous; only batch stride is - // parameterized in the params struct. - if (state_scale.has_value()) { - auto const& ss = state_scale.value(); - CHECK_CUDA(ss); - CHECK_DIM(3, ss); - FLASHINFER_CHECK(ss.size(0) == state_cache_size, "state_scale.size(0)=", ss.size(0), - " must equal state_cache_size=", state_cache_size); - FLASHINFER_CHECK(ss.size(1) == nheads, "state_scale.size(1)=", ss.size(1), - " must equal nheads=", nheads); - FLASHINFER_CHECK(ss.size(2) == dim, "state_scale.size(2)=", ss.size(2), - " must equal dim=", dim); - FLASHINFER_CHECK(ss.stride(2) == 1, "state_scale.stride(2) must be 1, got ", ss.stride(2)); - FLASHINFER_CHECK(ss.stride(1) == dim, "state_scale.stride(1)=", ss.stride(1), - " must equal dim=", dim, " ((nheads, dim) must be contiguous)"); - } - - // ── Dtype consistency ── - // input_dtype = x.dtype; all activation tensors (B, C, output, z, old_x, - // old_B) and the state cache's "input-side" mirrors must match it. - // weight_dtype = D.dtype = dt_bias.dtype (kernel template sees one - // weight_t for both). - // Cache scalar tensors have fixed dtypes hardcoded in the kernel. - { - auto input_dtype = x.dtype(); - FLASHINFER_CHECK(B.dtype() == input_dtype, "B.dtype must match x.dtype"); - FLASHINFER_CHECK(C.dtype() == input_dtype, "C.dtype must match x.dtype"); - FLASHINFER_CHECK(output.dtype() == input_dtype, "output.dtype must match x.dtype"); - FLASHINFER_CHECK(old_x.dtype() == input_dtype, "old_x.dtype must match x.dtype"); - FLASHINFER_CHECK(old_B.dtype() == input_dtype, "old_B.dtype must match x.dtype"); - if (z.has_value()) { - FLASHINFER_CHECK(z.value().dtype() == input_dtype, "z.dtype must match x.dtype"); + else + { + FLASHINFER_CHECK(x.size(0) == batch, "x.size(0)=", x.size(0), " must equal batch=", batch); + FLASHINFER_CHECK(x.size(1) == npredicted, "x.size(1)=", x.size(1), " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(x.size(2) == nheads, "x.size(2)=", x.size(2), " must equal nheads=", nheads); + FLASHINFER_CHECK(x.size(3) == dim, "x.size(3)=", x.size(3), " must equal dim=", dim); + CHECK_LAST_DIM_CONTIGUOUS(x); + FLASHINFER_CHECK(x.stride(2) == dim, "x.stride(2)=", x.stride(2), " must equal dim=", dim, + " ((nheads, dim) must be contiguous)"); + + // In varlen, all per-token tensors share the flattened token axis — use + // x.size(1) as the canonical total_tokens and cross-check the others below. + int64_t const total_tokens = is_varlen ? x.size(1) : 0; + + // ── Validate dt ── + CHECK_CUDA(dt); + CHECK_DIM(4, dt); + if (is_varlen) + { + FLASHINFER_CHECK(dt.size(0) == 1, "varlen: dt.size(0)=", dt.size(0), " must be 1"); + FLASHINFER_CHECK( + dt.size(1) == total_tokens, "varlen: dt.size(1)=", dt.size(1), " must equal x.size(1)=", total_tokens); } - if (D.has_value() && dt_bias.has_value()) { - FLASHINFER_CHECK(D.value().dtype() == dt_bias.value().dtype(), - "D.dtype must equal dt_bias.dtype (kernel uses a single weight_t)"); + else + { + FLASHINFER_CHECK(dt.size(0) == batch, "dt.size(0)=", dt.size(0), " must equal batch=", batch); + FLASHINFER_CHECK(dt.size(1) == npredicted, "dt.size(1)=", dt.size(1), " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(dt.size(2) == nheads, "dt.size(2)=", dt.size(2), " must equal nheads=", nheads); + FLASHINFER_CHECK(dt.size(3) == dim, "dt.size(3)=", dt.size(3), " must equal dim=", dim); + FLASHINFER_CHECK(dt.stride(2) == 1, "dt.stride(2) must be 1 (tie_hdim), got ", dt.stride(2)); + FLASHINFER_CHECK(dt.stride(3) == 0, "dt.stride(3) must be 0 (tie_hdim), got ", dt.stride(3)); + + // ── Validate A: (nheads, dim, dstate) tie_hdim ── + CHECK_CUDA(A); + CHECK_DIM(3, A); + FLASHINFER_CHECK(A.size(0) == nheads, "A.size(0)=", A.size(0), " must equal nheads=", nheads); + FLASHINFER_CHECK(A.size(1) == dim, "A.size(1)=", A.size(1), " must equal dim=", dim); + FLASHINFER_CHECK(A.size(2) == dstate, "A.size(2)=", A.size(2), " must equal dstate=", dstate); + FLASHINFER_CHECK(A.stride(0) == 1, "A.stride(0) must be 1, got ", A.stride(0)); + FLASHINFER_CHECK(A.stride(1) == 0, "A.stride(1) must be 0 (tie_hdim), got ", A.stride(1)); + FLASHINFER_CHECK(A.stride(2) == 0, "A.stride(2) must be 0 (tie_hdim), got ", A.stride(2)); + + // ── Validate B ── + CHECK_CUDA(B); + CHECK_DIM(4, B); + if (is_varlen) + { + FLASHINFER_CHECK(B.size(0) == 1, "varlen: B.size(0)=", B.size(0), " must be 1"); + FLASHINFER_CHECK( + B.size(1) == total_tokens, "varlen: B.size(1)=", B.size(1), " must equal x.size(1)=", total_tokens); + } + else + { + FLASHINFER_CHECK(B.size(0) == batch, "B.size(0)=", B.size(0), " must equal batch=", batch); + FLASHINFER_CHECK(B.size(1) == npredicted, "B.size(1)=", B.size(1), " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(B.size(3) == dstate, "B.size(3)=", B.size(3), " must equal dstate=", dstate); + CHECK_LAST_DIM_CONTIGUOUS(B); + FLASHINFER_CHECK(B.stride(2) == dstate, "B.stride(2)=", B.stride(2), " must equal dstate=", dstate, + " ((ngroups, dstate) must be contiguous)"); + FLASHINFER_CHECK(nheads % ngroups == 0, "nheads=", nheads, " must be divisible by ngroups=", ngroups); + + // ── Validate C ── + CHECK_CUDA(C); + CHECK_DIM(4, C); + if (is_varlen) + { + FLASHINFER_CHECK(C.size(0) == 1, "varlen: C.size(0)=", C.size(0), " must be 1"); + FLASHINFER_CHECK( + C.size(1) == total_tokens, "varlen: C.size(1)=", C.size(1), " must equal x.size(1)=", total_tokens); + } + else + { + FLASHINFER_CHECK(C.size(0) == batch, "C.size(0)=", C.size(0), " must equal batch=", batch); + FLASHINFER_CHECK(C.size(1) == npredicted, "C.size(1)=", C.size(1), " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(C.size(2) == ngroups, "C.size(2)=", C.size(2), " must equal ngroups=", ngroups); + FLASHINFER_CHECK(C.size(3) == dstate, "C.size(3)=", C.size(3), " must equal dstate=", dstate); + CHECK_LAST_DIM_CONTIGUOUS(C); + FLASHINFER_CHECK(C.stride(2) == dstate, "C.stride(2)=", C.stride(2), " must equal dstate=", dstate, + " ((ngroups, dstate) must be contiguous)"); + + // ── Validate output ── + CHECK_CUDA(output); + CHECK_DIM(4, output); + if (is_varlen) + { + FLASHINFER_CHECK(output.size(0) == 1, "varlen: output.size(0)=", output.size(0), " must be 1"); + FLASHINFER_CHECK(output.size(1) == total_tokens, "varlen: output.size(1)=", output.size(1), + " must equal x.size(1)=", total_tokens); + } + else + { + FLASHINFER_CHECK(output.size(0) == batch, "output.size(0)=", output.size(0), " must equal batch=", batch); + FLASHINFER_CHECK( + output.size(1) == npredicted, "output.size(1)=", output.size(1), " must equal npredicted=", npredicted); } - // old_dt / old_cumAdt are produced by this same kernel in f32 and - // consumed back in f32 on the next call. - FLASHINFER_CHECK(old_dt.dtype().code == kDLFloat && old_dt.dtype().bits == 32, - "old_dt must be float32"); - FLASHINFER_CHECK(old_cumAdt.dtype().code == kDLFloat && old_cumAdt.dtype().bits == 32, - "old_cumAdt must be float32"); - // Index tensors used by the kernel as int32 scalars. - FLASHINFER_CHECK(cache_buf_idx.dtype().code == kDLInt && cache_buf_idx.dtype().bits == 32, - "cache_buf_idx must be int32"); + FLASHINFER_CHECK(output.size(2) == nheads, "output.size(2)=", output.size(2), " must equal nheads=", nheads); + FLASHINFER_CHECK(output.size(3) == dim, "output.size(3)=", output.size(3), " must equal dim=", dim); + CHECK_LAST_DIM_CONTIGUOUS(output); + FLASHINFER_CHECK(output.stride(2) == dim, "output.stride(2)=", output.stride(2), " must equal dim=", dim, + " ((nheads, dim) must be contiguous)"); + + // ── Validate cache tensors ── + // old_x: kernel uses `head * DIM + d_tile_off` → (nheads, dim) contig. + CHECK_CUDA(old_x); + CHECK_DIM(4, old_x); // (state_cache_size, MAX_WINDOW, nheads, dim) + FLASHINFER_CHECK(old_x.size(0) == state_cache_size, "old_x.size(0)=", old_x.size(0), + " must equal state_cache_size=", state_cache_size); + FLASHINFER_CHECK( + old_x.size(1) == max_window, "old_x.size(1)=", old_x.size(1), " must equal max_window=", max_window); + FLASHINFER_CHECK(old_x.size(2) == nheads, "old_x.size(2)=", old_x.size(2), " must equal nheads=", nheads); + FLASHINFER_CHECK(old_x.size(3) == dim, "old_x.size(3)=", old_x.size(3), " must equal dim=", dim); + CHECK_LAST_DIM_CONTIGUOUS(old_x); + FLASHINFER_CHECK(old_x.stride(2) == dim, "old_x.stride(2)=", old_x.stride(2), " must equal dim=", dim, + " ((nheads, dim) must be contiguous)"); + + // old_B: kernel uses `group_idx * DSTATE` → (ngroups, dstate) contig. + CHECK_CUDA(old_B); + CHECK_DIM(5, old_B); // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) + FLASHINFER_CHECK(old_B.size(0) == state_cache_size, "old_B.size(0)=", old_B.size(0), + " must equal state_cache_size=", state_cache_size); + FLASHINFER_CHECK(old_B.size(1) == 2, "old_B.size(1) must be 2 (double-buffered), got ", old_B.size(1)); + FLASHINFER_CHECK( + old_B.size(2) == max_window, "old_B.size(2)=", old_B.size(2), " must equal max_window=", max_window); + FLASHINFER_CHECK(old_B.size(3) == ngroups, "old_B.size(3)=", old_B.size(3), " must equal ngroups=", ngroups); + FLASHINFER_CHECK(old_B.size(4) == dstate, "old_B.size(4)=", old_B.size(4), " must equal dstate=", dstate); + CHECK_LAST_DIM_CONTIGUOUS(old_B); + FLASHINFER_CHECK(old_B.stride(3) == dstate, "old_B.stride(3)=", old_B.stride(3), " must equal dstate=", dstate, + " ((ngroups, dstate) must be contiguous)"); + + // old_dt: kernel only assumes last-dim contig (head row). + CHECK_CUDA(old_dt); + CHECK_DIM(4, old_dt); // (state_cache_size, 2, nheads, MAX_WINDOW) + FLASHINFER_CHECK(old_dt.size(0) == state_cache_size, "old_dt.size(0)=", old_dt.size(0), + " must equal state_cache_size=", state_cache_size); + FLASHINFER_CHECK(old_dt.size(1) == 2, "old_dt.size(1) must be 2, got ", old_dt.size(1)); + FLASHINFER_CHECK(old_dt.size(2) == nheads, "old_dt.size(2)=", old_dt.size(2), " must equal nheads=", nheads); + FLASHINFER_CHECK( + old_dt.size(3) == max_window, "old_dt.size(3)=", old_dt.size(3), " must equal max_window=", max_window); + CHECK_LAST_DIM_CONTIGUOUS(old_dt); + + // old_cumAdt: same as old_dt. + CHECK_CUDA(old_cumAdt); + CHECK_DIM(4, old_cumAdt); // (state_cache_size, 2, nheads, MAX_WINDOW) + FLASHINFER_CHECK(old_cumAdt.size(0) == state_cache_size, "old_cumAdt.size(0)=", old_cumAdt.size(0), + " must equal state_cache_size=", state_cache_size); + FLASHINFER_CHECK(old_cumAdt.size(1) == 2, "old_cumAdt.size(1) must be 2, got ", old_cumAdt.size(1)); FLASHINFER_CHECK( - prev_num_accepted.dtype().code == kDLInt && prev_num_accepted.dtype().bits == 32, - "prev_num_accepted must be int32"); - if (state_batch_indices.has_value()) { - auto sbi_dt = state_batch_indices.value().dtype(); - FLASHINFER_CHECK(sbi_dt.code == kDLInt && (sbi_dt.bits == 32 || sbi_dt.bits == 64), - "state_batch_indices must be int32 or int64"); + old_cumAdt.size(2) == nheads, "old_cumAdt.size(2)=", old_cumAdt.size(2), " must equal nheads=", nheads); + FLASHINFER_CHECK(old_cumAdt.size(3) == max_window, "old_cumAdt.size(3)=", old_cumAdt.size(3), + " must equal max_window=", max_window); + CHECK_LAST_DIM_CONTIGUOUS(old_cumAdt); + + CHECK_CUDA(cache_buf_idx); + CHECK_DIM(1, cache_buf_idx); + FLASHINFER_CHECK(cache_buf_idx.size(0) == state_cache_size, "cache_buf_idx.size(0)=", cache_buf_idx.size(0), + " must equal state_cache_size=", state_cache_size); + CHECK_CONTIGUOUS(cache_buf_idx); + + CHECK_CUDA(prev_num_accepted); + CHECK_DIM(1, prev_num_accepted); + FLASHINFER_CHECK(prev_num_accepted.size(0) == state_cache_size, + "prev_num_accepted.size(0)=", prev_num_accepted.size(0), " must equal state_cache_size=", state_cache_size); + CHECK_CONTIGUOUS(prev_num_accepted); + + // ── Validate optional D ── + if (D.has_value()) + { + auto& Dv = D.value(); + CHECK_CUDA(Dv); + CHECK_DIM(2, Dv); + FLASHINFER_CHECK(Dv.size(0) == nheads, "D.size(0)=", Dv.size(0), " must equal nheads=", nheads); + FLASHINFER_CHECK(Dv.size(1) == dim, "D.size(1)=", Dv.size(1), " must equal dim=", dim); + FLASHINFER_CHECK(Dv.stride(0) == 1, "D.stride(0) must be 1 (tie_hdim), got ", Dv.stride(0)); + FLASHINFER_CHECK(Dv.stride(1) == 0, "D.stride(1) must be 0 (tie_hdim), got ", Dv.stride(1)); + } + + // ── Validate optional dt_bias ── + if (dt_bias.has_value()) + { + auto& db = dt_bias.value(); + CHECK_CUDA(db); + CHECK_DIM(2, db); + FLASHINFER_CHECK(db.size(0) == nheads, "dt_bias.size(0)=", db.size(0), " must equal nheads=", nheads); + FLASHINFER_CHECK(db.size(1) == dim, "dt_bias.size(1)=", db.size(1), " must equal dim=", dim); + FLASHINFER_CHECK(db.stride(0) == 1, "dt_bias.stride(0) must be 1 (tie_hdim), got ", db.stride(0)); + FLASHINFER_CHECK(db.stride(1) == 0, "dt_bias.stride(1) must be 0 (tie_hdim), got ", db.stride(1)); } - if (state_scale.has_value()) { - auto ss_dt = state_scale.value().dtype(); - FLASHINFER_CHECK(ss_dt.code == kDLFloat && ss_dt.bits == 32, "state_scale must be float32"); + + // ── Validate optional z: same layout/contig rules as x ── + if (z.has_value()) + { + auto& zv = z.value(); + CHECK_CUDA(zv); + CHECK_DIM(4, zv); + if (is_varlen) + { + FLASHINFER_CHECK(zv.size(0) == 1, "varlen: z.size(0)=", zv.size(0), " must be 1"); + FLASHINFER_CHECK( + zv.size(1) == total_tokens, "varlen: z.size(1)=", zv.size(1), " must equal x.size(1)=", total_tokens); + } + else + { + FLASHINFER_CHECK(zv.size(0) == batch, "z.size(0)=", zv.size(0), " must equal batch=", batch); + FLASHINFER_CHECK(zv.size(1) == npredicted, "z.size(1)=", zv.size(1), " must equal npredicted=", npredicted); + } + FLASHINFER_CHECK(zv.size(2) == nheads, "z.size(2)=", zv.size(2), " must equal nheads=", nheads); + FLASHINFER_CHECK(zv.size(3) == dim, "z.size(3)=", zv.size(3), " must equal dim=", dim); + CHECK_LAST_DIM_CONTIGUOUS(zv); + FLASHINFER_CHECK(zv.stride(2) == dim, "z.stride(2)=", zv.stride(2), " must equal dim=", dim, + " ((nheads, dim) must be contiguous)"); } - // Quantized state dtypes (int8, fp8_e4m3fn, ...) require a state_scale - // tensor; non-quantized dtypes must not pass one. Mirrors the Python - // wrapper assertion and matches the kernel's compile-time - // `state_scale_t == void` gating. + + // ── Validate optional state_batch_indices ── + if (state_batch_indices.has_value()) { - auto sd = state.dtype(); - bool const is_int8 = (sd.code == kDLInt && sd.bits == 8); - bool const is_fp8 = (sd.code == kDLFloat8_e4m3fn && sd.bits == 8); - bool const is_quantized_state = is_int8 || is_fp8; - if (is_quantized_state) { - FLASHINFER_CHECK(state_scale.has_value(), - "Quantized state.dtype (int8/fp8_e4m3fn) requires a state_scale tensor " - "of shape (state_cache_size, nheads, dim) and dtype float32"); - // The 8-bit replay path uses Layout<_4, _1> (M-shard per warp) which - // needs per-warp M = D_PER_CTA / 4 >= 16 (m16n8 atom M). This forces - // D_PER_CTA >= 64, i.e. d_split == 1. + auto& sbi = state_batch_indices.value(); + CHECK_CUDA(sbi); + CHECK_DIM(1, sbi); FLASHINFER_CHECK( - d_split == 1, - "Quantized state.dtype (int8/fp8_e4m3fn) requires d_split=1 (got d_split=", d_split, - "); the M-shard-per-warp replay layout needs D_PER_CTA / 4 >= 16."); - } else { - FLASHINFER_CHECK(!state_scale.has_value(), - "state_scale must be None for non-quantized state.dtype " - "(allowed quantized dtypes: {int8, fp8_e4m3fn})"); - } + sbi.size(0) == batch, "state_batch_indices.size(0)=", sbi.size(0), " must equal batch=", batch); + CHECK_CONTIGUOUS(sbi); + } + + // ── Validate optional state_scale: (state_cache_size, nheads, dim) ── + // Inner two dims (nheads, dim) must be contiguous; only batch stride is + // parameterized in the params struct. + if (state_scale.has_value()) + { + auto const& ss = state_scale.value(); + CHECK_CUDA(ss); + CHECK_DIM(3, ss); + FLASHINFER_CHECK(ss.size(0) == state_cache_size, "state_scale.size(0)=", ss.size(0), + " must equal state_cache_size=", state_cache_size); + FLASHINFER_CHECK(ss.size(1) == nheads, "state_scale.size(1)=", ss.size(1), " must equal nheads=", nheads); + FLASHINFER_CHECK(ss.size(2) == dim, "state_scale.size(2)=", ss.size(2), " must equal dim=", dim); + FLASHINFER_CHECK(ss.stride(2) == 1, "state_scale.stride(2) must be 1, got ", ss.stride(2)); + FLASHINFER_CHECK(ss.stride(1) == dim, "state_scale.stride(1)=", ss.stride(1), " must equal dim=", dim, + " ((nheads, dim) must be contiguous)"); } - } - - // ── Populate params ── - CheckpointingSsuParams p; - - // ── Validate d_split (v12 §59) ── - // Allowed for v12: {1, 2}. d_split=4 deferred to v12.x (needs warp-count - // restructure — output MMA `_1×4` layout requires D_PER_CTA ≥ 32). - FLASHINFER_CHECK(d_split == 1 || d_split == 2, "d_split=", d_split, - " must be one of {1, 2} (d_split=4 is deferred to v12.x)"); - FLASHINFER_CHECK(dim % d_split == 0, "dim=", dim, " must be divisible by d_split=", d_split); - FLASHINFER_CHECK(dim / d_split >= 32, "d_split=", d_split, " gives D_PER_CTA=", dim / d_split, - " < 32 (output MMA m16n8 atom floor with _1×4 warp layout)"); - - p.batch = batch; - p.nheads = nheads; - p.dim = dim; - p.dstate = dstate; - p.ngroups = ngroups; - p.state_cache_size = state_cache_size; - p.npredicted = npredicted; - p.max_window = max_window; - p.pad_slot_id = pad_slot_id; - p.d_split = static_cast(d_split); - p.dt_softplus = dt_softplus; - - // Pointers - p.state = state.data_ptr(); - p.x = const_cast(x.data_ptr()); - p.dt = const_cast(dt.data_ptr()); - p.A = const_cast(A.data_ptr()); - p.B = const_cast(B.data_ptr()); - p.C = const_cast(C.data_ptr()); - p.output = output.data_ptr(); - - p.old_x = old_x.data_ptr(); - p.old_B = const_cast(old_B.data_ptr()); - p.old_dt = const_cast(old_dt.data_ptr()); - p.old_cumAdt = const_cast(old_cumAdt.data_ptr()); - p.cache_buf_idx = const_cast(cache_buf_idx.data_ptr()); - p.prev_num_accepted = const_cast(prev_num_accepted.data_ptr()); - - if (D.has_value()) p.D = const_cast(D.value().data_ptr()); - if (z.has_value()) { - p.z = const_cast(z.value().data_ptr()); - // Same seq-dim selection as the rest of the batch-side tensors below. - p.z_stride_seq = z.value().stride(is_varlen ? 1 : 0); - p.z_stride_token = z.value().stride(1); - } - if (dt_bias.has_value()) p.dt_bias = const_cast(dt_bias.value().data_ptr()); - if (state_batch_indices.has_value()) - p.state_batch_indices = const_cast(state_batch_indices.value().data_ptr()); - if (is_varlen) { - p.cu_seqlens = const_cast(cu_seqlens.value().data_ptr()); - } - if (state_scale.has_value()) { - p.state_scale = state_scale.value().data_ptr(); - p.state_scale_stride_seq = state_scale.value().stride(0); - } - if (rand_seed.has_value()) { - auto const& rs = rand_seed.value(); - CHECK_CUDA(rs); - FLASHINFER_CHECK(rs.numel() == 1, "rand_seed must be single-element, got numel=", rs.numel()); - FLASHINFER_CHECK(rs.dtype().code == kDLInt && rs.dtype().bits == 64, "rand_seed must be int64"); - p.rand_seed = static_cast(rs.data_ptr()); - } - - // Strides - p.state_stride_seq = state.stride(0); - - // `*_stride_seq` is the outer iteration stride. Non-varlen iterates over - // dim 0 (per-batch), varlen iterates over dim 1 (per-token) — sequences - // are packed into a single batch in the (1, total_tokens, ...) layout. - // The kernel uses one formula `seq * *_stride_seq` for both modes. - int const seq_dim = is_varlen ? 1 : 0; - p.x_stride_seq = x.stride(seq_dim); - p.x_stride_token = x.stride(1); - p.dt_stride_seq = dt.stride(seq_dim); - p.dt_stride_token = dt.stride(1); - p.B_stride_seq = B.stride(seq_dim); - p.B_stride_token = B.stride(1); - p.C_stride_seq = C.stride(seq_dim); - p.C_stride_token = C.stride(1); - p.out_stride_seq = output.stride(seq_dim); - p.out_stride_token = output.stride(1); - - p.old_x_stride_seq = old_x.stride(0); - p.old_x_stride_token = old_x.stride(1); - p.old_B_stride_seq = old_B.stride(0); - p.old_B_stride_dbuf = old_B.stride(1); - p.old_B_stride_token = old_B.stride(2); - p.old_dt_stride_seq = old_dt.stride(0); - p.old_dt_stride_dbuf = old_dt.stride(1); - p.old_dt_stride_head = old_dt.stride(2); - p.old_cumAdt_stride_seq = old_cumAdt.stride(0); - p.old_cumAdt_stride_dbuf = old_cumAdt.stride(1); - p.old_cumAdt_stride_head = old_cumAdt.stride(2); - - // Launch - ffi::CUDADeviceGuard device_guard(state.device().device_id); - const cudaStream_t stream = get_stream(state.device()); - - launchCheckpointingSsu( - p, stream); + + // ── Dtype consistency ── + // input_dtype = x.dtype; all activation tensors (B, C, output, z, old_x, + // old_B) and the state cache's "input-side" mirrors must match it. + // weight_dtype = D.dtype = dt_bias.dtype (kernel template sees one + // weight_t for both). + // Cache scalar tensors have fixed dtypes hardcoded in the kernel. + { + auto input_dtype = x.dtype(); + FLASHINFER_CHECK(B.dtype() == input_dtype, "B.dtype must match x.dtype"); + FLASHINFER_CHECK(C.dtype() == input_dtype, "C.dtype must match x.dtype"); + FLASHINFER_CHECK(output.dtype() == input_dtype, "output.dtype must match x.dtype"); + FLASHINFER_CHECK(old_x.dtype() == input_dtype, "old_x.dtype must match x.dtype"); + FLASHINFER_CHECK(old_B.dtype() == input_dtype, "old_B.dtype must match x.dtype"); + if (z.has_value()) + { + FLASHINFER_CHECK(z.value().dtype() == input_dtype, "z.dtype must match x.dtype"); + } + if (D.has_value() && dt_bias.has_value()) + { + FLASHINFER_CHECK(D.value().dtype() == dt_bias.value().dtype(), + "D.dtype must equal dt_bias.dtype (kernel uses a single weight_t)"); + } + // old_dt / old_cumAdt are produced by this same kernel in f32 and + // consumed back in f32 on the next call. + FLASHINFER_CHECK(old_dt.dtype().code == kDLFloat && old_dt.dtype().bits == 32, "old_dt must be float32"); + FLASHINFER_CHECK( + old_cumAdt.dtype().code == kDLFloat && old_cumAdt.dtype().bits == 32, "old_cumAdt must be float32"); + // Index tensors used by the kernel as int32 scalars. + FLASHINFER_CHECK( + cache_buf_idx.dtype().code == kDLInt && cache_buf_idx.dtype().bits == 32, "cache_buf_idx must be int32"); + FLASHINFER_CHECK(prev_num_accepted.dtype().code == kDLInt && prev_num_accepted.dtype().bits == 32, + "prev_num_accepted must be int32"); + if (state_batch_indices.has_value()) + { + auto sbi_dt = state_batch_indices.value().dtype(); + FLASHINFER_CHECK(sbi_dt.code == kDLInt && (sbi_dt.bits == 32 || sbi_dt.bits == 64), + "state_batch_indices must be int32 or int64"); + } + if (state_scale.has_value()) + { + auto ss_dt = state_scale.value().dtype(); + FLASHINFER_CHECK(ss_dt.code == kDLFloat && ss_dt.bits == 32, "state_scale must be float32"); + } + // Quantized state dtypes (int8, fp8_e4m3fn, ...) require a state_scale + // tensor; non-quantized dtypes must not pass one. Mirrors the Python + // wrapper assertion and matches the kernel's compile-time + // `state_scale_t == void` gating. + { + auto sd = state.dtype(); + bool const is_int8 = (sd.code == kDLInt && sd.bits == 8); + bool const is_fp8 = (sd.code == kDLFloat8_e4m3fn && sd.bits == 8); + bool const is_quantized_state = is_int8 || is_fp8; + if (is_quantized_state) + { + FLASHINFER_CHECK(state_scale.has_value(), + "Quantized state.dtype (int8/fp8_e4m3fn) requires a state_scale tensor " + "of shape (state_cache_size, nheads, dim) and dtype float32"); + // The 8-bit replay path uses Layout<_4, _1> (M-shard per warp) which + // needs per-warp M = D_PER_CTA / 4 >= 16 (m16n8 atom M). This forces + // D_PER_CTA >= 64, i.e. d_split == 1. + FLASHINFER_CHECK(d_split == 1, + "Quantized state.dtype (int8/fp8_e4m3fn) requires d_split=1 (got d_split=", d_split, + "); the M-shard-per-warp replay layout needs D_PER_CTA / 4 >= 16."); + } + else + { + FLASHINFER_CHECK(!state_scale.has_value(), + "state_scale must be None for non-quantized state.dtype " + "(allowed quantized dtypes: {int8, fp8_e4m3fn})"); + } + } + } + + // ── Populate params ── + CheckpointingSsuParams p; + + // ── Validate d_split (v12 §59) ── + // Allowed for v12: {1, 2}. d_split=4 deferred to v12.x (needs warp-count + // restructure — output MMA `_1×4` layout requires D_PER_CTA ≥ 32). + FLASHINFER_CHECK( + d_split == 1 || d_split == 2, "d_split=", d_split, " must be one of {1, 2} (d_split=4 is deferred to v12.x)"); + FLASHINFER_CHECK(dim % d_split == 0, "dim=", dim, " must be divisible by d_split=", d_split); + FLASHINFER_CHECK(dim / d_split >= 32, "d_split=", d_split, " gives D_PER_CTA=", dim / d_split, + " < 32 (output MMA m16n8 atom floor with _1×4 warp layout)"); + + p.batch = batch; + p.nheads = nheads; + p.dim = dim; + p.dstate = dstate; + p.ngroups = ngroups; + p.state_cache_size = state_cache_size; + p.npredicted = npredicted; + p.max_window = max_window; + p.pad_slot_id = pad_slot_id; + p.d_split = static_cast(d_split); + p.dt_softplus = dt_softplus; + + // Pointers + p.state = state.data_ptr(); + p.x = const_cast(x.data_ptr()); + p.dt = const_cast(dt.data_ptr()); + p.A = const_cast(A.data_ptr()); + p.B = const_cast(B.data_ptr()); + p.C = const_cast(C.data_ptr()); + p.output = output.data_ptr(); + + p.old_x = old_x.data_ptr(); + p.old_B = const_cast(old_B.data_ptr()); + p.old_dt = const_cast(old_dt.data_ptr()); + p.old_cumAdt = const_cast(old_cumAdt.data_ptr()); + p.cache_buf_idx = const_cast(cache_buf_idx.data_ptr()); + p.prev_num_accepted = const_cast(prev_num_accepted.data_ptr()); + + if (D.has_value()) + p.D = const_cast(D.value().data_ptr()); + if (z.has_value()) + { + p.z = const_cast(z.value().data_ptr()); + // Same seq-dim selection as the rest of the batch-side tensors below. + p.z_stride_seq = z.value().stride(is_varlen ? 1 : 0); + p.z_stride_token = z.value().stride(1); + } + if (dt_bias.has_value()) + p.dt_bias = const_cast(dt_bias.value().data_ptr()); + if (state_batch_indices.has_value()) + p.state_batch_indices = const_cast(state_batch_indices.value().data_ptr()); + if (is_varlen) + { + p.cu_seqlens = const_cast(cu_seqlens.value().data_ptr()); + } + if (state_scale.has_value()) + { + p.state_scale = state_scale.value().data_ptr(); + p.state_scale_stride_seq = state_scale.value().stride(0); + } + if (rand_seed.has_value()) + { + auto const& rs = rand_seed.value(); + CHECK_CUDA(rs); + FLASHINFER_CHECK(rs.numel() == 1, "rand_seed must be single-element, got numel=", rs.numel()); + FLASHINFER_CHECK(rs.dtype().code == kDLInt && rs.dtype().bits == 64, "rand_seed must be int64"); + p.rand_seed = static_cast(rs.data_ptr()); + } + + // Strides + p.state_stride_seq = state.stride(0); + + // `*_stride_seq` is the outer iteration stride. Non-varlen iterates over + // dim 0 (per-batch), varlen iterates over dim 1 (per-token) — sequences + // are packed into a single batch in the (1, total_tokens, ...) layout. + // The kernel uses one formula `seq * *_stride_seq` for both modes. + int const seq_dim = is_varlen ? 1 : 0; + p.x_stride_seq = x.stride(seq_dim); + p.x_stride_token = x.stride(1); + p.dt_stride_seq = dt.stride(seq_dim); + p.dt_stride_token = dt.stride(1); + p.B_stride_seq = B.stride(seq_dim); + p.B_stride_token = B.stride(1); + p.C_stride_seq = C.stride(seq_dim); + p.C_stride_token = C.stride(1); + p.out_stride_seq = output.stride(seq_dim); + p.out_stride_token = output.stride(1); + + p.old_x_stride_seq = old_x.stride(0); + p.old_x_stride_token = old_x.stride(1); + p.old_B_stride_seq = old_B.stride(0); + p.old_B_stride_dbuf = old_B.stride(1); + p.old_B_stride_token = old_B.stride(2); + p.old_dt_stride_seq = old_dt.stride(0); + p.old_dt_stride_dbuf = old_dt.stride(1); + p.old_dt_stride_head = old_dt.stride(2); + p.old_cumAdt_stride_seq = old_cumAdt.stride(0); + p.old_cumAdt_stride_dbuf = old_cumAdt.stride(1); + p.old_cumAdt_stride_head = old_cumAdt.stride(2); + + // Launch + ffi::CUDADeviceGuard device_guard(state.device().device_id); + const cudaStream_t stream = get_stream(state.device()); + + launchCheckpointingSsu(p, stream); } -} // namespace flashinfer::mamba::checkpointing +} // namespace flashinfer::mamba::checkpointing diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu index 89034be5c9cc..508c52666633 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu @@ -17,37 +17,35 @@ using tvm::ffi::Optional; -namespace flashinfer::mamba::checkpointing { +namespace flashinfer::mamba::checkpointing +{ -void checkpointing_ssu( - TensorView state, // (cache, nheads, dim, dstate) - TensorView x, // 4D (batch, T, nheads, dim) or 4D (1, total_tokens, nheads, dim) under varlen - TensorView - dt, // (batch, T, nheads, dim) tie_hdim / (1, total_tokens, nheads, dim) under varlen - TensorView A, // (nheads, dim, dstate) tie_hdim - TensorView B, // (batch, T, ngroups, dstate) / (1, total_tokens, ngroups, dstate) under varlen - TensorView C, // same as B - TensorView output, // same layout as x +void checkpointing_ssu(TensorView state, // (cache, nheads, dim, dstate) + TensorView x, // 4D (batch, T, nheads, dim) or 4D (1, total_tokens, nheads, dim) under varlen + TensorView dt, // (batch, T, nheads, dim) tie_hdim / (1, total_tokens, nheads, dim) under varlen + TensorView A, // (nheads, dim, dstate) tie_hdim + TensorView B, // (batch, T, ngroups, dstate) / (1, total_tokens, ngroups, dstate) under varlen + TensorView C, // same as B + TensorView output, // same layout as x // Cache tensors - TensorView old_x, // (cache, T, nheads, dim) - TensorView old_B, // (cache, 2, T, ngroups, dstate) - TensorView old_dt, // (cache, 2, nheads, T) f32 - TensorView old_cumAdt, // (cache, 2, nheads, T) f32 - TensorView cache_buf_idx, // (cache,) int32 - TensorView prev_num_accepted, // (cache,) int32 + TensorView old_x, // (cache, T, nheads, dim) + TensorView old_B, // (cache, 2, T, ngroups, dstate) + TensorView old_dt, // (cache, 2, nheads, T) f32 + TensorView old_cumAdt, // (cache, 2, nheads, T) f32 + TensorView cache_buf_idx, // (cache,) int32 + TensorView prev_num_accepted, // (cache,) int32 // Optional tensors - Optional D, // (nheads, dim) - Optional z, // same layout as x - Optional dt_bias, // (nheads, dim) tie_hdim + Optional D, // (nheads, dim) + Optional z, // same layout as x + Optional dt_bias, // (nheads, dim) tie_hdim bool dt_softplus, - Optional state_batch_indices, // (batch,) int32 + Optional state_batch_indices, // (batch,) int32 int64_t pad_slot_id, - Optional state_scale, // (cache, nheads, dim) f32 - Optional rand_seed, // single int64 - int64_t d_split, // v12 §59: per-head DIM split factor (1, 2, or 4) - Optional cu_seqlens); // (batch+1,) int32 — varlen mode + Optional state_scale, // (cache, nheads, dim) f32 + Optional rand_seed, // single int64 + int64_t d_split, // v12 §59: per-head DIM split factor (1, 2, or 4) + Optional cu_seqlens); // (batch+1,) int32 — varlen mode -} // namespace flashinfer::mamba::checkpointing +} // namespace flashinfer::mamba::checkpointing -TVM_FFI_DLL_EXPORT_TYPED_FUNC(checkpointing_ssu, - flashinfer::mamba::checkpointing::checkpointing_ssu); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(checkpointing_ssu, flashinfer::mamba::checkpointing::checkpointing_ssu); diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu index 07eaa5e597a5..cbb945631093 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu @@ -2,11 +2,13 @@ #include "checkpointing_ssu_config.inc" #include #include + // clang-format on -namespace flashinfer::mamba::checkpointing { +namespace flashinfer::mamba::checkpointing +{ -template void launchCheckpointingSsu(CheckpointingSsuParams&, cudaStream_t); +template void launchCheckpointingSsu( + CheckpointingSsuParams&, cudaStream_t); -} // namespace flashinfer::mamba::checkpointing +} // namespace flashinfer::mamba::checkpointing diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h index aaaa2b5b3e51..4521f68c5086 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h @@ -23,84 +23,104 @@ #define FLASHINFER_ERROR(message) throw flashinfer::Error(__FUNCTION__, __FILE__, __LINE__, message) // Base case for empty arguments -inline void write_to_stream(std::ostringstream& oss) { - // No-op for empty arguments +inline void write_to_stream(std::ostringstream& oss) +{ + // No-op for empty arguments } template -void write_to_stream(std::ostringstream& oss, T&& val) { - oss << std::forward(val); +void write_to_stream(std::ostringstream& oss, T&& val) +{ + oss << std::forward(val); } template -void write_to_stream(std::ostringstream& oss, T&& val, Args&&... args) { - oss << std::forward(val) << " "; - write_to_stream(oss, std::forward(args)...); +void write_to_stream(std::ostringstream& oss, T&& val, Args&&... args) +{ + oss << std::forward(val) << " "; + write_to_stream(oss, std::forward(args)...); } // Helper macro to handle empty __VA_ARGS__ -#define FLASHINFER_CHECK_IMPL(condition, message) \ - if (!(condition)) { \ - FLASHINFER_ERROR(message); \ - } +#define FLASHINFER_CHECK_IMPL(condition, message) \ + if (!(condition)) \ + { \ + FLASHINFER_ERROR(message); \ + } // Main macro that handles both cases -#define FLASHINFER_CHECK(condition, ...) \ - do { \ - if (!(condition)) { \ - std::ostringstream oss; \ - write_to_stream(oss, ##__VA_ARGS__); \ - std::string msg = oss.str(); \ - if (msg.empty()) { \ - msg = "Check failed: " #condition; \ - } \ - FLASHINFER_ERROR(msg); \ - } \ - } while (0) +#define FLASHINFER_CHECK(condition, ...) \ + do \ + { \ + if (!(condition)) \ + { \ + std::ostringstream oss; \ + write_to_stream(oss, ##__VA_ARGS__); \ + std::string msg = oss.str(); \ + if (msg.empty()) \ + { \ + msg = "Check failed: " #condition; \ + } \ + FLASHINFER_ERROR(msg); \ + } \ + } while (0) // Warning macro -#define FLASHINFER_WARN(...) \ - do { \ - std::ostringstream oss; \ - write_to_stream(oss, ##__VA_ARGS__); \ - std::string msg = oss.str(); \ - if (msg.empty()) { \ - msg = "Warning triggered"; \ - } \ - flashinfer::Warning(__FUNCTION__, __FILE__, __LINE__, msg).emit(); \ - } while (0) - -namespace flashinfer { -class Error : public std::exception { - private: - std::string message_; - - public: - Error(const std::string& func, const std::string& file, int line, const std::string& message) { - std::ostringstream oss; - oss << "Error in function '" << func << "' " - << "at " << file << ":" << line << ": " << message; - message_ = oss.str(); - } - - virtual const char* what() const noexcept override { return message_.c_str(); } +#define FLASHINFER_WARN(...) \ + do \ + { \ + std::ostringstream oss; \ + write_to_stream(oss, ##__VA_ARGS__); \ + std::string msg = oss.str(); \ + if (msg.empty()) \ + { \ + msg = "Warning triggered"; \ + } \ + flashinfer::Warning(__FUNCTION__, __FILE__, __LINE__, msg).emit(); \ + } while (0) + +namespace flashinfer +{ +class Error : public std::exception +{ +private: + std::string message_; + +public: + Error(std::string const& func, std::string const& file, int line, std::string const& message) + { + std::ostringstream oss; + oss << "Error in function '" << func << "' " + << "at " << file << ":" << line << ": " << message; + message_ = oss.str(); + } + + virtual char const* what() const noexcept override + { + return message_.c_str(); + } }; -class Warning { - private: - std::string message_; - - public: - Warning(const std::string& func, const std::string& file, int line, const std::string& message) { - std::ostringstream oss; - oss << "Warning in function '" << func << "' " - << "at " << file << ":" << line << ": " << message; - message_ = oss.str(); - } - - void emit() const { std::cerr << message_ << std::endl; } +class Warning +{ +private: + std::string message_; + +public: + Warning(std::string const& func, std::string const& file, int line, std::string const& message) + { + std::ostringstream oss; + oss << "Warning in function '" << func << "' " + << "at " << file << ":" << line << ": " << message; + message_ = oss.str(); + } + + void emit() const + { + std::cerr << message_ << std::endl; + } }; -} // namespace flashinfer +} // namespace flashinfer -#endif // FLASHINFER_EXCEPTION_H_ +#endif // FLASHINFER_EXCEPTION_H_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh index d763d05135c6..c77e2afcf8ab 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh @@ -18,127 +18,126 @@ #include -namespace flashinfer::mamba::checkpointing { - -struct CheckpointingSsuParams { - uint32_t batch{}, nheads{}, dim{}, dstate{}, ngroups{}; - uint32_t state_cache_size{}; - uint32_t npredicted{}; - uint32_t max_window{}; - int32_t pad_slot_id{-1}; - - // v12 §59: per-head DIM split factor. Must be one of {1, 2, 4}. The host - // launcher dispatches to a kernel template specialized on this value; the - // kernel cross-checks via assert(params.d_split == D_SPLIT). - int32_t d_split{1}; - - bool dt_softplus{false}; - - // Note: Programmatic Dependent Launch is JIT-stamped via the `ENABLE_PDL` - // constexpr (see checkpointing_ssu_customize_config.jinja). Each .so has - // its PDL mode baked in; no runtime field needed. - - // ── Tensor pointers ── - void* __restrict__ state{nullptr}; // (state_cache_size, nheads, dim, dstate) - void* __restrict__ x{nullptr}; // (batch, NPREDICTED, nheads, dim) - void* __restrict__ dt{nullptr}; // (batch, NPREDICTED, nheads, dim) tie_hdim - void* __restrict__ A{nullptr}; // (nheads, dim, dstate) tie_hdim - void* __restrict__ B{nullptr}; // (batch, NPREDICTED, ngroups, dstate) - void* __restrict__ C{nullptr}; // (batch, NPREDICTED, ngroups, dstate) - void* __restrict__ D{nullptr}; // (nheads, dim), optional - void* __restrict__ z{nullptr}; // (batch, NPREDICTED, nheads, dim), optional - void* __restrict__ dt_bias{nullptr}; // (nheads, dim) tie_hdim, optional - void* __restrict__ output{nullptr}; // (batch, NPREDICTED, nheads, dim) - - // ── Cache tensors for incremental replay ── - void* __restrict__ old_x{nullptr}; // (state_cache_size, MAX_WINDOW, nheads, dim) single-buffered - void* __restrict__ old_B{ - nullptr}; // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) double-buffered - void* __restrict__ old_dt{ - nullptr}; // (state_cache_size, 2, nheads, MAX_WINDOW) double-buffered, f32 - void* __restrict__ old_cumAdt{ - nullptr}; // (state_cache_size, 2, nheads, MAX_WINDOW) double-buffered, f32 - void* __restrict__ cache_buf_idx{nullptr}; // (state_cache_size,) int32 - void* __restrict__ prev_num_accepted{nullptr}; // (state_cache_size,) int32 - - // ── Index tensors ── - void* __restrict__ state_batch_indices{nullptr}; // (batch,) optional - - // ── Varlen (v20): packed inputs ── - // When non-null, `x/dt/B/C/z/out` are laid out as - // `(1, total_tokens, nheads, dim)` / `(1, total_tokens, ngroups, dstate)` - // and `cu_seqlens[i]` gives the token-axis base of sequence i. - // `seq_len_i = cu_seqlens[i+1] - cu_seqlens[i]`. Kernel dispatch on - // `cu_seqlens != nullptr` selects a `VARLEN=true` template. - // - // The `*_stride_seq` fields below already encode the outer iteration - // stride for both modes — the wrapper sets them to: - // non-varlen: `tensor.stride(0)` (per-batch) - // varlen : `tensor.stride(1)` (per-token, since sequences are packed - // into a single batch of total_tokens) - // so the kernel uses one formula `seq * *_stride_seq` regardless of mode. - void* __restrict__ cu_seqlens{nullptr}; // (batch+1,) int32, optional - - // ── Block-scale decode factors for quantized state ── - void* __restrict__ state_scale{nullptr}; // float32: (state_cache_size, nheads, dim) - - // ── Philox PRNG seed for stochastic rounding ── - const int64_t* rand_seed{nullptr}; - - // ── Strides ── - // state: (state_cache_size, nheads, dim, dstate) — inner 3 dims contiguous - int64_t state_stride_seq{}; - - // For the six batch-side tensors (x, dt, B, C, out, z), `*_stride_seq` - // is the outer iteration stride — per-batch in non-varlen, per-token in - // varlen. `*_stride_token` is the inner per-row (T-axis) stride, same - // in both modes. - - // x: (batch, NPREDICTED, nheads, dim) [non-varlen] / (1, total_tokens, nheads, dim) [varlen] - int64_t x_stride_seq{}; - int64_t x_stride_token{}; - - // dt: (batch, NPREDICTED, nheads, dim) — tie_hdim (stride_dim=0) - int64_t dt_stride_seq{}; - int64_t dt_stride_token{}; - - // B: (batch, NPREDICTED, ngroups, dstate) - int64_t B_stride_seq{}; - int64_t B_stride_token{}; - - // C: (batch, NPREDICTED, ngroups, dstate) - int64_t C_stride_seq{}; - int64_t C_stride_token{}; - - // output: (batch, NPREDICTED, nheads, dim) - int64_t out_stride_seq{}; - int64_t out_stride_token{}; - - // z: (batch, NPREDICTED, nheads, dim) - int64_t z_stride_seq{}; - int64_t z_stride_token{}; - - // old_x: (state_cache_size, MAX_WINDOW, nheads, dim) — single-buffered - int64_t old_x_stride_seq{}; - int64_t old_x_stride_token{}; - - // old_B: (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) — double-buffered - int64_t old_B_stride_seq{}; - int64_t old_B_stride_dbuf{}; - int64_t old_B_stride_token{}; - - // old_dt: (state_cache_size, 2, nheads, MAX_WINDOW) — double-buffered, MAX_WINDOW contiguous - int64_t old_dt_stride_seq{}; - int64_t old_dt_stride_dbuf{}; - int64_t old_dt_stride_head{}; - - // old_cumAdt: (state_cache_size, 2, nheads, MAX_WINDOW) — double-buffered, MAX_WINDOW contiguous - int64_t old_cumAdt_stride_seq{}; - int64_t old_cumAdt_stride_dbuf{}; - int64_t old_cumAdt_stride_head{}; - - // state_scale: (state_cache_size, nheads, dim) - int64_t state_scale_stride_seq{}; +namespace flashinfer::mamba::checkpointing +{ + +struct CheckpointingSsuParams +{ + uint32_t batch{}, nheads{}, dim{}, dstate{}, ngroups{}; + uint32_t state_cache_size{}; + uint32_t npredicted{}; + uint32_t max_window{}; + int32_t pad_slot_id{-1}; + + // v12 §59: per-head DIM split factor. Must be one of {1, 2, 4}. The host + // launcher dispatches to a kernel template specialized on this value; the + // kernel cross-checks via assert(params.d_split == D_SPLIT). + int32_t d_split{1}; + + bool dt_softplus{false}; + + // Note: Programmatic Dependent Launch is JIT-stamped via the `ENABLE_PDL` + // constexpr (see checkpointing_ssu_customize_config.jinja). Each .so has + // its PDL mode baked in; no runtime field needed. + + // ── Tensor pointers ── + void* __restrict__ state{nullptr}; // (state_cache_size, nheads, dim, dstate) + void* __restrict__ x{nullptr}; // (batch, NPREDICTED, nheads, dim) + void* __restrict__ dt{nullptr}; // (batch, NPREDICTED, nheads, dim) tie_hdim + void* __restrict__ A{nullptr}; // (nheads, dim, dstate) tie_hdim + void* __restrict__ B{nullptr}; // (batch, NPREDICTED, ngroups, dstate) + void* __restrict__ C{nullptr}; // (batch, NPREDICTED, ngroups, dstate) + void* __restrict__ D{nullptr}; // (nheads, dim), optional + void* __restrict__ z{nullptr}; // (batch, NPREDICTED, nheads, dim), optional + void* __restrict__ dt_bias{nullptr}; // (nheads, dim) tie_hdim, optional + void* __restrict__ output{nullptr}; // (batch, NPREDICTED, nheads, dim) + + // ── Cache tensors for incremental replay ── + void* __restrict__ old_x{nullptr}; // (state_cache_size, MAX_WINDOW, nheads, dim) single-buffered + void* __restrict__ old_B{nullptr}; // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) double-buffered + void* __restrict__ old_dt{nullptr}; // (state_cache_size, 2, nheads, MAX_WINDOW) double-buffered, f32 + void* __restrict__ old_cumAdt{nullptr}; // (state_cache_size, 2, nheads, MAX_WINDOW) double-buffered, f32 + void* __restrict__ cache_buf_idx{nullptr}; // (state_cache_size,) int32 + void* __restrict__ prev_num_accepted{nullptr}; // (state_cache_size,) int32 + + // ── Index tensors ── + void* __restrict__ state_batch_indices{nullptr}; // (batch,) optional + + // ── Varlen (v20): packed inputs ── + // When non-null, `x/dt/B/C/z/out` are laid out as + // `(1, total_tokens, nheads, dim)` / `(1, total_tokens, ngroups, dstate)` + // and `cu_seqlens[i]` gives the token-axis base of sequence i. + // `seq_len_i = cu_seqlens[i+1] - cu_seqlens[i]`. Kernel dispatch on + // `cu_seqlens != nullptr` selects a `VARLEN=true` template. + // + // The `*_stride_seq` fields below already encode the outer iteration + // stride for both modes — the wrapper sets them to: + // non-varlen: `tensor.stride(0)` (per-batch) + // varlen : `tensor.stride(1)` (per-token, since sequences are packed + // into a single batch of total_tokens) + // so the kernel uses one formula `seq * *_stride_seq` regardless of mode. + void* __restrict__ cu_seqlens{nullptr}; // (batch+1,) int32, optional + + // ── Block-scale decode factors for quantized state ── + void* __restrict__ state_scale{nullptr}; // float32: (state_cache_size, nheads, dim) + + // ── Philox PRNG seed for stochastic rounding ── + int64_t const* rand_seed{nullptr}; + + // ── Strides ── + // state: (state_cache_size, nheads, dim, dstate) — inner 3 dims contiguous + int64_t state_stride_seq{}; + + // For the six batch-side tensors (x, dt, B, C, out, z), `*_stride_seq` + // is the outer iteration stride — per-batch in non-varlen, per-token in + // varlen. `*_stride_token` is the inner per-row (T-axis) stride, same + // in both modes. + + // x: (batch, NPREDICTED, nheads, dim) [non-varlen] / (1, total_tokens, nheads, dim) [varlen] + int64_t x_stride_seq{}; + int64_t x_stride_token{}; + + // dt: (batch, NPREDICTED, nheads, dim) — tie_hdim (stride_dim=0) + int64_t dt_stride_seq{}; + int64_t dt_stride_token{}; + + // B: (batch, NPREDICTED, ngroups, dstate) + int64_t B_stride_seq{}; + int64_t B_stride_token{}; + + // C: (batch, NPREDICTED, ngroups, dstate) + int64_t C_stride_seq{}; + int64_t C_stride_token{}; + + // output: (batch, NPREDICTED, nheads, dim) + int64_t out_stride_seq{}; + int64_t out_stride_token{}; + + // z: (batch, NPREDICTED, nheads, dim) + int64_t z_stride_seq{}; + int64_t z_stride_token{}; + + // old_x: (state_cache_size, MAX_WINDOW, nheads, dim) — single-buffered + int64_t old_x_stride_seq{}; + int64_t old_x_stride_token{}; + + // old_B: (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) — double-buffered + int64_t old_B_stride_seq{}; + int64_t old_B_stride_dbuf{}; + int64_t old_B_stride_token{}; + + // old_dt: (state_cache_size, 2, nheads, MAX_WINDOW) — double-buffered, MAX_WINDOW contiguous + int64_t old_dt_stride_seq{}; + int64_t old_dt_stride_dbuf{}; + int64_t old_dt_stride_head{}; + + // old_cumAdt: (state_cache_size, 2, nheads, MAX_WINDOW) — double-buffered, MAX_WINDOW contiguous + int64_t old_cumAdt_stride_seq{}; + int64_t old_cumAdt_stride_dbuf{}; + int64_t old_cumAdt_stride_head{}; + + // state_scale: (state_cache_size, nheads, dim) + int64_t state_scale_stride_seq{}; }; // Forward declaration — defined in kernel_checkpointing_ssu.cuh. @@ -146,9 +145,9 @@ struct CheckpointingSsuParams { // `params.d_split` and routes to the matching `launchCheckpointingSsuImpl` // specialization (v12 §59). Caller side stays single-entry. template + typename stateIndex_t, typename state_scale_t> void launchCheckpointingSsu(CheckpointingSsuParams& params, cudaStream_t stream); -} // namespace flashinfer::mamba::checkpointing +} // namespace flashinfer::mamba::checkpointing -#endif // FLASHINFER_MAMBA_CHECKPOINTING_SSU_CUH_ +#endif // FLASHINFER_MAMBA_CHECKPOINTING_SSU_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh index c36883f523cb..476458a24c5a 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh @@ -24,7 +24,8 @@ #include -namespace flashinfer::mamba { +namespace flashinfer::mamba +{ constexpr unsigned warpSize = 32; @@ -33,57 +34,71 @@ constexpr unsigned warpSize = 32; // ============================================================================= // Largest power of 2 that divides v (i.e. v & -v). Returns 1 when v == 0. -inline constexpr unsigned largestPow2Divisor(unsigned v) { return v ? (v & (~v + 1)) : 1; } +inline constexpr unsigned largestPow2Divisor(unsigned v) +{ + return v ? (v & (~v + 1)) : 1; +} // Simple packed vector type for loading N elements of type T. // Alignment is the largest power-of-2 factor of the total byte size, // so it is always valid even when N * sizeof(T) is not a power of 2 (e.g. 3 × 2 = 6). template -struct alignas(largestPow2Divisor(N * sizeof(T))) PackedAligned { - T val[N]; - static constexpr int count = N; - using dtype = T; +struct alignas(largestPow2Divisor(N * sizeof(T))) PackedAligned +{ + T val[N]; + static constexpr int count = N; + using dtype = T; }; template -__device__ __forceinline__ auto make_zeros() -> load_t { - load_t ret{}; +__device__ __forceinline__ auto make_zeros() -> load_t +{ + load_t ret{}; #pragma unroll - for (int i = 0; i < ret.count; i++) - ret.val[i] = typename load_t::dtype{}; // default initialization - return ret; + for (int i = 0; i < ret.count; i++) + ret.val[i] = typename load_t::dtype{}; // default initialization + return ret; }; // Computes the vector load size that ensures full warp utilization. // Avoids cases like: dstate=64, load_t = sizeof(float4)/sizeof(f16), warpsize=32 (32 * 8 > 64) // in which case a part of the warp would be idle. template -inline constexpr auto getVectorLoadSizeForFullUtilization() -> unsigned { - static_assert(sizeof(float4) >= sizeof(T)); - constexpr unsigned maxHardwareLoadSize = sizeof(float4) / sizeof(T); - constexpr unsigned maxLogicalLoadSize = (unsigned)DSTATE / warpSize; - return maxHardwareLoadSize < maxLogicalLoadSize ? maxHardwareLoadSize : maxLogicalLoadSize; +inline constexpr auto getVectorLoadSizeForFullUtilization() -> unsigned +{ + static_assert(sizeof(float4) >= sizeof(T)); + constexpr unsigned maxHardwareLoadSize = sizeof(float4) / sizeof(T); + constexpr unsigned maxLogicalLoadSize = (unsigned) DSTATE / warpSize; + return maxHardwareLoadSize < maxLogicalLoadSize ? maxHardwareLoadSize : maxLogicalLoadSize; } -__device__ __forceinline__ float warpReduceSum(float val) { - for (int s = warpSize / 2; s > 0; s /= 2) { - val += __shfl_down_sync(UINT32_MAX, val, s); - } - return val; +__device__ __forceinline__ float warpReduceSum(float val) +{ + for (int s = warpSize / 2; s > 0; s /= 2) + { + val += __shfl_down_sync(UINT32_MAX, val, s); + } + return val; } -__device__ __forceinline__ float warpReduceMax(float val) { - for (int s = warpSize / 2; s > 0; s /= 2) { - val = max(val, __shfl_down_sync(UINT32_MAX, val, s)); - } - return val; +__device__ __forceinline__ float warpReduceMax(float val) +{ + for (int s = warpSize / 2; s > 0; s /= 2) + { + val = max(val, __shfl_down_sync(UINT32_MAX, val, s)); + } + return val; } -__forceinline__ __device__ float softplus(float x) { return __logf(1.f + __expf(x)); } +__forceinline__ __device__ float softplus(float x) +{ + return __logf(1.f + __expf(x)); +} -__device__ __forceinline__ float thresholded_softplus(float dt_value) { - constexpr float threshold = 20.f; - return (dt_value <= threshold) ? softplus(dt_value) : dt_value; +__device__ __forceinline__ float thresholded_softplus(float dt_value) +{ + constexpr float threshold = 20.f; + return (dt_value <= threshold) ? softplus(dt_value) : dt_value; } // ============================================================================= @@ -92,86 +107,95 @@ __device__ __forceinline__ float thresholded_softplus(float dt_value) { // Format an integer_sequence as a comma-separated string for error messages template -std::string format_sequence(std::integer_sequence) { - std::ostringstream oss; - bool first = true; - ((oss << (first ? (first = false, "") : ", ") << Values), ...); - return oss.str(); +std::string format_sequence(std::integer_sequence) +{ + std::ostringstream oss; + bool first = true; + ((oss << (first ? (first = false, "") : ", ") << Values), ...); + return oss.str(); } // Helper function to dispatch dim and dstate with a kernel launcher template void dispatchDimDstate(ParamsType& params, std::integer_sequence dims_seq, - std::integer_sequence dstates_seq, - KernelLauncher&& launcher) { - auto dispatch_dstate = [&]() { - auto try_dstate = [&]() { - if (params.dstate == DSTATE) { - launcher.template operator()(); - return true; - } - return false; + std::integer_sequence dstates_seq, KernelLauncher&& launcher) +{ + auto dispatch_dstate = [&]() + { + auto try_dstate = [&]() + { + if (params.dstate == DSTATE) + { + launcher.template operator()(); + return true; + } + return false; + }; + bool dispatched = (try_dstate.template operator()() || ...); + FLASHINFER_CHECK(dispatched, "Unsupported dstate value: ", params.dstate, + ".\nSupported values: ", format_sequence(dstates_seq)); }; - bool dispatched = (try_dstate.template operator()() || ...); - FLASHINFER_CHECK(dispatched, "Unsupported dstate value: ", params.dstate, - ".\nSupported values: ", format_sequence(dstates_seq)); - }; - - auto try_dim = [&]() { - if (params.dim == DIM) { - dispatch_dstate.template operator()(); - return true; - } - return false; - }; - bool dim_dispatched = (try_dim.template operator()() || ...); - FLASHINFER_CHECK(dim_dispatched, "Unsupported dim value: ", params.dim, - ".\nSupported values: ", format_sequence(dims_seq)); + auto try_dim = [&]() + { + if (params.dim == DIM) + { + dispatch_dstate.template operator()(); + return true; + } + return false; + }; + + bool dim_dispatched = (try_dim.template operator()() || ...); + FLASHINFER_CHECK( + dim_dispatched, "Unsupported dim value: ", params.dim, ".\nSupported values: ", format_sequence(dims_seq)); } // Helper function to dispatch ratio with a kernel launcher template -void dispatchRatio(ParamsType& params, std::integer_sequence ratios_seq, - KernelLauncher&& launcher) { - auto try_ratio = [&]() { - if (params.nheads / params.ngroups == RATIO) { - launcher.template operator()(); - return true; - } - return false; - }; +void dispatchRatio( + ParamsType& params, std::integer_sequence ratios_seq, KernelLauncher&& launcher) +{ + auto try_ratio = [&]() + { + if (params.nheads / params.ngroups == RATIO) + { + launcher.template operator()(); + return true; + } + return false; + }; - bool ratio_dispatched = (try_ratio.template operator()() || ...); - FLASHINFER_CHECK(ratio_dispatched, - "Unsupported nheads/ngroups ratio: ", params.nheads / params.ngroups, - ".\nSupported values: ", format_sequence(ratios_seq)); + bool ratio_dispatched = (try_ratio.template operator()() || ...); + FLASHINFER_CHECK(ratio_dispatched, "Unsupported nheads/ngroups ratio: ", params.nheads / params.ngroups, + ".\nSupported values: ", format_sequence(ratios_seq)); } // Helper function to dispatch dim, dstate, and ntokens_mtp with a kernel launcher // Reuses dispatchDimDstate by wrapping the launcher to add token dispatch -template -void dispatchDimDstateTokens(ParamsType& params, - std::integer_sequence dims_seq, - std::integer_sequence dstates_seq, - std::integer_sequence tokens_seq, - KernelLauncher&& launcher) { - // Wrap the launcher to add token dispatch as the innermost level - auto dim_dstate_launcher = [&]() { - auto try_tokens = [&]() { - if (params.ntokens_mtp == TOKENS_MTP) { - launcher.template operator()(); - return true; - } - return false; +template +void dispatchDimDstateTokens(ParamsType& params, std::integer_sequence dims_seq, + std::integer_sequence dstates_seq, std::integer_sequence tokens_seq, + KernelLauncher&& launcher) +{ + // Wrap the launcher to add token dispatch as the innermost level + auto dim_dstate_launcher = [&]() + { + auto try_tokens = [&]() + { + if (params.ntokens_mtp == TOKENS_MTP) + { + launcher.template operator()(); + return true; + } + return false; + }; + bool dispatched = (try_tokens.template operator()() || ...); + FLASHINFER_CHECK(dispatched, "Unsupported ntokens_mtp value: ", params.ntokens_mtp, + ".\nSupported values: ", format_sequence(tokens_seq)); }; - bool dispatched = (try_tokens.template operator()() || ...); - FLASHINFER_CHECK(dispatched, "Unsupported ntokens_mtp value: ", params.ntokens_mtp, - ".\nSupported values: ", format_sequence(tokens_seq)); - }; - dispatchDimDstate(params, dims_seq, dstates_seq, dim_dstate_launcher); + dispatchDimDstate(params, dims_seq, dstates_seq, dim_dstate_launcher); } // ============================================================================= @@ -181,28 +205,30 @@ void dispatchDimDstateTokens(ParamsType& params, // Check alignment for common input variables (x, z, B, C) // Works for both STP (SelectiveStateUpdateParams) and MTP (SelectiveStateMTPParams) template -void check_ptr_alignment_input_vars(const ParamsType& params) { - using load_input_t = PackedAligned; - FLASHINFER_CHECK(reinterpret_cast(params.x) % sizeof(load_input_t) == 0, - "x pointer must be aligned to ", sizeof(load_input_t), " bytes"); - FLASHINFER_CHECK((params.x_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, - "x batch stride must be aligned to ", sizeof(load_input_t), " bytes"); - if (params.z) { - FLASHINFER_CHECK(reinterpret_cast(params.z) % sizeof(load_input_t) == 0, - "z pointer must be aligned to ", sizeof(load_input_t), " bytes"); - FLASHINFER_CHECK((params.z_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, - "z batch stride must be aligned to ", sizeof(load_input_t), " bytes"); - } - FLASHINFER_CHECK(reinterpret_cast(params.B) % sizeof(load_input_t) == 0, - "B pointer must be aligned to ", sizeof(load_input_t), " bytes"); - FLASHINFER_CHECK(reinterpret_cast(params.C) % sizeof(load_input_t) == 0, - "C pointer must be aligned to ", sizeof(load_input_t), " bytes"); - FLASHINFER_CHECK((params.B_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, - "B batch stride must be aligned to ", sizeof(load_input_t), " bytes"); - FLASHINFER_CHECK((params.C_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, - "C batch stride must be aligned to ", sizeof(load_input_t), " bytes"); +void check_ptr_alignment_input_vars(ParamsType const& params) +{ + using load_input_t = PackedAligned; + FLASHINFER_CHECK(reinterpret_cast(params.x) % sizeof(load_input_t) == 0, "x pointer must be aligned to ", + sizeof(load_input_t), " bytes"); + FLASHINFER_CHECK((params.x_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, + "x batch stride must be aligned to ", sizeof(load_input_t), " bytes"); + if (params.z) + { + FLASHINFER_CHECK(reinterpret_cast(params.z) % sizeof(load_input_t) == 0, + "z pointer must be aligned to ", sizeof(load_input_t), " bytes"); + FLASHINFER_CHECK((params.z_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, + "z batch stride must be aligned to ", sizeof(load_input_t), " bytes"); + } + FLASHINFER_CHECK(reinterpret_cast(params.B) % sizeof(load_input_t) == 0, "B pointer must be aligned to ", + sizeof(load_input_t), " bytes"); + FLASHINFER_CHECK(reinterpret_cast(params.C) % sizeof(load_input_t) == 0, "C pointer must be aligned to ", + sizeof(load_input_t), " bytes"); + FLASHINFER_CHECK((params.B_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, + "B batch stride must be aligned to ", sizeof(load_input_t), " bytes"); + FLASHINFER_CHECK((params.C_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, + "C batch stride must be aligned to ", sizeof(load_input_t), " bytes"); } -} // namespace flashinfer::mamba +} // namespace flashinfer::mamba -#endif // FLASHINFER_MAMBA_COMMON_CUH_ +#endif // FLASHINFER_MAMBA_COMMON_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh index d8b46058b4a7..bab89e520ab3 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh @@ -6,90 +6,134 @@ #include #endif -namespace flashinfer::mamba::conversion { +namespace flashinfer::mamba::conversion +{ -inline __device__ float toFloat(float f) { return f; } +inline __device__ float toFloat(float f) +{ + return f; +} -inline __device__ float toFloat(__half h) { return __half2float(h); } +inline __device__ float toFloat(__half h) +{ + return __half2float(h); +} #ifdef FLASHINFER_ENABLE_BF16 -inline __device__ float toFloat(__nv_bfloat16 val) { return __bfloat162float(val); } +inline __device__ float toFloat(__nv_bfloat16 val) +{ + return __bfloat162float(val); +} #endif // No accuracy loss: int8_t / int16_t range fits exactly in float32 (24-bit // mantissa represents all integers up to 2^24 = 16M exactly). -inline __device__ float toFloat(int8_t val) { return static_cast(val); } -inline __device__ float toFloat(int16_t val) { return static_cast(val); } +inline __device__ float toFloat(int8_t val) +{ + return static_cast(val); +} + +inline __device__ float toFloat(int16_t val) +{ + return static_cast(val); +} // fp8 e4m3 → fp32. Goes via __half (cuda_fp8 library has the implicit // conversion that compiles to `cvt.rn.f16.e4m3` PTX on sm_89+), then // __half2float for the final step. No direct fp8→fp32 PTX op exists. -inline __device__ float toFloat(__nv_fp8_e4m3 val) { - return __half2float(static_cast<__half>(val)); +inline __device__ float toFloat(__nv_fp8_e4m3 val) +{ + return __half2float(static_cast<__half>(val)); } // Packed 2-element conversion: convert a packed pair to float2. // Uses native packed intrinsics for bf16/fp16 (fewer PRMT/SHF instructions). -inline __device__ float2 toFloat2(float2 packed) { return packed; } +inline __device__ float2 toFloat2(float2 packed) +{ + return packed; +} -inline __device__ float2 toFloat2(__half2 packed) { return __half22float2(packed); } +inline __device__ float2 toFloat2(__half2 packed) +{ + return __half22float2(packed); +} // Pointer-based overloads: read two consecutive elements and convert to float2. // Dispatches to the packed intrinsic for bf16/fp16 via the overloads above. -inline __device__ float2 toFloat2(float const* ptr) { return {ptr[0], ptr[1]}; } +inline __device__ float2 toFloat2(float const* ptr) +{ + return {ptr[0], ptr[1]}; +} -inline __device__ float2 toFloat2(__half const* ptr) { - return toFloat2(*reinterpret_cast<__half2 const*>(ptr)); +inline __device__ float2 toFloat2(__half const* ptr) +{ + return toFloat2(*reinterpret_cast<__half2 const*>(ptr)); } #ifdef FLASHINFER_ENABLE_BF16 // inline __device__ float2 toFloat2(__nv_bfloat162 packed) { return __bfloat1622float2(packed); } -inline __device__ float2 toFloat2(__nv_bfloat162 packed) { - // bf16 is the upper 16 bits of f32 — shift/mask is cheaper than PRMT byte permutation. - // NOTE: this ignores denormals - uint32_t bits = reinterpret_cast(packed); - float2 out; - out.x = __uint_as_float(bits << 16); // low bf16 → upper 16 bits of f32 - out.y = __uint_as_float(bits & 0xFFFF0000u); // high bf16 already in upper 16 bits - return out; +inline __device__ float2 toFloat2(__nv_bfloat162 packed) +{ + // bf16 is the upper 16 bits of f32 — shift/mask is cheaper than PRMT byte permutation. + // NOTE: this ignores denormals + uint32_t bits = reinterpret_cast(packed); + float2 out; + out.x = __uint_as_float(bits << 16); // low bf16 → upper 16 bits of f32 + out.y = __uint_as_float(bits & 0xFFFF0000u); // high bf16 already in upper 16 bits + return out; } -inline __device__ float2 toFloat2(__nv_bfloat16 const* ptr) { - return toFloat2(*reinterpret_cast<__nv_bfloat162 const*>(ptr)); +inline __device__ float2 toFloat2(__nv_bfloat16 const* ptr) +{ + return toFloat2(*reinterpret_cast<__nv_bfloat162 const*>(ptr)); } // Paired f32 → bf16 conversion: pack two f32 values into __nv_bfloat162. // Uses native cvt.rn.bf16x2.f32 — single instruction, round-to-nearest-even. -inline __device__ __nv_bfloat162 fromFloat2(float2 val) { - uint32_t result; - asm("cvt.rn.bf16x2.f32 %0, %1, %2;\n" : "=r"(result) : "f"(val.y), "f"(val.x)); - return reinterpret_cast<__nv_bfloat162 const&>(result); +inline __device__ __nv_bfloat162 fromFloat2(float2 val) +{ + uint32_t result; + asm("cvt.rn.bf16x2.f32 %0, %1, %2;\n" : "=r"(result) : "f"(val.y), "f"(val.x)); + return reinterpret_cast<__nv_bfloat162 const&>(result); } #endif -inline __device__ float2 toFloat2(int8_t const* ptr) { return {toFloat(ptr[0]), toFloat(ptr[1])}; } -inline __device__ float2 toFloat2(int16_t const* ptr) { return {toFloat(ptr[0]), toFloat(ptr[1])}; } +inline __device__ float2 toFloat2(int8_t const* ptr) +{ + return {toFloat(ptr[0]), toFloat(ptr[1])}; +} + +inline __device__ float2 toFloat2(int16_t const* ptr) +{ + return {toFloat(ptr[0]), toFloat(ptr[1])}; +} -inline __device__ void convertAndStore(float* output, float input) { *output = input; } +inline __device__ void convertAndStore(float* output, float input) +{ + *output = input; +} -inline __device__ void convertAndStore(__half* output, float input) { - *output = __float2half(input); +inline __device__ void convertAndStore(__half* output, float input) +{ + *output = __float2half(input); } #ifdef FLASHINFER_ENABLE_BF16 -inline __device__ void convertAndStore(__nv_bfloat16* output, float input) { - *output = __float2bfloat16(input); +inline __device__ void convertAndStore(__nv_bfloat16* output, float input) +{ + *output = __float2bfloat16(input); } #endif -inline __device__ void convertAndStore(int16_t* output, float input) { - // Symmetric clip: [-max, max] (not [-max-1, max]) so that negation is safe. - // Matches Triton reference which clips to [-32767, 32767] before storing. - constexpr float int16_max = static_cast(std::numeric_limits::max()); - input = fminf(fmaxf(input, -int16_max), int16_max); - *output = static_cast(__float2int_rn(input)); +inline __device__ void convertAndStore(int16_t* output, float input) +{ + // Symmetric clip: [-max, max] (not [-max-1, max]) so that negation is safe. + // Matches Triton reference which clips to [-32767, 32767] before storing. + constexpr float int16_max = static_cast(std::numeric_limits::max()); + input = fminf(fmaxf(input, -int16_max), int16_max); + *output = static_cast(__float2int_rn(input)); } // ============================================================================= @@ -104,34 +148,36 @@ inline __device__ void convertAndStore(int16_t* output, float input) { // caches where `cache_slot * stride` exceeds 2^32. // All four outputs (c0..c3) are independent and uniformly distributed. template -__device__ __forceinline__ void philox_randint4x(int64_t seed, int64_t offset, uint32_t& r0, - uint32_t& r1, uint32_t& r2, uint32_t& r3) { - constexpr uint32_t PHILOX_KEY_A = 0x9E3779B9u; - constexpr uint32_t PHILOX_KEY_B = 0xBB67AE85u; - constexpr uint32_t PHILOX_ROUND_A = 0xD2511F53u; - constexpr uint32_t PHILOX_ROUND_B = 0xCD9E8D57u; - - uint32_t k0 = static_cast(static_cast(seed)); - uint32_t k1 = static_cast(static_cast(seed) >> 32); - uint64_t uoffset = static_cast(offset); - uint32_t c0 = static_cast(uoffset); - uint32_t c1 = static_cast(uoffset >> 32); - uint32_t c2 = 0, c3 = 0; +__device__ __forceinline__ void philox_randint4x( + int64_t seed, int64_t offset, uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3) +{ + constexpr uint32_t PHILOX_KEY_A = 0x9E3779B9u; + constexpr uint32_t PHILOX_KEY_B = 0xBB67AE85u; + constexpr uint32_t PHILOX_ROUND_A = 0xD2511F53u; + constexpr uint32_t PHILOX_ROUND_B = 0xCD9E8D57u; + + uint32_t k0 = static_cast(static_cast(seed)); + uint32_t k1 = static_cast(static_cast(seed) >> 32); + uint64_t uoffset = static_cast(offset); + uint32_t c0 = static_cast(uoffset); + uint32_t c1 = static_cast(uoffset >> 32); + uint32_t c2 = 0, c3 = 0; #pragma unroll - for (int i = 0; i < n_rounds; i++) { - uint32_t _c0 = c0, _c2 = c2; - c0 = __umulhi(PHILOX_ROUND_B, _c2) ^ c1 ^ k0; - c2 = __umulhi(PHILOX_ROUND_A, _c0) ^ c3 ^ k1; - c1 = PHILOX_ROUND_B * _c2; - c3 = PHILOX_ROUND_A * _c0; - k0 += PHILOX_KEY_A; - k1 += PHILOX_KEY_B; - } - r0 = c0; - r1 = c1; - r2 = c2; - r3 = c3; + for (int i = 0; i < n_rounds; i++) + { + uint32_t _c0 = c0, _c2 = c2; + c0 = __umulhi(PHILOX_ROUND_B, _c2) ^ c1 ^ k0; + c2 = __umulhi(PHILOX_ROUND_A, _c0) ^ c3 ^ k1; + c1 = PHILOX_ROUND_B * _c2; + c3 = PHILOX_ROUND_A * _c0; + k0 += PHILOX_KEY_A; + k1 += PHILOX_KEY_B; + } + r0 = c0; + r1 = c1; + r2 = c2; + r3 = c3; } // Generates a pseudorandom uint32 from (seed, offset) using the Philox-4x32 algorithm. @@ -141,10 +187,11 @@ __device__ __forceinline__ void philox_randint4x(int64_t seed, int64_t offset, u // NOTE: This discards 3 of the 4 Philox outputs. For better throughput, use // philox_randint4x to get all 4 outputs from a single Philox invocation. template -__device__ __forceinline__ uint32_t philox_randint(int64_t seed, int64_t offset) { - uint32_t r0, r1, r2, r3; - philox_randint4x(seed, offset, r0, r1, r2, r3); - return r0; +__device__ __forceinline__ uint32_t philox_randint(int64_t seed, int64_t offset) +{ + uint32_t r0, r1, r2, r3; + philox_randint4x(seed, offset, r0, r1, r2, r3); + return r0; } // ============================================================================= @@ -154,33 +201,41 @@ __device__ __forceinline__ uint32_t philox_randint(int64_t seed, int64_t offset) // Software stochastic rounding: convert one fp32 value to fp16 using 13 random bits. // Adds random noise at the sub-fp16-mantissa position, then truncates. // rand13: 13-bit random value in bits [12:0]. -__device__ __forceinline__ uint16_t cvt_rs_f16_sw(float x, uint32_t rand13) { - uint32_t bits = __float_as_uint(x); - uint32_t sign = bits & 0x80000000u; - uint32_t abs_bits = bits & 0x7FFFFFFFu; - - // fp32 has 23 mantissa bits, fp16 has 10. The 13 LSBs are the remainder. - // Add 13-bit random noise at bits [12:0]. Carry into bit 13 → round up. - abs_bits += (rand13 & 0x1FFFu); - - // Convert to fp16 by truncation. - uint32_t f32_exp = (abs_bits >> 23) & 0xFFu; - uint32_t f32_mantissa = abs_bits & 0x7FFFFFu; - - uint16_t f16_bits; - if (f32_exp == 0xFF) { - f16_bits = (f32_mantissa != 0) ? 0x7E00u : 0x7C00u; // NaN or Inf - } else if (f32_exp > 142) { // 127 + 15 = 142 → overflow to Inf - f16_bits = 0x7C00u; - } else if (f32_exp < 113) { // 127 - 14 = 113 → underflow to zero - f16_bits = 0; - } else { - uint16_t f16_exp = static_cast(f32_exp - 112); // rebias: 127→15 - uint16_t f16_mantissa = static_cast(f32_mantissa >> 13); - f16_bits = (f16_exp << 10) | f16_mantissa; - } - - return static_cast(sign >> 16) | f16_bits; +__device__ __forceinline__ uint16_t cvt_rs_f16_sw(float x, uint32_t rand13) +{ + uint32_t bits = __float_as_uint(x); + uint32_t sign = bits & 0x80000000u; + uint32_t abs_bits = bits & 0x7FFFFFFFu; + + // fp32 has 23 mantissa bits, fp16 has 10. The 13 LSBs are the remainder. + // Add 13-bit random noise at bits [12:0]. Carry into bit 13 → round up. + abs_bits += (rand13 & 0x1FFFu); + + // Convert to fp16 by truncation. + uint32_t f32_exp = (abs_bits >> 23) & 0xFFu; + uint32_t f32_mantissa = abs_bits & 0x7FFFFFu; + + uint16_t f16_bits; + if (f32_exp == 0xFF) + { + f16_bits = (f32_mantissa != 0) ? 0x7E00u : 0x7C00u; // NaN or Inf + } + else if (f32_exp > 142) + { // 127 + 15 = 142 → overflow to Inf + f16_bits = 0x7C00u; + } + else if (f32_exp < 113) + { // 127 - 14 = 113 → underflow to zero + f16_bits = 0; + } + else + { + uint16_t f16_exp = static_cast(f32_exp - 112); // rebias: 127→15 + uint16_t f16_mantissa = static_cast(f32_mantissa >> 13); + f16_bits = (f16_exp << 10) | f16_mantissa; + } + + return static_cast(sign >> 16) | f16_bits; } // Forward declaration (defined below, after cvt_rs_f16x2_f32). @@ -189,15 +244,16 @@ __device__ __forceinline__ uint32_t cvt_rs_f16x2_f32(float a, float b, uint32_t // Stochastic rounding: convert one fp32 value to fp16 using 13 random bits. // On sm_100a+: uses PTX cvt.rs.f16x2.f32 with a dummy zero second input. // On other archs: software emulation. -__device__ __forceinline__ __half cvt_rs_f16_f32(float x, uint32_t rand13) { +__device__ __forceinline__ __half cvt_rs_f16_f32(float x, uint32_t rand13) +{ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL) - // Pack rand13 into rbits[12:0] (for PTX operand b → low half → our x). - // High half gets zero noise for the dummy input. - uint32_t rbits = rand13 & 0x1FFFu; - uint32_t packed = cvt_rs_f16x2_f32(x, 0.0f, rbits); - return __ushort_as_half(static_cast(packed & 0xFFFFu)); + // Pack rand13 into rbits[12:0] (for PTX operand b → low half → our x). + // High half gets zero noise for the dummy input. + uint32_t rbits = rand13 & 0x1FFFu; + uint32_t packed = cvt_rs_f16x2_f32(x, 0.0f, rbits); + return __ushort_as_half(static_cast(packed & 0xFFFFu)); #else - return __ushort_as_half(cvt_rs_f16_sw(x, rand13)); + return __ushort_as_half(cvt_rs_f16_sw(x, rand13)); #endif } @@ -213,19 +269,20 @@ __device__ __forceinline__ __half cvt_rs_f16_f32(float x, uint32_t rand13) { // // Our asm maps: %1→C++ a→PTX b→d[15:0], %2→C++ b→PTX a→d[31:16] // So: C++ a uses rbits[12:0], C++ b uses rbits[28:16]. -__device__ __forceinline__ uint32_t cvt_rs_f16x2_f32(float a, float b, uint32_t rbits) { +__device__ __forceinline__ uint32_t cvt_rs_f16x2_f32(float a, float b, uint32_t rbits) +{ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL) - uint32_t packed; - asm("cvt.rs.f16x2.f32 %0, %2, %1, %3;" - : "=r"(packed) - : "r"(__float_as_uint(a)), "r"(__float_as_uint(b)), "r"(rbits)); - return packed; + uint32_t packed; + asm("cvt.rs.f16x2.f32 %0, %2, %1, %3;" + : "=r"(packed) + : "r"(__float_as_uint(a)), "r"(__float_as_uint(b)), "r"(rbits)); + return packed; #else - uint32_t rand_a = rbits & 0x1FFFu; // bits [12:0] → C++ a (PTX b → low half) - uint32_t rand_b = (rbits >> 16) & 0x1FFFu; // bits [28:16] → C++ b (PTX a → high half) - uint16_t a_fp16 = __half_as_ushort(cvt_rs_f16_f32(a, rand_a)); - uint16_t b_fp16 = __half_as_ushort(cvt_rs_f16_f32(b, rand_b)); - return static_cast(a_fp16) | (static_cast(b_fp16) << 16); + uint32_t rand_a = rbits & 0x1FFFu; // bits [12:0] → C++ a (PTX b → low half) + uint32_t rand_b = (rbits >> 16) & 0x1FFFu; // bits [28:16] → C++ b (PTX a → high half) + uint16_t a_fp16 = __half_as_ushort(cvt_rs_f16_f32(a, rand_a)); + uint16_t b_fp16 = __half_as_ushort(cvt_rs_f16_f32(b, rand_b)); + return static_cast(a_fp16) | (static_cast(b_fp16) << 16); #endif } @@ -234,10 +291,10 @@ __device__ __forceinline__ uint32_t cvt_rs_f16x2_f32(float a, float b, uint32_t // seed: Philox seed (from params.rand_seed). // offset: unique per-element offset (e.g. d * DSTATE + i) for deterministic randomness. template -inline __device__ void convertSRAndStore(__half* output, float input, int64_t seed, - uint32_t offset) { - uint32_t rand = philox_randint(seed, offset); - *output = cvt_rs_f16_f32(input, rand & 0x1FFFu); +inline __device__ void convertSRAndStore(__half* output, float input, int64_t seed, uint32_t offset) +{ + uint32_t rand = philox_randint(seed, offset); + *output = cvt_rs_f16_f32(input, rand & 0x1FFFu); } // ============================================================================= @@ -255,7 +312,10 @@ inline __device__ void convertSRAndStore(__half* output, float input, int64_t se // 32-bit reverse; we right-shift by 16 to land the 16 LSBs of input in the 16 // LSBs of output. Total: 2 SASS instructions (brev + shf/shr), vs the prior // software 4-step mask/shift/OR chain (~12-16 inst). -__device__ __forceinline__ uint32_t bitrev16(uint32_t b) { return __brev(b) >> 16; } +__device__ __forceinline__ uint32_t bitrev16(uint32_t b) +{ + return __brev(b) >> 16; +} // Software stochastic rounding: convert one fp32 value to e4m3 (FN, satfinite) using 16 random // bits. @@ -276,81 +336,101 @@ __device__ __forceinline__ uint32_t bitrev16(uint32_t b) { return __brev(b) >> 1 // Verified bitwise against HW (cvt.rs.satfinite.e4m3x4.f32 on sm_100a) across 22528 // inputs spanning subnormal, normal, and saturation regions during the SR // reverse-engineering effort — see .plans/e4m3_stochastic_rounding.md. -__device__ __forceinline__ uint8_t cvt_rs_e4m3_sw(float x, uint32_t rand16) { - uint32_t bits = __float_as_uint(x); - uint32_t sign = (bits >> 31) & 1u; - uint32_t abs_bits = bits & 0x7FFFFFFFu; - uint32_t f32_exp = (abs_bits >> 23) & 0xFFu; - uint32_t f32_mant = abs_bits & 0x7FFFFFu; - - // NaN / Inf - if (f32_exp == 0xFFu) { - if (f32_mant != 0) { - return static_cast(0x7Fu | (sign << 7)); // canonical e4m3 NaN - } else { - return static_cast(0x7Eu | (sign << 7)); // Inf → ±max finite +__device__ __forceinline__ uint8_t cvt_rs_e4m3_sw(float x, uint32_t rand16) +{ + uint32_t bits = __float_as_uint(x); + uint32_t sign = (bits >> 31) & 1u; + uint32_t abs_bits = bits & 0x7FFFFFFFu; + uint32_t f32_exp = (abs_bits >> 23) & 0xFFu; + uint32_t f32_mant = abs_bits & 0x7FFFFFu; + + // NaN / Inf + if (f32_exp == 0xFFu) + { + if (f32_mant != 0) + { + return static_cast(0x7Fu | (sign << 7)); // canonical e4m3 NaN + } + else + { + return static_cast(0x7Eu | (sign << 7)); // Inf → ±max finite + } } - } - - // fp32 zero / denormal → e4m3 zero (with sign). - if (f32_exp == 0u) { - return static_cast(sign << 7); - } - - int unbiased = static_cast(f32_exp) - 127; - uint64_t mant24 = 0x800000u | f32_mant; // implicit-1 + mantissa, 24-bit - int shift_truncate = (unbiased >= -6) ? 20 : (14 - unbiased); - int rand_shift = shift_truncate - 16; - uint64_t rand_contrib; - if (rand_shift < 0) { - // Defensive: shift_truncate < 16 shouldn't happen for valid normal/subnormal e4m3. - rand_contrib = static_cast(rand16 & 0xFFFFu) >> (-rand_shift); - } else if (rand_shift < 56) { - rand_contrib = static_cast(rand16 & 0xFFFFu) << rand_shift; - } else { - rand_contrib = 0; - } - uint64_t total = mant24 + rand_contrib; - // Guard the shift: shift_truncate reaches 140 for tiny-normal fp32 - // (unbiased < -49), which is UB for a uint64_t shift. Mathematically the - // result is 0 there (the value is far below e4m3's smallest subnormal), - // so flush to int_part = 0. The downstream subnormal branch rounds it to ±0. - uint32_t int_part = (shift_truncate >= 64) ? 0u : static_cast(total >> shift_truncate); - - if (unbiased >= -6) { - // Started in normal binade. int_part ∈ [8, 15] normally; can overflow to 16+ if rand - // bumped the exponent. - int e4m3_exp = unbiased + 7; - while (int_part >= 16u) { - int_part >>= 1; - e4m3_exp += 1; + + // fp32 zero / denormal → e4m3 zero (with sign). + if (f32_exp == 0u) + { + return static_cast(sign << 7); } - if (e4m3_exp > 15 || (e4m3_exp == 15 && (int_part & 0x7u) == 7u)) { - return static_cast(0x7Eu | (sign << 7)); + + int unbiased = static_cast(f32_exp) - 127; + uint64_t mant24 = 0x800000u | f32_mant; // implicit-1 + mantissa, 24-bit + int shift_truncate = (unbiased >= -6) ? 20 : (14 - unbiased); + int rand_shift = shift_truncate - 16; + uint64_t rand_contrib; + if (rand_shift < 0) + { + // Defensive: shift_truncate < 16 shouldn't happen for valid normal/subnormal e4m3. + rand_contrib = static_cast(rand16 & 0xFFFFu) >> (-rand_shift); } - return static_cast((sign << 7) | (e4m3_exp << 3) | (int_part & 0x7u)); - } else { - // Started in subnormal/underflow. int_part: - // 0 → zero - // 1..7 → subnormal e4m3 - // 8..15 → smallest normal binade (e4m3_exp = 1) - // 16+ → higher normal binades (rare; only if rand pushed up multiple binades) - if (int_part == 0u) { - return static_cast(sign << 7); + else if (rand_shift < 56) + { + rand_contrib = static_cast(rand16 & 0xFFFFu) << rand_shift; } - if (int_part <= 7u) { - return static_cast((sign << 7) | int_part); + else + { + rand_contrib = 0; } - int e4m3_exp = 1; - while (int_part >= 16u) { - int_part >>= 1; - e4m3_exp += 1; + uint64_t total = mant24 + rand_contrib; + // Guard the shift: shift_truncate reaches 140 for tiny-normal fp32 + // (unbiased < -49), which is UB for a uint64_t shift. Mathematically the + // result is 0 there (the value is far below e4m3's smallest subnormal), + // so flush to int_part = 0. The downstream subnormal branch rounds it to ±0. + uint32_t int_part = (shift_truncate >= 64) ? 0u : static_cast(total >> shift_truncate); + + if (unbiased >= -6) + { + // Started in normal binade. int_part ∈ [8, 15] normally; can overflow to 16+ if rand + // bumped the exponent. + int e4m3_exp = unbiased + 7; + while (int_part >= 16u) + { + int_part >>= 1; + e4m3_exp += 1; + } + if (e4m3_exp > 15 || (e4m3_exp == 15 && (int_part & 0x7u) == 7u)) + { + return static_cast(0x7Eu | (sign << 7)); + } + return static_cast((sign << 7) | (e4m3_exp << 3) | (int_part & 0x7u)); } - if (e4m3_exp > 15 || (e4m3_exp == 15 && (int_part & 0x7u) == 7u)) { - return static_cast(0x7Eu | (sign << 7)); + else + { + // Started in subnormal/underflow. int_part: + // 0 → zero + // 1..7 → subnormal e4m3 + // 8..15 → smallest normal binade (e4m3_exp = 1) + // 16+ → higher normal binades (rare; only if rand pushed up multiple binades) + if (int_part == 0u) + { + return static_cast(sign << 7); + } + if (int_part <= 7u) + { + return static_cast((sign << 7) | int_part); + } + int e4m3_exp = 1; + while (int_part >= 16u) + { + int_part >>= 1; + e4m3_exp += 1; + } + if (e4m3_exp > 15 || (e4m3_exp == 15 && (int_part & 0x7u) == 7u)) + { + return static_cast(0x7Eu | (sign << 7)); + } + return static_cast((sign << 7) | (e4m3_exp << 3) | (int_part & 0x7u)); } - return static_cast((sign << 7) | (e4m3_exp << 3) | (int_part & 0x7u)); - } } // Stochastic rounding: convert four fp32 values to packed fp8x4 e4m3 using random bits. @@ -379,24 +459,24 @@ __device__ __forceinline__ uint8_t cvt_rs_e4m3_sw(float x, uint32_t rand16) { // e4m3(a_i) into byte i of d. We want byte 0 = e4m3(a), so the source-vector // ordering is {d, c, b, a} = {%4, %3, %2, %1}. See PTX ISA: // https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt -__device__ __forceinline__ uint32_t cvt_rs_e4m3x4_f32(float a, float b, float c, float d, - uint32_t rbits) { +__device__ __forceinline__ uint32_t cvt_rs_e4m3x4_f32(float a, float b, float c, float d, uint32_t rbits) +{ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL) - uint32_t packed; - asm("cvt.rs.satfinite.e4m3x4.f32 %0, {%4, %3, %2, %1}, %5;" - : "=r"(packed) - : "r"(__float_as_uint(a)), "r"(__float_as_uint(b)), "r"(__float_as_uint(c)), - "r"(__float_as_uint(d)), "r"(rbits)); - return packed; + uint32_t packed; + asm("cvt.rs.satfinite.e4m3x4.f32 %0, {%4, %3, %2, %1}, %5;" + : "=r"(packed) + : "r"(__float_as_uint(a)), "r"(__float_as_uint(b)), "r"(__float_as_uint(c)), "r"(__float_as_uint(d)), + "r"(rbits)); + return packed; #else - uint32_t low_chunk = rbits & 0xFFFFu; - uint32_t high_chunk = (rbits >> 16) & 0xFFFFu; - uint8_t pa = cvt_rs_e4m3_sw(a, low_chunk); // PTX f → byte 0 - uint8_t pb = cvt_rs_e4m3_sw(b, bitrev16(low_chunk)); // PTX e → byte 1 - uint8_t pc = cvt_rs_e4m3_sw(c, high_chunk); // PTX b → byte 2 - uint8_t pd = cvt_rs_e4m3_sw(d, bitrev16(high_chunk)); // PTX a → byte 3 - return static_cast(pa) | (static_cast(pb) << 8) | - (static_cast(pc) << 16) | (static_cast(pd) << 24); + uint32_t low_chunk = rbits & 0xFFFFu; + uint32_t high_chunk = (rbits >> 16) & 0xFFFFu; + uint8_t pa = cvt_rs_e4m3_sw(a, low_chunk); // PTX f → byte 0 + uint8_t pb = cvt_rs_e4m3_sw(b, bitrev16(low_chunk)); // PTX e → byte 1 + uint8_t pc = cvt_rs_e4m3_sw(c, high_chunk); // PTX b → byte 2 + uint8_t pd = cvt_rs_e4m3_sw(d, bitrev16(high_chunk)); // PTX a → byte 3 + return static_cast(pa) | (static_cast(pb) << 8) | (static_cast(pc) << 16) + | (static_cast(pd) << 24); #endif } @@ -408,13 +488,14 @@ __device__ __forceinline__ uint32_t cvt_rs_e4m3x4_f32(float a, float b, float c, // the F2I.S32 + VIMNMX(min 127) + VIMNMX(max -127) chain. // Saturates to [-128, 127]. Callers using encode_scale = 127/amax // guarantee |input| ≤ 127.0, so -128 is never produced. -__device__ __forceinline__ int8_t cvt_rni_sat_s8(float x) { +__device__ __forceinline__ int8_t cvt_rni_sat_s8(float x) +{ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - int32_t result; - asm("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(result) : "f"(x)); - return static_cast(result); + int32_t result; + asm("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(result) : "f"(x)); + return static_cast(result); #else - return static_cast(max(-128, min(127, __float2int_rn(x)))); + return static_cast(max(-128, min(127, __float2int_rn(x)))); #endif } @@ -426,13 +507,13 @@ __device__ __forceinline__ int8_t cvt_rni_sat_s8(float x) { // Matches Triton's `floor(scaled_value + rand01)` where // `rand01 = (rand & 0x00FFFFFF) * (1.0 / (1 << 24))`. // Saturates to [-127, 127] (symmetric, matching encode_scale = 127/amax). -__device__ __forceinline__ int8_t cvt_rs_sat_s8(float x, uint32_t rand_bits) { - float const rand01 = - static_cast(rand_bits & 0x00FFFFFFu) * (1.0f / static_cast(1 << 24)); - // `__float2int_rd` (round toward -infinity) fuses `floorf` + `__float2int_rz` - // into a single `cvt.rmi.s32.f32` SASS instruction, saving one FRND per call. - int32_t const clamped = max(-127, min(127, __float2int_rd(x + rand01))); - return static_cast(clamped); +__device__ __forceinline__ int8_t cvt_rs_sat_s8(float x, uint32_t rand_bits) +{ + float const rand01 = static_cast(rand_bits & 0x00FFFFFFu) * (1.0f / static_cast(1 << 24)); + // `__float2int_rd` (round toward -infinity) fuses `floorf` + `__float2int_rz` + // into a single `cvt.rmi.s32.f32` SASS instruction, saving one FRND per call. + int32_t const clamped = max(-127, min(127, __float2int_rd(x + rand01))); + return static_cast(clamped); } // Stochastic rounding: convert four fp32 values to packed s8x4 using a single @@ -449,26 +530,26 @@ __device__ __forceinline__ int8_t cvt_rs_sat_s8(float x, uint32_t rand_bits) { // Amortization: 1 random u32 → 4 SR int8s. A single Philox call (4 u32s) // covers 16 int8 conversions, a 4× reduction in PRNG cost vs the scalar // cvt_rs_sat_s8 path. -__device__ __forceinline__ uint32_t cvt_rs_sat_s8x4_f32(float a, float b, float c, float d, - uint32_t rbits) { - uint32_t const low_chunk = rbits & 0xFFFFu; - uint32_t const high_chunk = (rbits >> 16) & 0xFFFFu; - constexpr float kInv16 = 1.0f / static_cast(1u << 16); - - float const r_a = static_cast(low_chunk) * kInv16; - float const r_b = static_cast(bitrev16(low_chunk)) * kInv16; - float const r_c = static_cast(high_chunk) * kInv16; - float const r_d = static_cast(bitrev16(high_chunk)) * kInv16; - - // `__float2int_rd` (round toward -infinity) emits a single `cvt.rmi.s32.f32` - // SASS op, fusing the `floorf` + `__float2int_rz` chain into one instruction. - int32_t const pa = max(-127, min(127, __float2int_rd(a + r_a))); - int32_t const pb = max(-127, min(127, __float2int_rd(b + r_b))); - int32_t const pc = max(-127, min(127, __float2int_rd(c + r_c))); - int32_t const pd = max(-127, min(127, __float2int_rd(d + r_d))); - - return (static_cast(pa) & 0xFFu) | ((static_cast(pb) & 0xFFu) << 8) | - ((static_cast(pc) & 0xFFu) << 16) | ((static_cast(pd) & 0xFFu) << 24); +__device__ __forceinline__ uint32_t cvt_rs_sat_s8x4_f32(float a, float b, float c, float d, uint32_t rbits) +{ + uint32_t const low_chunk = rbits & 0xFFFFu; + uint32_t const high_chunk = (rbits >> 16) & 0xFFFFu; + constexpr float kInv16 = 1.0f / static_cast(1u << 16); + + float const r_a = static_cast(low_chunk) * kInv16; + float const r_b = static_cast(bitrev16(low_chunk)) * kInv16; + float const r_c = static_cast(high_chunk) * kInv16; + float const r_d = static_cast(bitrev16(high_chunk)) * kInv16; + + // `__float2int_rd` (round toward -infinity) emits a single `cvt.rmi.s32.f32` + // SASS op, fusing the `floorf` + `__float2int_rz` chain into one instruction. + int32_t const pa = max(-127, min(127, __float2int_rd(a + r_a))); + int32_t const pb = max(-127, min(127, __float2int_rd(b + r_b))); + int32_t const pc = max(-127, min(127, __float2int_rd(c + r_c))); + int32_t const pd = max(-127, min(127, __float2int_rd(d + r_d))); + + return (static_cast(pa) & 0xFFu) | ((static_cast(pb) & 0xFFu) << 8) + | ((static_cast(pc) & 0xFFu) << 16) | ((static_cast(pd) & 0xFFu) << 24); } -} // namespace flashinfer::mamba::conversion +} // namespace flashinfer::mamba::conversion diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh index c12ae8b5841b..1cf1211bcd8f 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh @@ -52,7 +52,8 @@ #include "kernel_checkpointing_ssu_common.cuh" -namespace flashinfer::mamba::checkpointing { +namespace flashinfer::mamba::checkpointing +{ // ============================================================================= // Shared memory layout @@ -73,112 +74,111 @@ namespace flashinfer::mamba::checkpointing { // cols up to the swizzle atom width. The cp.async only fills the first // D_PER_CTA cols; the padded tail is unused but keeps the swizzle layout // well-formed. Cost: 1 KB per [NPREDICTED_PAD_MMA_M, 64] buffer at D_PER_CTA=32. -template -struct CheckpointingSsuStorage { - // Re-export the two T/W axis sizes so helpers that only see SmemT can - // recover them. - static constexpr int NPREDICTED = NPREDICTED_; - static constexpr int MAX_WINDOW = MAX_WINDOW_; - // Swizzle atom width for input_t (= 64 cols for 2-byte types). - static constexpr int D_SMEM_COLS = next_multiple_of::ATOM_COLS>(D_PER_CTA); - // M-dim of the output MMAs (C, x, z, CB_scaled): always m16-tiled (keyed - // off NPREDICTED, the new-tokens count). - static constexpr int NPREDICTED_PAD_MMA_M = next_multiple_of(NPREDICTED); - // N-dim of the precompute-CB MMA (matmul-1: C @ B^T). B's row count is - // the matmul N-axis → padded to MMA::N=8. When NPREDICTED ≤ 8 only warp - // 0 has valid B rows; warp 1 zero-fills its CB slice. - static constexpr int NPREDICTED_PAD_MMA_N = next_multiple_of(NPREDICTED); - // K-dim of the replay MMA (matmul-2: old_x^T @ dB_scaled). Padded to - // the small atom's K (== the LDSM unit for 2-byte elements). When - // MAX_WINDOW ≤ MMA::K_SMALL=8, replay picks the small atom (1 K-tile, - // smaller smem, +1 CTA/SM occupancy); otherwise the big atom. Assumes - // MAX_WINDOW ≤ MMA::K_BIG (asserted in the wrapper). - static constexpr int MAX_WINDOW_PAD_MMA_K = next_multiple_of(MAX_WINDOW); - // Row count for buffers padded only to the input-type swizzle atom's row - // extent (8 for 2-byte, 4 for 4-byte) — used by C and z, which alias the - // second m-tile back onto the first via `make_aliased_swizzled_layout_rc`. - // Keyed off NPREDICTED. - static constexpr int NPREDICTED_SWIZZLE_R = - next_multiple_of::ATOM_ROWS>(NPREDICTED); - - // All 2D smem buffers below are stored as flat 1D arrays — the actual - // physical layout is determined by `make_swizzled_layout_rc<...>` at each - // access site, which scrambles (row, col) → physical offset via the - // Swizzle XOR. Declaring them as `T[ROWS][COLS]` would falsely suggest a - // row-major C-array layout that nobody ever uses; the only thing that - // matters here is total byte count and 16-byte alignment. - - // CB_scaled — logical (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE) Swizzle<3,3,3>. - // CB_ROW_STRIDE pads each row to one bank cycle (128 B = 32 banks × 4 B) - // worth of `input_t` so LDSM reads in matmul-4's A operand are - // conflict-free. Equals the swizzle atom's col extent for `input_t` - // (64 for 2-byte, 32 for 4-byte). Logical CB matrix is - // (NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M); trailing cols are padding. - static constexpr int CB_ROW_STRIDE = SmemSwizzle::ATOM_COLS; - alignas(16) input_t CB_scaled[NPREDICTED_PAD_MMA_M * CB_ROW_STRIDE]; - - // B — logical (NPREDICTED_PAD_MMA_N, DSTATE). Row count is matmul-1's - // N-axis (since matmul-1 = C @ B^T). Padding rows inside [NPREDICTED, - // NPREDICTED_PAD_MMA_N) contain garbage — valid output uses only - // [0, NPREDICTED). Warp-1 of compute_CB_scaled_2warp reads rows ≥ 8 of a - // 16-row view; those reads spill into C/old_B smem but are masked to 0 by - // the (j < NPREDICTED) CB-store predicate since j ≥ 8 ≥ NPREDICTED when - // NPREDICTED_PAD_MMA_N == 8. - alignas(16) input_t B[NPREDICTED_PAD_MMA_N * DSTATE]; - - // C — physical (next_multiple_of(NPREDICTED), DSTATE). Padded - // only to the swizzle atom's row extent (8 for 2-byte, 4 for 4-byte), not - // to MMA_prop::M=16. cp.async writes to this exact extent (CShape's first - // dim shrunk to match — see load_data). The MMA still views it as - // NPREDICTED_PAD_MMA_M=16 rows via `make_aliased_swizzled_layout_rc`, - // which aliases the second m-tile back onto the first via stride-0 - // row-tile mode. Garbage feeds output rows ≥ NPREDICTED — predicated - // out at gmem store. Saves up to 2 KB of smem at NPREDICTED ≤ ATOM_ROWS, - // no-op when NPREDICTED > ATOM_ROWS. - alignas(16) input_t C[NPREDICTED_SWIZZLE_R * DSTATE]; - - // x — logical (NPREDICTED_PAD_MMA_M, D_SMEM_COLS). Cols padded to - // D_SMEM_COLS for swizzle atom alignment; cp.async only fills cols - // [0, D_PER_CTA), the tail is unused. - alignas(16) input_t x[NPREDICTED_PAD_MMA_M * D_SMEM_COLS]; - - // z — physical (next_multiple_of(NPREDICTED), D_SMEM_COLS). - // Padded only to the swizzle atom's row extent (8 for 2-byte, 4 for - // 4-byte), not to MMA_prop::M=16. z is never an MMA operand — the - // z-gating epilogue reads it via `partition_C` of the m16n8 c-frag, so - // the MMA still views it as NPREDICTED_PAD_MMA_M=16 rows via - // `make_aliased_swizzled_layout_rc`, which aliases the second m-tile back - // onto the first via stride-0 row-tile mode. Garbage feeds output rows - // ≥ NPREDICTED — predicated out at gmem store. Saves up to 1 KB of smem - // at NPREDICTED ≤ ATOM_ROWS, no-op when NPREDICTED > ATOM_ROWS. - alignas(16) input_t z[NPREDICTED_SWIZZLE_R * D_SMEM_COLS]; - - // Old cache data loaded in Phase 0 (consumed in Phase 1 replay). - // old_x — logical (MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS); ldmatrix.trans - // feeds replay MMA A-operand (only the first D_PER_CTA cols are valid - // data). - alignas(16) input_t old_x[MAX_WINDOW_PAD_MMA_K * D_SMEM_COLS]; - - // old_B — logical (MAX_WINDOW_PAD_MMA_K, DSTATE) Swizzle<3,3,3>. Replay - // MMA reads via ldmatrix.trans (LDSM_T) + register scaling. Padding - // rows zero-filled via cp.async ZFILL. - alignas(16) input_t old_B[MAX_WINDOW_PAD_MMA_K * DSTATE]; - - float old_dt[MAX_WINDOW]; - float old_cumAdt[MAX_WINDOW]; - - // Processed dt for new tokens (Phase 1a uses this for CB_scaled + cumAdt) - float dt_proc[NPREDICTED]; - - // Cumulative A*dt — computed once by warp 0, read by all warps after sync - float cumAdt[NPREDICTED]; - - // state — logical (D_PER_CTA, DSTATE) in `state_t` (native dtype). The - // MMA path reinterprets 2-byte state as bf16 for LDSM; f32 state is loaded - // via UniversalCopy and converted to bf16 in registers inside - // add_init_out. - alignas(16) state_t state[D_PER_CTA * DSTATE]; +template +struct CheckpointingSsuStorage +{ + // Re-export the two T/W axis sizes so helpers that only see SmemT can + // recover them. + static constexpr int NPREDICTED = NPREDICTED_; + static constexpr int MAX_WINDOW = MAX_WINDOW_; + // Swizzle atom width for input_t (= 64 cols for 2-byte types). + static constexpr int D_SMEM_COLS = next_multiple_of::ATOM_COLS>(D_PER_CTA); + // M-dim of the output MMAs (C, x, z, CB_scaled): always m16-tiled (keyed + // off NPREDICTED, the new-tokens count). + static constexpr int NPREDICTED_PAD_MMA_M = next_multiple_of(NPREDICTED); + // N-dim of the precompute-CB MMA (matmul-1: C @ B^T). B's row count is + // the matmul N-axis → padded to MMA::N=8. When NPREDICTED ≤ 8 only warp + // 0 has valid B rows; warp 1 zero-fills its CB slice. + static constexpr int NPREDICTED_PAD_MMA_N = next_multiple_of(NPREDICTED); + // K-dim of the replay MMA (matmul-2: old_x^T @ dB_scaled). Padded to + // the small atom's K (== the LDSM unit for 2-byte elements). When + // MAX_WINDOW ≤ MMA::K_SMALL=8, replay picks the small atom (1 K-tile, + // smaller smem, +1 CTA/SM occupancy); otherwise the big atom. Assumes + // MAX_WINDOW ≤ MMA::K_BIG (asserted in the wrapper). + static constexpr int MAX_WINDOW_PAD_MMA_K = next_multiple_of(MAX_WINDOW); + // Row count for buffers padded only to the input-type swizzle atom's row + // extent (8 for 2-byte, 4 for 4-byte) — used by C and z, which alias the + // second m-tile back onto the first via `make_aliased_swizzled_layout_rc`. + // Keyed off NPREDICTED. + static constexpr int NPREDICTED_SWIZZLE_R = next_multiple_of::ATOM_ROWS>(NPREDICTED); + + // All 2D smem buffers below are stored as flat 1D arrays — the actual + // physical layout is determined by `make_swizzled_layout_rc<...>` at each + // access site, which scrambles (row, col) → physical offset via the + // Swizzle XOR. Declaring them as `T[ROWS][COLS]` would falsely suggest a + // row-major C-array layout that nobody ever uses; the only thing that + // matters here is total byte count and 16-byte alignment. + + // CB_scaled — logical (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE) Swizzle<3,3,3>. + // CB_ROW_STRIDE pads each row to one bank cycle (128 B = 32 banks × 4 B) + // worth of `input_t` so LDSM reads in matmul-4's A operand are + // conflict-free. Equals the swizzle atom's col extent for `input_t` + // (64 for 2-byte, 32 for 4-byte). Logical CB matrix is + // (NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M); trailing cols are padding. + static constexpr int CB_ROW_STRIDE = SmemSwizzle::ATOM_COLS; + alignas(16) input_t CB_scaled[NPREDICTED_PAD_MMA_M * CB_ROW_STRIDE]; + + // B — logical (NPREDICTED_PAD_MMA_N, DSTATE). Row count is matmul-1's + // N-axis (since matmul-1 = C @ B^T). Padding rows inside [NPREDICTED, + // NPREDICTED_PAD_MMA_N) contain garbage — valid output uses only + // [0, NPREDICTED). Warp-1 of compute_CB_scaled_2warp reads rows ≥ 8 of a + // 16-row view; those reads spill into C/old_B smem but are masked to 0 by + // the (j < NPREDICTED) CB-store predicate since j ≥ 8 ≥ NPREDICTED when + // NPREDICTED_PAD_MMA_N == 8. + alignas(16) input_t B[NPREDICTED_PAD_MMA_N * DSTATE]; + + // C — physical (next_multiple_of(NPREDICTED), DSTATE). Padded + // only to the swizzle atom's row extent (8 for 2-byte, 4 for 4-byte), not + // to MMA_prop::M=16. cp.async writes to this exact extent (CShape's first + // dim shrunk to match — see load_data). The MMA still views it as + // NPREDICTED_PAD_MMA_M=16 rows via `make_aliased_swizzled_layout_rc`, + // which aliases the second m-tile back onto the first via stride-0 + // row-tile mode. Garbage feeds output rows ≥ NPREDICTED — predicated + // out at gmem store. Saves up to 2 KB of smem at NPREDICTED ≤ ATOM_ROWS, + // no-op when NPREDICTED > ATOM_ROWS. + alignas(16) input_t C[NPREDICTED_SWIZZLE_R * DSTATE]; + + // x — logical (NPREDICTED_PAD_MMA_M, D_SMEM_COLS). Cols padded to + // D_SMEM_COLS for swizzle atom alignment; cp.async only fills cols + // [0, D_PER_CTA), the tail is unused. + alignas(16) input_t x[NPREDICTED_PAD_MMA_M * D_SMEM_COLS]; + + // z — physical (next_multiple_of(NPREDICTED), D_SMEM_COLS). + // Padded only to the swizzle atom's row extent (8 for 2-byte, 4 for + // 4-byte), not to MMA_prop::M=16. z is never an MMA operand — the + // z-gating epilogue reads it via `partition_C` of the m16n8 c-frag, so + // the MMA still views it as NPREDICTED_PAD_MMA_M=16 rows via + // `make_aliased_swizzled_layout_rc`, which aliases the second m-tile back + // onto the first via stride-0 row-tile mode. Garbage feeds output rows + // ≥ NPREDICTED — predicated out at gmem store. Saves up to 1 KB of smem + // at NPREDICTED ≤ ATOM_ROWS, no-op when NPREDICTED > ATOM_ROWS. + alignas(16) input_t z[NPREDICTED_SWIZZLE_R * D_SMEM_COLS]; + + // Old cache data loaded in Phase 0 (consumed in Phase 1 replay). + // old_x — logical (MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS); ldmatrix.trans + // feeds replay MMA A-operand (only the first D_PER_CTA cols are valid + // data). + alignas(16) input_t old_x[MAX_WINDOW_PAD_MMA_K * D_SMEM_COLS]; + + // old_B — logical (MAX_WINDOW_PAD_MMA_K, DSTATE) Swizzle<3,3,3>. Replay + // MMA reads via ldmatrix.trans (LDSM_T) + register scaling. Padding + // rows zero-filled via cp.async ZFILL. + alignas(16) input_t old_B[MAX_WINDOW_PAD_MMA_K * DSTATE]; + + float old_dt[MAX_WINDOW]; + float old_cumAdt[MAX_WINDOW]; + + // Processed dt for new tokens (Phase 1a uses this for CB_scaled + cumAdt) + float dt_proc[NPREDICTED]; + + // Cumulative A*dt — computed once by warp 0, read by all warps after sync + float cumAdt[NPREDICTED]; + + // state — logical (D_PER_CTA, DSTATE) in `state_t` (native dtype). The + // MMA path reinterprets 2-byte state as bf16 for LDSM; f32 state is loaded + // via UniversalCopy and converted to bf16 in registers inside + // add_init_out. + alignas(16) state_t state[D_PER_CTA * DSTATE]; }; // ============================================================================= @@ -191,15 +191,16 @@ struct CheckpointingSsuStorage { // cvt_rs gets its own dedicated 32-bit randint. // ============================================================================= template -__device__ __forceinline__ uint32_t -stochastic_round_pair_with_philox_refresh(float a, float b, int pair_idx, int64_t rand_seed, - int64_t philox_off, uint32_t (&rand_idx)[4]) { - int const rand_pos = pair_idx & 3; - if (rand_pos == 0) { - conversion::philox_randint4x(rand_seed, philox_off, rand_idx[0], rand_idx[1], - rand_idx[2], rand_idx[3]); - } - return conversion::cvt_rs_f16x2_f32(a, b, rand_idx[rand_pos]); +__device__ __forceinline__ uint32_t stochastic_round_pair_with_philox_refresh( + float a, float b, int pair_idx, int64_t rand_seed, int64_t philox_off, uint32_t (&rand_idx)[4]) +{ + int const rand_pos = pair_idx & 3; + if (rand_pos == 0) + { + conversion::philox_randint4x( + rand_seed, philox_off, rand_idx[0], rand_idx[1], rand_idx[2], rand_idx[3]); + } + return conversion::cvt_rs_f16x2_f32(a, b, rand_idx[rand_pos]); } // ============================================================================= @@ -215,41 +216,45 @@ stochastic_round_pair_with_philox_refresh(float a, float b, int pair_idx, int64_ // iter covers BOTH passes' data via cross-lane participation. // ============================================================================= template -__device__ __forceinline__ void exchange_ntile_state_store_global( - state_t* __restrict__ state_w_base, int np, int lane, - uint32_t const (&my_packed)[2][PAIRS_PER_PASS], IdPart const& id_part) { - using namespace cute; - static_assert(sizeof(state_t) == 2, - "exchange_ntile_state_store_global requires 2-byte state_t for STG.64 alignment"); - int const n_base_p0 = np * N_PER_PASS; - int const n_base_p1 = (np + 1) * N_PER_PASS; +__device__ __forceinline__ void exchange_ntile_state_store_global(state_t* __restrict__ state_w_base, int np, int lane, + uint32_t const (&my_packed)[2][PAIRS_PER_PASS], IdPart const& id_part) +{ + using namespace cute; + static_assert( + sizeof(state_t) == 2, "exchange_ntile_state_store_global requires 2-byte state_t for STG.64 alignment"); + int const n_base_p0 = np * N_PER_PASS; + int const n_base_p1 = (np + 1) * N_PER_PASS; #pragma unroll - for (int p = 0; p < PAIRS_PER_PASS; ++p) { - int const i = p * 2; - // xor mask = 1 swaps neighbor lanes: lane 0 <-> lane 1, lane 2 <-> lane 3, ... - uint32_t const peer_p0 = __shfl_xor_sync(constants::MASK_ALL_LANES, my_packed[0][p], 1); - uint32_t const peer_p1 = __shfl_xor_sync(constants::MASK_ALL_LANES, my_packed[1][p], 1); - - int const row = get<0>(id_part(i)); - int const col_p0 = get<1>(id_part(i)) + n_base_p0; - int const col_p1 = get<1>(id_part(i)) + n_base_p1; - - uint64_t combined; - int32_t gmem_off; - if ((lane & 1) == 0) { - // Even lane: store PASS n0 — my (lower col) in low, peer in high. - combined = static_cast(my_packed[0][p]) | - (static_cast(peer_p0) << constants::num_bits_uint32); - gmem_off = row * DSTATE + col_p0; - } else { - // Odd lane: store PASS n1 — peer (lower col) in low, my in high. - // STG addr = gmem[row*DSTATE + (peer's col base)] = col_p1 - 2. - combined = static_cast(peer_p1) | - (static_cast(my_packed[1][p]) << constants::num_bits_uint32); - gmem_off = row * DSTATE + (col_p1 - 2); + for (int p = 0; p < PAIRS_PER_PASS; ++p) + { + int const i = p * 2; + // xor mask = 1 swaps neighbor lanes: lane 0 <-> lane 1, lane 2 <-> lane 3, ... + uint32_t const peer_p0 = __shfl_xor_sync(constants::MASK_ALL_LANES, my_packed[0][p], 1); + uint32_t const peer_p1 = __shfl_xor_sync(constants::MASK_ALL_LANES, my_packed[1][p], 1); + + int const row = get<0>(id_part(i)); + int const col_p0 = get<1>(id_part(i)) + n_base_p0; + int const col_p1 = get<1>(id_part(i)) + n_base_p1; + + uint64_t combined; + int32_t gmem_off; + if ((lane & 1) == 0) + { + // Even lane: store PASS n0 — my (lower col) in low, peer in high. + combined = static_cast(my_packed[0][p]) + | (static_cast(peer_p0) << constants::num_bits_uint32); + gmem_off = row * DSTATE + col_p0; + } + else + { + // Odd lane: store PASS n1 — peer (lower col) in low, my in high. + // STG addr = gmem[row*DSTATE + (peer's col base)] = col_p1 - 2. + combined = static_cast(peer_p1) + | (static_cast(my_packed[1][p]) << constants::num_bits_uint32); + gmem_off = row * DSTATE + (col_p1 - 2); + } + *reinterpret_cast(&state_w_base[gmem_off]) = combined; } - *reinterpret_cast(&state_w_base[gmem_off]) = combined; - } } // ============================================================================= @@ -274,220 +279,218 @@ __device__ __forceinline__ void exchange_ntile_state_store_global( // state_gmem_off) to 2 regs (just the base), and the per-pair STG.32 uses an // i32 element offset inside the chunk. Use this instead of separately // holding params.state-ptr and state_gmem_off. -template -__device__ __forceinline__ void replay_state_mma(SmemT& smem, CheckpointingSsuParams const& params, - int warp, int lane, int prev_k, int d_tile, - int64_t state_ptr_offset, state_t* state_w_base, - int64_t rand_seed, bool must_checkpoint) { - using namespace cute; - static_assert(sizeof(input_t) == 2, "replay_state_mma requires 2-byte input type"); - static_assert(D_PER_CTA % 16 == 0, "D_PER_CTA must be divisible by 16 (m16n8 atom)"); - static_assert(D_PER_CTA >= 16, "D_PER_CTA must be at least 16"); - - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; // 8 or 16 - int const tid = warp * warpSize + lane; - - // Atom K matches the cache-window tile (MAX_WINDOW_PAD_MMA_K). - // K == MMA_prop::K_BIG (16) → m16n8k16 + x4/x2 ldmatrix.trans - // K == MMA_prop::K_SMALL (8) → m16n8k8 + x2/x1 ldmatrix.trans - using MmaAtomType = std::conditional_t; - using LdsmA = std::conditional_t; - using LdsmB = std::conditional_t; - - // 4 warps along N=DSTATE; each warp covers full M (D_PER_CTA/16 m-atoms). - auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(tid); - - // Per-pass output tile is (D_PER_CTA, N_PER_PASS). N_PER_PASS = 4 warps × n8 = 32 cols. - constexpr int N_PER_PASS = 4 * MMA_prop::N; - static_assert(DSTATE % N_PER_PASS == 0, - "DSTATE must be divisible by 4 * MMA_prop::N for _1x4 warp layout"); - constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; - - float total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - float total_decay = (prev_k > 0) ? __expf(total_cumAdt) : 1.f; - - // ── A operand: old_x [MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS] Swizzle<3,3,3>, transposed - // view [M=D_SMEM_COLS, K=MAX_WINDOW_PAD_MMA_K]. D_SMEM_COLS may be padded above - // D_PER_CTA when D_PER_CTA < swizzle atom; local_tile to D_PER_CTA - // restricts the LDSM to the valid sub-tile. Each warp loads the FULL M (4× - // redundant across warps). See header comment for traffic accounting. ── - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - auto layout_A_full = - make_swizzled_layout_rc_transpose(); - Tensor smem_A_full = make_tensor( - make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); - Tensor smem_A = local_tile(smem_A_full, make_shape(Int{}, Int{}), - make_coord(_0{}, _0{})); - - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); - Tensor frag_A = thr_mma.partition_fragment_A(make_tensor( - (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); - Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); - - cute::copy(s2r_A, smem_A_s2r, frag_A_view); - // old_x is input_t == MMA_prop::operand_t (bf16) — no conversion needed. - - // ── B operand: old_B [MAX_WINDOW_PAD_MMA_K, DSTATE] swizzled, transposed view - // [N=DSTATE, K=MAX_WINDOW_PAD_MMA_K]. Per pass loads N_PER_PASS=32 cols across - // 4 warps; partition_S splits — each warp gets its disjoint 8-col slice. ── - auto layout_B = make_swizzled_layout_rc_transpose(); - Tensor smem_B_full = make_tensor( - make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B); - - auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B = s2r_B.get_slice(tid); - - // ── State: per-CTA swizzle layout [D_PER_CTA, DSTATE]. ── - auto layout_state_swz = make_swizzled_layout_rc(); - state_t* state_base = reinterpret_cast(smem.state); - - // ── Per-pass identity for (row, col) coords ── - // partition_C of an identity tensor of the per-pass output shape gives this - // thread's (row, col) at every C-frag position, including warp-N offset. - // Frag size per thread = (M_atoms=D_PER_CTA/16) × (N_atoms_per_warp=1) × 4 elts. - auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_tile); - // Linear order from CuTe's column-major partition_C with m16n8 atom: - // i=0,1: same row (= row_lo of M-atom 0), adjacent cols (col_off, col_off+1) - // i=2,3: same row (= row_hi of M-atom 0), adjacent cols - // i=4,5: same row (= row_lo of M-atom 1) - // ... (V index 0..3 inside each m16n8, then M-atoms in M-major order) - // Pair load at (i, i+1) covers two consecutive bf16 elts → one 32-bit LDS. - - // Precompute dB coefficients once — depend only on K (lane), not on N. - constexpr int LANES_PER_N_COL = warpSize / MMA_prop::N; // = 4 for m16n8k_ - constexpr int DB_COEFFS_PER_LANE = MAX_WINDOW_PAD_MMA_K / LANES_PER_N_COL; - float dB_coeff[DB_COEFFS_PER_LANE]; - precompute_dB_coeff(dB_coeff, smem, total_cumAdt, prev_k, lane); - - using pair_t = Pair; - - // Philox state amortized across 4 consecutive pair conversions: each call - // returns 4 randints, all 4 get consumed before the next refresh (vs. 1-of-4 - // in the Triton-bit-equal layout — see writeback loop below). Compile-time - // pair_idx (n-loop and i-loop both unrolled) keeps `rand_idx[pair_idx & 3]` - // as a known register access — no local-memory spill. - constexpr bool kPhiloxF16 = (PHILOX_ROUNDS > 0) && std::is_same_v; - [[maybe_unused]] uint32_t rand_idx[4]; - // state_w_base is the pre-combined (params.state + state_gmem_off) base - // pointer — see the function header. No separate state_w / state_gmem_off - // alive in this scope. - - // ── Vectorized state writeback (cross-pass STG.64 fusion) ────────── - // smem always gets nearest-even f32→state_t (consumed by matmul 3 — must - // match Triton's f32→bf16 path as closely as possible). Gmem cache, when - // PHILOX_ROUNDS > 0 and state_t == __half, gets PTX cvt.rs.f16x2.f32 - // stochastic rounding direct from registers via cross-pass STG.64; the - // smem→gmem `store_state` is gated off in compute_and_store_output. - // - // Cross-pass STG fusion: do PASS n0 and PASS n1 back-to-back, buffering - // the post-cvt_rs packed u32s of n0 across n1's HMMA + cvt_rs. Then issue - // ONE STG.64 instruction per pair iter, all 32 lanes active: - // - even lane stores PASS n0 data at the warp's n0 column slice - // - odd lane stores PASS n1 data at the warp's n1 column slice - // Halves the STG instruction count vs per-pass writeback (16 STG.64/thread - // per 2 passes vs 16 + 16 = 32 STG.64/thread previously — same byte volume). - // - // Randint amortization: rand_idx[4] refreshed every 4 pairs; each pair's - // cvt_rs uses one of the 4 randints. Triton bit-equality is intentionally - // given up; unbiasedness still holds. - constexpr int PAIRS_PER_PASS = D_PER_CTA / 8; // = (D_PER_CTA/16) × 2 row-pair iters - static_assert(NUM_N_PASSES % 2 == 0, "Cross-pass STG fusion requires even NUM_N_PASSES"); +template +__device__ __forceinline__ void replay_state_mma(SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, + int prev_k, int d_tile, int64_t state_ptr_offset, state_t* state_w_base, int64_t rand_seed, bool must_checkpoint) +{ + using namespace cute; + static_assert(sizeof(input_t) == 2, "replay_state_mma requires 2-byte input type"); + static_assert(D_PER_CTA % 16 == 0, "D_PER_CTA must be divisible by 16 (m16n8 atom)"); + static_assert(D_PER_CTA >= 16, "D_PER_CTA must be at least 16"); + + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; // 8 or 16 + int const tid = warp * warpSize + lane; + + // Atom K matches the cache-window tile (MAX_WINDOW_PAD_MMA_K). + // K == MMA_prop::K_BIG (16) → m16n8k16 + x4/x2 ldmatrix.trans + // K == MMA_prop::K_SMALL (8) → m16n8k8 + x2/x1 ldmatrix.trans + using MmaAtomType + = std::conditional_t; + using LdsmA = std::conditional_t; + using LdsmB = std::conditional_t; + + // 4 warps along N=DSTATE; each warp covers full M (D_PER_CTA/16 m-atoms). + auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(tid); + + // Per-pass output tile is (D_PER_CTA, N_PER_PASS). N_PER_PASS = 4 warps × n8 = 32 cols. + constexpr int N_PER_PASS = 4 * MMA_prop::N; + static_assert(DSTATE % N_PER_PASS == 0, "DSTATE must be divisible by 4 * MMA_prop::N for _1x4 warp layout"); + constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; + + float total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + float total_decay = (prev_k > 0) ? __expf(total_cumAdt) : 1.f; + + // ── A operand: old_x [MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS] Swizzle<3,3,3>, transposed + // view [M=D_SMEM_COLS, K=MAX_WINDOW_PAD_MMA_K]. D_SMEM_COLS may be padded above + // D_PER_CTA when D_PER_CTA < swizzle atom; local_tile to D_PER_CTA + // restricts the LDSM to the valid sub-tile. Each warp loads the FULL M (4× + // redundant across warps). See header comment for traffic accounting. ── + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + auto layout_A_full = make_swizzled_layout_rc_transpose(); + Tensor smem_A_full + = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); + Tensor smem_A + = local_tile(smem_A_full, make_shape(Int{}, Int{}), make_coord(_0{}, _0{})); + + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); + Tensor frag_A = thr_mma.partition_fragment_A( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); + + cute::copy(s2r_A, smem_A_s2r, frag_A_view); + // old_x is input_t == MMA_prop::operand_t (bf16) — no conversion needed. + + // ── B operand: old_B [MAX_WINDOW_PAD_MMA_K, DSTATE] swizzled, transposed view + // [N=DSTATE, K=MAX_WINDOW_PAD_MMA_K]. Per pass loads N_PER_PASS=32 cols across + // 4 warps; partition_S splits — each warp gets its disjoint 8-col slice. ── + auto layout_B = make_swizzled_layout_rc_transpose(); + Tensor smem_B_full = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B); + + auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B = s2r_B.get_slice(tid); + + // ── State: per-CTA swizzle layout [D_PER_CTA, DSTATE]. ── + auto layout_state_swz = make_swizzled_layout_rc(); + state_t* state_base = reinterpret_cast(smem.state); + + // ── Per-pass identity for (row, col) coords ── + // partition_C of an identity tensor of the per-pass output shape gives this + // thread's (row, col) at every C-frag position, including warp-N offset. + // Frag size per thread = (M_atoms=D_PER_CTA/16) × (N_atoms_per_warp=1) × 4 elts. + auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_tile); + // Linear order from CuTe's column-major partition_C with m16n8 atom: + // i=0,1: same row (= row_lo of M-atom 0), adjacent cols (col_off, col_off+1) + // i=2,3: same row (= row_hi of M-atom 0), adjacent cols + // i=4,5: same row (= row_lo of M-atom 1) + // ... (V index 0..3 inside each m16n8, then M-atoms in M-major order) + // Pair load at (i, i+1) covers two consecutive bf16 elts → one 32-bit LDS. + + // Precompute dB coefficients once — depend only on K (lane), not on N. + constexpr int LANES_PER_N_COL = warpSize / MMA_prop::N; // = 4 for m16n8k_ + constexpr int DB_COEFFS_PER_LANE = MAX_WINDOW_PAD_MMA_K / LANES_PER_N_COL; + float dB_coeff[DB_COEFFS_PER_LANE]; + precompute_dB_coeff(dB_coeff, smem, total_cumAdt, prev_k, lane); + + using pair_t = Pair; + + // Philox state amortized across 4 consecutive pair conversions: each call + // returns 4 randints, all 4 get consumed before the next refresh (vs. 1-of-4 + // in the Triton-bit-equal layout — see writeback loop below). Compile-time + // pair_idx (n-loop and i-loop both unrolled) keeps `rand_idx[pair_idx & 3]` + // as a known register access — no local-memory spill. + constexpr bool kPhiloxF16 = (PHILOX_ROUNDS > 0) && std::is_same_v; + [[maybe_unused]] uint32_t rand_idx[4]; + // state_w_base is the pre-combined (params.state + state_gmem_off) base + // pointer — see the function header. No separate state_w / state_gmem_off + // alive in this scope. + + // ── Vectorized state writeback (cross-pass STG.64 fusion) ────────── + // smem always gets nearest-even f32→state_t (consumed by matmul 3 — must + // match Triton's f32→bf16 path as closely as possible). Gmem cache, when + // PHILOX_ROUNDS > 0 and state_t == __half, gets PTX cvt.rs.f16x2.f32 + // stochastic rounding direct from registers via cross-pass STG.64; the + // smem→gmem `store_state` is gated off in compute_and_store_output. + // + // Cross-pass STG fusion: do PASS n0 and PASS n1 back-to-back, buffering + // the post-cvt_rs packed u32s of n0 across n1's HMMA + cvt_rs. Then issue + // ONE STG.64 instruction per pair iter, all 32 lanes active: + // - even lane stores PASS n0 data at the warp's n0 column slice + // - odd lane stores PASS n1 data at the warp's n1 column slice + // Halves the STG instruction count vs per-pass writeback (16 STG.64/thread + // per 2 passes vs 16 + 16 = 32 STG.64/thread previously — same byte volume). + // + // Randint amortization: rand_idx[4] refreshed every 4 pairs; each pair's + // cvt_rs uses one of the 4 randints. Triton bit-equality is intentionally + // given up; unbiasedness still holds. + constexpr int PAIRS_PER_PASS = D_PER_CTA / 8; // = (D_PER_CTA/16) × 2 row-pair iters + static_assert(NUM_N_PASSES % 2 == 0, "Cross-pass STG fusion requires even NUM_N_PASSES"); #pragma unroll - for (int np = 0; np < NUM_N_PASSES; np += 2) { - // Buffer of post-cvt_rs packed u32s for both passes (philox path only). - [[maybe_unused]] uint32_t my_packed[2][PAIRS_PER_PASS]; + for (int np = 0; np < NUM_N_PASSES; np += 2) + { + // Buffer of post-cvt_rs packed u32s for both passes (philox path only). + [[maybe_unused]] uint32_t my_packed[2][PAIRS_PER_PASS]; #pragma unroll - for (int local_n = 0; local_n < 2; ++local_n) { - int const n = np + local_n; - int const n_base = n * N_PER_PASS; + for (int local_n = 0; local_n < 2; ++local_n) + { + int const n = np + local_n; + int const n_base = n * N_PER_PASS; - // ── Allocate per-pass C-frag (4 × M_atoms fp32 elts/thread) ── - Tensor frag_h = thr_mma.partition_fragment_C( - make_tensor((float*)0x0, make_shape(Int{}, Int{}))); + // ── Allocate per-pass C-frag (4 × M_atoms fp32 elts/thread) ── + Tensor frag_h = thr_mma.partition_fragment_C( + make_tensor((float*) 0x0, make_shape(Int{}, Int{}))); - // ── Load state × total_decay into frag_h. ── + // ── Load state × total_decay into frag_h. ── #pragma unroll - for (int i = 0; i < size(frag_h); i += 2) { - int const row = get<0>(id_part(i)); - int const col = get<1>(id_part(i)) + n_base; - int const off = layout_state_swz(row, col); - pair_t const p = *reinterpret_cast(&state_base[off]); - frag_h(i) = toFloat(p[cute::Int<0>{}]) * total_decay; - frag_h(i + 1) = toFloat(p[cute::Int<1>{}]) * total_decay; - } + for (int i = 0; i < size(frag_h); i += 2) + { + int const row = get<0>(id_part(i)); + int const col = get<1>(id_part(i)) + n_base; + int const off = layout_state_swz(row, col); + pair_t const p = *reinterpret_cast(&state_base[off]); + frag_h(i) = toFloat(p[cute::Int<0>{}]) * total_decay; + frag_h(i + 1) = toFloat(p[cute::Int<1>{}]) * total_decay; + } - // ── LDSM.T per-pass B (per warp = 1 atom of 8 cols of N) ── - Tensor smem_B_n = - local_tile(smem_B_full, make_tile(Int{}, Int{}), - make_coord(n, _0{})); - auto smem_B_s2r_n = s2r_thr_B.partition_S(smem_B_n); + // ── LDSM.T per-pass B (per warp = 1 atom of 8 cols of N) ── + Tensor smem_B_n = local_tile( + smem_B_full, make_tile(Int{}, Int{}), make_coord(n, _0{})); + auto smem_B_s2r_n = s2r_thr_B.partition_S(smem_B_n); - Tensor frag_B = thr_mma.partition_fragment_B(make_tensor( - (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); - auto frag_B_view = s2r_thr_B.retile_D(frag_B); + Tensor frag_B = thr_mma.partition_fragment_B( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + auto frag_B_view = s2r_thr_B.retile_D(frag_B); - cute::copy(s2r_B, smem_B_s2r_n, frag_B_view); + cute::copy(s2r_B, smem_B_s2r_n, frag_B_view); - compute_dB_scaling(frag_B, dB_coeff); + compute_dB_scaling(frag_B, dB_coeff); - // ── HMMA: frag_h += frag_A @ frag_B ── - cute::gemm(tiled_mma, frag_h, frag_A, frag_B, frag_h); + // ── HMMA: frag_h += frag_A @ frag_B ── + cute::gemm(tiled_mma, frag_h, frag_A, frag_B, frag_h); - // ── Smem write (always) + cvt_rs into my_packed (philox path) ── + // ── Smem write (always) + cvt_rs into my_packed (philox path) ── #pragma unroll - for (int i = 0; i < size(frag_h); i += 2) { - int const row = get<0>(id_part(i)); - int const col = get<1>(id_part(i)) + n_base; - int const off = layout_state_swz(row, col); - - // Smem write — always nearest-even (output's matmul 3 reads this). - pair_t const q = pack_float2(make_float2(frag_h(i), frag_h(i + 1))); - *reinterpret_cast(&state_base[off]) = q; - - if constexpr (kPhiloxF16) { - static_assert(sizeof(state_t) == 2, "STG.64 cooperative path requires 2-byte state_t"); - int const pair_idx = n * PAIRS_PER_PASS + i / 2; - // Per-lane philox_off is unique per (thread, refresh group) — each - // pair gets its own randint bits. Always computed; only consumed - // by the refresh branch inside the helper. - int64_t const philox_off = - state_ptr_offset + (int64_t)(d_tile * D_PER_CTA + row) * DSTATE + col; - // Buffer the SR'd packed u32 — store happens after BOTH passes. - my_packed[local_n][i / 2] = stochastic_round_pair_with_philox_refresh( - frag_h(i), frag_h(i + 1), pair_idx, rand_seed, philox_off, rand_idx); + for (int i = 0; i < size(frag_h); i += 2) + { + int const row = get<0>(id_part(i)); + int const col = get<1>(id_part(i)) + n_base; + int const off = layout_state_swz(row, col); + + // Smem write — always nearest-even (output's matmul 3 reads this). + pair_t const q = pack_float2(make_float2(frag_h(i), frag_h(i + 1))); + *reinterpret_cast(&state_base[off]) = q; + + if constexpr (kPhiloxF16) + { + static_assert(sizeof(state_t) == 2, "STG.64 cooperative path requires 2-byte state_t"); + int const pair_idx = n * PAIRS_PER_PASS + i / 2; + // Per-lane philox_off is unique per (thread, refresh group) — each + // pair gets its own randint bits. Always computed; only consumed + // by the refresh branch inside the helper. + int64_t const philox_off = state_ptr_offset + (int64_t) (d_tile * D_PER_CTA + row) * DSTATE + col; + // Buffer the SR'd packed u32 — store happens after BOTH passes. + my_packed[local_n][i / 2] = stochastic_round_pair_with_philox_refresh( + frag_h(i), frag_h(i + 1), pair_idx, rand_seed, philox_off, rand_idx); + } + } } - } - } - // ── Cross-pass STG.64: all 32 lanes active. ───────────────────────── - // m16n8 lane layout: lane k → row k/4, cols (k%4)*2..(k%4)*2+1. Lanes - // (2k, 2k+1) hold adjacent col-pairs of the same row. After shfl_xor, - // the even/odd lane each has a 4-col contiguous block (in different - // bit-orders). Even lane STG.64s the n0-pass block at its own col - // base; odd lane STG.64s the n1-pass block at the peer's (lower) col - // — both 8-byte aligned for state_t = f16. - // Runtime-gated on must_checkpoint: non-checkpoint steps skip the gmem - // STGs entirely (state HBM remains the prior checkpoint). The cvt_rs - // SR + philox refresh above still ran — only the STGs are elided — - // because skipping them would require routing must_checkpoint into the - // pair_idx amortization logic, which lives across the n-loop. - if constexpr (kPhiloxF16) { - if (must_checkpoint) { - exchange_ntile_state_store_global( - state_w_base, np, lane, my_packed, id_part); - } + // ── Cross-pass STG.64: all 32 lanes active. ───────────────────────── + // m16n8 lane layout: lane k → row k/4, cols (k%4)*2..(k%4)*2+1. Lanes + // (2k, 2k+1) hold adjacent col-pairs of the same row. After shfl_xor, + // the even/odd lane each has a 4-col contiguous block (in different + // bit-orders). Even lane STG.64s the n0-pass block at its own col + // base; odd lane STG.64s the n1-pass block at the peer's (lower) col + // — both 8-byte aligned for state_t = f16. + // Runtime-gated on must_checkpoint: non-checkpoint steps skip the gmem + // STGs entirely (state HBM remains the prior checkpoint). The cvt_rs + // SR + philox refresh above still ran — only the STGs are elided — + // because skipping them would require routing must_checkpoint into the + // pair_idx amortization logic, which lives across the n-loop. + if constexpr (kPhiloxF16) + { + if (must_checkpoint) + { + exchange_ntile_state_store_global( + state_w_base, np, lane, my_packed, id_part); + } + } } - } } // ── Orchestrator: compute_and_store_output ───────────────────────────── @@ -495,178 +498,178 @@ __device__ __forceinline__ void replay_state_mma(SmemT& smem, CheckpointingSsuPa // All operations on register-resident frag_y — no smem round-trip. // Result converted f32 → input_t in registers and stored directly to gmem // via partition_C of the global output tensor (like CUTLASS sgemm_sm80 epilogue). -template -__device__ __forceinline__ void compute_and_store_output(SmemT& smem, - CheckpointingSsuParams const& params, - int warp, int lane, int d_tile, - int64_t out_seq_base, int head, - int64_t cache_slot, float D_val, - bool must_checkpoint, int seq_len) { - using namespace cute; - static_assert(sizeof(input_t) == 2, "compute_and_store_output requires 2-byte input type"); - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - int const tid = warp * warpSize + lane; - - // ── TiledMMA: 128 threads, covers [16, 32] output per step ── - auto tiled_mma = - make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(tid); - - // ── Swizzled smem views ── - // When D_PER_CTA < swizzle atom (= 64 for bf16), the underlying - // smem buffer is padded to D_SMEM_COLS so the swizzle layout is well-formed. - // Per-pass MMA loops only iterate D_PER_CTA / N_TILE tiles → never touch - // the padded tail. - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - - // x: swizzled [NPREDICTED_PAD_MMA_M, D_SMEM_COLS] - auto layout_x_swz = make_swizzled_layout_rc(); - Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), - layout_x_swz); - auto layout_x_trans_swz = - make_swizzled_layout_rc_transpose(); - Tensor smem_x_trans = make_tensor( - make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); - - // z: aliased swizzled [NPREDICTED_PAD_MMA_M, D_SMEM_COLS] — physical buffer - // is only next_multiple_of(NPREDICTED) rows tall; second m-tile - // aliases first. Ghost rows feed predicated-out output rows. - auto layout_z_swz = - make_aliased_swizzled_layout_rc(); - Tensor smem_z = - make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); - - // ── S2R copies ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B = s2r_B.get_slice(tid); - auto s2r_B_trans = - make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); - - // ── Load CB_scaled A operand from smem (precomputed by warps 0,1 between syncs) ── - // Row stride matches the buffer's padded width (one swizzle atom of `input_t`). - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - auto layout_cb_swz = - make_swizzled_layout_rc(); - Tensor smem_CB = make_tensor( - make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); - auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); - Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); - auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); - cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); - - // Decay broadcast: cumAdt[t] → [NPREDICTED_PAD_MMA_M, N_TILE] with stride-0 on N. - constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); - Tensor decay_bcast = make_tensor( - make_smem_ptr(smem.cumAdt), - make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); - Tensor decay_part = thr_mma.partition_C(decay_bcast); - - // ── Gmem output: partition_C for direct register → gmem store ── - auto* __restrict__ output_ptr = reinterpret_cast(params.output); - // out_base lands on this CTA's D-slice within the head. - int64_t const out_base = out_seq_base + (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; - - // Row predicate for padding. The epilogue store loop iterates i in steps - // of 2 and only consults pred(0) and pred(2) — m16n8k16 C-frag per thread - // has 4 elts at rows {t/4, t/4, t/4+8, t/4+8}, so there are only 2 unique - // row predicates. Compute them once and skip the 4-wide pred tensor. - auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_tile); - bool const pred_row_lo = get<0>(id_part(0)) < seq_len; - bool const pred_row_hi = get<0>(id_part(2)) < seq_len; - - // Number of output N-tiles per pass = D_PER_CTA / N_TILE. - // D_SPLIT=1, D_PER_CTA=64, N_TILE=32 → NUM_N_TILES = 2 (current behavior). - // D_SPLIT=2, D_PER_CTA=32 → NUM_N_TILES = 1 (uses _n1 variant). - constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; - static_assert(NUM_N_TILES == 1 || NUM_N_TILES == 2, - "Output epilogue supports NUM_N_TILES = D_PER_CTA / N_TILE in {1, 2}"); - - // ── Epilogue lambda (defined once; called per N-tile from each branch) ── - auto epilogue = [&](auto& frag_y, int n) { - // Decay: frag_y *= exp(cumAdt[t]) +template +__device__ __forceinline__ void compute_and_store_output(SmemT& smem, CheckpointingSsuParams const& params, int warp, + int lane, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, bool must_checkpoint, + int seq_len) +{ + using namespace cute; + static_assert(sizeof(input_t) == 2, "compute_and_store_output requires 2-byte input type"); + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + int const tid = warp * warpSize + lane; + + // ── TiledMMA: 128 threads, covers [16, 32] output per step ── + auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(tid); + + // ── Swizzled smem views ── + // When D_PER_CTA < swizzle atom (= 64 for bf16), the underlying + // smem buffer is padded to D_SMEM_COLS so the swizzle layout is well-formed. + // Per-pass MMA loops only iterate D_PER_CTA / N_TILE tiles → never touch + // the padded tail. + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + + // x: swizzled [NPREDICTED_PAD_MMA_M, D_SMEM_COLS] + auto layout_x_swz = make_swizzled_layout_rc(); + Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_swz); + auto layout_x_trans_swz = make_swizzled_layout_rc_transpose(); + Tensor smem_x_trans + = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); + + // z: aliased swizzled [NPREDICTED_PAD_MMA_M, D_SMEM_COLS] — physical buffer + // is only next_multiple_of(NPREDICTED) rows tall; second m-tile + // aliases first. Ghost rows feed predicated-out output rows. + auto layout_z_swz = make_aliased_swizzled_layout_rc(); + Tensor smem_z = make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); + + // ── S2R copies ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B = s2r_B.get_slice(tid); + auto s2r_B_trans = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); + + // ── Load CB_scaled A operand from smem (precomputed by warps 0,1 between syncs) ── + // Row stride matches the buffer's padded width (one swizzle atom of `input_t`). + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + auto layout_cb_swz = make_swizzled_layout_rc(); + Tensor smem_CB + = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); + auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); + Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); + auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); + cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); + + // Decay broadcast: cumAdt[t] → [NPREDICTED_PAD_MMA_M, N_TILE] with stride-0 on N. + constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); + Tensor decay_bcast = make_tensor(make_smem_ptr(smem.cumAdt), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); + Tensor decay_part = thr_mma.partition_C(decay_bcast); + + // ── Gmem output: partition_C for direct register → gmem store ── + auto* __restrict__ output_ptr = reinterpret_cast(params.output); + // out_base lands on this CTA's D-slice within the head. + int64_t const out_base = out_seq_base + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; + + // Row predicate for padding. The epilogue store loop iterates i in steps + // of 2 and only consults pred(0) and pred(2) — m16n8k16 C-frag per thread + // has 4 elts at rows {t/4, t/4, t/4+8, t/4+8}, so there are only 2 unique + // row predicates. Compute them once and skip the 4-wide pred tensor. + auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_tile); + bool const pred_row_lo = get<0>(id_part(0)) < seq_len; + bool const pred_row_hi = get<0>(id_part(2)) < seq_len; + + // Number of output N-tiles per pass = D_PER_CTA / N_TILE. + // D_SPLIT=1, D_PER_CTA=64, N_TILE=32 → NUM_N_TILES = 2 (current behavior). + // D_SPLIT=2, D_PER_CTA=32 → NUM_N_TILES = 1 (uses _n1 variant). + constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; + static_assert( + NUM_N_TILES == 1 || NUM_N_TILES == 2, "Output epilogue supports NUM_N_TILES = D_PER_CTA / N_TILE in {1, 2}"); + + // ── Epilogue lambda (defined once; called per N-tile from each branch) ── + auto epilogue = [&](auto& frag_y, int n) + { + // Decay: frag_y *= exp(cumAdt[t]) #pragma unroll - for (int i = 0; i < size(frag_y); ++i) { - frag_y(i) *= __expf(decay_part(i)); - } + for (int i = 0; i < size(frag_y); ++i) + { + frag_y(i) *= __expf(decay_part(i)); + } - // frag_y += CB_scaled @ x (CB from smem LDSM, x from smem via ldmatrix.trans) - add_cb_x( - frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); - - // frag_y += D * x[t, d] - add_D_skip(frag_y, smem_x, thr_mma, D_val, n); - - // frag_y *= z * sigmoid(z) - compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); - - // Store frag_y directly to gmem (register → gmem, no smem round-trip). - auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), - make_layout(make_shape(Int{}, Int{}), - make_stride(params.out_stride_token, _1{}))); - auto gOut_part = thr_mma.partition_C(gOut_tile); - // Vectorized pair store: elements i and i+1 are same-row, consecutive columns - // in the m16n8k16 partition_C layout, so &gOut_part(i+1) == &gOut_part(i) + 1. - // Address is naturally aligned to sizeof(Pair) since MMA column - // index = (lane%4)*2 → even. pack_float2 dispatches to the native packed - // cvt (e.g. cvt.rn.bf16x2.f32 for bf16) — one instruction for the pair. + // frag_y += CB_scaled @ x (CB from smem LDSM, x from smem via ldmatrix.trans) + add_cb_x( + frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); + + // frag_y += D * x[t, d] + add_D_skip(frag_y, smem_x, thr_mma, D_val, n); + + // frag_y *= z * sigmoid(z) + compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); + + // Store frag_y directly to gmem (register → gmem, no smem round-trip). + auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), + make_layout( + make_shape(Int{}, Int{}), make_stride(params.out_stride_token, _1{}))); + auto gOut_part = thr_mma.partition_C(gOut_tile); + // Vectorized pair store: elements i and i+1 are same-row, consecutive columns + // in the m16n8k16 partition_C layout, so &gOut_part(i+1) == &gOut_part(i) + 1. + // Address is naturally aligned to sizeof(Pair) since MMA column + // index = (lane%4)*2 → even. pack_float2 dispatches to the native packed + // cvt (e.g. cvt.rn.bf16x2.f32 for bf16) — one instruction for the pair. #pragma unroll - for (int i = 0; i < size(frag_y); i += 2) { - // Bit 1 of i toggles between the two row groups of the m16n8k16 - // C-frag: i∈{0,1} → row t/4, i∈{2,3} → row t/4+8 (repeats per M-atom). - bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; - if (pred_i) { - *reinterpret_cast*>(&gOut_part(i)) = - pack_float2(make_float2(frag_y(i), frag_y(i + 1))); - } - } - }; - - // Skip the smem→gmem state copy when philox+f16: `replay_state_mma` - // already did the gmem store with stochastic rounding direct from registers. - constexpr bool kSkipSmemToGmemState = (PHILOX_ROUNDS > 0) && std::is_same_v; - - // ── Matmul 3 + store_state + epilogue, dispatching on NUM_N_TILES ── - // (NumNTiles is deduced from the variadic frag_y... pack in `add_init_out`.) - if constexpr (NUM_N_TILES == 2) { - Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); - Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); - add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0, - frag_y_1); - // State writeback hoisted here — after matmul 3 has finished consuming - // smem.state, before matmul 4 which reads only smem.x / smem.CB_scaled - // / smem.z. STGs fire-and-forget alongside the epilogue (matmul 4 + - // D*x + z-gate + output STG). Runtime-gated on must_checkpoint: - // non-checkpoint steps leave the prior state HBM intact (saving - // bandwidth — that's the perf win of the checkpointing design). - if constexpr (!kSkipSmemToGmemState) { - if (must_checkpoint) { - store_state(smem, params, warp, lane, d_tile, - head, cache_slot); - } + for (int i = 0; i < size(frag_y); i += 2) + { + // Bit 1 of i toggles between the two row groups of the m16n8k16 + // C-frag: i∈{0,1} → row t/4, i∈{2,3} → row t/4+8 (repeats per M-atom). + bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; + if (pred_i) + { + *reinterpret_cast*>(&gOut_part(i)) + = pack_float2(make_float2(frag_y(i), frag_y(i + 1))); + } + } + }; + + // Skip the smem→gmem state copy when philox+f16: `replay_state_mma` + // already did the gmem store with stochastic rounding direct from registers. + constexpr bool kSkipSmemToGmemState = (PHILOX_ROUNDS > 0) && std::is_same_v; + + // ── Matmul 3 + store_state + epilogue, dispatching on NUM_N_TILES ── + // (NumNTiles is deduced from the variadic frag_y... pack in `add_init_out`.) + if constexpr (NUM_N_TILES == 2) + { + Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); + Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); + add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0, frag_y_1); + // State writeback hoisted here — after matmul 3 has finished consuming + // smem.state, before matmul 4 which reads only smem.x / smem.CB_scaled + // / smem.z. STGs fire-and-forget alongside the epilogue (matmul 4 + + // D*x + z-gate + output STG). Runtime-gated on must_checkpoint: + // non-checkpoint steps leave the prior state HBM intact (saving + // bandwidth — that's the perf win of the checkpointing design). + if constexpr (!kSkipSmemToGmemState) + { + if (must_checkpoint) + { + store_state( + smem, params, warp, lane, d_tile, head, cache_slot); + } + } + epilogue(frag_y_0, 0); + epilogue(frag_y_1, 1); } - epilogue(frag_y_0, 0); - epilogue(frag_y_1, 1); - } else { // NUM_N_TILES == 1 (D_SPLIT = 2 path) - Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); - add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0); - // No sync needed before store_state: the post-replay __syncthreads() - // in the kernel already established cross-warp visibility of replay's - // writes to smem.state, and nothing after that point writes to it - // (add_init_out is read-only on smem.state). - if constexpr (!kSkipSmemToGmemState) { - if (must_checkpoint) { - store_state(smem, params, warp, lane, d_tile, - head, cache_slot); - } + else + { // NUM_N_TILES == 1 (D_SPLIT = 2 path) + Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); + add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0); + // No sync needed before store_state: the post-replay __syncthreads() + // in the kernel already established cross-warp visibility of replay's + // writes to smem.state, and nothing after that point writes to it + // (add_init_out is read-only on smem.state). + if constexpr (!kSkipSmemToGmemState) + { + if (must_checkpoint) + { + store_state( + smem, params, warp, lane, d_tile, head, cache_slot); + } + } + epilogue(frag_y_0, 0); } - epilogue(frag_y_0, 0); - } } // ── Orchestrator: compute_no_write_output (must_checkpoint == false path) ── @@ -684,434 +687,428 @@ __device__ __forceinline__ void compute_and_store_output(SmemT& smem, // CB_old is populated by `compute_CB_old_2warp` on warps 2,3 before the // `__syncthreads()` in the no-write dispatcher. It lives in the CB_scaled // buffer at cols [NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M + MAX_WINDOW_PAD_MMA_K). -template -__device__ __forceinline__ void compute_no_write_output( - SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int prev_k, int d_tile, - int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) { - using namespace cute; - static_assert(sizeof(input_t) == 2, "compute_no_write_output requires 2-byte input type"); - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - int const tid = warp * warpSize + lane; - - // ── TiledMMA for matmul-3 + matmul-4-new (K=NPREDICTED_PAD_MMA_M=16 fits K_BIG). ── - auto tiled_mma = - make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(tid); - - // ── TiledMMA for matmul-4-old: K = MAX_WINDOW_PAD_MMA_K ∈ {8, 16} → atom dispatch. ── - using MmaAtomOld = std::conditional_t; - using LdsmAOld = std::conditional_t; - using LdsmBOld = std::conditional_t; - auto tiled_mma_old = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_old = tiled_mma_old.get_slice(tid); - - // ── Swizzled smem views ── - auto layout_x_swz = make_swizzled_layout_rc(); - Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), - layout_x_swz); - auto layout_x_trans_swz = - make_swizzled_layout_rc_transpose(); - Tensor smem_x_trans = make_tensor( - make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); - - auto layout_old_x_trans_swz = - make_swizzled_layout_rc_transpose(); - Tensor smem_old_x_trans = - make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), - layout_old_x_trans_swz); - - auto layout_z_swz = - make_aliased_swizzled_layout_rc(); - Tensor smem_z = - make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); - - // ── S2R copies (matmul-4-new) ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - auto s2r_B_trans = - make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); - - // ── S2R copies (matmul-4-old, K-dispatched atoms) ── - auto s2r_A_old = make_tiled_copy_A(Copy_Atom{}, tiled_mma_old); - auto s2r_thr_A_old = s2r_A_old.get_slice(tid); - auto s2r_B_old_trans = - make_tiled_copy_B(Copy_Atom{}, tiled_mma_old); - auto s2r_thr_B_old_trans = s2r_B_old_trans.get_slice(tid); - - // ── Load CB_scaled A operand (cols [0, NPREDICTED_PAD_MMA_M)) ── - auto layout_cb_swz = - make_swizzled_layout_rc(); - Tensor smem_CB = make_tensor( - make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); - auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); - Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); - auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); - cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); - - // ── Load CB_old A operand (cols [NPREDICTED_PAD_MMA_M, +MAX_WINDOW_PAD_MMA_K)) ── - // Use the full physical (T_pad, CB_ROW_STRIDE) padded swizzle view — byte- - // compatible with both the CB_scaled (T_pad, T_pad, CB_ROW_STRIDE) write - // layout and compute_CB_old_2warp's wide write layout (inner offset - // r*CB_ROW_STRIDE + c is identical across the three views). - auto layout_cb_full = make_swizzled_layout_rc(); - Tensor smem_CB_full = make_tensor( - make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); - Tensor smem_CB_old = - local_tile(smem_CB_full, make_tile(Int{}, Int{}), - make_coord(_0{}, NPREDICTED_PAD_MMA_M / MAX_WINDOW_PAD_MMA_K)); - auto smem_CB_old_s2r = s2r_thr_A_old.partition_S(smem_CB_old); - Tensor frag_CB_old_A = thr_mma_old.partition_fragment_A(smem_CB_old); - auto frag_CB_old_A_view = s2r_thr_A_old.retile_D(frag_CB_old_A); - cute::copy(s2r_A_old, smem_CB_old_s2r, frag_CB_old_A_view); - - // ── Decay broadcast: cumAdt[t] (per-T scalar) with stride-0 on N. ── - constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); - Tensor decay_bcast = make_tensor( - make_smem_ptr(smem.cumAdt), - make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); - Tensor decay_part = thr_mma.partition_C(decay_bcast); - - // ── β extra factor: exp(total_old_cumAdt) — uniform constant across (t, d) ── - float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - float const beta_extra = __expf(total_old_cumAdt); - - // ── Gmem output base ── - auto* __restrict__ output_ptr = reinterpret_cast(params.output); - int64_t const out_base = out_seq_base + (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; - - // ── Row predicate (same pattern as compute_and_store_output) ── - auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_tile); - bool const pred_row_lo = get<0>(id_part(0)) < seq_len; - bool const pred_row_hi = get<0>(id_part(2)) < seq_len; - - constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; - static_assert(NUM_N_TILES == 1 || NUM_N_TILES == 2, - "compute_no_write_output supports NUM_N_TILES in {1, 2}"); - - // ── Epilogue per N-tile ── - auto epilogue = [&](auto& frag_y, int n) { - // 1. β-scale: frag_y(t, d) *= exp(total_old_cumAdt + cumAdt[t]). - // Matmul-3 produced u^T = C @ s_0^T; this scales the u term by β - // BEFORE matmul-4 adds the CB·x and CB_old·old_x contributions. +template +__device__ __forceinline__ void compute_no_write_output(SmemT& smem, CheckpointingSsuParams const& params, int warp, + int lane, int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) +{ + using namespace cute; + static_assert(sizeof(input_t) == 2, "compute_no_write_output requires 2-byte input type"); + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + int const tid = warp * warpSize + lane; + + // ── TiledMMA for matmul-3 + matmul-4-new (K=NPREDICTED_PAD_MMA_M=16 fits K_BIG). ── + auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(tid); + + // ── TiledMMA for matmul-4-old: K = MAX_WINDOW_PAD_MMA_K ∈ {8, 16} → atom dispatch. ── + using MmaAtomOld = std::conditional_t; + using LdsmAOld = std::conditional_t; + using LdsmBOld = std::conditional_t; + auto tiled_mma_old = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_old = tiled_mma_old.get_slice(tid); + + // ── Swizzled smem views ── + auto layout_x_swz = make_swizzled_layout_rc(); + Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_swz); + auto layout_x_trans_swz = make_swizzled_layout_rc_transpose(); + Tensor smem_x_trans + = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); + + auto layout_old_x_trans_swz = make_swizzled_layout_rc_transpose(); + Tensor smem_old_x_trans + = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), layout_old_x_trans_swz); + + auto layout_z_swz = make_aliased_swizzled_layout_rc(); + Tensor smem_z = make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); + + // ── S2R copies (matmul-4-new) ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + auto s2r_B_trans = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); + + // ── S2R copies (matmul-4-old, K-dispatched atoms) ── + auto s2r_A_old = make_tiled_copy_A(Copy_Atom{}, tiled_mma_old); + auto s2r_thr_A_old = s2r_A_old.get_slice(tid); + auto s2r_B_old_trans = make_tiled_copy_B(Copy_Atom{}, tiled_mma_old); + auto s2r_thr_B_old_trans = s2r_B_old_trans.get_slice(tid); + + // ── Load CB_scaled A operand (cols [0, NPREDICTED_PAD_MMA_M)) ── + auto layout_cb_swz = make_swizzled_layout_rc(); + Tensor smem_CB + = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); + auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); + Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); + auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); + cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); + + // ── Load CB_old A operand (cols [NPREDICTED_PAD_MMA_M, +MAX_WINDOW_PAD_MMA_K)) ── + // Use the full physical (T_pad, CB_ROW_STRIDE) padded swizzle view — byte- + // compatible with both the CB_scaled (T_pad, T_pad, CB_ROW_STRIDE) write + // layout and compute_CB_old_2warp's wide write layout (inner offset + // r*CB_ROW_STRIDE + c is identical across the three views). + auto layout_cb_full = make_swizzled_layout_rc(); + Tensor smem_CB_full + = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); + Tensor smem_CB_old = local_tile(smem_CB_full, make_tile(Int{}, Int{}), + make_coord(_0{}, NPREDICTED_PAD_MMA_M / MAX_WINDOW_PAD_MMA_K)); + auto smem_CB_old_s2r = s2r_thr_A_old.partition_S(smem_CB_old); + Tensor frag_CB_old_A = thr_mma_old.partition_fragment_A(smem_CB_old); + auto frag_CB_old_A_view = s2r_thr_A_old.retile_D(frag_CB_old_A); + cute::copy(s2r_A_old, smem_CB_old_s2r, frag_CB_old_A_view); + + // ── Decay broadcast: cumAdt[t] (per-T scalar) with stride-0 on N. ── + constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); + Tensor decay_bcast = make_tensor(make_smem_ptr(smem.cumAdt), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); + Tensor decay_part = thr_mma.partition_C(decay_bcast); + + // ── β extra factor: exp(total_old_cumAdt) — uniform constant across (t, d) ── + float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + float const beta_extra = __expf(total_old_cumAdt); + + // ── Gmem output base ── + auto* __restrict__ output_ptr = reinterpret_cast(params.output); + int64_t const out_base = out_seq_base + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; + + // ── Row predicate (same pattern as compute_and_store_output) ── + auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_tile); + bool const pred_row_lo = get<0>(id_part(0)) < seq_len; + bool const pred_row_hi = get<0>(id_part(2)) < seq_len; + + constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; + static_assert(NUM_N_TILES == 1 || NUM_N_TILES == 2, "compute_no_write_output supports NUM_N_TILES in {1, 2}"); + + // ── Epilogue per N-tile ── + auto epilogue = [&](auto& frag_y, int n) + { + // 1. β-scale: frag_y(t, d) *= exp(total_old_cumAdt + cumAdt[t]). + // Matmul-3 produced u^T = C @ s_0^T; this scales the u term by β + // BEFORE matmul-4 adds the CB·x and CB_old·old_x contributions. #pragma unroll - for (int i = 0; i < size(frag_y); ++i) { - frag_y(i) *= beta_extra * __expf(decay_part(i)); - } + for (int i = 0; i < size(frag_y); ++i) + { + frag_y(i) *= beta_extra * __expf(decay_part(i)); + } - // 2. frag_y += CB_scaled @ x (matmul-4 over new tokens). - add_cb_x( - frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); + // 2. frag_y += CB_scaled @ x (matmul-4 over new tokens). + add_cb_x( + frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); - // 3. frag_y += CB_old @ old_x (matmul-4 over old tokens — NEW). - add_cb_old_x( - frag_y, frag_CB_old_A, smem_old_x_trans, s2r_B_old_trans, s2r_thr_B_old_trans, thr_mma_old, - tiled_mma_old, n); + // 3. frag_y += CB_old @ old_x (matmul-4 over old tokens — NEW). + add_cb_old_x(frag_y, frag_CB_old_A, + smem_old_x_trans, s2r_B_old_trans, s2r_thr_B_old_trans, thr_mma_old, tiled_mma_old, n); - // 4. frag_y += D · x[t, d]. - add_D_skip(frag_y, smem_x, thr_mma, D_val, n); + // 4. frag_y += D · x[t, d]. + add_D_skip(frag_y, smem_x, thr_mma, D_val, n); - // 5. frag_y *= z · sigmoid(z). - compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); + // 5. frag_y *= z · sigmoid(z). + compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); - // 6. Store frag_y → gmem via partition_C (same pattern as compute_and_store_output). - auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), - make_layout(make_shape(Int{}, Int{}), - make_stride(params.out_stride_token, _1{}))); - auto gOut_part = thr_mma.partition_C(gOut_tile); + // 6. Store frag_y → gmem via partition_C (same pattern as compute_and_store_output). + auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), + make_layout( + make_shape(Int{}, Int{}), make_stride(params.out_stride_token, _1{}))); + auto gOut_part = thr_mma.partition_C(gOut_tile); #pragma unroll - for (int i = 0; i < size(frag_y); i += 2) { - bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; - if (pred_i) { - *reinterpret_cast*>(&gOut_part(i)) = - pack_float2(make_float2(frag_y(i), frag_y(i + 1))); - } + for (int i = 0; i < size(frag_y); i += 2) + { + bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; + if (pred_i) + { + *reinterpret_cast*>(&gOut_part(i)) + = pack_float2(make_float2(frag_y(i), frag_y(i + 1))); + } + } + }; + + // ── Matmul-3: frag_y = C @ s_0^T (smem.state retains s_0 since replay skipped) ── + if constexpr (NUM_N_TILES == 2) + { + Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); + Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); + add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0, frag_y_1); + epilogue(frag_y_0, 0); + epilogue(frag_y_1, 1); + } + else + { // NUM_N_TILES == 1 (D_SPLIT = 2 path) + Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); + add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0); + epilogue(frag_y_0, 0); } - }; - - // ── Matmul-3: frag_y = C @ s_0^T (smem.state retains s_0 since replay skipped) ── - if constexpr (NUM_N_TILES == 2) { - Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); - Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); - add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0, - frag_y_1); - epilogue(frag_y_0, 0); - epilogue(frag_y_1, 1); - } else { // NUM_N_TILES == 1 (D_SPLIT = 2 path) - Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); - add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0); - epilogue(frag_y_0, 0); - } } // ── Per-path dispatchers (called from checkpointing_ssu_kernel) ── // ssu_checkpoint: replay → sync → output (today's body). // ssu_nocheckpoint: sync → no-write output (skips replay). -template -__device__ __forceinline__ void ssu_checkpoint(SmemT& smem, CheckpointingSsuParams const& params, - int warp, int lane, int prev_k, int d_tile, - int64_t out_seq_base, int head, int64_t cache_slot, - float D_val, int seq_len) { - // ── DO NOT HOIST `rand_seed` ── see kernel preamble for the perf rationale. - int64_t const rand_seed = (PHILOX_ROUNDS > 0) ? *params.rand_seed : 0; - // `state_ptr_offset` is int64 — matches Triton's `base_rand = - // cache_batch_idx * stride_state_batch + ...` (cache_batch_idx is .to(int64)). - // Full 64 bits flow through `philox_randint4x`, which splits low/high - // across Philox c0/c1. No collision risk at large serving cache sizes. - int64_t const state_ptr_offset = - cache_slot * params.state_stride_seq + (int64_t)head * DIM * DSTATE; - state_t* const state_w_base = reinterpret_cast(params.state) + - cache_slot * params.state_stride_seq + - (int64_t)head * DIM * DSTATE + (int64_t)d_tile * D_PER_CTA * DSTATE; - replay_state_mma( - smem, params, warp, lane, prev_k, d_tile, state_ptr_offset, state_w_base, rand_seed, - /*must_checkpoint=*/true); - - __syncthreads(); - - compute_and_store_output(smem, params, warp, lane, d_tile, out_seq_base, head, - cache_slot, D_val, /*must_checkpoint=*/true, seq_len); +template +__device__ __forceinline__ void ssu_checkpoint(SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, + int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) +{ + // ── DO NOT HOIST `rand_seed` ── see kernel preamble for the perf rationale. + int64_t const rand_seed = (PHILOX_ROUNDS > 0) ? *params.rand_seed : 0; + // `state_ptr_offset` is int64 — matches Triton's `base_rand = + // cache_batch_idx * stride_state_batch + ...` (cache_batch_idx is .to(int64)). + // Full 64 bits flow through `philox_randint4x`, which splits low/high + // across Philox c0/c1. No collision risk at large serving cache sizes. + int64_t const state_ptr_offset = cache_slot * params.state_stride_seq + (int64_t) head * DIM * DSTATE; + state_t* const state_w_base = reinterpret_cast(params.state) + cache_slot * params.state_stride_seq + + (int64_t) head * DIM * DSTATE + (int64_t) d_tile * D_PER_CTA * DSTATE; + replay_state_mma(smem, params, warp, lane, prev_k, d_tile, + state_ptr_offset, state_w_base, rand_seed, + /*must_checkpoint=*/true); + + __syncthreads(); + + compute_and_store_output( + smem, params, warp, lane, d_tile, out_seq_base, head, cache_slot, D_val, /*must_checkpoint=*/true, seq_len); } -template -__device__ __forceinline__ void ssu_nocheckpoint(SmemT& smem, CheckpointingSsuParams const& params, - int warp, int lane, int prev_k, int d_tile, - int64_t out_seq_base, int head, int64_t cache_slot, - float D_val, int seq_len) { - // Sync makes warps 0,1's CB_scaled writes and warps 2,3's CB_old writes - // visible to all warps before matmul-4 reads CB_scaled + CB_old. Also - // covers smem.x (warp 2-loaded) and smem.z (warp 3-loaded) for Phase 2. - __syncthreads(); - - compute_no_write_output(smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, - cache_slot, D_val, seq_len); +template +__device__ __forceinline__ void ssu_nocheckpoint(SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, + int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) +{ + // Sync makes warps 0,1's CB_scaled writes and warps 2,3's CB_old writes + // visible to all warps before matmul-4 reads CB_scaled + CB_old. Also + // covers smem.x (warp 2-loaded) and smem.z (warp 3-loaded) for Phase 2. + __syncthreads(); + + compute_no_write_output( + smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); } // ============================================================================= // Kernel // ============================================================================= template -__global__ void checkpointing_ssu_kernel(CheckpointingSsuParams params) { - // Per-head DIM is sharded across `D_SPLIT` CTAs (D_PER_CTA each). - static_assert(DIM % D_SPLIT == 0, "DIM must be divisible by D_SPLIT"); - constexpr int D_PER_CTA = DIM / D_SPLIT; - static_assert(D_PER_CTA >= 32, - "D_PER_CTA must be >= 32 (output MMA m16n8 with _1×4 warp layout). " - "D_SPLIT=4 (D_PER_CTA=16) needs warp-count restructure."); - static_assert(NPREDICTED <= MAX_WINDOW, - "NPREDICTED must be <= MAX_WINDOW (new tokens must fit in cache)"); - static_assert(MAX_WINDOW <= MMA_prop::K_BIG, - "MAX_WINDOW must be <= MMA::K_BIG=16 (single replay K-tile assumption)"); - // Cross-check: host launcher must dispatch the template specialization - // matching the runtime params.d_split it stamped into the struct. - assert(params.d_split == D_SPLIT); - using SmemT = - CheckpointingSsuStorage; - extern __shared__ __align__(128) char smem_buf[]; - auto& smem = *reinterpret_cast(smem_buf); - - // Grid layout (D_SPLIT, batch, nheads). - int const d_tile = blockIdx.x; - int const seq = blockIdx.y; - int const head = blockIdx.z; - int const lane = threadIdx.x; - int const warp = threadIdx.y; - int const group_idx = head / HEADS_PER_GROUP; - - // ── Resolve cache slot ── - auto const* __restrict__ sbi = reinterpret_cast(params.state_batch_indices); - int64_t const cache_slot = sbi ? static_cast(sbi[seq]) : seq; - if (cache_slot == params.pad_slot_id) return; - - // ── Double-buffer index ── - auto const* __restrict__ buf_idx_ptr = reinterpret_cast(params.cache_buf_idx); - int const buf_read = __ldg(&buf_idx_ptr[cache_slot]); - - // ── prev_num_accepted_tokens ── - auto const* __restrict__ prev_ptr = reinterpret_cast(params.prev_num_accepted); - int const prev_k = prev_ptr[cache_slot]; - - // ── Varlen vs non-varlen prologue. See checkpointing_ssu_kernel_8bit for - // the rationale: `seq_len` flows downstream as a constexpr-foldable - // NPREDICTED in non-varlen, runtime int in varlen. - // - // Uniform gmem-base formula: `outer * *_stride_seq` where - // non-varlen: outer = seq (= blockIdx.y), stride_seq = x.stride(0). - // varlen : outer = cu_seqlens[seq], stride_seq = x.stride(1). - // The wrapper picks the right stride_seq value; the kernel only branches - // on whether to load cu_seqlens. - int seq_len; - int64_t outer; - if constexpr (VARLEN) { - auto const* __restrict__ cu_seqlens = reinterpret_cast(params.cu_seqlens); - // Two LDG.E.32 (not one LDG.E.64): cu_seqlens is only 4-byte aligned - // at `&cu_seqlens[seq]` when seq is odd, and PTX - // `ld.global.v2.b32` faults on a 4-byte-aligned address. ptxas emits - // the two scalar loads back-to-back; latency is hidden against the - // following ALU work. - int const bos = __ldg(&cu_seqlens[seq]); - int const eos = __ldg(&cu_seqlens[seq + 1]); - seq_len = eos - bos; - if (seq_len <= 0) return; - outer = (int64_t)bos; - } else { - seq_len = NPREDICTED; - outer = (int64_t)seq; - } - // x/B/C bases are computed inside `load_post_pdl_wait_data` from `outer` - // so the products don't get pinned in registers across `gdc_wait` (asm - // volatile blocks rematerialization; cost was ~6 extra regs). dt/z bases - // are only consumed pre-wait, and out_base only post-replay — fine to - // precompute. - int64_t const dt_seq_base = outer * params.dt_stride_seq + head; - int64_t const z_seq_base = outer * params.z_stride_seq; - int64_t const out_seq_base = outer * params.out_stride_seq; - - // ── Per-CTA implicit checkpoint criterion ── - // When the new tokens would overflow the cache buffer, we must checkpoint: - // replay [0, prev_k) into state, write state to HBM, write the new tokens - // to the **staging** buffer (1 - buf_read) at offset 0. Otherwise, we - // append the new tokens to the **active** buffer (buf_read) at offset - // prev_k and skip the state HBM write entirely. Cache writes always - // happen — only their target buffer + offset depends on must_checkpoint. - bool const must_checkpoint = (prev_k + seq_len > MAX_WINDOW); - int const buf_write = must_checkpoint ? (1 - buf_read) : buf_read; - int const write_offset = must_checkpoint ? 0 : prev_k; - - // ── Load A (scalar, tie_hdim), dt_bias, and D (hoisted to hide gmem latency) ── - auto const* __restrict__ A_ptr = reinterpret_cast(params.A); - auto const* __restrict__ dt_bias_ptr = reinterpret_cast(params.dt_bias); - auto const* __restrict__ D_ptr = reinterpret_cast(params.D); - float const A_val = toFloat(A_ptr[head]); - float const dt_bias_val = dt_bias_ptr ? toFloat(dt_bias_ptr[head]) : 0.f; - float const D_val = D_ptr ? toFloat(D_ptr[head]) : 0.f; - - // ════════════════════════════════════════════════════════════════════════ - // Phase 0: Load all data into smem (per-warp ownership) - // ════════════════════════════════════════════════════════════════════════ - // Two-phase load around the PDL barrier: - // 1. Issue cp.async for cache (state, old_B, old_x) and in_proj-derived - // data (z); run scalar LDGs (old_dt, old_cumAdt, dt → dt_proc) and the - // cumAdt warp scan. None of these depend on conv1d, so they overlap - // with the upstream's tail. - // 2. `gdc_wait()` — wait for the upstream conv1d to signal (no-op when - // the kernel isn't launched with the PDL attribute). - // 3. Issue cp.async for conv1d outputs (x, B, C), then __pipeline_commit - // + __pipeline_wait_prior(0) + __syncwarp drains BOTH halves' cp.async - // (they share the per-thread async group). - // - // Each warp sees its own cp.async via __syncwarp. Cross-warp visibility - // is established by the post-replay __syncthreads below — replay reads of - // state are safe because (a) replay's frag_h initial load sees only the - // current warp's lane positions, and (b) the actual _1×4 cross-warp - // dependency is on writes that haven't happened yet at this point. - // ENABLE_PDL is JIT-stamped (see checkpointing_ssu_customize_config.jinja). - // `if constexpr` keeps only the chosen branch in the binary — no register - // pressure leak from the unused path. - if constexpr (ENABLE_PDL) { - load_pre_pdl_wait_data(smem, params, lane, warp, d_tile, head, group_idx, cache_slot, - buf_read, A_val, dt_bias_val, dt_seq_base, z_seq_base, - seq_len); - gdc_wait(); - load_post_pdl_wait_data( - smem, params, lane, warp, d_tile, head, group_idx, outer, seq_len); - } else { - load_data( - smem, params, lane, warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, - outer, seq_len); - } - - // old_B writeback hoisted ahead of Phase 1. Source (smem.B) is consumed - // only by Phase 1a CB; the STGs fire-and-forget onto the memory subsystem - // and complete in parallel with all subsequent compute. Only W0, W1 hold - // valid smem.B at this point (they're the ones that cp.async'd B). Gate - // accordingly — store halves its thread count but B is small (4 KB) so - // still cheap. old_B is D-independent (per-group, full DSTATE) — only - // d_tile == 0 writes; other d_tiles would emit identical payloads. - if (d_tile == 0 && warp < 2) { - store_old_B( - smem, params, warp, lane, head, group_idx, cache_slot, buf_write, write_offset, seq_len); - } - - // CB precompute (4-warp split): warps 0,1 compute CB_scaled (new tokens); - // warps 2,3 compute CB_old (old tokens) in the no-write path only. Both - // halves write to disjoint col ranges of the same swizzled smem.CB_scaled - // buffer. In the checkpoint path, warps 2,3 stay idle here and pick up - // work below in `ssu_checkpoint`'s replay matmul. - if (warp < 2) { - compute_CB_scaled_2warp(smem, warp, lane, seq_len); - } else if (!must_checkpoint) { - compute_CB_old_2warp(smem, warp, lane, prev_k, - seq_len); - } - - // ════════════════════════════════════════════════════════════════════════ - // Phase 1b + 2: Per-path dispatch - // ════════════════════════════════════════════════════════════════════════ - // Checkpoint path: replay + sync + compute_and_store_output (today's body). - // No-write path : sync + compute_no_write_output (skips replay; matmul-3 - // reads s_0 directly from smem.state, matmul-4 extends with - // the CB_old @ old_x contribution over [0, prev_k)). - // must_checkpoint is uniform across the CTA (derived from broadcast prev_k - // + compile-time NPREDICTED + MAX_WINDOW), so both branches contain a - // __syncthreads and divergence is balanced. - if (must_checkpoint) { - ssu_checkpoint( - smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); - } else { - ssu_nocheckpoint( - smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); - } - - // ── PDL: signal downstream that `output` is written. The cache writes - // below target tensors that only the next SSU step reads, not the - // immediate downstream kernel, so we can signal before issuing them. - if constexpr (ENABLE_PDL) { - gdc_launch_dependents(); - } - - // ── Phase 3: Store to global memory ── - // (old_B hoisted to pre-Phase-1; state hoisted into compute_and_store_output.) - - // Cache writes — old_x uses all warps (vectorized), dt/cumAdt one warp each. - // Each writes the new NPREDICTED tokens at gmem offset `write_offset` into - // buffer `buf_write` (computed above from must_checkpoint). - store_old_x(smem, params, warp, lane, d_tile, head, - cache_slot, write_offset, seq_len); - // dt_proc / cumAdt are D-independent — only d_tile == 0 writes. - if (d_tile == 0 && warp == 0 && lane < seq_len) { - auto* __restrict__ old_dt_w = reinterpret_cast(params.old_dt); - int64_t const dt_w_base = cache_slot * params.old_dt_stride_seq + - buf_write * params.old_dt_stride_dbuf + - head * params.old_dt_stride_head; - old_dt_w[dt_w_base + write_offset + lane] = smem.dt_proc[lane]; - } - if (d_tile == 0 && warp == 1 && lane < seq_len) { - auto* __restrict__ old_cumAdt_w = reinterpret_cast(params.old_cumAdt); - int64_t const ca_w_base = cache_slot * params.old_cumAdt_stride_seq + - buf_write * params.old_cumAdt_stride_dbuf + - head * params.old_cumAdt_stride_head; - old_cumAdt_w[ca_w_base + write_offset + lane] = smem.cumAdt[lane]; - } + typename stateIndex_t, typename state_scale_t, int NPREDICTED, int MAX_WINDOW, int DIM, int DSTATE, + int HEADS_PER_GROUP, int PHILOX_ROUNDS, int NUM_WARPS, int D_SPLIT = 1, bool VARLEN = false> +__global__ void checkpointing_ssu_kernel(CheckpointingSsuParams params) +{ + // Per-head DIM is sharded across `D_SPLIT` CTAs (D_PER_CTA each). + static_assert(DIM % D_SPLIT == 0, "DIM must be divisible by D_SPLIT"); + constexpr int D_PER_CTA = DIM / D_SPLIT; + static_assert(D_PER_CTA >= 32, + "D_PER_CTA must be >= 32 (output MMA m16n8 with _1×4 warp layout). " + "D_SPLIT=4 (D_PER_CTA=16) needs warp-count restructure."); + static_assert(NPREDICTED <= MAX_WINDOW, "NPREDICTED must be <= MAX_WINDOW (new tokens must fit in cache)"); + static_assert( + MAX_WINDOW <= MMA_prop::K_BIG, "MAX_WINDOW must be <= MMA::K_BIG=16 (single replay K-tile assumption)"); + // Cross-check: host launcher must dispatch the template specialization + // matching the runtime params.d_split it stamped into the struct. + assert(params.d_split == D_SPLIT); + using SmemT = CheckpointingSsuStorage; + extern __shared__ __align__(128) char smem_buf[]; + auto& smem = *reinterpret_cast(smem_buf); + + // Grid layout (D_SPLIT, batch, nheads). + int const d_tile = blockIdx.x; + int const seq = blockIdx.y; + int const head = blockIdx.z; + int const lane = threadIdx.x; + int const warp = threadIdx.y; + int const group_idx = head / HEADS_PER_GROUP; + + // ── Resolve cache slot ── + auto const* __restrict__ sbi = reinterpret_cast(params.state_batch_indices); + int64_t const cache_slot = sbi ? static_cast(sbi[seq]) : seq; + if (cache_slot == params.pad_slot_id) + return; + + // ── Double-buffer index ── + auto const* __restrict__ buf_idx_ptr = reinterpret_cast(params.cache_buf_idx); + int const buf_read = __ldg(&buf_idx_ptr[cache_slot]); + + // ── prev_num_accepted_tokens ── + auto const* __restrict__ prev_ptr = reinterpret_cast(params.prev_num_accepted); + int const prev_k = prev_ptr[cache_slot]; + + // ── Varlen vs non-varlen prologue. See checkpointing_ssu_kernel_8bit for + // the rationale: `seq_len` flows downstream as a constexpr-foldable + // NPREDICTED in non-varlen, runtime int in varlen. + // + // Uniform gmem-base formula: `outer * *_stride_seq` where + // non-varlen: outer = seq (= blockIdx.y), stride_seq = x.stride(0). + // varlen : outer = cu_seqlens[seq], stride_seq = x.stride(1). + // The wrapper picks the right stride_seq value; the kernel only branches + // on whether to load cu_seqlens. + int seq_len; + int64_t outer; + if constexpr (VARLEN) + { + auto const* __restrict__ cu_seqlens = reinterpret_cast(params.cu_seqlens); + // Two LDG.E.32 (not one LDG.E.64): cu_seqlens is only 4-byte aligned + // at `&cu_seqlens[seq]` when seq is odd, and PTX + // `ld.global.v2.b32` faults on a 4-byte-aligned address. ptxas emits + // the two scalar loads back-to-back; latency is hidden against the + // following ALU work. + int const bos = __ldg(&cu_seqlens[seq]); + int const eos = __ldg(&cu_seqlens[seq + 1]); + seq_len = eos - bos; + if (seq_len <= 0) + return; + outer = (int64_t) bos; + } + else + { + seq_len = NPREDICTED; + outer = (int64_t) seq; + } + // x/B/C bases are computed inside `load_post_pdl_wait_data` from `outer` + // so the products don't get pinned in registers across `gdc_wait` (asm + // volatile blocks rematerialization; cost was ~6 extra regs). dt/z bases + // are only consumed pre-wait, and out_base only post-replay — fine to + // precompute. + int64_t const dt_seq_base = outer * params.dt_stride_seq + head; + int64_t const z_seq_base = outer * params.z_stride_seq; + int64_t const out_seq_base = outer * params.out_stride_seq; + + // ── Per-CTA implicit checkpoint criterion ── + // When the new tokens would overflow the cache buffer, we must checkpoint: + // replay [0, prev_k) into state, write state to HBM, write the new tokens + // to the **staging** buffer (1 - buf_read) at offset 0. Otherwise, we + // append the new tokens to the **active** buffer (buf_read) at offset + // prev_k and skip the state HBM write entirely. Cache writes always + // happen — only their target buffer + offset depends on must_checkpoint. + bool const must_checkpoint = (prev_k + seq_len > MAX_WINDOW); + int const buf_write = must_checkpoint ? (1 - buf_read) : buf_read; + int const write_offset = must_checkpoint ? 0 : prev_k; + + // ── Load A (scalar, tie_hdim), dt_bias, and D (hoisted to hide gmem latency) ── + auto const* __restrict__ A_ptr = reinterpret_cast(params.A); + auto const* __restrict__ dt_bias_ptr = reinterpret_cast(params.dt_bias); + auto const* __restrict__ D_ptr = reinterpret_cast(params.D); + float const A_val = toFloat(A_ptr[head]); + float const dt_bias_val = dt_bias_ptr ? toFloat(dt_bias_ptr[head]) : 0.f; + float const D_val = D_ptr ? toFloat(D_ptr[head]) : 0.f; + + // ════════════════════════════════════════════════════════════════════════ + // Phase 0: Load all data into smem (per-warp ownership) + // ════════════════════════════════════════════════════════════════════════ + // Two-phase load around the PDL barrier: + // 1. Issue cp.async for cache (state, old_B, old_x) and in_proj-derived + // data (z); run scalar LDGs (old_dt, old_cumAdt, dt → dt_proc) and the + // cumAdt warp scan. None of these depend on conv1d, so they overlap + // with the upstream's tail. + // 2. `gdc_wait()` — wait for the upstream conv1d to signal (no-op when + // the kernel isn't launched with the PDL attribute). + // 3. Issue cp.async for conv1d outputs (x, B, C), then __pipeline_commit + // + __pipeline_wait_prior(0) + __syncwarp drains BOTH halves' cp.async + // (they share the per-thread async group). + // + // Each warp sees its own cp.async via __syncwarp. Cross-warp visibility + // is established by the post-replay __syncthreads below — replay reads of + // state are safe because (a) replay's frag_h initial load sees only the + // current warp's lane positions, and (b) the actual _1×4 cross-warp + // dependency is on writes that haven't happened yet at this point. + // ENABLE_PDL is JIT-stamped (see checkpointing_ssu_customize_config.jinja). + // `if constexpr` keeps only the chosen branch in the binary — no register + // pressure leak from the unused path. + if constexpr (ENABLE_PDL) + { + load_pre_pdl_wait_data(smem, + params, lane, warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, dt_seq_base, + z_seq_base, seq_len); + gdc_wait(); + load_post_pdl_wait_data( + smem, params, lane, warp, d_tile, head, group_idx, outer, seq_len); + } + else + { + load_data(smem, params, lane, + warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, outer, seq_len); + } + + // old_B writeback hoisted ahead of Phase 1. Source (smem.B) is consumed + // only by Phase 1a CB; the STGs fire-and-forget onto the memory subsystem + // and complete in parallel with all subsequent compute. Only W0, W1 hold + // valid smem.B at this point (they're the ones that cp.async'd B). Gate + // accordingly — store halves its thread count but B is small (4 KB) so + // still cheap. old_B is D-independent (per-group, full DSTATE) — only + // d_tile == 0 writes; other d_tiles would emit identical payloads. + if (d_tile == 0 && warp < 2) + { + store_old_B( + smem, params, warp, lane, head, group_idx, cache_slot, buf_write, write_offset, seq_len); + } + + // CB precompute (4-warp split): warps 0,1 compute CB_scaled (new tokens); + // warps 2,3 compute CB_old (old tokens) in the no-write path only. Both + // halves write to disjoint col ranges of the same swizzled smem.CB_scaled + // buffer. In the checkpoint path, warps 2,3 stay idle here and pick up + // work below in `ssu_checkpoint`'s replay matmul. + if (warp < 2) + { + compute_CB_scaled_2warp(smem, warp, lane, seq_len); + } + else if (!must_checkpoint) + { + compute_CB_old_2warp(smem, warp, lane, prev_k, seq_len); + } + + // ════════════════════════════════════════════════════════════════════════ + // Phase 1b + 2: Per-path dispatch + // ════════════════════════════════════════════════════════════════════════ + // Checkpoint path: replay + sync + compute_and_store_output (today's body). + // No-write path : sync + compute_no_write_output (skips replay; matmul-3 + // reads s_0 directly from smem.state, matmul-4 extends with + // the CB_old @ old_x contribution over [0, prev_k)). + // must_checkpoint is uniform across the CTA (derived from broadcast prev_k + // + compile-time NPREDICTED + MAX_WINDOW), so both branches contain a + // __syncthreads and divergence is balanced. + if (must_checkpoint) + { + ssu_checkpoint( + smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); + } + else + { + ssu_nocheckpoint( + smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); + } + + // ── PDL: signal downstream that `output` is written. The cache writes + // below target tensors that only the next SSU step reads, not the + // immediate downstream kernel, so we can signal before issuing them. + if constexpr (ENABLE_PDL) + { + gdc_launch_dependents(); + } + + // ── Phase 3: Store to global memory ── + // (old_B hoisted to pre-Phase-1; state hoisted into compute_and_store_output.) + + // Cache writes — old_x uses all warps (vectorized), dt/cumAdt one warp each. + // Each writes the new NPREDICTED tokens at gmem offset `write_offset` into + // buffer `buf_write` (computed above from must_checkpoint). + store_old_x( + smem, params, warp, lane, d_tile, head, cache_slot, write_offset, seq_len); + // dt_proc / cumAdt are D-independent — only d_tile == 0 writes. + if (d_tile == 0 && warp == 0 && lane < seq_len) + { + auto* __restrict__ old_dt_w = reinterpret_cast(params.old_dt); + int64_t const dt_w_base = cache_slot * params.old_dt_stride_seq + buf_write * params.old_dt_stride_dbuf + + head * params.old_dt_stride_head; + old_dt_w[dt_w_base + write_offset + lane] = smem.dt_proc[lane]; + } + if (d_tile == 0 && warp == 1 && lane < seq_len) + { + auto* __restrict__ old_cumAdt_w = reinterpret_cast(params.old_cumAdt); + int64_t const ca_w_base = cache_slot * params.old_cumAdt_stride_seq + buf_write * params.old_cumAdt_stride_dbuf + + head * params.old_cumAdt_stride_head; + old_cumAdt_w[ca_w_base + write_offset + lane] = smem.cumAdt[lane]; + } } -} // namespace flashinfer::mamba::checkpointing +} // namespace flashinfer::mamba::checkpointing -#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_CUH_ +#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh index 6733919d477b..3322e626ad9d 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh @@ -21,7 +21,8 @@ #include "kernel_checkpointing_ssu_common.cuh" -namespace flashinfer::mamba::checkpointing { +namespace flashinfer::mamba::checkpointing +{ // ============================================================================= // 8-bit chain-rewrite storage (sibling of CheckpointingSsuStorage) @@ -40,65 +41,64 @@ namespace flashinfer::mamba::checkpointing { // state) are byte-for-byte identical to the generic struct — Phase 0/1 // helpers (`compute_CB_scaled_2warp`, B/C/x/z loaders, etc.) are templated on // `SmemT` and read these by name, so they work unchanged. -template -struct CheckpointingSsuStorage8bit { - using state_t = state_t_; - static_assert(sizeof(state_t) == 1, "CheckpointingSsuStorage8bit requires a 1-byte state_t"); - - static constexpr int NPREDICTED = NPREDICTED_; - static constexpr int MAX_WINDOW = MAX_WINDOW_; - static constexpr int D_SMEM_COLS = next_multiple_of::ATOM_COLS>(D_PER_CTA); - static constexpr int NPREDICTED_PAD_MMA_M = next_multiple_of(NPREDICTED); - static constexpr int NPREDICTED_PAD_MMA_N = next_multiple_of(NPREDICTED); - static constexpr int MAX_WINDOW_PAD_MMA_K = next_multiple_of(MAX_WINDOW); - static constexpr int NPREDICTED_SWIZZLE_R = - next_multiple_of::ATOM_ROWS>(NPREDICTED); - static constexpr int CB_ROW_STRIDE = SmemSwizzle::ATOM_COLS; - - // Shared Phase 0/1 buffers (same shape/swizzle as `CheckpointingSsuStorage`). - alignas(16) input_t CB_scaled[NPREDICTED_PAD_MMA_M * CB_ROW_STRIDE]; - alignas(16) input_t B[NPREDICTED_PAD_MMA_N * DSTATE]; - alignas(16) input_t C[NPREDICTED_SWIZZLE_R * DSTATE]; - alignas(16) input_t x[NPREDICTED_PAD_MMA_M * D_SMEM_COLS]; - alignas(16) input_t z[NPREDICTED_SWIZZLE_R * D_SMEM_COLS]; - alignas(16) input_t old_x[MAX_WINDOW_PAD_MMA_K * D_SMEM_COLS]; - alignas(16) input_t old_B[MAX_WINDOW_PAD_MMA_K * DSTATE]; - - float old_dt[MAX_WINDOW]; - float old_cumAdt[MAX_WINDOW]; - float dt_proc[NPREDICTED]; - float cumAdt[NPREDICTED]; - // decay[t] = exp(cumAdt[t]) — precomputed at Phase 0 alongside cumAdt so the - // output decay broadcast in `compute_output_8bit` is a plain LDS instead of a - // per-element __expf. Each of the 4 warps redundantly writes the same values - // (same pattern as cumAdt); cross-warp visibility comes via the kernel's - // existing __syncthreads before compute_output_8bit. - float decay[NPREDICTED]; - - // state — int8 input, only LDS'd in the single replay pass. After replay - // completes its dequant + matmul into the C-frag, smem.state is dead — so - // `output_transpose` could in principle alias it (8 KB int8 vs 2 KB bf16 - // overlap easily), but for clarity we keep them separate; alias is a - // Phase-4 micro-optimization. - alignas(16) state_t state[D_PER_CTA * DSTATE]; - - // output_transpose — physical (NPREDICTED_PAD_MMA_M, OUTPUT_TRANSPOSE_ROW_STRIDE) - // input_t scratch buffer with PADDED row stride for bank-conflict-free per-thread - // STS + 16-byte-aligned cooperative LDS.128. Used by `compute_output_int8` to - // flip the per-warp `frag_y_DxT[D, T]` register layout into `(T, D)` gmem order. - // - // Row stride: D_PER_CTA + 8 = 72 bf16 elts = 144 bytes. The 8-elt (16-byte) pad - // gives: - // - 144 % 16 == 0 → LDS.128 / STG.128 stays 16-byte aligned across all rows. - // - 144 / 4 % 32 == 4 → adjacent t-rows shift bank assignment by 4 banks. - // For the m16n8 partition_C STS pattern (per-elt: 4 lanes write at fixed d, - // t ∈ {0, 2, 4, 6} → banks {0, 4, 8, 12} on the padded layout — all distinct, - // no conflicts), the padded layout cuts STS bank conflicts from ~63% of - // wavefronts (NCU v16.0) down to 0%. - // Volume: 16 × 72 × 2 B = 2.25 KB (vs unswizzled 2 KB; +256 B). - static constexpr int OUTPUT_TRANSPOSE_ROW_STRIDE = D_PER_CTA + 8; - alignas(16) input_t output_transpose[NPREDICTED_PAD_MMA_M * OUTPUT_TRANSPOSE_ROW_STRIDE]; +template +struct CheckpointingSsuStorage8bit +{ + using state_t = state_t_; + static_assert(sizeof(state_t) == 1, "CheckpointingSsuStorage8bit requires a 1-byte state_t"); + + static constexpr int NPREDICTED = NPREDICTED_; + static constexpr int MAX_WINDOW = MAX_WINDOW_; + static constexpr int D_SMEM_COLS = next_multiple_of::ATOM_COLS>(D_PER_CTA); + static constexpr int NPREDICTED_PAD_MMA_M = next_multiple_of(NPREDICTED); + static constexpr int NPREDICTED_PAD_MMA_N = next_multiple_of(NPREDICTED); + static constexpr int MAX_WINDOW_PAD_MMA_K = next_multiple_of(MAX_WINDOW); + static constexpr int NPREDICTED_SWIZZLE_R = next_multiple_of::ATOM_ROWS>(NPREDICTED); + static constexpr int CB_ROW_STRIDE = SmemSwizzle::ATOM_COLS; + + // Shared Phase 0/1 buffers (same shape/swizzle as `CheckpointingSsuStorage`). + alignas(16) input_t CB_scaled[NPREDICTED_PAD_MMA_M * CB_ROW_STRIDE]; + alignas(16) input_t B[NPREDICTED_PAD_MMA_N * DSTATE]; + alignas(16) input_t C[NPREDICTED_SWIZZLE_R * DSTATE]; + alignas(16) input_t x[NPREDICTED_PAD_MMA_M * D_SMEM_COLS]; + alignas(16) input_t z[NPREDICTED_SWIZZLE_R * D_SMEM_COLS]; + alignas(16) input_t old_x[MAX_WINDOW_PAD_MMA_K * D_SMEM_COLS]; + alignas(16) input_t old_B[MAX_WINDOW_PAD_MMA_K * DSTATE]; + + float old_dt[MAX_WINDOW]; + float old_cumAdt[MAX_WINDOW]; + float dt_proc[NPREDICTED]; + float cumAdt[NPREDICTED]; + // decay[t] = exp(cumAdt[t]) — precomputed at Phase 0 alongside cumAdt so the + // output decay broadcast in `compute_output_8bit` is a plain LDS instead of a + // per-element __expf. Each of the 4 warps redundantly writes the same values + // (same pattern as cumAdt); cross-warp visibility comes via the kernel's + // existing __syncthreads before compute_output_8bit. + float decay[NPREDICTED]; + + // state — int8 input, only LDS'd in the single replay pass. After replay + // completes its dequant + matmul into the C-frag, smem.state is dead — so + // `output_transpose` could in principle alias it (8 KB int8 vs 2 KB bf16 + // overlap easily), but for clarity we keep them separate; alias is a + // Phase-4 micro-optimization. + alignas(16) state_t state[D_PER_CTA * DSTATE]; + + // output_transpose — physical (NPREDICTED_PAD_MMA_M, OUTPUT_TRANSPOSE_ROW_STRIDE) + // input_t scratch buffer with PADDED row stride for bank-conflict-free per-thread + // STS + 16-byte-aligned cooperative LDS.128. Used by `compute_output_int8` to + // flip the per-warp `frag_y_DxT[D, T]` register layout into `(T, D)` gmem order. + // + // Row stride: D_PER_CTA + 8 = 72 bf16 elts = 144 bytes. The 8-elt (16-byte) pad + // gives: + // - 144 % 16 == 0 → LDS.128 / STG.128 stays 16-byte aligned across all rows. + // - 144 / 4 % 32 == 4 → adjacent t-rows shift bank assignment by 4 banks. + // For the m16n8 partition_C STS pattern (per-elt: 4 lanes write at fixed d, + // t ∈ {0, 2, 4, 6} → banks {0, 4, 8, 12} on the padded layout — all distinct, + // no conflicts), the padded layout cuts STS bank conflicts from ~63% of + // wavefronts (NCU v16.0) down to 0%. + // Volume: 16 × 72 × 2 B = 2.25 KB (vs unswizzled 2 KB; +256 B). + static constexpr int OUTPUT_TRANSPOSE_ROW_STRIDE = D_PER_CTA + 8; + alignas(16) input_t output_transpose[NPREDICTED_PAD_MMA_M * OUTPUT_TRANSPOSE_ROW_STRIDE]; }; // ============================================================================= @@ -110,43 +110,52 @@ struct CheckpointingSsuStorage8bit { // bound, and (c) packing/unpacking the byte from a u16 `Pair`. template -__device__ __forceinline__ uint8_t state_byte_of(state_t v) { - if constexpr (std::is_same_v) { - return static_cast(static_cast(v)); - } else { - static_assert(std::is_same_v, - "8-bit state_t must be int8_t or __nv_fp8_e4m3"); - return reinterpret_cast<__nv_fp8_storage_t const&>(v); - } +__device__ __forceinline__ uint8_t state_byte_of(state_t v) +{ + if constexpr (std::is_same_v) + { + return static_cast(static_cast(v)); + } + else + { + static_assert(std::is_same_v, "8-bit state_t must be int8_t or __nv_fp8_e4m3"); + return reinterpret_cast<__nv_fp8_storage_t const&>(v); + } } // fp32 → state_t with RN + saturate. Single-element scalar — the kernel's // smem layout writes pairs as u16, so per-element conversion fits the // per-thread fragment topology directly. template -__device__ __forceinline__ state_t encode_rn_8bit(float x) { - if constexpr (std::is_same_v) { - return conversion::cvt_rni_sat_s8(x); - } else { - static_assert(std::is_same_v, - "8-bit state_t must be int8_t or __nv_fp8_e4m3"); - // cuda_fp8 ctor compiles to `cvt.rn.satfinite.e4m3.f32` on sm_89+. - return __nv_fp8_e4m3(x); - } +__device__ __forceinline__ state_t encode_rn_8bit(float x) +{ + if constexpr (std::is_same_v) + { + return conversion::cvt_rni_sat_s8(x); + } + else + { + static_assert(std::is_same_v, "8-bit state_t must be int8_t or __nv_fp8_e4m3"); + // cuda_fp8 ctor compiles to `cvt.rn.satfinite.e4m3.f32` on sm_89+. + return __nv_fp8_e4m3(x); + } } // Per-state-dtype symmetric clip / encode-scale denominator. // int8: ±127 (matches Triton reference, leaves -128 unused) // fp8_e4m3fn: ±448 (max finite e4m3 value) template -__device__ __forceinline__ constexpr float quant_max_8bit() { - if constexpr (std::is_same_v) { - return 127.0f; - } else { - static_assert(std::is_same_v, - "8-bit state_t must be int8_t or __nv_fp8_e4m3"); - return 448.0f; - } +__device__ __forceinline__ constexpr float quant_max_8bit() +{ + if constexpr (std::is_same_v) + { + return 127.0f; + } + else + { + static_assert(std::is_same_v, "8-bit state_t must be int8_t or __nv_fp8_e4m3"); + return 448.0f; + } } // SM80 m16n8k16 C-frag → A-frag layout reshape for chained mma (state → @@ -168,23 +177,20 @@ __device__ __forceinline__ constexpr float quant_max_8bit() { // Output layout: ((2, 2, 2), MMA_M, MMA_N / 2) — m16n8k16 A-frag, // MMA_K = MMA_N / 2 template -__forceinline__ __device__ auto convert_layout_acc_Aregs_sm80(Layout acc_layout) { - using namespace cute; - using X = Underscore; - static_assert(decltype(size<0, 0>(acc_layout))::value == 2, - "C-frag inner mode must be (col_pair=2, row_pair=2)"); - static_assert(decltype(size<0, 1>(acc_layout))::value == 2, - "C-frag inner mode must be (col_pair=2, row_pair=2)"); - static_assert(decltype(rank(acc_layout))::value == 3, - "C-frag must be rank-3 ((C0,C1), MMA_M, MMA_N)"); - static_assert(decltype(rank(get<0>(acc_layout)))::value == 2, - "SM80 m16n8 C-frag inner is rank-2 (no inner stride mode like SM90)"); - // logical_divide the outer MMA_N axis by 2 → ((2, 2), MMA_M, (2, MMA_N/2)) - auto l = logical_divide(acc_layout, Shape{}); - return make_layout( - make_layout(get<0, 0>(l), get<0, 1>(l), get<2, 0>(l)), // ((col_pair, row_pair, k_half)) - get<1>(l), // MMA_M - get<2, 1>(l)); // MMA_K = MMA_N / 2 +__forceinline__ __device__ auto convert_layout_acc_Aregs_sm80(Layout acc_layout) +{ + using namespace cute; + using X = Underscore; + static_assert(decltype(size<0, 0>(acc_layout))::value == 2, "C-frag inner mode must be (col_pair=2, row_pair=2)"); + static_assert(decltype(size<0, 1>(acc_layout))::value == 2, "C-frag inner mode must be (col_pair=2, row_pair=2)"); + static_assert(decltype(rank(acc_layout))::value == 3, "C-frag must be rank-3 ((C0,C1), MMA_M, MMA_N)"); + static_assert(decltype(rank(get<0>(acc_layout)))::value == 2, + "SM80 m16n8 C-frag inner is rank-2 (no inner stride mode like SM90)"); + // logical_divide the outer MMA_N axis by 2 → ((2, 2), MMA_M, (2, MMA_N/2)) + auto l = logical_divide(acc_layout, Shape{}); + return make_layout(make_layout(get<0, 0>(l), get<0, 1>(l), get<2, 0>(l)), // ((col_pair, row_pair, k_half)) + get<1>(l), // MMA_M + get<2, 1>(l)); // MMA_K = MMA_N / 2 } // ============================================================================= @@ -256,273 +262,264 @@ __forceinline__ __device__ auto convert_layout_acc_Aregs_sm80(Layout acc_layout) // so chain matmul-3 sees each warp's own data without cross-warp sync. // The caller's single __syncthreads between all replay passes and // compute_output_8bit provides smem.CB_scaled / smem.x / smem.z visibility. -template -__device__ __forceinline__ void replay_state_mma_8bit_chain( - SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int prev_k, int d_tile, - int64_t cache_slot, int head, bool must_checkpoint, FragYDxT& frag_y_DxT, - float (&encode_scale_per_row_out)[2], float (&total_scale_out)[2]) { - using namespace cute; - static_assert(sizeof(input_t) == 2, "replay_state_mma_8bit_chain requires 2-byte input_t"); - static_assert(sizeof(state_t) == 1, - "replay_state_mma_8bit_chain is for 1-byte state_t (int8/fp8) only"); - static_assert(D_PER_CTA == 64, - "replay_state_mma_8bit_chain requires D_PER_CTA == 64 (M-shard, per-warp M=16)."); - - constexpr int NUM_WARPS = 4; - constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; // 16 - static_assert(M_PER_WARP == MMA_prop::M, "Per-warp M must equal m16n8 atom M (=16)"); - - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - int const tid = warp * warpSize + lane; - - // Atom-K dispatch: K_BIG=16 (default), K_SMALL=8 if MAX_WINDOW ≤ 8. - using MmaAtomReplayType = std::conditional_t; - using LdsmA = std::conditional_t; - using LdsmB = std::conditional_t; - - // Replay TiledMma: M-shard, 4 warps along M, 1 along N. Output is - // ((2,2), 1, NUM_N_PASSES) per thread of fp32 (or bf16 view for new_state). - auto tiled_mma_replay = - make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_replay = tiled_mma_replay.get_slice(tid); - - // Chain TiledMma: m16n8k16 (always K_BIG=16 since K=DSTATE/16 atoms ≥ 1), - // same M-shard layout as replay. M_per_warp=16 (1 m-atom), - // N=NPREDICTED_PAD_MMA_M (T_pad, ≤ 16 = up to 2 n-atoms per warp), K=DSTATE. - auto tiled_mma_chain = - make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_chain = tiled_mma_chain.get_slice(tid); - - constexpr int N_PER_PASS = MMA_prop::N; // 8 - constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; // 16 - constexpr int FRAG_SIZE = 4; - constexpr int D_ROWS_PER_THREAD = 2; - constexpr float QUANT_MAX = quant_max_8bit(); - - float const total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - float const total_decay = (prev_k > 0) ? __expf(total_cumAdt) : 1.f; - - int const lane_d = lane / 4; - int const warp_d_base = warp * M_PER_WARP; - - // ── Per-row decode_scale for state init. - auto const* __restrict__ state_scale_ptr = reinterpret_cast(params.state_scale); - int64_t const state_scale_base = cache_slot * params.state_scale_stride_seq + - (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; - float decode_scale_in[D_ROWS_PER_THREAD]; - decode_scale_in[0] = state_scale_ptr[state_scale_base + warp_d_base + lane_d]; - decode_scale_in[1] = state_scale_ptr[state_scale_base + warp_d_base + lane_d + 8]; - float total_scale[D_ROWS_PER_THREAD]; - total_scale[0] = decode_scale_in[0] * total_decay; - total_scale[1] = decode_scale_in[1] * total_decay; - - // ── A operand (replay): old_x [MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS] → LDSM_T. - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - auto layout_A_full = - make_swizzled_layout_rc_transpose(); - Tensor smem_A_full = make_tensor( - make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); - Tensor smem_A = local_tile(smem_A_full, make_shape(Int{}, Int{}), - make_coord(_0{}, _0{})); - - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma_replay); - auto s2r_thr_A = s2r_A.get_slice(tid); - Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); - Tensor frag_A_replay = thr_mma_replay.partition_fragment_A(make_tensor( - (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); - Tensor frag_A_replay_view = s2r_thr_A.retile_D(frag_A_replay); - cute::copy(s2r_A, smem_A_s2r, frag_A_replay_view); - - // ── Bake dB coefficients into frag_A once (8 scale ops), replacing 16× - // per-N-pass compute_dB_scaling on frag_B (64 scale ops). - // dB coefficients c[k] baked into frag_A once, replacing per-N-pass B scaling. - apply_dA_coeff(frag_A_replay, smem, total_cumAdt, prev_k, lane); - - // ── B operand (replay): old_B per-pass. - auto layout_B_replay = make_swizzled_layout_rc_transpose(); - Tensor smem_B_full = make_tensor( - make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B_replay); - auto s2r_B_replay = make_tiled_copy_B(Copy_Atom{}, tiled_mma_replay); - auto s2r_thr_B_replay = s2r_B_replay.get_slice(tid); - - // ── State: 1-byte input pointer + manual swizzle offsets (read in BOTH passes). - // Drop bf16 new_state staging — replay's fp32 frag flows directly into the - // register-resident `new_state` tensor below. - state_t* state_base = reinterpret_cast(smem.state); - - // Manual swizzle offsets for m16n8 C-fragment layout (1-byte Swizzle<3,4,3>). - // off = row * 128 + (col ^ ((row & 7) << 4)). - // row_hi = row_lo + 8; (row+8)&7 == row&7 ⇒ off_hi = off_lo + 1024. - // Fragment col within each N_PER_PASS=8 tile: (lane % 4) * 2. - int const row_lo = warp_d_base + lane_d; - int const frag_col_base = (lane & 3) << 1; - int const state_base_lo = row_lo << 7; // row_lo * DSTATE - int const state_xor = (row_lo & 7) << 4; - - float per_thread_amax[D_ROWS_PER_THREAD] = {0.f, 0.f}; - - // No __syncthreads here — smem.C is redundantly loaded by all 4 warps - // (each warp sees its own cp.async via __syncwarp in load_data). Cross-warp - // visibility for smem.CB_scaled / smem.x / smem.z is established by the - // caller's __syncthreads between this function and compute_output_8bit. - - // ── smem.C view + B-operand TiledCopy for chain matmul-3 (hoisted before - // the loop; same view per K-pair, B sliced per K-atom inside the loop). - auto layout_C_swz = - make_aliased_swizzled_layout_rc(); - Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), - layout_C_swz); - auto s2r_B_chain = - make_tiled_copy_B(Copy_Atom{}, tiled_mma_chain); - auto s2r_thr_B_chain = s2r_B_chain.get_slice(tid); - - constexpr int NUM_K_PAIRS = NUM_N_PASSES / 2; // 8 for DSTATE=128 - static_assert(NUM_N_PASSES % 2 == 0, "Per-K-pair fusion requires NUM_N_PASSES to be even"); - static_assert(MMA_prop::K_BIG == 16, "Chain mma assumes m16n8k16 K-atom = 16"); - - // ════════════════════════════════════════════════════════════════════════ - // PASS 1 — fused replay + chain matmul-3 (per-K-pair): - // For each kpair ∈ [0, NUM_K_PAIRS): - // - Run 2 replay HMMAs (N-passes 2*kpair, 2*kpair+1) → fp32 frag_h × 2. - // - Update per-thread amax (bit-exact fp32). - // - Pack each pair's 4 fp32 → 4 bf16 into a tiny K-atom-sized A frag - // (`a_kpair` shape ((2,2,2), 1, 1) of bf16 = 8 elts/thread = 4 - // 32-bit regs). Linear positions [local_n*4 .. local_n*4+3] - // within `a_kpair` map to the m16n8k16 A operand's (kh=local_n) - // slice — proven by the linear-index identity in the deleted - // `new_state`-tensor comment above. - // - LDS one K-atom of B (smem.C[T_pad, kpair*16..+16]) into a - // similarly small `b_kpair` frag (4 32-bit regs / thread). - // - `cute::gemm` accumulates one K-atom into `frag_y_DxT`. - // - Both `a_kpair` and `b_kpair` go out of scope at iter end → the - // compiler frees those ~8 32-bit regs/thread for the next iter. - // Net: register footprint drops from the 32 regs of the old register- - // resident `new_state` array (held across the whole loop) to ~8 regs in - // flight. Frees ~24 regs/thread → potentially +1-2 blocks/SM occupancy. - // ════════════════════════════════════════════════════════════════════════ +template +__device__ __forceinline__ void replay_state_mma_8bit_chain(SmemT& smem, CheckpointingSsuParams const& params, int warp, + int lane, int prev_k, int d_tile, int64_t cache_slot, int head, bool must_checkpoint, FragYDxT& frag_y_DxT, + float (&encode_scale_per_row_out)[2], float (&total_scale_out)[2]) +{ + using namespace cute; + static_assert(sizeof(input_t) == 2, "replay_state_mma_8bit_chain requires 2-byte input_t"); + static_assert(sizeof(state_t) == 1, "replay_state_mma_8bit_chain is for 1-byte state_t (int8/fp8) only"); + static_assert(D_PER_CTA == 64, "replay_state_mma_8bit_chain requires D_PER_CTA == 64 (M-shard, per-warp M=16)."); + + constexpr int NUM_WARPS = 4; + constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; // 16 + static_assert(M_PER_WARP == MMA_prop::M, "Per-warp M must equal m16n8 atom M (=16)"); + + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + int const tid = warp * warpSize + lane; + + // Atom-K dispatch: K_BIG=16 (default), K_SMALL=8 if MAX_WINDOW ≤ 8. + using MmaAtomReplayType + = std::conditional_t; + using LdsmA = std::conditional_t; + using LdsmB = std::conditional_t; + + // Replay TiledMma: M-shard, 4 warps along M, 1 along N. Output is + // ((2,2), 1, NUM_N_PASSES) per thread of fp32 (or bf16 view for new_state). + auto tiled_mma_replay = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_replay = tiled_mma_replay.get_slice(tid); + + // Chain TiledMma: m16n8k16 (always K_BIG=16 since K=DSTATE/16 atoms ≥ 1), + // same M-shard layout as replay. M_per_warp=16 (1 m-atom), + // N=NPREDICTED_PAD_MMA_M (T_pad, ≤ 16 = up to 2 n-atoms per warp), K=DSTATE. + auto tiled_mma_chain = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_chain = tiled_mma_chain.get_slice(tid); + + constexpr int N_PER_PASS = MMA_prop::N; // 8 + constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; // 16 + constexpr int FRAG_SIZE = 4; + constexpr int D_ROWS_PER_THREAD = 2; + constexpr float QUANT_MAX = quant_max_8bit(); + + float const total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + float const total_decay = (prev_k > 0) ? __expf(total_cumAdt) : 1.f; + + int const lane_d = lane / 4; + int const warp_d_base = warp * M_PER_WARP; + + // ── Per-row decode_scale for state init. + auto const* __restrict__ state_scale_ptr = reinterpret_cast(params.state_scale); + int64_t const state_scale_base + = cache_slot * params.state_scale_stride_seq + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; + float decode_scale_in[D_ROWS_PER_THREAD]; + decode_scale_in[0] = state_scale_ptr[state_scale_base + warp_d_base + lane_d]; + decode_scale_in[1] = state_scale_ptr[state_scale_base + warp_d_base + lane_d + 8]; + float total_scale[D_ROWS_PER_THREAD]; + total_scale[0] = decode_scale_in[0] * total_decay; + total_scale[1] = decode_scale_in[1] * total_decay; + + // ── A operand (replay): old_x [MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS] → LDSM_T. + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + auto layout_A_full = make_swizzled_layout_rc_transpose(); + Tensor smem_A_full + = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); + Tensor smem_A + = local_tile(smem_A_full, make_shape(Int{}, Int{}), make_coord(_0{}, _0{})); + + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma_replay); + auto s2r_thr_A = s2r_A.get_slice(tid); + Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); + Tensor frag_A_replay = thr_mma_replay.partition_fragment_A( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + Tensor frag_A_replay_view = s2r_thr_A.retile_D(frag_A_replay); + cute::copy(s2r_A, smem_A_s2r, frag_A_replay_view); + + // ── Bake dB coefficients into frag_A once (8 scale ops), replacing 16× + // per-N-pass compute_dB_scaling on frag_B (64 scale ops). + // dB coefficients c[k] baked into frag_A once, replacing per-N-pass B scaling. + apply_dA_coeff(frag_A_replay, smem, total_cumAdt, prev_k, lane); + + // ── B operand (replay): old_B per-pass. + auto layout_B_replay = make_swizzled_layout_rc_transpose(); + Tensor smem_B_full + = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B_replay); + auto s2r_B_replay = make_tiled_copy_B(Copy_Atom{}, tiled_mma_replay); + auto s2r_thr_B_replay = s2r_B_replay.get_slice(tid); + + // ── State: 1-byte input pointer + manual swizzle offsets (read in BOTH passes). + // Drop bf16 new_state staging — replay's fp32 frag flows directly into the + // register-resident `new_state` tensor below. + state_t* state_base = reinterpret_cast(smem.state); + + // Manual swizzle offsets for m16n8 C-fragment layout (1-byte Swizzle<3,4,3>). + // off = row * 128 + (col ^ ((row & 7) << 4)). + // row_hi = row_lo + 8; (row+8)&7 == row&7 ⇒ off_hi = off_lo + 1024. + // Fragment col within each N_PER_PASS=8 tile: (lane % 4) * 2. + int const row_lo = warp_d_base + lane_d; + int const frag_col_base = (lane & 3) << 1; + int const state_base_lo = row_lo << 7; // row_lo * DSTATE + int const state_xor = (row_lo & 7) << 4; + + float per_thread_amax[D_ROWS_PER_THREAD] = {0.f, 0.f}; + + // No __syncthreads here — smem.C is redundantly loaded by all 4 warps + // (each warp sees its own cp.async via __syncwarp in load_data). Cross-warp + // visibility for smem.CB_scaled / smem.x / smem.z is established by the + // caller's __syncthreads between this function and compute_output_8bit. + + // ── smem.C view + B-operand TiledCopy for chain matmul-3 (hoisted before + // the loop; same view per K-pair, B sliced per K-atom inside the loop). + auto layout_C_swz = make_aliased_swizzled_layout_rc(); + Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C_swz); + auto s2r_B_chain = make_tiled_copy_B(Copy_Atom{}, tiled_mma_chain); + auto s2r_thr_B_chain = s2r_B_chain.get_slice(tid); + + constexpr int NUM_K_PAIRS = NUM_N_PASSES / 2; // 8 for DSTATE=128 + static_assert(NUM_N_PASSES % 2 == 0, "Per-K-pair fusion requires NUM_N_PASSES to be even"); + static_assert(MMA_prop::K_BIG == 16, "Chain mma assumes m16n8k16 K-atom = 16"); + + // ════════════════════════════════════════════════════════════════════════ + // PASS 1 — fused replay + chain matmul-3 (per-K-pair): + // For each kpair ∈ [0, NUM_K_PAIRS): + // - Run 2 replay HMMAs (N-passes 2*kpair, 2*kpair+1) → fp32 frag_h × 2. + // - Update per-thread amax (bit-exact fp32). + // - Pack each pair's 4 fp32 → 4 bf16 into a tiny K-atom-sized A frag + // (`a_kpair` shape ((2,2,2), 1, 1) of bf16 = 8 elts/thread = 4 + // 32-bit regs). Linear positions [local_n*4 .. local_n*4+3] + // within `a_kpair` map to the m16n8k16 A operand's (kh=local_n) + // slice — proven by the linear-index identity in the deleted + // `new_state`-tensor comment above. + // - LDS one K-atom of B (smem.C[T_pad, kpair*16..+16]) into a + // similarly small `b_kpair` frag (4 32-bit regs / thread). + // - `cute::gemm` accumulates one K-atom into `frag_y_DxT`. + // - Both `a_kpair` and `b_kpair` go out of scope at iter end → the + // compiler frees those ~8 32-bit regs/thread for the next iter. + // Net: register footprint drops from the 32 regs of the old register- + // resident `new_state` array (held across the whole loop) to ~8 regs in + // flight. Frees ~24 regs/thread → potentially +1-2 blocks/SM occupancy. + // ════════════════════════════════════════════════════════════════════════ #pragma unroll - for (int kpair = 0; kpair < NUM_K_PAIRS; ++kpair) { - // K-atom-sized A frag for chain matmul-3 (filled across the 2 N-passes). - Tensor a_kpair = thr_mma_chain.partition_fragment_A(make_tensor( - (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); - static_assert(decltype(size(a_kpair))::value == 8, - "a_kpair must hold 1 m16n8k16 K-atom of A = 8 bf16/thread"); + for (int kpair = 0; kpair < NUM_K_PAIRS; ++kpair) + { + // K-atom-sized A frag for chain matmul-3 (filled across the 2 N-passes). + Tensor a_kpair = thr_mma_chain.partition_fragment_A( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + static_assert(decltype(size(a_kpair))::value == 8, "a_kpair must hold 1 m16n8k16 K-atom of A = 8 bf16/thread"); #pragma unroll - for (int local_n = 0; local_n < 2; ++local_n) { - int const n = kpair * 2 + local_n; - int const n_base = n * N_PER_PASS; - - Tensor frag_h = thr_mma_replay.partition_fragment_C( - make_tensor((float*)0x0, make_shape(Int{}, Int{}))); - static_assert(decltype(size(frag_h))::value == FRAG_SIZE, - "FRAG_SIZE must match the partitioned C-fragment size"); - - // Zero-init accumulator — MMA from scratch, state added after. - clear(frag_h); - - // Replay B operand load. - Tensor smem_B_n = - local_tile(smem_B_full, make_tile(Int{}, Int{}), - make_coord(n, _0{})); - auto smem_B_s2r_n = s2r_thr_B_replay.partition_S(smem_B_n); - Tensor frag_B_replay = thr_mma_replay.partition_fragment_B(make_tensor( - (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); - auto frag_B_replay_view = s2r_thr_B_replay.retile_D(frag_B_replay); - cute::copy(s2r_B_replay, smem_B_s2r_n, frag_B_replay_view); - - // Replay HMMA: frag_h = frag_A_scaled @ frag_B (c[k] baked into A). - cute::gemm(tiled_mma_replay, frag_h, frag_A_replay, frag_B_replay, frag_h); - - { - int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); - Pair const p0 = *reinterpret_cast const*>(&state_base[off_lo]); - frag_h(0) += toFloat(p0[Int<0>{}]) * total_scale[0]; - frag_h(1) += toFloat(p0[Int<1>{}]) * total_scale[0]; - Pair const p1 = - *reinterpret_cast const*>(&state_base[off_lo + 1024]); - frag_h(2) += toFloat(p1[Int<0>{}]) * total_scale[1]; - frag_h(3) += toFloat(p1[Int<1>{}]) * total_scale[1]; - } - - // Update amax (fp32, bit-exact) AND pack 4 fp32 → 4 bf16 into a_kpair - // at offset local_n*4 (matches A-frag's (kh=local_n) slice). + for (int local_n = 0; local_n < 2; ++local_n) + { + int const n = kpair * 2 + local_n; + int const n_base = n * N_PER_PASS; + + Tensor frag_h = thr_mma_replay.partition_fragment_C( + make_tensor((float*) 0x0, make_shape(Int{}, Int{}))); + static_assert( + decltype(size(frag_h))::value == FRAG_SIZE, "FRAG_SIZE must match the partitioned C-fragment size"); + + // Zero-init accumulator — MMA from scratch, state added after. + clear(frag_h); + + // Replay B operand load. + Tensor smem_B_n = local_tile( + smem_B_full, make_tile(Int{}, Int{}), make_coord(n, _0{})); + auto smem_B_s2r_n = s2r_thr_B_replay.partition_S(smem_B_n); + Tensor frag_B_replay = thr_mma_replay.partition_fragment_B( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + auto frag_B_replay_view = s2r_thr_B_replay.retile_D(frag_B_replay); + cute::copy(s2r_B_replay, smem_B_s2r_n, frag_B_replay_view); + + // Replay HMMA: frag_h = frag_A_scaled @ frag_B (c[k] baked into A). + cute::gemm(tiled_mma_replay, frag_h, frag_A_replay, frag_B_replay, frag_h); + + { + int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); + Pair const p0 = *reinterpret_cast const*>(&state_base[off_lo]); + frag_h(0) += toFloat(p0[Int<0>{}]) * total_scale[0]; + frag_h(1) += toFloat(p0[Int<1>{}]) * total_scale[0]; + Pair const p1 = *reinterpret_cast const*>(&state_base[off_lo + 1024]); + frag_h(2) += toFloat(p1[Int<0>{}]) * total_scale[1]; + frag_h(3) += toFloat(p1[Int<1>{}]) * total_scale[1]; + } + + // Update amax (fp32, bit-exact) AND pack 4 fp32 → 4 bf16 into a_kpair + // at offset local_n*4 (matches A-frag's (kh=local_n) slice). #pragma unroll - for (int i = 0; i < FRAG_SIZE; i += 2) { - int const d_idx = i / 2; - float const a0 = fabsf(frag_h(i)); - float const a1 = fabsf(frag_h(i + 1)); - per_thread_amax[d_idx] = fmaxf(per_thread_amax[d_idx], fmaxf(a0, a1)); - - Pair const q = - pack_float2(make_float2(frag_h(i), frag_h(i + 1))); - *reinterpret_cast*>(&a_kpair(local_n * FRAG_SIZE + i)) = q; - } + for (int i = 0; i < FRAG_SIZE; i += 2) + { + int const d_idx = i / 2; + float const a0 = fabsf(frag_h(i)); + float const a1 = fabsf(frag_h(i + 1)); + per_thread_amax[d_idx] = fmaxf(per_thread_amax[d_idx], fmaxf(a0, a1)); + + Pair const q + = pack_float2(make_float2(frag_h(i), frag_h(i + 1))); + *reinterpret_cast*>(&a_kpair(local_n * FRAG_SIZE + i)) = q; + } + } + + // ── B operand for chain matmul-3 K-atom: smem.C[T_pad, kpair*16..+16] ── + Tensor smem_C_k = local_tile( + smem_C, make_tile(Int{}, Int{}), make_coord(_0{}, kpair)); + auto smem_C_k_s2r = s2r_thr_B_chain.partition_S(smem_C_k); + Tensor b_kpair = thr_mma_chain.partition_fragment_B( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + auto b_kpair_view = s2r_thr_B_chain.retile_D(b_kpair); + cute::copy(s2r_B_chain, smem_C_k_s2r, b_kpair_view); + + // Single K-atom chain matmul-3: frag_y_DxT += a_kpair @ b_kpair + // frag_y_DxT (pre-zeroed by caller) accumulates across all 8 K-atoms. + cute::gemm(tiled_mma_chain, frag_y_DxT, a_kpair, b_kpair, frag_y_DxT); } - // ── B operand for chain matmul-3 K-atom: smem.C[T_pad, kpair*16..+16] ── - Tensor smem_C_k = - local_tile(smem_C, make_tile(Int{}, Int{}), - make_coord(_0{}, kpair)); - auto smem_C_k_s2r = s2r_thr_B_chain.partition_S(smem_C_k); - Tensor b_kpair = thr_mma_chain.partition_fragment_B( - make_tensor((MMA_prop::operand_t*)0x0, - make_shape(Int{}, Int{}))); - auto b_kpair_view = s2r_thr_B_chain.retile_D(b_kpair); - cute::copy(s2r_B_chain, smem_C_k_s2r, b_kpair_view); - - // Single K-atom chain matmul-3: frag_y_DxT += a_kpair @ b_kpair - // frag_y_DxT (pre-zeroed by caller) accumulates across all 8 K-atoms. - cute::gemm(tiled_mma_chain, frag_y_DxT, a_kpair, b_kpair, frag_y_DxT); - } - - // ── Warp-local amax reduce (Layout<_4,_1> → fully warp-local; no atomics). + // ── Warp-local amax reduce (Layout<_4,_1> → fully warp-local; no atomics). #pragma unroll - for (int i = 0; i < D_ROWS_PER_THREAD; ++i) { - per_thread_amax[i] = fmaxf(per_thread_amax[i], - __shfl_xor_sync(constants::MASK_ALL_LANES, per_thread_amax[i], 1)); - per_thread_amax[i] = fmaxf(per_thread_amax[i], - __shfl_xor_sync(constants::MASK_ALL_LANES, per_thread_amax[i], 2)); - } - - // ── encode scale (Triton fall-through for amax==0). - // decode_scale = 1 / encode_scale (mathematically: decode = amax/QUANT_MAX, - // encode = QUANT_MAX/amax, so decode = 1/encode; and when amax==0 both fall - // through to 1.f → 1/1 == 1). Computed inline at the STG below — keeping - // only `encode_scale_per_row` in regs saves 2 fp32 regs across PASS 2. - float encode_scale_per_row[D_ROWS_PER_THREAD]; + for (int i = 0; i < D_ROWS_PER_THREAD; ++i) + { + per_thread_amax[i] + = fmaxf(per_thread_amax[i], __shfl_xor_sync(constants::MASK_ALL_LANES, per_thread_amax[i], 1)); + per_thread_amax[i] + = fmaxf(per_thread_amax[i], __shfl_xor_sync(constants::MASK_ALL_LANES, per_thread_amax[i], 2)); + } + + // ── encode scale (Triton fall-through for amax==0). + // decode_scale = 1 / encode_scale (mathematically: decode = amax/QUANT_MAX, + // encode = QUANT_MAX/amax, so decode = 1/encode; and when amax==0 both fall + // through to 1.f → 1/1 == 1). Computed inline at the STG below — keeping + // only `encode_scale_per_row` in regs saves 2 fp32 regs across PASS 2. + float encode_scale_per_row[D_ROWS_PER_THREAD]; #pragma unroll - for (int i = 0; i < D_ROWS_PER_THREAD; ++i) { - float const a = per_thread_amax[i]; - encode_scale_per_row[i] = (a == 0.f) ? 1.f : (QUANT_MAX / a); - } - - // ── STG decode_scale (one writer per (cache, head, d_row)). - if (must_checkpoint && (lane & 3) == 0) { - auto* __restrict__ state_scale_w = reinterpret_cast(params.state_scale); + for (int i = 0; i < D_ROWS_PER_THREAD; ++i) + { + float const a = per_thread_amax[i]; + encode_scale_per_row[i] = (a == 0.f) ? 1.f : (QUANT_MAX / a); + } + + // ── STG decode_scale (one writer per (cache, head, d_row)). + if (must_checkpoint && (lane & 3) == 0) + { + auto* __restrict__ state_scale_w = reinterpret_cast(params.state_scale); #pragma unroll - for (int i = 0; i < D_ROWS_PER_THREAD; ++i) { - int const d_row_in_atom = lane_d + (i & 1) * 8; - int const d_row = warp_d_base + d_row_in_atom; - state_scale_w[state_scale_base + d_row] = 1.f / encode_scale_per_row[i]; + for (int i = 0; i < D_ROWS_PER_THREAD; ++i) + { + int const d_row_in_atom = lane_d + (i & 1) * 8; + int const d_row = warp_d_base + d_row_in_atom; + state_scale_w[state_scale_base + d_row] = 1.f / encode_scale_per_row[i]; + } } - } - - // Hand `encode_scale_per_row` AND `total_scale` (= OLD decode_scale_in × - // total_decay) to the caller so PASS 2 (encode replay-again) can: - // - dequantize the OLD int8 state with the OLD decode_scale (NOT the NEW - // one we just STG'd above — re-reading params.state_scale in PASS 2 - // would pick up the new value and corrupt the encode), and - // - encode the NEW state with the right encode_scale = 127 / amax. - encode_scale_per_row_out[0] = encode_scale_per_row[0]; - encode_scale_per_row_out[1] = encode_scale_per_row[1]; - total_scale_out[0] = total_scale[0]; - total_scale_out[1] = total_scale[1]; + + // Hand `encode_scale_per_row` AND `total_scale` (= OLD decode_scale_in × + // total_decay) to the caller so PASS 2 (encode replay-again) can: + // - dequantize the OLD int8 state with the OLD decode_scale (NOT the NEW + // one we just STG'd above — re-reading params.state_scale in PASS 2 + // would pick up the new value and corrupt the encode), and + // - encode the NEW state with the right encode_scale = 127 / amax. + encode_scale_per_row_out[0] = encode_scale_per_row[0]; + encode_scale_per_row_out[1] = encode_scale_per_row[1]; + total_scale_out[0] = total_scale[0]; + total_scale_out[1] = total_scale[1]; } // ───────────────────────────────────────────────────────────────────────── @@ -542,182 +539,181 @@ __device__ __forceinline__ void replay_state_mma_8bit_chain( // The setup (TiledMma, frag_A_replay, smem layouts) is duplicated // from `replay_state_mma_8bit_chain` — separate stack frame keeps register // allocation simple and avoids cross-function lifetime tracking. -template -__device__ __forceinline__ void encode_state_replay_8bit( - SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int prev_k, int d_tile, - int64_t cache_slot, int head, float const (&encode_scale_per_row)[2], - float const (&total_scale)[2], int64_t rand_seed, int64_t state_ptr_offset) { - using namespace cute; - static_assert(sizeof(input_t) == 2, "encode_state_replay_8bit requires 2-byte input_t"); - static_assert(sizeof(state_t) == 1, - "encode_state_replay_8bit is for 1-byte state_t (int8/fp8) only"); - static_assert(D_PER_CTA == 64, - "encode_state_replay_8bit requires D_PER_CTA == 64 (M-shard, per-warp M=16)."); - - constexpr int NUM_WARPS = 4; - constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; - static_assert(M_PER_WARP == MMA_prop::M, "Per-warp M must equal m16n8 atom M (=16)"); - - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - int const tid = warp * warpSize + lane; - - using MmaAtomReplayType = std::conditional_t; - using LdsmA = std::conditional_t; - using LdsmB = std::conditional_t; - - auto tiled_mma_replay = - make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_replay = tiled_mma_replay.get_slice(tid); - - constexpr int N_PER_PASS = MMA_prop::N; - constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; - constexpr int FRAG_SIZE = 4; - constexpr int D_ROWS_PER_THREAD = 2; - - float const total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - - // total_scale (= OLD decode_scale_in × total_decay) was computed in PASS 1 - // and is passed in by reference. We MUST NOT re-load decode_scale_in from - // params.state_scale here — by the time PASS 2 runs, PASS 1 has already - // STG'd the NEW decode_scale to that same gmem location. - - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - auto layout_A_full = - make_swizzled_layout_rc_transpose(); - Tensor smem_A_full = make_tensor( - make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); - Tensor smem_A = local_tile(smem_A_full, make_shape(Int{}, Int{}), - make_coord(_0{}, _0{})); - - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma_replay); - auto s2r_thr_A = s2r_A.get_slice(tid); - Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); - Tensor frag_A_replay = thr_mma_replay.partition_fragment_A(make_tensor( - (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); - Tensor frag_A_replay_view = s2r_thr_A.retile_D(frag_A_replay); - cute::copy(s2r_A, smem_A_s2r, frag_A_replay_view); - - // dB coefficients baked into frag_A (same identity as PASS 1). - apply_dA_coeff(frag_A_replay, smem, total_cumAdt, prev_k, lane); - - auto layout_B_replay = make_swizzled_layout_rc_transpose(); - Tensor smem_B_full = make_tensor( - make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B_replay); - auto s2r_B_replay = make_tiled_copy_B(Copy_Atom{}, tiled_mma_replay); - auto s2r_thr_B_replay = s2r_B_replay.get_slice(tid); - - state_t* state_base = reinterpret_cast(smem.state); - - // Manual swizzle offsets (same derivation as replay_state_mma_8bit_chain). - int const lane_d = lane / 4; - int const warp_d_base = warp * M_PER_WARP; - int const row_lo = warp_d_base + lane_d; - int const frag_col_base = (lane & 3) << 1; - int const state_base_lo = row_lo << 7; - int const state_xor = (row_lo & 7) << 4; - - // Philox state for SR — one refresh every 4 n-passes (cvt_rs_sat_s8x4_f32 - // packs 4 int8s per u32 of randomness, so 1 Philox call covers 16 int8s). - [[maybe_unused]] uint32_t rand_idx[4]; +template +__device__ __forceinline__ void encode_state_replay_8bit(SmemT& smem, CheckpointingSsuParams const& params, int warp, + int lane, int prev_k, int d_tile, int64_t cache_slot, int head, float const (&encode_scale_per_row)[2], + float const (&total_scale)[2], int64_t rand_seed, int64_t state_ptr_offset) +{ + using namespace cute; + static_assert(sizeof(input_t) == 2, "encode_state_replay_8bit requires 2-byte input_t"); + static_assert(sizeof(state_t) == 1, "encode_state_replay_8bit is for 1-byte state_t (int8/fp8) only"); + static_assert(D_PER_CTA == 64, "encode_state_replay_8bit requires D_PER_CTA == 64 (M-shard, per-warp M=16)."); + + constexpr int NUM_WARPS = 4; + constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; + static_assert(M_PER_WARP == MMA_prop::M, "Per-warp M must equal m16n8 atom M (=16)"); + + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + int const tid = warp * warpSize + lane; + + using MmaAtomReplayType + = std::conditional_t; + using LdsmA = std::conditional_t; + using LdsmB = std::conditional_t; + + auto tiled_mma_replay = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_replay = tiled_mma_replay.get_slice(tid); + + constexpr int N_PER_PASS = MMA_prop::N; + constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; + constexpr int FRAG_SIZE = 4; + constexpr int D_ROWS_PER_THREAD = 2; + + float const total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + + // total_scale (= OLD decode_scale_in × total_decay) was computed in PASS 1 + // and is passed in by reference. We MUST NOT re-load decode_scale_in from + // params.state_scale here — by the time PASS 2 runs, PASS 1 has already + // STG'd the NEW decode_scale to that same gmem location. + + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + auto layout_A_full = make_swizzled_layout_rc_transpose(); + Tensor smem_A_full + = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); + Tensor smem_A + = local_tile(smem_A_full, make_shape(Int{}, Int{}), make_coord(_0{}, _0{})); + + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma_replay); + auto s2r_thr_A = s2r_A.get_slice(tid); + Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); + Tensor frag_A_replay = thr_mma_replay.partition_fragment_A( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + Tensor frag_A_replay_view = s2r_thr_A.retile_D(frag_A_replay); + cute::copy(s2r_A, smem_A_s2r, frag_A_replay_view); + + // dB coefficients baked into frag_A (same identity as PASS 1). + apply_dA_coeff(frag_A_replay, smem, total_cumAdt, prev_k, lane); + + auto layout_B_replay = make_swizzled_layout_rc_transpose(); + Tensor smem_B_full + = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B_replay); + auto s2r_B_replay = make_tiled_copy_B(Copy_Atom{}, tiled_mma_replay); + auto s2r_thr_B_replay = s2r_B_replay.get_slice(tid); + + state_t* state_base = reinterpret_cast(smem.state); + + // Manual swizzle offsets (same derivation as replay_state_mma_8bit_chain). + int const lane_d = lane / 4; + int const warp_d_base = warp * M_PER_WARP; + int const row_lo = warp_d_base + lane_d; + int const frag_col_base = (lane & 3) << 1; + int const state_base_lo = row_lo << 7; + int const state_xor = (row_lo & 7) << 4; + + // Philox state for SR — one refresh every 4 n-passes (cvt_rs_sat_s8x4_f32 + // packs 4 int8s per u32 of randomness, so 1 Philox call covers 16 int8s). + [[maybe_unused]] uint32_t rand_idx[4]; #pragma unroll - for (int n = 0; n < NUM_N_PASSES; ++n) { - int const n_base = n * N_PER_PASS; - - Tensor frag_h = thr_mma_replay.partition_fragment_C( - make_tensor((float*)0x0, make_shape(Int{}, Int{}))); - - // Zero-init accumulator — MMA from scratch, state added after. - clear(frag_h); - - Tensor smem_B_n = - local_tile(smem_B_full, make_tile(Int{}, Int{}), - make_coord(n, _0{})); - auto smem_B_s2r_n = s2r_thr_B_replay.partition_S(smem_B_n); - Tensor frag_B_replay = thr_mma_replay.partition_fragment_B(make_tensor( - (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); - auto frag_B_replay_view = s2r_thr_B_replay.retile_D(frag_B_replay); - cute::copy(s2r_B_replay, smem_B_s2r_n, frag_B_replay_view); - - // HMMA: frag_h = frag_A_scaled @ frag_B (c[k] baked into A). - cute::gemm(tiled_mma_replay, frag_h, frag_A_replay, frag_B_replay, frag_h); - - { - int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); - Pair const p0 = *reinterpret_cast const*>(&state_base[off_lo]); - frag_h(0) += toFloat(p0[Int<0>{}]) * total_scale[0]; - frag_h(1) += toFloat(p0[Int<1>{}]) * total_scale[0]; - Pair const p1 = *reinterpret_cast const*>(&state_base[off_lo + 1024]); - frag_h(2) += toFloat(p1[Int<0>{}]) * total_scale[1]; - frag_h(3) += toFloat(p1[Int<1>{}]) * total_scale[1]; - } - - // ── Encode + in-place STS to smem.state at cols [n*8, n*8+8) ── - // Overwrites the OLD 8-bit input *for this n-pass's cols only*. The next - // n-pass dequants from a DIFFERENT col band [(n+1)*8, +8) — still OLD — - // so no read-after-write hazard. Each warp writes only to its own - // M-shard rows; cross-warp visibility is established by the - // caller's __syncthreads before the cooperative store_state call. + for (int n = 0; n < NUM_N_PASSES; ++n) { - int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); - float const e0 = encode_scale_per_row[0]; - float const e1 = encode_scale_per_row[1]; - - if constexpr (PHILOX_ROUNDS > 0) { - // One Philox4x call yields 4 independent u32s — enough for 4 n-passes - // when each pass packs all 4 of its 8-bit outputs into one u32 via the - // dtype-specific x4 cvt_rs (int8: `cvt_rs_sat_s8x4_f32`, 16-bit - // randomness/elt via bitrev16 trick; fp8 e4m3: `cvt_rs_e4m3x4_f32`, - // native PTX `cvt.rs.satfinite.e4m3x4.f32` on sm_100a+ with SW fallback). - int const rand_pos = n & 3; - if (rand_pos == 0) { - int64_t const philox_off = - state_ptr_offset + (int64_t)row_lo * DSTATE + (frag_col_base + n_base); - conversion::philox_randint4x(rand_seed, philox_off, rand_idx[0], - rand_idx[1], rand_idx[2], rand_idx[3]); + int const n_base = n * N_PER_PASS; + + Tensor frag_h = thr_mma_replay.partition_fragment_C( + make_tensor((float*) 0x0, make_shape(Int{}, Int{}))); + + // Zero-init accumulator — MMA from scratch, state added after. + clear(frag_h); + + Tensor smem_B_n + = local_tile(smem_B_full, make_tile(Int{}, Int{}), make_coord(n, _0{})); + auto smem_B_s2r_n = s2r_thr_B_replay.partition_S(smem_B_n); + Tensor frag_B_replay = thr_mma_replay.partition_fragment_B( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + auto frag_B_replay_view = s2r_thr_B_replay.retile_D(frag_B_replay); + cute::copy(s2r_B_replay, smem_B_s2r_n, frag_B_replay_view); + + // HMMA: frag_h = frag_A_scaled @ frag_B (c[k] baked into A). + cute::gemm(tiled_mma_replay, frag_h, frag_A_replay, frag_B_replay, frag_h); + + { + int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); + Pair const p0 = *reinterpret_cast const*>(&state_base[off_lo]); + frag_h(0) += toFloat(p0[Int<0>{}]) * total_scale[0]; + frag_h(1) += toFloat(p0[Int<1>{}]) * total_scale[0]; + Pair const p1 = *reinterpret_cast const*>(&state_base[off_lo + 1024]); + frag_h(2) += toFloat(p1[Int<0>{}]) * total_scale[1]; + frag_h(3) += toFloat(p1[Int<1>{}]) * total_scale[1]; } - // Packed layout: byte 0 = q0_lo, byte 1 = q1_lo (→ row_lo store at off_lo) - // byte 2 = q0_hi, byte 3 = q1_hi (→ row_hi store at off_lo + 1024) - uint32_t packed; - if constexpr (std::is_same_v) { - packed = conversion::cvt_rs_sat_s8x4_f32(frag_h(0) * e0, frag_h(1) * e0, frag_h(2) * e1, - frag_h(3) * e1, rand_idx[rand_pos]); - } else { - static_assert(std::is_same_v, - "8-bit SR supports state_t in {int8_t, __nv_fp8_e4m3}"); - packed = conversion::cvt_rs_e4m3x4_f32(frag_h(0) * e0, frag_h(1) * e0, frag_h(2) * e1, - frag_h(3) * e1, rand_idx[rand_pos]); + + // ── Encode + in-place STS to smem.state at cols [n*8, n*8+8) ── + // Overwrites the OLD 8-bit input *for this n-pass's cols only*. The next + // n-pass dequants from a DIFFERENT col band [(n+1)*8, +8) — still OLD — + // so no read-after-write hazard. Each warp writes only to its own + // M-shard rows; cross-warp visibility is established by the + // caller's __syncthreads before the cooperative store_state call. + { + int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); + float const e0 = encode_scale_per_row[0]; + float const e1 = encode_scale_per_row[1]; + + if constexpr (PHILOX_ROUNDS > 0) + { + // One Philox4x call yields 4 independent u32s — enough for 4 n-passes + // when each pass packs all 4 of its 8-bit outputs into one u32 via the + // dtype-specific x4 cvt_rs (int8: `cvt_rs_sat_s8x4_f32`, 16-bit + // randomness/elt via bitrev16 trick; fp8 e4m3: `cvt_rs_e4m3x4_f32`, + // native PTX `cvt.rs.satfinite.e4m3x4.f32` on sm_100a+ with SW fallback). + int const rand_pos = n & 3; + if (rand_pos == 0) + { + int64_t const philox_off = state_ptr_offset + (int64_t) row_lo * DSTATE + (frag_col_base + n_base); + conversion::philox_randint4x( + rand_seed, philox_off, rand_idx[0], rand_idx[1], rand_idx[2], rand_idx[3]); + } + // Packed layout: byte 0 = q0_lo, byte 1 = q1_lo (→ row_lo store at off_lo) + // byte 2 = q0_hi, byte 3 = q1_hi (→ row_hi store at off_lo + 1024) + uint32_t packed; + if constexpr (std::is_same_v) + { + packed = conversion::cvt_rs_sat_s8x4_f32( + frag_h(0) * e0, frag_h(1) * e0, frag_h(2) * e1, frag_h(3) * e1, rand_idx[rand_pos]); + } + else + { + static_assert( + std::is_same_v, "8-bit SR supports state_t in {int8_t, __nv_fp8_e4m3}"); + packed = conversion::cvt_rs_e4m3x4_f32( + frag_h(0) * e0, frag_h(1) * e0, frag_h(2) * e1, frag_h(3) * e1, rand_idx[rand_pos]); + } + Pair q_lo, q_hi; + q_lo.raw = static_cast(packed & 0xFFFFu); + q_hi.raw = static_cast(packed >> 16); + *reinterpret_cast*>(&state_base[off_lo]) = q_lo; + *reinterpret_cast*>(&state_base[off_lo + 1024]) = q_hi; + } + else + { + // d_idx=0: row_lo + state_t const q0_lo = encode_rn_8bit(frag_h(0) * e0); + state_t const q1_lo = encode_rn_8bit(frag_h(1) * e0); + Pair q_lo; + q_lo.raw + = static_cast(state_byte_of(q0_lo)) | (static_cast(state_byte_of(q1_lo)) << 8); + *reinterpret_cast*>(&state_base[off_lo]) = q_lo; + // d_idx=1: row_hi = row_lo + 8, off_hi = off_lo + 1024 + state_t const q0_hi = encode_rn_8bit(frag_h(2) * e1); + state_t const q1_hi = encode_rn_8bit(frag_h(3) * e1); + Pair q_hi; + q_hi.raw + = static_cast(state_byte_of(q0_hi)) | (static_cast(state_byte_of(q1_hi)) << 8); + *reinterpret_cast*>(&state_base[off_lo + 1024]) = q_hi; + } } - Pair q_lo, q_hi; - q_lo.raw = static_cast(packed & 0xFFFFu); - q_hi.raw = static_cast(packed >> 16); - *reinterpret_cast*>(&state_base[off_lo]) = q_lo; - *reinterpret_cast*>(&state_base[off_lo + 1024]) = q_hi; - } else { - // d_idx=0: row_lo - state_t const q0_lo = encode_rn_8bit(frag_h(0) * e0); - state_t const q1_lo = encode_rn_8bit(frag_h(1) * e0); - Pair q_lo; - q_lo.raw = static_cast(state_byte_of(q0_lo)) | - (static_cast(state_byte_of(q1_lo)) << 8); - *reinterpret_cast*>(&state_base[off_lo]) = q_lo; - // d_idx=1: row_hi = row_lo + 8, off_hi = off_lo + 1024 - state_t const q0_hi = encode_rn_8bit(frag_h(2) * e1); - state_t const q1_hi = encode_rn_8bit(frag_h(3) * e1); - Pair q_hi; - q_hi.raw = static_cast(state_byte_of(q0_hi)) | - (static_cast(state_byte_of(q1_hi)) << 8); - *reinterpret_cast*>(&state_base[off_lo + 1024]) = q_hi; - } } - } - // No __syncthreads or cooperative STG here — the caller's single sync - // provides cross-warp smem.state visibility, then calls store_state. + // No __syncthreads or cooperative STG here — the caller's single sync + // provides cross-warp smem.state visibility, then calls store_state. } // ──────────────────────────────────────────────────────────────────────── @@ -741,180 +737,180 @@ __device__ __forceinline__ void encode_state_replay_8bit( // Cross-warp dependencies (smem.x, smem.z, smem.CB_scaled) are already // visible because the caller's __syncthreads fires between all replay // passes and this function. -template -__device__ __forceinline__ void compute_output_8bit(SmemT& smem, - CheckpointingSsuParams const& params, int warp, - int lane, int d_tile, int64_t out_seq_base, - int head, int64_t cache_slot, float D_val, - int seq_len, FragYDxT& frag_y_DxT) { - using namespace cute; - static_assert(sizeof(input_t) == 2, "compute_output_8bit requires 2-byte input_t"); - static_assert(D_PER_CTA == 64, "compute_output_8bit requires D_PER_CTA == 64"); - static_assert(NUM_WARPS == 4, "compute_output_8bit requires 4 warps"); - - int const tid = warp * warpSize + lane; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; // 16 - - // Same TiledMma as replay_state_mma_int8_chain (M-shard, m16n8k16). - auto tiled_mma_chain = - make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_chain = tiled_mma_chain.get_slice(tid); - - // ── Smem views ── - // x_trans: x physically stored at (T, D); transposed view at (D, T). - // Used as the A operand of the chain matmul-4. - auto layout_x_trans = - make_swizzled_layout_rc_transpose(); - Tensor smem_x_trans = make_tensor( - make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans); - Tensor smem_x_trans_tile = - local_tile(smem_x_trans, make_shape(Int{}, Int{}), - make_coord(_0{}, _0{})); - - // x natural (T, D) view — for D-skip + z-gate per-element scalar LDS. - auto layout_x = make_swizzled_layout_rc(); - - // z natural (T, D) view (aliased so padded rows alias valid rows). - auto layout_z = make_aliased_swizzled_layout_rc(); - - // CB_scaled (T, T_pad) within (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE). - auto layout_cb = - make_swizzled_layout_rc(); - Tensor smem_CB = make_tensor( - make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb); - - // ── Per-thread (d, t) coord lookup for epilogue scalar reads + smem-transpose write ── - auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma_chain.partition_C(id_tile); - - // ── 1. Decay broadcast: frag_y(i) *= exp(cumAdt[t]) ── - // Reads `smem.decay[t]` (= exp(cumAdt[t])) precomputed in Phase 0 by - // `compute_cumAdt` — fused EX2 with the cumsum write, replacing ~512 per-CTA - // __expf calls in this inner loop with a single LDS per element. - // For padded T-cols (t >= NPREDICTED), the read returns garbage but the STG - // at the end is predicated on t < NPREDICTED, so the garbage never reaches gmem. +template +__device__ __forceinline__ void compute_output_8bit(SmemT& smem, CheckpointingSsuParams const& params, int warp, + int lane, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len, + FragYDxT& frag_y_DxT) +{ + using namespace cute; + static_assert(sizeof(input_t) == 2, "compute_output_8bit requires 2-byte input_t"); + static_assert(D_PER_CTA == 64, "compute_output_8bit requires D_PER_CTA == 64"); + static_assert(NUM_WARPS == 4, "compute_output_8bit requires 4 warps"); + + int const tid = warp * warpSize + lane; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; // 16 + + // Same TiledMma as replay_state_mma_int8_chain (M-shard, m16n8k16). + auto tiled_mma_chain = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_chain = tiled_mma_chain.get_slice(tid); + + // ── Smem views ── + // x_trans: x physically stored at (T, D); transposed view at (D, T). + // Used as the A operand of the chain matmul-4. + auto layout_x_trans = make_swizzled_layout_rc_transpose(); + Tensor smem_x_trans + = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans); + Tensor smem_x_trans_tile + = local_tile(smem_x_trans, make_shape(Int{}, Int{}), make_coord(_0{}, _0{})); + + // x natural (T, D) view — for D-skip + z-gate per-element scalar LDS. + auto layout_x = make_swizzled_layout_rc(); + + // z natural (T, D) view (aliased so padded rows alias valid rows). + auto layout_z = make_aliased_swizzled_layout_rc(); + + // CB_scaled (T, T_pad) within (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE). + auto layout_cb = make_swizzled_layout_rc(); + Tensor smem_CB + = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb); + + // ── Per-thread (d, t) coord lookup for epilogue scalar reads + smem-transpose write ── + auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma_chain.partition_C(id_tile); + + // ── 1. Decay broadcast: frag_y(i) *= exp(cumAdt[t]) ── + // Reads `smem.decay[t]` (= exp(cumAdt[t])) precomputed in Phase 0 by + // `compute_cumAdt` — fused EX2 with the cumsum write, replacing ~512 per-CTA + // __expf calls in this inner loop with a single LDS per element. + // For padded T-cols (t >= NPREDICTED), the read returns garbage but the STG + // at the end is predicated on t < NPREDICTED, so the garbage never reaches gmem. #pragma unroll - for (int i = 0; i < size(frag_y_DxT); ++i) { - int const t = get<1>(id_part(i)); - if (t < seq_len) { - frag_y_DxT(i) *= smem.decay[t]; + for (int i = 0; i < size(frag_y_DxT); ++i) + { + int const t = get<1>(id_part(i)); + if (t < seq_len) + { + frag_y_DxT(i) *= smem.decay[t]; + } } - } - - // ── 2. Chain matmul-4: frag_y_DxT += x^T @ CB^T ── - // A operand: smem.x physically (T, D); transposed view (D, T) used as - // A(M=D, K=T). The transposed view has D-stride=1, T-stride=D — same - // pattern as replay's A from old_x — so use LDSM_T to produce row-major - // A from this column-wise smem source. - auto s2r_A_x = - make_tiled_copy_A(Copy_Atom{}, tiled_mma_chain); - auto s2r_thr_A_x = s2r_A_x.get_slice(tid); - auto smem_x_s2r = s2r_thr_A_x.partition_S(smem_x_trans_tile); - Tensor frag_A_x = thr_mma_chain.partition_fragment_A(make_tensor( - (MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); - auto frag_A_x_view = s2r_thr_A_x.retile_D(frag_A_x); - cute::copy(s2r_A_x, smem_x_s2r, frag_A_x_view); - - // B operand for chain matmul-4 = CB^T. smem.CB natural view shape (T, T) - // already has T_inner stride 1 = K-major. Use LDSM_N (no transpose). - auto s2r_B_CB = - make_tiled_copy_B(Copy_Atom{}, tiled_mma_chain); - auto s2r_thr_B_CB = s2r_B_CB.get_slice(tid); - auto smem_CB_s2r = s2r_thr_B_CB.partition_S(smem_CB); - Tensor frag_B_CB = thr_mma_chain.partition_fragment_B( - make_tensor((MMA_prop::operand_t*)0x0, - make_shape(Int{}, Int{}))); - auto frag_B_CB_view = s2r_thr_B_CB.retile_D(frag_B_CB); - cute::copy(s2r_B_CB, smem_CB_s2r, frag_B_CB_view); - - cute::gemm(tiled_mma_chain, frag_y_DxT, frag_A_x, frag_B_CB, frag_y_DxT); - - // ── 3. D*x skip: frag_y(d, t) += D_val * x[t, d] (scalar LDS per element) ── - if (D_val != 0.f) { - auto* __restrict__ smem_x_base = reinterpret_cast(smem.x); + + // ── 2. Chain matmul-4: frag_y_DxT += x^T @ CB^T ── + // A operand: smem.x physically (T, D); transposed view (D, T) used as + // A(M=D, K=T). The transposed view has D-stride=1, T-stride=D — same + // pattern as replay's A from old_x — so use LDSM_T to produce row-major + // A from this column-wise smem source. + auto s2r_A_x = make_tiled_copy_A(Copy_Atom{}, tiled_mma_chain); + auto s2r_thr_A_x = s2r_A_x.get_slice(tid); + auto smem_x_s2r = s2r_thr_A_x.partition_S(smem_x_trans_tile); + Tensor frag_A_x = thr_mma_chain.partition_fragment_A( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + auto frag_A_x_view = s2r_thr_A_x.retile_D(frag_A_x); + cute::copy(s2r_A_x, smem_x_s2r, frag_A_x_view); + + // B operand for chain matmul-4 = CB^T. smem.CB natural view shape (T, T) + // already has T_inner stride 1 = K-major. Use LDSM_N (no transpose). + auto s2r_B_CB = make_tiled_copy_B(Copy_Atom{}, tiled_mma_chain); + auto s2r_thr_B_CB = s2r_B_CB.get_slice(tid); + auto smem_CB_s2r = s2r_thr_B_CB.partition_S(smem_CB); + Tensor frag_B_CB = thr_mma_chain.partition_fragment_B( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + auto frag_B_CB_view = s2r_thr_B_CB.retile_D(frag_B_CB); + cute::copy(s2r_B_CB, smem_CB_s2r, frag_B_CB_view); + + cute::gemm(tiled_mma_chain, frag_y_DxT, frag_A_x, frag_B_CB, frag_y_DxT); + + // ── 3. D*x skip: frag_y(d, t) += D_val * x[t, d] (scalar LDS per element) ── + if (D_val != 0.f) + { + auto* __restrict__ smem_x_base = reinterpret_cast(smem.x); #pragma unroll - for (int i = 0; i < size(frag_y_DxT); ++i) { - int const d = get<0>(id_part(i)); - int const t = get<1>(id_part(i)); - if (t < seq_len) { - int const off = layout_x(t, d); - frag_y_DxT(i) += D_val * toFloat(smem_x_base[off]); - } + for (int i = 0; i < size(frag_y_DxT); ++i) + { + int const d = get<0>(id_part(i)); + int const t = get<1>(id_part(i)); + if (t < seq_len) + { + int const off = layout_x(t, d); + frag_y_DxT(i) += D_val * toFloat(smem_x_base[off]); + } + } } - } - // ── 4. z-gate: frag_y *= z * sigmoid(z) (scalar LDS per element) ── - if (params.z != nullptr) { - auto* __restrict__ smem_z_base = reinterpret_cast(smem.z); + // ── 4. z-gate: frag_y *= z * sigmoid(z) (scalar LDS per element) ── + if (params.z != nullptr) + { + auto* __restrict__ smem_z_base = reinterpret_cast(smem.z); #pragma unroll - for (int i = 0; i < size(frag_y_DxT); ++i) { - int const d = get<0>(id_part(i)); - int const t = get<1>(id_part(i)); - if (t < seq_len) { - int const off = layout_z(t, d); - float const z = toFloat(smem_z_base[off]); - frag_y_DxT(i) *= z * __fdividef(1.f, (1.f + __expf(-z))); - } + for (int i = 0; i < size(frag_y_DxT); ++i) + { + int const d = get<0>(id_part(i)); + int const t = get<1>(id_part(i)); + if (t < seq_len) + { + int const off = layout_z(t, d); + float const z = toFloat(smem_z_base[off]); + frag_y_DxT(i) *= z * __fdividef(1.f, (1.f + __expf(-z))); + } + } } - } - - // ── 5. Pack fp32 → input_t per element + 6. STS to smem.output_transpose (T, D) ── - // Padded row stride (D_PER_CTA + 8 = 72 bf16 = 144 bytes) gives: - // - 16-byte-aligned LDS.128 / STG.128 across all rows. - // - 4-bank shift per row → m16n8 STS pattern hits {bank 0, 4, 8, 12} for the - // 4 t-rows of an elt → bank-conflict-free (vs 4-way conflict at stride 64). - // See CheckpointingSsuStorage8bit::OUTPUT_TRANSPOSE_ROW_STRIDE for derivation. - constexpr int kSmemRowStride = SmemT::OUTPUT_TRANSPOSE_ROW_STRIDE; // 72 bf16 elts - auto* __restrict__ smem_out_base = reinterpret_cast(smem.output_transpose); + + // ── 5. Pack fp32 → input_t per element + 6. STS to smem.output_transpose (T, D) ── + // Padded row stride (D_PER_CTA + 8 = 72 bf16 = 144 bytes) gives: + // - 16-byte-aligned LDS.128 / STG.128 across all rows. + // - 4-bank shift per row → m16n8 STS pattern hits {bank 0, 4, 8, 12} for the + // 4 t-rows of an elt → bank-conflict-free (vs 4-way conflict at stride 64). + // See CheckpointingSsuStorage8bit::OUTPUT_TRANSPOSE_ROW_STRIDE for derivation. + constexpr int kSmemRowStride = SmemT::OUTPUT_TRANSPOSE_ROW_STRIDE; // 72 bf16 elts + auto* __restrict__ smem_out_base = reinterpret_cast(smem.output_transpose); #pragma unroll - for (int i = 0; i < size(frag_y_DxT); ++i) { - int const d = get<0>(id_part(i)); - int const t = get<1>(id_part(i)); - if (t < seq_len) { - // Pack via pack_float2(f, 0.f) and take low elt — emits a single cvt - // (compiler folds the dummy into a no-op for the discarded high half). - smem_out_base[t * kSmemRowStride + d] = - pack_float2(make_float2(frag_y_DxT(i), 0.f))[Int<0>{}]; + for (int i = 0; i < size(frag_y_DxT); ++i) + { + int const d = get<0>(id_part(i)); + int const t = get<1>(id_part(i)); + if (t < seq_len) + { + // Pack via pack_float2(f, 0.f) and take low elt — emits a single cvt + // (compiler folds the dummy into a no-op for the discarded high half). + smem_out_base[t * kSmemRowStride + d] = pack_float2(make_float2(frag_y_DxT(i), 0.f))[Int<0>{}]; + } } - } - - // ── 7. Warp sync for cross-lane STS→LDS ordering ── - __syncwarp(); - - // ── 8. Warp-local cooperative STG.128: 32 lanes → one warp's 16 D-rows ── - // Each warp's data: 16 D-rows × T_pad=16 cols × 2 B = 512 B. - // Re-tile 32 lanes: (t = lane%16, d_group = lane/16 ∈ {0, 1}) → covers - // T_pad × 2 D-groups = 32 slots, each STG.128 = 8 D-cols × 2 B = 16 B. - // No cross-warp coordination → no __syncthreads. - constexpr int kElsPerSTG = 16 / sizeof(input_t); // 8 bf16 elts per STG.128 - constexpr int kDGroupsPerWarp = M_PER_WARP / kElsPerSTG; // = 16 / 8 = 2 - static_assert(NPREDICTED_PAD_MMA_M * kDGroupsPerWarp == 32, - "warp-local STG re-tile: T_pad × dGroupsPerWarp must equal warpSize"); - int const stg_t = lane % NPREDICTED_PAD_MMA_M; - int const stg_d_group = lane / NPREDICTED_PAD_MMA_M; - int const warp_d_base = warp * M_PER_WARP; - int const stg_d = warp_d_base + stg_d_group * kElsPerSTG; + // ── 7. Warp sync for cross-lane STS→LDS ordering ── + __syncwarp(); + + // ── 8. Warp-local cooperative STG.128: 32 lanes → one warp's 16 D-rows ── + // Each warp's data: 16 D-rows × T_pad=16 cols × 2 B = 512 B. + // Re-tile 32 lanes: (t = lane%16, d_group = lane/16 ∈ {0, 1}) → covers + // T_pad × 2 D-groups = 32 slots, each STG.128 = 8 D-cols × 2 B = 16 B. + // No cross-warp coordination → no __syncthreads. + constexpr int kElsPerSTG = 16 / sizeof(input_t); // 8 bf16 elts per STG.128 + constexpr int kDGroupsPerWarp = M_PER_WARP / kElsPerSTG; // = 16 / 8 = 2 + static_assert(NPREDICTED_PAD_MMA_M * kDGroupsPerWarp == 32, + "warp-local STG re-tile: T_pad × dGroupsPerWarp must equal warpSize"); + + int const stg_t = lane % NPREDICTED_PAD_MMA_M; + int const stg_d_group = lane / NPREDICTED_PAD_MMA_M; + int const warp_d_base = warp * M_PER_WARP; + int const stg_d = warp_d_base + stg_d_group * kElsPerSTG; + + if (stg_t < seq_len) + { + int const smem_off = stg_t * kSmemRowStride + stg_d; - if (stg_t < seq_len) { - int const smem_off = stg_t * kSmemRowStride + stg_d; + auto* __restrict__ output_ptr = reinterpret_cast(params.output); + int64_t const out_base = out_seq_base + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; + int64_t const gmem_off = out_base + (int64_t) stg_t * params.out_stride_token + stg_d; - auto* __restrict__ output_ptr = reinterpret_cast(params.output); - int64_t const out_base = out_seq_base + (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; - int64_t const gmem_off = out_base + (int64_t)stg_t * params.out_stride_token + stg_d; - - // 128-bit copy. smem_off * 2 B = (t * 144 + d * 2) is 16-byte aligned - // for any t when d % 8 == 0 (here d_offset_within_warp = 0 or 8). - using Vec = uint4; - *reinterpret_cast(&output_ptr[gmem_off]) = - *reinterpret_cast(&smem_out_base[smem_off]); - } + // 128-bit copy. smem_off * 2 B = (t * 144 + d * 2) is 16-byte aligned + // for any t when d % 8 == 0 (here d_offset_within_warp = 0 or 8). + using Vec = uint4; + *reinterpret_cast(&output_ptr[gmem_off]) = *reinterpret_cast(&smem_out_base[smem_off]); + } } // ============================================================================= @@ -941,90 +937,90 @@ __device__ __forceinline__ void compute_output_8bit(SmemT& smem, // This eliminates the per-cell `... * scale` FMUL chain (was the // long_scoreboard hotspot at line 1066) at the cost of 2 extra FMUL/elt in // the post-matmul C-frag scale (net 1792× fewer FMUL per warp). -template -__device__ __forceinline__ void add_init_out_8bit(SmemT const& smem, int warp, int lane, - TiledMma const& tiled_mma, ThrMma const& thr_mma, - int tid, FragY&... frag_y) { - using namespace cute; - static_assert(sizeof(state_t) == 1, "add_init_out_8bit requires 1-byte state"); - static_assert(D_PER_CTA == 64, "add_init_out_8bit requires D_PER_CTA == 64 (8-bit D_SPLIT=1)"); - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int K_TILE = MMA_prop::K_BIG; // 16 - constexpr int NUM_K_TILES = DSTATE / K_TILE; // 8 - constexpr int N_TILE = cute::tile_size<1>(TiledMma{}); // 32 (4 warps × MMA::N=8) - constexpr int NUM_N_TILES = sizeof...(FragY); // 2 (D_PER_CTA / N_TILE) - static_assert(NUM_N_TILES * N_TILE == D_PER_CTA, "FragY count must match D_PER_CTA / N_TILE"); - - // ── Per-thread coords ── - int const t = lane & 3; // K-pair index within K-atom - int const lane_d = lane >> 2; // gID = lane/4; selects N-col within atom - int const warp_d_base = warp * MMA_prop::N; // warp's 8-col offset within an N-tile - - // ── A operand (C): swizzled (T_pad, DSTATE), K-tiled per K-loop iter ── - auto layout_C_swz = - make_aliased_swizzled_layout_rc(); - Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), - layout_C_swz); - Tensor smem_C_ktiled = local_tile(smem_C, make_tile(Int{}, Int{}), - make_coord(_0{}, _)); - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - auto smem_A_s2r = s2r_thr_A.partition_S(smem_C_ktiled); - Tensor frag_A = thr_mma.partition_fragment_A(smem_C_ktiled(_, _, _0{})); - auto frag_A_view = s2r_thr_A.retile_D(frag_A); - - // ── B-frag (one m16n8k16 atom of B; 4 bf16 elts/lane) ── - Tensor frag_B = thr_mma.partition_fragment_B( - make_tensor((MMA_prop::operand_t*)0x0, make_shape(Int{}, Int{}))); - static_assert(decltype(size(frag_B))::value == 4, "B-frag must be 4 elts/lane for m16n8k16"); - - // ── Smem state base (1-byte) + Swizzle<3,4,3> XOR formula: - // off = d_row * DSTATE + (K XOR ((d_row & 7) << 4)) - // K within the same swizzle row-group {0..7 mod 8} shares the same XOR mask. ── - state_t const* state_base = reinterpret_cast(smem.state); - - // Pre-clear accumulators (caller doesn't pre-zero — matches bf16 add_init_out). - (clear(frag_y), ...); - - // Parameter-pack indexing via pointer array (same pattern as pipelined_kloop_gemm). - using FragY0 = std::tuple_element_t<0, std::tuple>; - FragY0* frag_y_p[NUM_N_TILES] = {(&frag_y)...}; - - // ── K-loop ── +template +__device__ __forceinline__ void add_init_out_8bit( + SmemT const& smem, int warp, int lane, TiledMma const& tiled_mma, ThrMma const& thr_mma, int tid, FragY&... frag_y) +{ + using namespace cute; + static_assert(sizeof(state_t) == 1, "add_init_out_8bit requires 1-byte state"); + static_assert(D_PER_CTA == 64, "add_init_out_8bit requires D_PER_CTA == 64 (8-bit D_SPLIT=1)"); + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int K_TILE = MMA_prop::K_BIG; // 16 + constexpr int NUM_K_TILES = DSTATE / K_TILE; // 8 + constexpr int N_TILE = cute::tile_size<1>(TiledMma{}); // 32 (4 warps × MMA::N=8) + constexpr int NUM_N_TILES = sizeof...(FragY); // 2 (D_PER_CTA / N_TILE) + static_assert(NUM_N_TILES * N_TILE == D_PER_CTA, "FragY count must match D_PER_CTA / N_TILE"); + + // ── Per-thread coords ── + int const t = lane & 3; // K-pair index within K-atom + int const lane_d = lane >> 2; // gID = lane/4; selects N-col within atom + int const warp_d_base = warp * MMA_prop::N; // warp's 8-col offset within an N-tile + + // ── A operand (C): swizzled (T_pad, DSTATE), K-tiled per K-loop iter ── + auto layout_C_swz = make_aliased_swizzled_layout_rc(); + Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C_swz); + Tensor smem_C_ktiled + = local_tile(smem_C, make_tile(Int{}, Int{}), make_coord(_0{}, _)); + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + auto smem_A_s2r = s2r_thr_A.partition_S(smem_C_ktiled); + Tensor frag_A = thr_mma.partition_fragment_A(smem_C_ktiled(_, _, _0{})); + auto frag_A_view = s2r_thr_A.retile_D(frag_A); + + // ── B-frag (one m16n8k16 atom of B; 4 bf16 elts/lane) ── + Tensor frag_B = thr_mma.partition_fragment_B( + make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); + static_assert(decltype(size(frag_B))::value == 4, "B-frag must be 4 elts/lane for m16n8k16"); + + // ── Smem state base (1-byte) + Swizzle<3,4,3> XOR formula: + // off = d_row * DSTATE + (K XOR ((d_row & 7) << 4)) + // K within the same swizzle row-group {0..7 mod 8} shares the same XOR mask. ── + state_t const* state_base = reinterpret_cast(smem.state); + + // Pre-clear accumulators (caller doesn't pre-zero — matches bf16 add_init_out). + (clear(frag_y), ...); + + // Parameter-pack indexing via pointer array (same pattern as pipelined_kloop_gemm). + using FragY0 = std::tuple_element_t<0, std::tuple>; + FragY0* frag_y_p[NUM_N_TILES] = {(&frag_y)...}; + + // ── K-loop ── #pragma unroll - for (int k = 0; k < NUM_K_TILES; ++k) { - int const K_base = k * K_TILE; - - // Load A K-tile via LDSM (shared across all N-tiles within this K-tile). - cute::copy(s2r_A, smem_A_s2r(_, _, _, k), frag_A_view); - - CUTE_UNROLL - for (int n = 0; n < NUM_N_TILES; ++n) { - int const d_row = n * N_TILE + warp_d_base + lane_d; - int const state_base_lo = d_row << 7; // d_row * DSTATE (DSTATE=128 → <<7) - int const state_xor = (d_row & 7) << 4; // Swizzle<3,4,3> - int const off_lo = state_base_lo + ((K_base + (t << 1)) ^ state_xor); - int const off_hi = state_base_lo + ((K_base + (t << 1) + 8) ^ state_xor); - - Pair const p_lo = *reinterpret_cast const*>(&state_base[off_lo]); - Pair const p_hi = *reinterpret_cast const*>(&state_base[off_hi]); - - // Pure int8/fp8 → bf16 cast. decode_scale is applied post-matmul in - // the caller's β-scale loop. - Pair const b_lo = pack_float2( - make_float2(toFloat(p_lo[Int<0>{}]), toFloat(p_lo[Int<1>{}]))); - Pair const b_hi = pack_float2( - make_float2(toFloat(p_hi[Int<0>{}]), toFloat(p_hi[Int<1>{}]))); - - // frag_B(0,1) = K-pair at {K_base+2t, K_base+2t+1}; (2,3) = at {+8, +9}. - *reinterpret_cast*>(&frag_B(0)) = b_lo; - *reinterpret_cast*>(&frag_B(2)) = b_hi; - - cute::gemm(tiled_mma, *frag_y_p[n], frag_A, frag_B, *frag_y_p[n]); + for (int k = 0; k < NUM_K_TILES; ++k) + { + int const K_base = k * K_TILE; + + // Load A K-tile via LDSM (shared across all N-tiles within this K-tile). + cute::copy(s2r_A, smem_A_s2r(_, _, _, k), frag_A_view); + + CUTE_UNROLL + for (int n = 0; n < NUM_N_TILES; ++n) + { + int const d_row = n * N_TILE + warp_d_base + lane_d; + int const state_base_lo = d_row << 7; // d_row * DSTATE (DSTATE=128 → <<7) + int const state_xor = (d_row & 7) << 4; // Swizzle<3,4,3> + int const off_lo = state_base_lo + ((K_base + (t << 1)) ^ state_xor); + int const off_hi = state_base_lo + ((K_base + (t << 1) + 8) ^ state_xor); + + Pair const p_lo = *reinterpret_cast const*>(&state_base[off_lo]); + Pair const p_hi = *reinterpret_cast const*>(&state_base[off_hi]); + + // Pure int8/fp8 → bf16 cast. decode_scale is applied post-matmul in + // the caller's β-scale loop. + Pair const b_lo + = pack_float2(make_float2(toFloat(p_lo[Int<0>{}]), toFloat(p_lo[Int<1>{}]))); + Pair const b_hi + = pack_float2(make_float2(toFloat(p_hi[Int<0>{}]), toFloat(p_hi[Int<1>{}]))); + + // frag_B(0,1) = K-pair at {K_base+2t, K_base+2t+1}; (2,3) = at {+8, +9}. + *reinterpret_cast*>(&frag_B(0)) = b_lo; + *reinterpret_cast*>(&frag_B(2)) = b_hi; + + cute::gemm(tiled_mma, *frag_y_p[n], frag_A, frag_B, *frag_y_p[n]); + } } - } } // ============================================================================= @@ -1043,180 +1039,169 @@ __device__ __forceinline__ void add_init_out_8bit(SmemT const& smem, int warp, i // // β(t) = exp(total_old_cumAdt + cumAdt[t]) where total_old_cumAdt = // smem.old_cumAdt[prev_k − 1] (= 0 when prev_k == 0). -template -__device__ __forceinline__ void compute_no_write_output_8bit( - SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int prev_k, int d_tile, - int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) { - using namespace cute; - static_assert(sizeof(input_t) == 2, "compute_no_write_output_8bit requires 2-byte input_t"); - static_assert(sizeof(state_t) == 1, "compute_no_write_output_8bit is for 1-byte state"); - static_assert(D_PER_CTA == 64, "compute_no_write_output_8bit requires D_PER_CTA == 64"); - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - int const tid = warp * warpSize + lane; - - // ── TiledMMA for matmul-3 + matmul-4-new (m16n8k16) ── - auto tiled_mma = - make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(tid); - - // ── TiledMMA for matmul-4-old: K = MAX_WINDOW_PAD_MMA_K ∈ {8, 16} → atom dispatch ── - using MmaAtomOld = std::conditional_t; - using LdsmAOld = std::conditional_t; - using LdsmBOld = std::conditional_t; - auto tiled_mma_old = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_old = tiled_mma_old.get_slice(tid); - - // ── Swizzled smem views ── - auto layout_x_swz = make_swizzled_layout_rc(); - Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), - layout_x_swz); - auto layout_x_trans_swz = - make_swizzled_layout_rc_transpose(); - Tensor smem_x_trans = make_tensor( - make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); - - auto layout_old_x_trans_swz = - make_swizzled_layout_rc_transpose(); - Tensor smem_old_x_trans = - make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), - layout_old_x_trans_swz); - - auto layout_z_swz = - make_aliased_swizzled_layout_rc(); - Tensor smem_z = - make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); - - // ── S2R copies (matmul-4-new) ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - auto s2r_B_trans = - make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); - - // ── S2R copies (matmul-4-old) ── - auto s2r_A_old = make_tiled_copy_A(Copy_Atom{}, tiled_mma_old); - auto s2r_thr_A_old = s2r_A_old.get_slice(tid); - auto s2r_B_old_trans = - make_tiled_copy_B(Copy_Atom{}, tiled_mma_old); - auto s2r_thr_B_old_trans = s2r_B_old_trans.get_slice(tid); - - // ── Load CB_scaled A operand (cols [0, NPREDICTED_PAD_MMA_M)) ── - auto layout_cb_swz = - make_swizzled_layout_rc(); - Tensor smem_CB = make_tensor( - make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); - auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); - Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); - auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); - cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); - - // ── Load CB_old A operand (cols [NPREDICTED_PAD_MMA_M, +MAX_WINDOW_PAD_MMA_K)) ── - auto layout_cb_full = make_swizzled_layout_rc(); - Tensor smem_CB_full = make_tensor( - make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); - Tensor smem_CB_old = - local_tile(smem_CB_full, make_tile(Int{}, Int{}), - make_coord(_0{}, NPREDICTED_PAD_MMA_M / MAX_WINDOW_PAD_MMA_K)); - auto smem_CB_old_s2r = s2r_thr_A_old.partition_S(smem_CB_old); - Tensor frag_CB_old_A = thr_mma_old.partition_fragment_A(smem_CB_old); - auto frag_CB_old_A_view = s2r_thr_A_old.retile_D(frag_CB_old_A); - cute::copy(s2r_A_old, smem_CB_old_s2r, frag_CB_old_A_view); - - // ── Decay broadcast (per-T scalar, stride-0 on N) ── - constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); - Tensor decay_bcast = make_tensor( - make_smem_ptr(smem.cumAdt), - make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); - Tensor decay_part = thr_mma.partition_C(decay_bcast); - - float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - float const beta_extra = __expf(total_old_cumAdt); - - // ── Post-matmul-3 state decode_scale (factored OUT of the inner K-product). ── - // y[t, d] = decode_scale[d] · raw_y[t, d] where raw_y = C @ (bf16)(state_byte)^T. - // Per lane, the m16n8 C-frag holds 4 elts spanning 2 unique d-cols (d_lo = 2t, - // d_hi = 2t+1) per N-tile. Indexed by `i & 1` in the epilogue scale loop. - auto const* __restrict__ state_scale_ptr = reinterpret_cast(params.state_scale); - int64_t const state_scale_base = cache_slot * params.state_scale_stride_seq + - (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; - constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; - static_assert(NUM_N_TILES == 2, - "compute_no_write_output_8bit assumes NUM_N_TILES == 2 (D_PER_CTA=64, N_TILE=32)"); - int const t_col = lane & 3; // 2t and 2t+1 are this lane's two C-frag d-cols - float decode_scale[NUM_N_TILES][2]; - CUTE_UNROLL - for (int n = 0; n < NUM_N_TILES; ++n) { - int const d_lo = n * N_TILE + warp * MMA_prop::N + (t_col << 1); - decode_scale[n][0] = state_scale_ptr[state_scale_base + d_lo]; - decode_scale[n][1] = state_scale_ptr[state_scale_base + d_lo + 1]; - } - - // ── Gmem output base ── - auto* __restrict__ output_ptr = reinterpret_cast(params.output); - int64_t const out_base = out_seq_base + (int64_t)head * DIM + (int64_t)d_tile * D_PER_CTA; - - // ── Row predicate ── - auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_tile); - bool const pred_row_lo = get<0>(id_part(0)) < seq_len; - bool const pred_row_hi = get<0>(id_part(2)) < seq_len; - - auto epilogue = [&](auto& frag_y, int n) { - // β-scale + state decode_scale fused: frag_y(i) *= β · exp(cumAdt[t]) · decode_scale[d]. - // The decode_scale[d] absorbs the per-row state quant factor (was previously - // multiplied into each B-element during dequant). -#pragma unroll - for (int i = 0; i < size(frag_y); ++i) { - int const d_idx = i & 1; // i=0,2 → d_lo; i=1,3 → d_hi - frag_y(i) *= beta_extra * __expf(decay_part(i)) * decode_scale[n][d_idx]; +template +__device__ __forceinline__ void compute_no_write_output_8bit(SmemT& smem, CheckpointingSsuParams const& params, + int warp, int lane, int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, + int seq_len) +{ + using namespace cute; + static_assert(sizeof(input_t) == 2, "compute_no_write_output_8bit requires 2-byte input_t"); + static_assert(sizeof(state_t) == 1, "compute_no_write_output_8bit is for 1-byte state"); + static_assert(D_PER_CTA == 64, "compute_no_write_output_8bit requires D_PER_CTA == 64"); + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + int const tid = warp * warpSize + lane; + + // ── TiledMMA for matmul-3 + matmul-4-new (m16n8k16) ── + auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(tid); + + // ── TiledMMA for matmul-4-old: K = MAX_WINDOW_PAD_MMA_K ∈ {8, 16} → atom dispatch ── + using MmaAtomOld = std::conditional_t; + using LdsmAOld = std::conditional_t; + using LdsmBOld = std::conditional_t; + auto tiled_mma_old = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_old = tiled_mma_old.get_slice(tid); + + // ── Swizzled smem views ── + auto layout_x_swz = make_swizzled_layout_rc(); + Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_swz); + auto layout_x_trans_swz = make_swizzled_layout_rc_transpose(); + Tensor smem_x_trans + = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); + + auto layout_old_x_trans_swz = make_swizzled_layout_rc_transpose(); + Tensor smem_old_x_trans + = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), layout_old_x_trans_swz); + + auto layout_z_swz = make_aliased_swizzled_layout_rc(); + Tensor smem_z = make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); + + // ── S2R copies (matmul-4-new) ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + auto s2r_B_trans = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); + + // ── S2R copies (matmul-4-old) ── + auto s2r_A_old = make_tiled_copy_A(Copy_Atom{}, tiled_mma_old); + auto s2r_thr_A_old = s2r_A_old.get_slice(tid); + auto s2r_B_old_trans = make_tiled_copy_B(Copy_Atom{}, tiled_mma_old); + auto s2r_thr_B_old_trans = s2r_B_old_trans.get_slice(tid); + + // ── Load CB_scaled A operand (cols [0, NPREDICTED_PAD_MMA_M)) ── + auto layout_cb_swz = make_swizzled_layout_rc(); + Tensor smem_CB + = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); + auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); + Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); + auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); + cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); + + // ── Load CB_old A operand (cols [NPREDICTED_PAD_MMA_M, +MAX_WINDOW_PAD_MMA_K)) ── + auto layout_cb_full = make_swizzled_layout_rc(); + Tensor smem_CB_full + = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); + Tensor smem_CB_old = local_tile(smem_CB_full, make_tile(Int{}, Int{}), + make_coord(_0{}, NPREDICTED_PAD_MMA_M / MAX_WINDOW_PAD_MMA_K)); + auto smem_CB_old_s2r = s2r_thr_A_old.partition_S(smem_CB_old); + Tensor frag_CB_old_A = thr_mma_old.partition_fragment_A(smem_CB_old); + auto frag_CB_old_A_view = s2r_thr_A_old.retile_D(frag_CB_old_A); + cute::copy(s2r_A_old, smem_CB_old_s2r, frag_CB_old_A_view); + + // ── Decay broadcast (per-T scalar, stride-0 on N) ── + constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); + Tensor decay_bcast = make_tensor(make_smem_ptr(smem.cumAdt), + make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); + Tensor decay_part = thr_mma.partition_C(decay_bcast); + + float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + float const beta_extra = __expf(total_old_cumAdt); + + // ── Post-matmul-3 state decode_scale (factored OUT of the inner K-product). ── + // y[t, d] = decode_scale[d] · raw_y[t, d] where raw_y = C @ (bf16)(state_byte)^T. + // Per lane, the m16n8 C-frag holds 4 elts spanning 2 unique d-cols (d_lo = 2t, + // d_hi = 2t+1) per N-tile. Indexed by `i & 1` in the epilogue scale loop. + auto const* __restrict__ state_scale_ptr = reinterpret_cast(params.state_scale); + int64_t const state_scale_base + = cache_slot * params.state_scale_stride_seq + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; + constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; + static_assert(NUM_N_TILES == 2, "compute_no_write_output_8bit assumes NUM_N_TILES == 2 (D_PER_CTA=64, N_TILE=32)"); + int const t_col = lane & 3; // 2t and 2t+1 are this lane's two C-frag d-cols + float decode_scale[NUM_N_TILES][2]; + CUTE_UNROLL + for (int n = 0; n < NUM_N_TILES; ++n) + { + int const d_lo = n * N_TILE + warp * MMA_prop::N + (t_col << 1); + decode_scale[n][0] = state_scale_ptr[state_scale_base + d_lo]; + decode_scale[n][1] = state_scale_ptr[state_scale_base + d_lo + 1]; } - // matmul-4-new (CB_scaled @ x). - add_cb_x( - frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); + // ── Gmem output base ── + auto* __restrict__ output_ptr = reinterpret_cast(params.output); + int64_t const out_base = out_seq_base + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; + + // ── Row predicate ── + auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_tile); + bool const pred_row_lo = get<0>(id_part(0)) < seq_len; + bool const pred_row_hi = get<0>(id_part(2)) < seq_len; + + auto epilogue = [&](auto& frag_y, int n) + { + // β-scale + state decode_scale fused: frag_y(i) *= β · exp(cumAdt[t]) · decode_scale[d]. + // The decode_scale[d] absorbs the per-row state quant factor (was previously + // multiplied into each B-element during dequant). +#pragma unroll + for (int i = 0; i < size(frag_y); ++i) + { + int const d_idx = i & 1; // i=0,2 → d_lo; i=1,3 → d_hi + frag_y(i) *= beta_extra * __expf(decay_part(i)) * decode_scale[n][d_idx]; + } - // matmul-4-old (CB_old @ old_x). - add_cb_old_x( - frag_y, frag_CB_old_A, smem_old_x_trans, s2r_B_old_trans, s2r_thr_B_old_trans, thr_mma_old, - tiled_mma_old, n); + // matmul-4-new (CB_scaled @ x). + add_cb_x( + frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); - // D·x. - add_D_skip(frag_y, smem_x, thr_mma, D_val, n); + // matmul-4-old (CB_old @ old_x). + add_cb_old_x(frag_y, frag_CB_old_A, + smem_old_x_trans, s2r_B_old_trans, s2r_thr_B_old_trans, thr_mma_old, tiled_mma_old, n); - // z-gate. - compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); + // D·x. + add_D_skip(frag_y, smem_x, thr_mma, D_val, n); - // Direct partition_C STG. - auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), - make_layout(make_shape(Int{}, Int{}), - make_stride(params.out_stride_token, _1{}))); - auto gOut_part = thr_mma.partition_C(gOut_tile); + // z-gate. + compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); + + // Direct partition_C STG. + auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), + make_layout( + make_shape(Int{}, Int{}), make_stride(params.out_stride_token, _1{}))); + auto gOut_part = thr_mma.partition_C(gOut_tile); #pragma unroll - for (int i = 0; i < size(frag_y); i += 2) { - bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; - if (pred_i) { - *reinterpret_cast*>(&gOut_part(i)) = - pack_float2(make_float2(frag_y(i), frag_y(i + 1))); - } - } - }; - - // ── Matmul-3: frag_y = C @ (bf16)(state_byte)^T (smem.state retains s_0 since - // replay skipped; decode_scale[d] applied post-matmul in the epilogue). ── - Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); - Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); - add_init_out_8bit(smem, warp, lane, tiled_mma, thr_mma, - tid, frag_y_0, frag_y_1); - epilogue(frag_y_0, 0); - epilogue(frag_y_1, 1); + for (int i = 0; i < size(frag_y); i += 2) + { + bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; + if (pred_i) + { + *reinterpret_cast*>(&gOut_part(i)) + = pack_float2(make_float2(frag_y(i), frag_y(i + 1))); + } + } + }; + + // ── Matmul-3: frag_y = C @ (bf16)(state_byte)^T (smem.state retains s_0 since + // replay skipped; decode_scale[d] applied post-matmul in the epilogue). ── + Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); + Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); + add_init_out_8bit( + smem, warp, lane, tiled_mma, thr_mma, tid, frag_y_0, frag_y_1); + epilogue(frag_y_0, 0); + epilogue(frag_y_1, 1); } // ============================================================================= @@ -1225,15 +1210,14 @@ __device__ __forceinline__ void compute_no_write_output_8bit( // Sync makes warps 0,1's CB_scaled writes AND warps 2,3's CB_old writes // visible to all warps before matmul-3 and matmul-4 read smem.{CB_scaled, // CB_old, x, z}. Matches the bf16 path's `ssu_nocheckpoint`. -template -__device__ __forceinline__ void ssu_nocheckpoint_8bit( - SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int prev_k, int d_tile, - int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) { - __syncthreads(); - compute_no_write_output_8bit(smem, params, warp, lane, prev_k, d_tile, out_seq_base, - head, cache_slot, D_val, seq_len); +template +__device__ __forceinline__ void ssu_nocheckpoint_8bit(SmemT& smem, CheckpointingSsuParams const& params, int warp, + int lane, int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) +{ + __syncthreads(); + compute_no_write_output_8bit( + smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); } // ============================================================================= @@ -1247,66 +1231,58 @@ __device__ __forceinline__ void ssu_nocheckpoint_8bit( // Pulled out of `checkpointing_ssu_kernel_8bit` to mirror the bf16 path's // `ssu_checkpoint` and make the kernel-body dispatch on // must_checkpoint readable. -template -__device__ __forceinline__ void ssu_checkpoint_8bit(SmemT& smem, - CheckpointingSsuParams const& params, int warp, - int lane, int prev_k, int d_tile, - int64_t out_seq_base, int head, - int64_t cache_slot, float D_val, int seq_len) { - using namespace cute; - int const tid = warp * warpSize + lane; - - // ── Allocate per-warp frag_y_DxT (chain mma C-frag, fp32) ── - // Layout ((2, 2), MMA_M=1, MMA_N=NPREDICTED_PAD_MMA_M/8) per thread. - // Caller must zero before chain matmul-3 accumulates. - auto tiled_mma_chain = - make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_chain = tiled_mma_chain.get_slice(tid); - auto id_DxT = - make_identity_tensor(make_shape(Int{}, Int{})); - Tensor frag_y_DxT = thr_mma_chain.partition_fragment_C(id_DxT); - cute::clear(frag_y_DxT); - - // ── Phase 1b: replay + amax + chain matmul-3 → frag_y_DxT (init_out^T). - // `encode_scale_per_row[]` is computed at the end of PASS 1 (from the warp- - // reduced amax) and consumed by `encode_state_replay_8bit` further below — - // *after* `compute_output_8bit` consumes `frag_y_DxT` and STGs the output. - float encode_scale_per_row[2]; - float total_scale[2]; - replay_state_mma_8bit_chain( - smem, params, warp, lane, prev_k, d_tile, cache_slot, head, /*must_checkpoint=*/true, - frag_y_DxT, encode_scale_per_row, total_scale); - - // ── Philox seed for stochastic rounding (deferred to reduce register pressure) ── - [[maybe_unused]] int64_t const rand_seed = (PHILOX_ROUNDS > 0) ? *params.rand_seed : 0; - // `state_ptr_offset` is int64 — matches Triton's `base_rand = - // cache_batch_idx * stride_state_batch + ...` (cache_batch_idx is .to(int64)). - // Full 64 bits flow through `philox_randint4x`, which splits low/high - // across Philox c0/c1. No collision risk at large serving cache sizes. - int64_t const state_ptr_offset = - cache_slot * params.state_stride_seq + (int64_t)head * DIM * DSTATE; - - // ── PASS 2 (replay-again): re-run replay HMMA, encode fp32 → int8 to - // smem.state. Runs BEFORE the sync so both replay passes overlap with - // warps 0,1's CB precompute — one fewer __syncthreads in the kernel. - // frag_y_DxT stays live through PASS 2 (extra register pressure accepted). ── - encode_state_replay_8bit( - smem, params, warp, lane, prev_k, d_tile, cache_slot, head, encode_scale_per_row, total_scale, - rand_seed, state_ptr_offset); - - // ── Single sync: cross-warp visibility for smem.CB_scaled (warps 0,1) / - // smem.x (warp 2) / smem.z (warp 3) / smem.state (all warps' M-shards). ── - __syncthreads(); - - // ── Cooperative STG.128 for encoded state (after sync for cross-warp - // smem.state visibility). Fire-and-forget before compute_output_8bit. ── - store_state(smem, params, warp, lane, d_tile, head, - cache_slot); - - // ── Phase 2: transposed matmul-4 + epilogue + smem-transpose STG ── - compute_output_8bit( - smem, params, warp, lane, d_tile, out_seq_base, head, cache_slot, D_val, seq_len, frag_y_DxT); +template +__device__ __forceinline__ void ssu_checkpoint_8bit(SmemT& smem, CheckpointingSsuParams const& params, int warp, + int lane, int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) +{ + using namespace cute; + int const tid = warp * warpSize + lane; + + // ── Allocate per-warp frag_y_DxT (chain mma C-frag, fp32) ── + // Layout ((2, 2), MMA_M=1, MMA_N=NPREDICTED_PAD_MMA_M/8) per thread. + // Caller must zero before chain matmul-3 accumulates. + auto tiled_mma_chain = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma_chain = tiled_mma_chain.get_slice(tid); + auto id_DxT = make_identity_tensor(make_shape(Int{}, Int{})); + Tensor frag_y_DxT = thr_mma_chain.partition_fragment_C(id_DxT); + cute::clear(frag_y_DxT); + + // ── Phase 1b: replay + amax + chain matmul-3 → frag_y_DxT (init_out^T). + // `encode_scale_per_row[]` is computed at the end of PASS 1 (from the warp- + // reduced amax) and consumed by `encode_state_replay_8bit` further below — + // *after* `compute_output_8bit` consumes `frag_y_DxT` and STGs the output. + float encode_scale_per_row[2]; + float total_scale[2]; + replay_state_mma_8bit_chain(smem, params, warp, lane, prev_k, d_tile, + cache_slot, head, /*must_checkpoint=*/true, frag_y_DxT, encode_scale_per_row, total_scale); + + // ── Philox seed for stochastic rounding (deferred to reduce register pressure) ── + [[maybe_unused]] int64_t const rand_seed = (PHILOX_ROUNDS > 0) ? *params.rand_seed : 0; + // `state_ptr_offset` is int64 — matches Triton's `base_rand = + // cache_batch_idx * stride_state_batch + ...` (cache_batch_idx is .to(int64)). + // Full 64 bits flow through `philox_randint4x`, which splits low/high + // across Philox c0/c1. No collision risk at large serving cache sizes. + int64_t const state_ptr_offset = cache_slot * params.state_stride_seq + (int64_t) head * DIM * DSTATE; + + // ── PASS 2 (replay-again): re-run replay HMMA, encode fp32 → int8 to + // smem.state. Runs BEFORE the sync so both replay passes overlap with + // warps 0,1's CB precompute — one fewer __syncthreads in the kernel. + // frag_y_DxT stays live through PASS 2 (extra register pressure accepted). ── + encode_state_replay_8bit(smem, params, warp, lane, prev_k, + d_tile, cache_slot, head, encode_scale_per_row, total_scale, rand_seed, state_ptr_offset); + + // ── Single sync: cross-warp visibility for smem.CB_scaled (warps 0,1) / + // smem.x (warp 2) / smem.z (warp 3) / smem.state (all warps' M-shards). ── + __syncthreads(); + + // ── Cooperative STG.128 for encoded state (after sync for cross-warp + // smem.state visibility). Fire-and-forget before compute_output_8bit. ── + store_state(smem, params, warp, lane, d_tile, head, cache_slot); + + // ── Phase 2: transposed matmul-4 + epilogue + smem-transpose STG ── + compute_output_8bit( + smem, params, warp, lane, d_tile, out_seq_base, head, cache_slot, D_val, seq_len, frag_y_DxT); } // ============================================================================= @@ -1326,170 +1302,180 @@ __device__ __forceinline__ void ssu_checkpoint_8bit(SmemT& smem, // touch shared smem fields that both storage structs expose by name. // template -__global__ void checkpointing_ssu_kernel_8bit(CheckpointingSsuParams params) { - using namespace cute; - static_assert(sizeof(state_t) == 1, - "checkpointing_ssu_kernel_8bit requires 1-byte state_t (int8 or fp8 e4m3)"); - static_assert(NPREDICTED <= MAX_WINDOW); - static_assert(MAX_WINDOW <= MMA_prop::K_BIG); - // int8 path uses M-shard layout (Layout<_4,_1>): per-warp M = 16 = m16n8 - // atom M. D_PER_CTA must equal DIM (D_SPLIT=1) to give 4×16=64 D-rows/CTA. - // The wrapper enforces d_split == 1 for int8. - constexpr int D_PER_CTA = DIM; - static_assert(D_PER_CTA == 64, "int8 chain kernel requires DIM == 64"); - assert(params.d_split == 1); - - using SmemT = - CheckpointingSsuStorage8bit; - extern __shared__ __align__(128) char smem_buf[]; - auto& smem = *reinterpret_cast(smem_buf); - - // Grid: (1, batch, nheads). D-tile is always 0 for int8 (D_SPLIT=1). - int const d_tile = blockIdx.x; - int const seq = blockIdx.y; - int const head = blockIdx.z; - int const lane = threadIdx.x; - int const warp = threadIdx.y; - int const group_idx = head / HEADS_PER_GROUP; - - // ── Resolve cache slot ── - auto const* __restrict__ sbi = reinterpret_cast(params.state_batch_indices); - int64_t const cache_slot = sbi ? static_cast(sbi[seq]) : seq; - if (cache_slot == params.pad_slot_id) return; - - auto const* __restrict__ buf_idx_ptr = reinterpret_cast(params.cache_buf_idx); - int const buf_read = __ldg(&buf_idx_ptr[cache_slot]); - - auto const* __restrict__ prev_ptr = reinterpret_cast(params.prev_num_accepted); - int const prev_k = prev_ptr[cache_slot]; - - // ── Varlen vs non-varlen prologue. The kernel branches once on the - // VARLEN template; downstream helpers receive `seq_len` (constexpr-foldable - // NPREDICTED in non-varlen, runtime in varlen) and pre-computed per-sequence - // gmem base offsets (`x_seq_base` etc.) — they're varlen-agnostic. - // - // Uniform gmem-base formula: `outer * *_stride_seq` where - // non-varlen: outer = seq (= blockIdx.y), stride_seq = x.stride(0). - // varlen : outer = cu_seqlens[seq], stride_seq = x.stride(1). - // The wrapper picks the right stride_seq value; the kernel only branches - // on whether to load cu_seqlens. - int seq_len; - int64_t outer; - if constexpr (VARLEN) { - auto const* __restrict__ cu_seqlens = reinterpret_cast(params.cu_seqlens); - // Two LDG.E.32 (not one LDG.E.64): cu_seqlens is only 4-byte aligned - // at `&cu_seqlens[seq]` when seq is odd, and PTX - // `ld.global.v2.b32` faults on a 4-byte-aligned address. ptxas emits - // the two scalar loads back-to-back; latency is hidden against the - // following ALU work. - int const bos = __ldg(&cu_seqlens[seq]); - int const eos = __ldg(&cu_seqlens[seq + 1]); - seq_len = eos - bos; - if (seq_len <= 0) return; - outer = (int64_t)bos; - } else { - seq_len = NPREDICTED; - outer = (int64_t)seq; - } - // x/B/C bases computed inside `load_post_pdl_wait_data` from `outer` — - // see generic kernel for rationale (avoid pinning 6 regs across gdc_wait). - int64_t const dt_seq_base = outer * params.dt_stride_seq + head; - int64_t const z_seq_base = outer * params.z_stride_seq; - int64_t const out_seq_base = outer * params.out_stride_seq; - - bool const must_checkpoint = (prev_k + seq_len > MAX_WINDOW); - int const buf_write = must_checkpoint ? (1 - buf_read) : buf_read; - int const write_offset = must_checkpoint ? 0 : prev_k; - - // ── Load scalars (A, dt_bias, D) ── - auto const* __restrict__ A_ptr = reinterpret_cast(params.A); - auto const* __restrict__ dt_bias_ptr = reinterpret_cast(params.dt_bias); - auto const* __restrict__ D_ptr = reinterpret_cast(params.D); - float const A_val = toFloat(A_ptr[head]); - float const dt_bias_val = dt_bias_ptr ? toFloat(dt_bias_ptr[head]) : 0.f; - float const D_val = D_ptr ? toFloat(D_ptr[head]) : 0.f; - - // ── Phase 0: two-phase load around the PDL barrier (see generic kernel - // for the full rationale). Pre-wait: state + old_* cache + in_proj - // outputs (dt, z) + scalar scans. Post-wait: x/B/C from conv1d. ── - // ENABLE_PDL is JIT-stamped; `if constexpr` keeps only one load path in - // the binary (no register pressure leak from the unused path). - if constexpr (ENABLE_PDL) { - load_pre_pdl_wait_data(smem, params, lane, warp, d_tile, head, group_idx, cache_slot, - buf_read, A_val, dt_bias_val, dt_seq_base, z_seq_base, - seq_len); - gdc_wait(); - load_post_pdl_wait_data( - smem, params, lane, warp, d_tile, head, group_idx, outer, seq_len); - } else { - load_data( - smem, params, lane, warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, - outer, seq_len); - } - - // ── store_old_B hoist (warps 0,1 only, d_tile == 0) ── - if (d_tile == 0 && warp < 2) { - store_old_B( - smem, params, warp, lane, head, group_idx, cache_slot, buf_write, write_offset, seq_len); - } - - // ── CB precompute (4-warp split): warps 0,1 compute CB_scaled (new tokens); - // warps 2,3 compute CB_old (old tokens) in the no-write path only. Mirrors - // the bf16 path's dispatch — warps 2,3 stay idle in checkpoint mode and - // pick up work below inside `ssu_checkpoint_8bit`'s replay. ── - if (warp < 2) { - compute_CB_scaled_2warp(smem, warp, lane, seq_len); - } else if (!must_checkpoint) { - compute_CB_old_2warp(smem, warp, lane, prev_k, - seq_len); - } - - // ── Phase 1b + 2: per-path dispatch ── - // Checkpoint: M-shard chain (PASS 1 + PASS 2 + sync + state STG + transposed - // matmul-4 with smem-transpose STG). - // No-write : N-shard matmul-3 from int8/fp8 state + matmul-4-new + matmul-4-old - // + direct partition_C STG (mirrors the bf16 no-write path). - // must_checkpoint is uniform across the CTA — both branches contain a - // __syncthreads so divergence is balanced. - if (must_checkpoint) { - ssu_checkpoint_8bit(smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, - cache_slot, D_val, seq_len); - } else { - ssu_nocheckpoint_8bit(smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, - cache_slot, D_val, seq_len); - } - - // ── PDL: signal downstream that `output` is written. Cache writes below - // target tensors only the next SSU step reads, not the immediate - // downstream kernel — safe to signal first. ── - if constexpr (ENABLE_PDL) { - gdc_launch_dependents(); - } - - // ── Phase 3: cache writes (old_x, dt_proc, cumAdt) ── - store_old_x(smem, params, warp, lane, d_tile, head, - cache_slot, write_offset, seq_len); - if (d_tile == 0 && warp == 0 && lane < seq_len) { - auto* __restrict__ old_dt_w = reinterpret_cast(params.old_dt); - int64_t const dt_w_base = cache_slot * params.old_dt_stride_seq + - buf_write * params.old_dt_stride_dbuf + - head * params.old_dt_stride_head; - old_dt_w[dt_w_base + write_offset + lane] = smem.dt_proc[lane]; - } - if (d_tile == 0 && warp == 1 && lane < seq_len) { - auto* __restrict__ old_cumAdt_w = reinterpret_cast(params.old_cumAdt); - int64_t const ca_w_base = cache_slot * params.old_cumAdt_stride_seq + - buf_write * params.old_cumAdt_stride_dbuf + - head * params.old_cumAdt_stride_head; - old_cumAdt_w[ca_w_base + write_offset + lane] = smem.cumAdt[lane]; - } + typename stateIndex_t, typename state_scale_t, int NPREDICTED, int MAX_WINDOW, int DIM, int DSTATE, + int HEADS_PER_GROUP, int PHILOX_ROUNDS, int NUM_WARPS, bool VARLEN = false> +__global__ void checkpointing_ssu_kernel_8bit(CheckpointingSsuParams params) +{ + using namespace cute; + static_assert(sizeof(state_t) == 1, "checkpointing_ssu_kernel_8bit requires 1-byte state_t (int8 or fp8 e4m3)"); + static_assert(NPREDICTED <= MAX_WINDOW); + static_assert(MAX_WINDOW <= MMA_prop::K_BIG); + // int8 path uses M-shard layout (Layout<_4,_1>): per-warp M = 16 = m16n8 + // atom M. D_PER_CTA must equal DIM (D_SPLIT=1) to give 4×16=64 D-rows/CTA. + // The wrapper enforces d_split == 1 for int8. + constexpr int D_PER_CTA = DIM; + static_assert(D_PER_CTA == 64, "int8 chain kernel requires DIM == 64"); + assert(params.d_split == 1); + + using SmemT = CheckpointingSsuStorage8bit; + extern __shared__ __align__(128) char smem_buf[]; + auto& smem = *reinterpret_cast(smem_buf); + + // Grid: (1, batch, nheads). D-tile is always 0 for int8 (D_SPLIT=1). + int const d_tile = blockIdx.x; + int const seq = blockIdx.y; + int const head = blockIdx.z; + int const lane = threadIdx.x; + int const warp = threadIdx.y; + int const group_idx = head / HEADS_PER_GROUP; + + // ── Resolve cache slot ── + auto const* __restrict__ sbi = reinterpret_cast(params.state_batch_indices); + int64_t const cache_slot = sbi ? static_cast(sbi[seq]) : seq; + if (cache_slot == params.pad_slot_id) + return; + + auto const* __restrict__ buf_idx_ptr = reinterpret_cast(params.cache_buf_idx); + int const buf_read = __ldg(&buf_idx_ptr[cache_slot]); + + auto const* __restrict__ prev_ptr = reinterpret_cast(params.prev_num_accepted); + int const prev_k = prev_ptr[cache_slot]; + + // ── Varlen vs non-varlen prologue. The kernel branches once on the + // VARLEN template; downstream helpers receive `seq_len` (constexpr-foldable + // NPREDICTED in non-varlen, runtime in varlen) and pre-computed per-sequence + // gmem base offsets (`x_seq_base` etc.) — they're varlen-agnostic. + // + // Uniform gmem-base formula: `outer * *_stride_seq` where + // non-varlen: outer = seq (= blockIdx.y), stride_seq = x.stride(0). + // varlen : outer = cu_seqlens[seq], stride_seq = x.stride(1). + // The wrapper picks the right stride_seq value; the kernel only branches + // on whether to load cu_seqlens. + int seq_len; + int64_t outer; + if constexpr (VARLEN) + { + auto const* __restrict__ cu_seqlens = reinterpret_cast(params.cu_seqlens); + // Two LDG.E.32 (not one LDG.E.64): cu_seqlens is only 4-byte aligned + // at `&cu_seqlens[seq]` when seq is odd, and PTX + // `ld.global.v2.b32` faults on a 4-byte-aligned address. ptxas emits + // the two scalar loads back-to-back; latency is hidden against the + // following ALU work. + int const bos = __ldg(&cu_seqlens[seq]); + int const eos = __ldg(&cu_seqlens[seq + 1]); + seq_len = eos - bos; + if (seq_len <= 0) + return; + outer = (int64_t) bos; + } + else + { + seq_len = NPREDICTED; + outer = (int64_t) seq; + } + // x/B/C bases computed inside `load_post_pdl_wait_data` from `outer` — + // see generic kernel for rationale (avoid pinning 6 regs across gdc_wait). + int64_t const dt_seq_base = outer * params.dt_stride_seq + head; + int64_t const z_seq_base = outer * params.z_stride_seq; + int64_t const out_seq_base = outer * params.out_stride_seq; + + bool const must_checkpoint = (prev_k + seq_len > MAX_WINDOW); + int const buf_write = must_checkpoint ? (1 - buf_read) : buf_read; + int const write_offset = must_checkpoint ? 0 : prev_k; + + // ── Load scalars (A, dt_bias, D) ── + auto const* __restrict__ A_ptr = reinterpret_cast(params.A); + auto const* __restrict__ dt_bias_ptr = reinterpret_cast(params.dt_bias); + auto const* __restrict__ D_ptr = reinterpret_cast(params.D); + float const A_val = toFloat(A_ptr[head]); + float const dt_bias_val = dt_bias_ptr ? toFloat(dt_bias_ptr[head]) : 0.f; + float const D_val = D_ptr ? toFloat(D_ptr[head]) : 0.f; + + // ── Phase 0: two-phase load around the PDL barrier (see generic kernel + // for the full rationale). Pre-wait: state + old_* cache + in_proj + // outputs (dt, z) + scalar scans. Post-wait: x/B/C from conv1d. ── + // ENABLE_PDL is JIT-stamped; `if constexpr` keeps only one load path in + // the binary (no register pressure leak from the unused path). + if constexpr (ENABLE_PDL) + { + load_pre_pdl_wait_data(smem, + params, lane, warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, dt_seq_base, + z_seq_base, seq_len); + gdc_wait(); + load_post_pdl_wait_data( + smem, params, lane, warp, d_tile, head, group_idx, outer, seq_len); + } + else + { + load_data(smem, params, lane, + warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, outer, seq_len); + } + + // ── store_old_B hoist (warps 0,1 only, d_tile == 0) ── + if (d_tile == 0 && warp < 2) + { + store_old_B( + smem, params, warp, lane, head, group_idx, cache_slot, buf_write, write_offset, seq_len); + } + + // ── CB precompute (4-warp split): warps 0,1 compute CB_scaled (new tokens); + // warps 2,3 compute CB_old (old tokens) in the no-write path only. Mirrors + // the bf16 path's dispatch — warps 2,3 stay idle in checkpoint mode and + // pick up work below inside `ssu_checkpoint_8bit`'s replay. ── + if (warp < 2) + { + compute_CB_scaled_2warp(smem, warp, lane, seq_len); + } + else if (!must_checkpoint) + { + compute_CB_old_2warp(smem, warp, lane, prev_k, seq_len); + } + + // ── Phase 1b + 2: per-path dispatch ── + // Checkpoint: M-shard chain (PASS 1 + PASS 2 + sync + state STG + transposed + // matmul-4 with smem-transpose STG). + // No-write : N-shard matmul-3 from int8/fp8 state + matmul-4-new + matmul-4-old + // + direct partition_C STG (mirrors the bf16 no-write path). + // must_checkpoint is uniform across the CTA — both branches contain a + // __syncthreads so divergence is balanced. + if (must_checkpoint) + { + ssu_checkpoint_8bit( + smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); + } + else + { + ssu_nocheckpoint_8bit( + smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); + } + + // ── PDL: signal downstream that `output` is written. Cache writes below + // target tensors only the next SSU step reads, not the immediate + // downstream kernel — safe to signal first. ── + if constexpr (ENABLE_PDL) + { + gdc_launch_dependents(); + } + + // ── Phase 3: cache writes (old_x, dt_proc, cumAdt) ── + store_old_x( + smem, params, warp, lane, d_tile, head, cache_slot, write_offset, seq_len); + if (d_tile == 0 && warp == 0 && lane < seq_len) + { + auto* __restrict__ old_dt_w = reinterpret_cast(params.old_dt); + int64_t const dt_w_base = cache_slot * params.old_dt_stride_seq + buf_write * params.old_dt_stride_dbuf + + head * params.old_dt_stride_head; + old_dt_w[dt_w_base + write_offset + lane] = smem.dt_proc[lane]; + } + if (d_tile == 0 && warp == 1 && lane < seq_len) + { + auto* __restrict__ old_cumAdt_w = reinterpret_cast(params.old_cumAdt); + int64_t const ca_w_base = cache_slot * params.old_cumAdt_stride_seq + buf_write * params.old_cumAdt_stride_dbuf + + head * params.old_cumAdt_stride_head; + old_cumAdt_w[ca_w_base + write_offset + lane] = smem.cumAdt[lane]; + } } -} // namespace flashinfer::mamba::checkpointing +} // namespace flashinfer::mamba::checkpointing -#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_8BIT_CUH_ +#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_8BIT_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh index 83f0d2f8bfee..c22e9c62c304 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh @@ -34,7 +34,8 @@ #include "cute/tensor.hpp" #include "ssu_mtp_common.cuh" -namespace flashinfer::mamba::checkpointing { +namespace flashinfer::mamba::checkpointing +{ using namespace conversion; @@ -44,10 +45,11 @@ using namespace conversion; // (2) state is the accumulator (C-frag), not an A/B operand — layout // remapping costs 8 shuffles + byte extractions, // (3) dynamic byte selection via SHF adds 15%+ short_scoreboard stalls. -namespace constants { +namespace constants +{ constexpr unsigned int MASK_ALL_LANES = 0xFFFFFFFFu; constexpr unsigned int num_bits_uint32 = 32u; -} // namespace constants +} // namespace constants // ── Programmatic Dependent Launch (PDL) helpers ──────────────────────────── // `gdc_wait` enforces no gmem access before the upstream PDL-paired kernel @@ -56,39 +58,48 @@ constexpr unsigned int num_bits_uint32 = 32u; // the launch-time `cudaLaunchAttributeProgrammaticStreamSerialization` // attribute, so the kernel can always emit them; the host-side `enable_pdl` // toggle is what flips the launch attribute. -__forceinline__ __device__ void gdc_wait() { +__forceinline__ __device__ void gdc_wait() +{ #if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + asm volatile("griddepcontrol.wait;"); #endif } -__forceinline__ __device__ void gdc_launch_dependents() { +__forceinline__ __device__ void gdc_launch_dependents() +{ #if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + asm volatile("griddepcontrol.launch_dependents;"); #endif } // Round x up to the next multiple of Y (Y must be a power of 2). template -constexpr int next_multiple_of(int x) { - static_assert(Y > 0 && (Y & (Y - 1)) == 0, "Y must be a power of 2"); - return (x + Y - 1) & ~(Y - 1); +constexpr int next_multiple_of(int x) +{ + static_assert(Y > 0 && (Y & (Y - 1)) == 0, "Y must be a power of 2"); + return (x + Y - 1) & ~(Y - 1); } // NativeOf::type: scalar T → 2-wide native CUDA vector type. template struct NativeOf; + template <> -struct NativeOf { - using type = float2; +struct NativeOf +{ + using type = float2; }; + template <> -struct NativeOf<__half> { - using type = __half2; +struct NativeOf<__half> +{ + using type = __half2; }; + template <> -struct NativeOf<__nv_bfloat16> { - using type = __nv_bfloat162; +struct NativeOf<__nv_bfloat16> +{ + using type = __nv_bfloat162; }; // Pair: thin wrapper over the native 2-wide type, adding compile-time @@ -96,16 +107,19 @@ struct NativeOf<__nv_bfloat16> { // layout as the native type — `=` compiles to one LDS.U32 / STS.U32 and the // pair stays in one register. template -struct Pair { - typename NativeOf::type raw; - template - __device__ __forceinline__ auto operator[](cute::Int) const { - static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); - if constexpr (I == 0) - return raw.x; - else - return raw.y; - } +struct Pair +{ + typename NativeOf::type raw; + + template + __device__ __forceinline__ auto operator[](cute::Int) const + { + static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); + if constexpr (I == 0) + return raw.x; + else + return raw.y; + } }; // Pair: explicit 16-bit packed specialization. A struct of two @@ -115,47 +129,59 @@ struct Pair { // extracting via shift+cast, we keep the two elements in one register // throughout the load → unpack → cast pipeline. template <> -struct Pair { - uint16_t raw; // [bits 7:0] = element 0, [bits 15:8] = element 1 - template - __device__ __forceinline__ int8_t operator[](cute::Int) const { - static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); - if constexpr (I == 0) - return static_cast(raw & 0xFFu); - else - return static_cast(raw >> 8); - } +struct Pair +{ + uint16_t raw; // [bits 7:0] = element 0, [bits 15:8] = element 1 + + template + __device__ __forceinline__ int8_t operator[](cute::Int) const + { + static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); + if constexpr (I == 0) + return static_cast(raw & 0xFFu); + else + return static_cast(raw >> 8); + } }; // Pair<__nv_fp8_e4m3>: same single-u16 backing as `Pair` — fp8 e4m3 // is also a 1-byte storage type, so the load → unpack → cast pipeline runs // through one 16-bit register. template <> -struct Pair<__nv_fp8_e4m3> { - uint16_t raw; - template - __device__ __forceinline__ __nv_fp8_e4m3 operator[](cute::Int) const { - static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); - __nv_fp8_storage_t const byte = (I == 0) ? static_cast<__nv_fp8_storage_t>(raw & 0xFFu) - : static_cast<__nv_fp8_storage_t>(raw >> 8); - return reinterpret_cast<__nv_fp8_e4m3 const&>(byte); - } +struct Pair<__nv_fp8_e4m3> +{ + uint16_t raw; + + template + __device__ __forceinline__ __nv_fp8_e4m3 operator[](cute::Int) const + { + static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); + __nv_fp8_storage_t const byte + = (I == 0) ? static_cast<__nv_fp8_storage_t>(raw & 0xFFu) : static_cast<__nv_fp8_storage_t>(raw >> 8); + return reinterpret_cast<__nv_fp8_e4m3 const&>(byte); + } }; // pack_float2: float2 → Pair, using packed hardware cvt when available. template __device__ __forceinline__ Pair pack_float2(float2 val); + template <> -__device__ __forceinline__ Pair pack_float2(float2 val) { - return {val}; +__device__ __forceinline__ Pair pack_float2(float2 val) +{ + return {val}; } + template <> -__device__ __forceinline__ Pair<__half> pack_float2<__half>(float2 val) { - return {__float22half2_rn(val)}; +__device__ __forceinline__ Pair<__half> pack_float2<__half>(float2 val) +{ + return {__float22half2_rn(val)}; } + template <> -__device__ __forceinline__ Pair<__nv_bfloat16> pack_float2<__nv_bfloat16>(float2 val) { - return {conversion::fromFloat2(val)}; +__device__ __forceinline__ Pair<__nv_bfloat16> pack_float2<__nv_bfloat16>(float2 val) +{ + return {conversion::fromFloat2(val)}; } // ============================================================================= @@ -164,10 +190,11 @@ __device__ __forceinline__ Pair<__nv_bfloat16> pack_float2<__nv_bfloat16>(float2 // 128-bit vector loads, shared by every gmem→smem copy in the kernel. The // ldmatrix unit has the same vector width so `vec_bytes` is derived from the // atom's source-register type and reused as the LDSM vector width. -struct Copy_prop { - using Atom = cute::SM80_CP_ASYNC_CACHEALWAYS; - using AtomZFill = cute::SM80_CP_ASYNC_CACHEALWAYS_ZFILL; - static constexpr int vec_bytes = sizeof(std::remove_extent_t); +struct Copy_prop +{ + using Atom = cute::SM80_CP_ASYNC_CACHEALWAYS; + using AtomZFill = cute::SM80_CP_ASYNC_CACHEALWAYS_ZFILL; + static constexpr int vec_bytes = sizeof(std::remove_extent_t); }; // ============================================================================= @@ -178,16 +205,17 @@ struct Copy_prop { // k=8 and k=16 atoms at compile time (MAX_WINDOW ≤ 8 picks K8 for smaller smem, // +1 CTA/SM); dims are pulled from MMA_Traits so they stay in sync with the // atom choice (e.g. m16n8k32 for int8 would just need AtomK16/K8 swapped). -struct MMA_prop { - using AtomK16 = cute::SM80_16x8x16_F32BF16BF16F32_TN; - using AtomK8 = cute::SM80_16x8x8_F32BF16BF16F32_TN; - // Operand dtype — matches the bf16 input of the atoms above. - using operand_t = __nv_bfloat16; - - static constexpr int M = cute::size<0>(typename cute::MMA_Traits::Shape_MNK{}); - static constexpr int N = cute::size<1>(typename cute::MMA_Traits::Shape_MNK{}); - static constexpr int K_BIG = cute::size<2>(typename cute::MMA_Traits::Shape_MNK{}); - static constexpr int K_SMALL = cute::size<2>(typename cute::MMA_Traits::Shape_MNK{}); +struct MMA_prop +{ + using AtomK16 = cute::SM80_16x8x16_F32BF16BF16F32_TN; + using AtomK8 = cute::SM80_16x8x8_F32BF16BF16F32_TN; + // Operand dtype — matches the bf16 input of the atoms above. + using operand_t = __nv_bfloat16; + + static constexpr int M = cute::size<0>(typename cute::MMA_Traits::Shape_MNK{}); + static constexpr int N = cute::size<1>(typename cute::MMA_Traits::Shape_MNK{}); + static constexpr int K_BIG = cute::size<2>(typename cute::MMA_Traits::Shape_MNK{}); + static constexpr int K_SMALL = cute::size<2>(typename cute::MMA_Traits::Shape_MNK{}); }; // ============================================================================= @@ -205,23 +233,25 @@ struct MMA_prop { // // The MMA operand element type dictates the smem buffer element type, which in // turn dictates the swizzle — so every call site passes its own element type. -constexpr int log2_pow2(int x) { - int r = 0; - while (x > 1) { - x >>= 1; - ++r; - } - return r; +constexpr int log2_pow2(int x) +{ + int r = 0; + while (x > 1) + { + x >>= 1; + ++r; + } + return r; } template -struct SmemSwizzle { - static_assert(Copy_prop::vec_bytes % sizeof(Elem) == 0, - "element size must divide LDSM atom (16 bytes)"); - static constexpr int ELEMS_PER_ATOM = Copy_prop::vec_bytes / sizeof(Elem); - using type = cute::Swizzle<3, log2_pow2(ELEMS_PER_ATOM), 3>; - static constexpr int ATOM_ROWS = 1 << type::num_bits; - static constexpr int ATOM_COLS = 1 << (type::num_base + type::num_shft); +struct SmemSwizzle +{ + static_assert(Copy_prop::vec_bytes % sizeof(Elem) == 0, "element size must divide LDSM atom (16 bytes)"); + static constexpr int ELEMS_PER_ATOM = Copy_prop::vec_bytes / sizeof(Elem); + using type = cute::Swizzle<3, log2_pow2(ELEMS_PER_ATOM), 3>; + static constexpr int ATOM_ROWS = 1 << type::num_bits; + static constexpr int ATOM_COLS = 1 << (type::num_base + type::num_shft); }; // Default (ROW_STRIDE == COLS): tile the swizzle atom into a (ROWS, COLS) @@ -233,22 +263,24 @@ struct SmemSwizzle { // "wasted padding" — the swizzle XOR scatters logical cells across the full // ROW_STRIDE, so the physical extent is what the bijection actually needs. template -__device__ __forceinline__ auto make_swizzled_layout_rc() { - using namespace cute; - using S = SmemSwizzle; - static_assert(ROWS % S::ATOM_ROWS == 0, "ROWS must be a multiple of the swizzle atom rows"); - static_assert(ROW_STRIDE % S::ATOM_COLS == 0, - "ROW_STRIDE must be a multiple of the swizzle atom cols"); - static_assert(ROW_STRIDE >= COLS, "ROW_STRIDE must be at least COLS"); - if constexpr (ROW_STRIDE == COLS) { - auto atom = composition(typename S::type{}, - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, _1{}))); - return tile_to_shape(atom, make_shape(Int{}, Int{})); - } else { - return composition(typename S::type{}, make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, _1{}))); - } +__device__ __forceinline__ auto make_swizzled_layout_rc() +{ + using namespace cute; + using S = SmemSwizzle; + static_assert(ROWS % S::ATOM_ROWS == 0, "ROWS must be a multiple of the swizzle atom rows"); + static_assert(ROW_STRIDE % S::ATOM_COLS == 0, "ROW_STRIDE must be a multiple of the swizzle atom cols"); + static_assert(ROW_STRIDE >= COLS, "ROW_STRIDE must be at least COLS"); + if constexpr (ROW_STRIDE == COLS) + { + auto atom = composition(typename S::type{}, + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); + return tile_to_shape(atom, make_shape(Int{}, Int{})); + } + else + { + return composition(typename S::type{}, + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); + } } // Aliased-row swizzled smem layout: logical (LOGICAL_ROWS, COLS) view over a @@ -265,29 +297,26 @@ __device__ __forceinline__ auto make_swizzled_layout_rc() { // and the alias factor collapses to 1 — this then degenerates to the same // layout that `make_swizzled_layout_rc` produces. template -__device__ __forceinline__ auto make_aliased_swizzled_layout_rc() { - using namespace cute; - using S = SmemSwizzle; - static_assert(LOGICAL_ROWS % S::ATOM_ROWS == 0, - "LOGICAL_ROWS must be a multiple of the swizzle atom rows"); - static_assert(COLS % S::ATOM_COLS == 0, "COLS must be a multiple of the swizzle atom cols"); - constexpr int PHYS_ROWS = next_multiple_of(VALID_ROWS); - constexpr int LOG_M_TILES = LOGICAL_ROWS / S::ATOM_ROWS; - constexpr int PHYS_M_TILES = PHYS_ROWS / S::ATOM_ROWS; - static_assert(LOG_M_TILES % PHYS_M_TILES == 0, - "LOGICAL_ROWS must be a multiple of PHYS_ROWS for clean alias"); - constexpr int ALIAS = LOG_M_TILES / PHYS_M_TILES; - constexpr int N_TILES = COLS / S::ATOM_COLS; - auto atom = composition(typename S::type{}, - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, _1{}))); - // Outer layout (in atom-units): row-tile mode = (PHYS_M_TILES, ALIAS) strides - // (1, 0); col-tile mode = N_TILES stride PHYS_M_TILES. blocked_product - // scales these by the atom cosize (= ATOM_ROWS * ATOM_COLS). - auto outer = - make_layout(make_shape(make_shape(Int{}, Int{}), Int{}), - make_stride(make_stride(_1{}, _0{}), Int{})); - return blocked_product(atom, outer); +__device__ __forceinline__ auto make_aliased_swizzled_layout_rc() +{ + using namespace cute; + using S = SmemSwizzle; + static_assert(LOGICAL_ROWS % S::ATOM_ROWS == 0, "LOGICAL_ROWS must be a multiple of the swizzle atom rows"); + static_assert(COLS % S::ATOM_COLS == 0, "COLS must be a multiple of the swizzle atom cols"); + constexpr int PHYS_ROWS = next_multiple_of(VALID_ROWS); + constexpr int LOG_M_TILES = LOGICAL_ROWS / S::ATOM_ROWS; + constexpr int PHYS_M_TILES = PHYS_ROWS / S::ATOM_ROWS; + static_assert(LOG_M_TILES % PHYS_M_TILES == 0, "LOGICAL_ROWS must be a multiple of PHYS_ROWS for clean alias"); + constexpr int ALIAS = LOG_M_TILES / PHYS_M_TILES; + constexpr int N_TILES = COLS / S::ATOM_COLS; + auto atom = composition(typename S::type{}, + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); + // Outer layout (in atom-units): row-tile mode = (PHYS_M_TILES, ALIAS) strides + // (1, 0); col-tile mode = N_TILES stride PHYS_M_TILES. blocked_product + // scales these by the atom cosize (= ATOM_ROWS * ATOM_COLS). + auto outer = make_layout(make_shape(make_shape(Int{}, Int{}), Int{}), + make_stride(make_stride(_1{}, _0{}), Int{})); + return blocked_product(atom, outer); } // Transposed swizzled smem layout: maps (col, row) → same physical offset as @@ -296,18 +325,19 @@ __device__ __forceinline__ auto make_aliased_swizzled_layout_rc() { // Built by swapping modes of the original inner layout (before swizzle), which // guarantees correct cross-atom offsets when both dimensions have multiple atoms. template -__device__ __forceinline__ auto make_swizzled_layout_rc_transpose() { - using namespace cute; - using S = SmemSwizzle; - static_assert(ROWS % S::ATOM_ROWS == 0, "ROWS must be a multiple of the swizzle atom rows"); - static_assert(COLS % S::ATOM_COLS == 0, "COLS must be a multiple of the swizzle atom cols"); - // Build the inner (un-swizzled) tiled layout for the original (ROWS, COLS) layout - auto inner = tile_to_shape(make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, _1{})), - make_shape(Int{}, Int{})); - // Swap modes to get true transpose: result(c, r) == original(r, c) - auto inner_T = make_layout(get<1>(inner), get<0>(inner)); - return composition(typename S::type{}, inner_T); +__device__ __forceinline__ auto make_swizzled_layout_rc_transpose() +{ + using namespace cute; + using S = SmemSwizzle; + static_assert(ROWS % S::ATOM_ROWS == 0, "ROWS must be a multiple of the swizzle atom rows"); + static_assert(COLS % S::ATOM_COLS == 0, "COLS must be a multiple of the swizzle atom cols"); + // Build the inner (un-swizzled) tiled layout for the original (ROWS, COLS) layout + auto inner = tile_to_shape( + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{})), + make_shape(Int{}, Int{})); + // Swap modes to get true transpose: result(c, r) == original(r, c) + auto inner_T = make_layout(get<1>(inner), get<0>(inner)); + return composition(typename S::type{}, inner_T); } // ============================================================================= @@ -333,49 +363,46 @@ __device__ __forceinline__ auto make_swizzled_layout_rc_transpose() { // the default is the constexpr `VALID_ROWS` so non-varlen call sites fold to // the same SASS as before. template -__device__ __forceinline__ void load_tile_async(input_t* __restrict__ smem_dst, - input_t const* __restrict__ gmem_src, - int gmem_row_stride, int lane, - int valid_rows_rt = VALID_ROWS) { - using namespace cute; - constexpr int ROWS_PAD = size<0>(SmemShape{}); - constexpr int VALID_COLS = size<1>(SmemShape{}); - // Smem cols are padded up to the swizzle atom width. We always use the - // wide thread layout (4 thread-rows × 8 thread-cols × 1×8 val = 4 rows × - // 64 cols/pass for bf16) so each thread-row covers all 8 vec-cols of the - // Swizzle atom — the design contract that makes cp.async writes - // bank-conflict-free (each row consumes one full bank cycle, rows - // serialize across cycles). A "narrow" layout (½-atom-width per row) - // would force adjacent rows to compete for the same 16 banks, costing - // ~3-4× replay (observed as 12-way LDGSTS conflicts in d_split=2 ncu). - // When VALID_COLS < SMEM_COLS (D_SPLIT > 1 path), cp.async ZFILL drops - // the predicated-out cols as zeros without touching gmem — same mechanism - // as ZFILL'ing rows ≥ VALID_ROWS. The padded smem cells are unused. - constexpr int SMEM_COLS = next_multiple_of::ATOM_COLS>(VALID_COLS); - Tensor s_full = - make_tensor(make_smem_ptr(smem_dst), make_swizzled_layout_rc()); - Tensor g_full = make_tensor(make_gmem_ptr(gmem_src), - make_layout(make_shape(Int{}, Int{}), - make_stride(gmem_row_stride, Int<1>{}))); - - constexpr int VAL_COLS_PER_THREAD = Copy_prop::vec_bytes / sizeof(input_t); - static_assert(SMEM_COLS % VAL_COLS_PER_THREAD == 0, - "SMEM_COLS must be divisible by VAL_COLS_PER_THREAD"); - using ThrLayout = Layout, Stride<_8, _1>>; - static_assert(size<1>(ThrLayout{}) * VAL_COLS_PER_THREAD == SmemSwizzle::ATOM_COLS, - "wide thread layout must cover one full swizzle atom width per row"); - auto g2s = make_tiled_copy(Copy_Atom{}, ThrLayout{}, - Layout>>{}); - auto thr = g2s.get_slice(lane); - - auto id = make_identity_tensor(make_shape(Int{}, Int{})); - auto thr_id = thr.partition_S(id); - auto pred = make_tensor(shape(thr_id)); - CUTE_UNROLL - for (int i = 0; i < size(pred); ++i) { - pred(i) = (get<0>(thr_id(i)) < valid_rows_rt) && (get<1>(thr_id(i)) < VALID_COLS); - } - copy_if(g2s, pred, thr.partition_S(g_full), thr.partition_D(s_full)); +__device__ __forceinline__ void load_tile_async(input_t* __restrict__ smem_dst, input_t const* __restrict__ gmem_src, + int gmem_row_stride, int lane, int valid_rows_rt = VALID_ROWS) +{ + using namespace cute; + constexpr int ROWS_PAD = size<0>(SmemShape{}); + constexpr int VALID_COLS = size<1>(SmemShape{}); + // Smem cols are padded up to the swizzle atom width. We always use the + // wide thread layout (4 thread-rows × 8 thread-cols × 1×8 val = 4 rows × + // 64 cols/pass for bf16) so each thread-row covers all 8 vec-cols of the + // Swizzle atom — the design contract that makes cp.async writes + // bank-conflict-free (each row consumes one full bank cycle, rows + // serialize across cycles). A "narrow" layout (½-atom-width per row) + // would force adjacent rows to compete for the same 16 banks, costing + // ~3-4× replay (observed as 12-way LDGSTS conflicts in d_split=2 ncu). + // When VALID_COLS < SMEM_COLS (D_SPLIT > 1 path), cp.async ZFILL drops + // the predicated-out cols as zeros without touching gmem — same mechanism + // as ZFILL'ing rows ≥ VALID_ROWS. The padded smem cells are unused. + constexpr int SMEM_COLS = next_multiple_of::ATOM_COLS>(VALID_COLS); + Tensor s_full = make_tensor(make_smem_ptr(smem_dst), make_swizzled_layout_rc()); + Tensor g_full = make_tensor(make_gmem_ptr(gmem_src), + make_layout(make_shape(Int{}, Int{}), make_stride(gmem_row_stride, Int<1>{}))); + + constexpr int VAL_COLS_PER_THREAD = Copy_prop::vec_bytes / sizeof(input_t); + static_assert(SMEM_COLS % VAL_COLS_PER_THREAD == 0, "SMEM_COLS must be divisible by VAL_COLS_PER_THREAD"); + using ThrLayout = Layout, Stride<_8, _1>>; + static_assert(size<1>(ThrLayout{}) * VAL_COLS_PER_THREAD == SmemSwizzle::ATOM_COLS, + "wide thread layout must cover one full swizzle atom width per row"); + auto g2s = make_tiled_copy( + Copy_Atom{}, ThrLayout{}, Layout>>{}); + auto thr = g2s.get_slice(lane); + + auto id = make_identity_tensor(make_shape(Int{}, Int{})); + auto thr_id = thr.partition_S(id); + auto pred = make_tensor(shape(thr_id)); + CUTE_UNROLL + for (int i = 0; i < size(pred); ++i) + { + pred(i) = (get<0>(thr_id(i)) < valid_rows_rt) && (get<1>(thr_id(i)) < VALID_COLS); + } + copy_if(g2s, pred, thr.partition_S(g_full), thr.partition_D(s_full)); } // State load — D_SPLIT-conditional dispatch: @@ -396,56 +423,50 @@ __device__ __forceinline__ void load_tile_async(input_t* __restrict__ smem_dst, // followed by a `local_tile` to the (D_PER_CTA, DSTATE) slice this CTA // owns — the swizzle outer-stride is invariant across D_SPLIT. template -__device__ __forceinline__ void load_state_per_warp(SmemT& smem, - state_t const* __restrict__ state_ptr, - int64_t state_base, int warp, int lane) { - using namespace cute; - static_assert(NUM_WARPS == 4, "Expected 4 warps"); - static_assert(D_PER_CTA % NUM_WARPS == 0, "D_PER_CTA must be divisible by NUM_WARPS"); - constexpr int DIM_PER_WARP = D_PER_CTA / NUM_WARPS; - - // Single-local_tile path — swizzle layout sized to this CTA's - // D_PER_CTA slice; one local_tile splits it directly per-warp. - Tensor sState_full = make_tensor(make_smem_ptr(reinterpret_cast(smem.state)), - make_swizzled_layout_rc()); - Tensor gState_full = make_tensor(make_gmem_ptr(state_ptr + state_base), - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - - Tensor sState = local_tile(sState_full, make_shape(Int{}, Int{}), - make_coord(warp, _0{})); - Tensor gState = local_tile(gState_full, make_shape(Int{}, Int{}), - make_coord(warp, _0{})); - - constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); - auto g2s = - make_tiled_copy(Copy_Atom{}, - Layout, Stride<_8, _1>>{}, Layout>>{}); - auto thr = g2s.get_slice(lane); - copy(g2s, thr.partition_S(gState), thr.partition_D(sState)); +__device__ __forceinline__ void load_state_per_warp( + SmemT& smem, state_t const* __restrict__ state_ptr, int64_t state_base, int warp, int lane) +{ + using namespace cute; + static_assert(NUM_WARPS == 4, "Expected 4 warps"); + static_assert(D_PER_CTA % NUM_WARPS == 0, "D_PER_CTA must be divisible by NUM_WARPS"); + constexpr int DIM_PER_WARP = D_PER_CTA / NUM_WARPS; + + // Single-local_tile path — swizzle layout sized to this CTA's + // D_PER_CTA slice; one local_tile splits it directly per-warp. + Tensor sState_full = make_tensor( + make_smem_ptr(reinterpret_cast(smem.state)), make_swizzled_layout_rc()); + Tensor gState_full = make_tensor(make_gmem_ptr(state_ptr + state_base), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{}))); + + Tensor sState = local_tile(sState_full, make_shape(Int{}, Int{}), make_coord(warp, _0{})); + Tensor gState = local_tile(gState_full, make_shape(Int{}, Int{}), make_coord(warp, _0{})); + + constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); + auto g2s = make_tiled_copy(Copy_Atom{}, Layout, Stride<_8, _1>>{}, + Layout>>{}); + auto thr = g2s.get_slice(lane); + copy(g2s, thr.partition_S(gState), thr.partition_D(sState)); } template -__device__ __forceinline__ void load_state_cta(SmemT& smem, state_t const* __restrict__ state_ptr, - int64_t state_base, int tid) { - using namespace cute; - static_assert(NUM_WARPS == 4, "Expected 4 warps"); - - Tensor sState = make_tensor(make_smem_ptr(reinterpret_cast(smem.state)), - make_swizzled_layout_rc()); - Tensor gState = make_tensor(make_gmem_ptr(state_ptr + state_base), - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - - constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); - using ThrLayout = Layout, Stride<_8, _1>>; - constexpr int THR_ROWS = decltype(size<0>(ThrLayout{}))::value; - static_assert(D_PER_CTA % THR_ROWS == 0, - "D_PER_CTA must be divisible by the thread layout's row count"); - auto g2s = make_tiled_copy(Copy_Atom{}, ThrLayout{}, - Layout>>{}); - auto thr = g2s.get_slice(tid); - copy(g2s, thr.partition_S(gState), thr.partition_D(sState)); +__device__ __forceinline__ void load_state_cta( + SmemT& smem, state_t const* __restrict__ state_ptr, int64_t state_base, int tid) +{ + using namespace cute; + static_assert(NUM_WARPS == 4, "Expected 4 warps"); + + Tensor sState = make_tensor( + make_smem_ptr(reinterpret_cast(smem.state)), make_swizzled_layout_rc()); + Tensor gState = make_tensor(make_gmem_ptr(state_ptr + state_base), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{}))); + + constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); + using ThrLayout = Layout, Stride<_8, _1>>; + constexpr int THR_ROWS = decltype(size<0>(ThrLayout{}))::value; + static_assert(D_PER_CTA % THR_ROWS == 0, "D_PER_CTA must be divisible by the thread layout's row count"); + auto g2s = make_tiled_copy(Copy_Atom{}, ThrLayout{}, Layout>>{}); + auto thr = g2s.get_slice(tid); + copy(g2s, thr.partition_S(gState), thr.partition_D(sState)); } // ============================================================================= @@ -459,27 +480,38 @@ __device__ __forceinline__ void load_state_cta(SmemT& smem, state_t const* __res // decay broadcast in `compute_output_8bit` becomes a plain LDS (no per-element // __expf). Detected via SFINAE so the 2/4-byte storage (no decay field) is // unaffected. -namespace detail { +namespace detail +{ template -struct has_decay : std::false_type {}; +struct has_decay : std::false_type +{ +}; + template -struct has_decay().decay[0])>> : std::true_type {}; -} // namespace detail +struct has_decay().decay[0])>> : std::true_type +{ +}; +} // namespace detail template -__device__ __forceinline__ void compute_cumAdt(SmemT& smem, int lane, float A_val) { - float val = (lane < NPREDICTED) ? A_val * smem.dt_proc[lane] : 0.f; - // Inclusive prefix sum (Hillis-Steele) - for (int offset = 1; offset < NPREDICTED; offset *= 2) { - float other = __shfl_up_sync(constants::MASK_ALL_LANES, val, offset); - if (lane >= offset) val += other; - } - if (lane < NPREDICTED) { - smem.cumAdt[lane] = val; - if constexpr (detail::has_decay::value) { - smem.decay[lane] = __expf(val); +__device__ __forceinline__ void compute_cumAdt(SmemT& smem, int lane, float A_val) +{ + float val = (lane < NPREDICTED) ? A_val * smem.dt_proc[lane] : 0.f; + // Inclusive prefix sum (Hillis-Steele) + for (int offset = 1; offset < NPREDICTED; offset *= 2) + { + float other = __shfl_up_sync(constants::MASK_ALL_LANES, val, offset); + if (lane >= offset) + val += other; + } + if (lane < NPREDICTED) + { + smem.cumAdt[lane] = val; + if constexpr (detail::has_decay::value) + { + smem.decay[lane] = __expf(val); + } } - } } // Load phase. Split into two halves around the PDL barrier (`gdc_wait`): @@ -519,126 +551,127 @@ __device__ __forceinline__ void compute_cumAdt(SmemT& smem, int lane, float A_va // non-varlen, runtime int in varlen). Used as the cp.async row predicate // and the dt/scalar lane predicate so trailing rows past `seq_len` ZFILL to // zero in smem. -template -__device__ __forceinline__ void load_pre_pdl_wait_data( - SmemT& smem, CheckpointingSsuParams const& params, int lane, int warp, int d_tile, int head, - int group_idx, int64_t cache_slot, int buf_read, float A_val, float dt_bias_val, - int64_t dt_seq_base, int64_t z_seq_base, int seq_len) { - constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 - static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); - static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); - - int const d_tile_off = d_tile * D_PER_CTA; - - auto const* __restrict__ z_ptr = reinterpret_cast(params.z); - auto const* __restrict__ old_x_ptr = reinterpret_cast(params.old_x); - auto const* __restrict__ old_B_ptr = reinterpret_cast(params.old_B); - auto const* __restrict__ old_dt_ptr = reinterpret_cast(params.old_dt); - auto const* __restrict__ old_cumAdt_ptr = reinterpret_cast(params.old_cumAdt); - auto const* __restrict__ dt_ptr = reinterpret_cast(params.dt); - - int64_t const ox_base = cache_slot * params.old_x_stride_seq + head * DIM + d_tile_off; - int64_t const oB_base = cache_slot * params.old_B_stride_seq + - buf_read * params.old_B_stride_dbuf + group_idx * DSTATE; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - // ZShape: shrunk to the swizzle atom's row extent (z is read via - // partition_C alias, the physical buffer only needs to be one swizzle - // row-atom tall). - using ZShape = cute::Shape, cute::Int>; - // old_B / old_x's smem row count = replay matmul K-axis = MAX_WINDOW_PAD_MMA_K. - using OldBShape = cute::Shape, cute::Int>; - using OxShape = cute::Shape, cute::Int>; - - // ── State: per-CTA D-slice ([D_PER_CTA, DSTATE]). Dispatch on D_SPLIT - // (= DIM == D_PER_CTA): per-warp coalesced load when one CTA owns the - // full head's D, cooperative 128-thread load when D is sharded. ── - { - auto const* __restrict__ state_ptr = reinterpret_cast(params.state); - int64_t const state_base = cache_slot * params.state_stride_seq + (int64_t)head * DIM * DSTATE + - (int64_t)d_tile_off * DSTATE; - if constexpr (DIM == D_PER_CTA) { - // D_SPLIT=1: per-warp partition (warp w loads contiguous DIM/4 D-rows). - load_state_per_warp(smem, state_ptr, state_base, warp, - lane); - } else { - // D_SPLIT>=2: 128-thread cooperative load (per-warp doesn't divide - // cleanly when D_PER_CTA/4 is too small for the (4,8) thread atom). - int const tid = warp * warpSize + lane; - load_state_cta(smem, state_ptr, state_base, tid); +template +__device__ __forceinline__ void load_pre_pdl_wait_data(SmemT& smem, CheckpointingSsuParams const& params, int lane, + int warp, int d_tile, int head, int group_idx, int64_t cache_slot, int buf_read, float A_val, float dt_bias_val, + int64_t dt_seq_base, int64_t z_seq_base, int seq_len) +{ + constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 + static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); + static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); + + int const d_tile_off = d_tile * D_PER_CTA; + + auto const* __restrict__ z_ptr = reinterpret_cast(params.z); + auto const* __restrict__ old_x_ptr = reinterpret_cast(params.old_x); + auto const* __restrict__ old_B_ptr = reinterpret_cast(params.old_B); + auto const* __restrict__ old_dt_ptr = reinterpret_cast(params.old_dt); + auto const* __restrict__ old_cumAdt_ptr = reinterpret_cast(params.old_cumAdt); + auto const* __restrict__ dt_ptr = reinterpret_cast(params.dt); + + int64_t const ox_base = cache_slot * params.old_x_stride_seq + head * DIM + d_tile_off; + int64_t const oB_base + = cache_slot * params.old_B_stride_seq + buf_read * params.old_B_stride_dbuf + group_idx * DSTATE; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + // ZShape: shrunk to the swizzle atom's row extent (z is read via + // partition_C alias, the physical buffer only needs to be one swizzle + // row-atom tall). + using ZShape = cute::Shape, cute::Int>; + // old_B / old_x's smem row count = replay matmul K-axis = MAX_WINDOW_PAD_MMA_K. + using OldBShape = cute::Shape, cute::Int>; + using OxShape = cute::Shape, cute::Int>; + + // ── State: per-CTA D-slice ([D_PER_CTA, DSTATE]). Dispatch on D_SPLIT + // (= DIM == D_PER_CTA): per-warp coalesced load when one CTA owns the + // full head's D, cooperative 128-thread load when D is sharded. ── + { + auto const* __restrict__ state_ptr = reinterpret_cast(params.state); + int64_t const state_base + = cache_slot * params.state_stride_seq + (int64_t) head * DIM * DSTATE + (int64_t) d_tile_off * DSTATE; + if constexpr (DIM == D_PER_CTA) + { + // D_SPLIT=1: per-warp partition (warp w loads contiguous DIM/4 D-rows). + load_state_per_warp(smem, state_ptr, state_base, warp, lane); + } + else + { + // D_SPLIT>=2: 128-thread cooperative load (per-warp doesn't divide + // cleanly when D_PER_CTA/4 is too small for the (4,8) thread atom). + int const tid = warp * warpSize + lane; + load_state_cta(smem, state_ptr, state_base, tid); + } + } + + // ── old_B: redundant on all 4 warps (each warp's replay consumes full + // DSTATE). Identical payloads to same smem dest — final bytes + // deterministic. VALID_ROWS = MAX_WINDOW (cache rows). ── + load_tile_async(smem.old_B, old_B_ptr + oB_base, params.old_B_stride_token, lane); + + // ── old_x: redundant on all 4 warps (small, simpler than partitioning). + // VALID_ROWS = MAX_WINDOW (cache rows). ── + load_tile_async(smem.old_x, old_x_ptr + ox_base, params.old_x_stride_token, lane); + + // ── z: W3 only (Phase-2 read, final __syncthreads makes it visible). + // Sourced from in_proj — not from conv1d — so safe to issue pre-wait. ── + if (warp == 3 && z_ptr) + { + int64_t const z_base = z_seq_base + head * DIM + d_tile_off; + load_tile_async(smem.z, z_ptr + z_base, params.z_stride_token, lane, seq_len); } - } - - // ── old_B: redundant on all 4 warps (each warp's replay consumes full - // DSTATE). Identical payloads to same smem dest — final bytes - // deterministic. VALID_ROWS = MAX_WINDOW (cache rows). ── - load_tile_async(smem.old_B, old_B_ptr + oB_base, params.old_B_stride_token, - lane); - - // ── old_x: redundant on all 4 warps (small, simpler than partitioning). - // VALID_ROWS = MAX_WINDOW (cache rows). ── - load_tile_async(smem.old_x, old_x_ptr + ox_base, params.old_x_stride_token, - lane); - - // ── z: W3 only (Phase-2 read, final __syncthreads makes it visible). - // Sourced from in_proj — not from conv1d — so safe to issue pre-wait. ── - if (warp == 3 && z_ptr) { - int64_t const z_base = z_seq_base + head * DIM + d_tile_off; - load_tile_async(smem.z, z_ptr + z_base, params.z_stride_token, lane, - seq_len); - } - - // Commit the cache cp.async group BEFORE the caller's `gdc_wait()` so the - // hardware actually issues the gmem→smem transfers while the wait is in - // flight (without commit, the operations sit pending and only kick off - // once the post-wait commit fires — no overlap). Placed immediately after - // the last cp.async (z); the synchronous LDGs + cumAdt scan below are not - // part of any pipeline group and run in parallel with the in-flight - // transfers. The post half issues a second group; `__pipeline_wait_prior(0)` - // there drains both. - __pipeline_commit(); - - // ── Scalar loads + cumAdt cumsum: redundant per warp. - // old_dt / old_cumAdt: load up to MAX_WINDOW lanes (cache scalars). - // dt_proc: load up to NPREDICTED lanes (new-token scalars from in_proj). - // Synchronous LDG + plain smem stores — no cp.async. Writes from 4 - // warps to the same slots are idempotent (same payloads). ── - static_assert(MAX_WINDOW <= warpSize, "MAX_WINDOW must fit in a single warp"); - if (lane < MAX_WINDOW) { - int64_t const dt_rd_base = cache_slot * params.old_dt_stride_seq + - buf_read * params.old_dt_stride_dbuf + - head * params.old_dt_stride_head; - smem.old_dt[lane] = old_dt_ptr[dt_rd_base + lane]; - - int64_t const ca_rd_base = cache_slot * params.old_cumAdt_stride_seq + - buf_read * params.old_cumAdt_stride_dbuf + - head * params.old_cumAdt_stride_head; - smem.old_cumAdt[lane] = old_cumAdt_ptr[ca_rd_base + lane]; - } - // dt → softplus → smem.dt_proc. Under varlen the active lane range is - // `[0, seq_len)`; lanes `[seq_len, NPREDICTED)` are left uninitialized — - // `compute_cumAdt` will scan over them and produce garbage in the - // `cumAdt[seq_len:NPREDICTED]` tail, but every downstream consumer - // (`compute_CB_scaled_2warp` mask, output STG, dt_proc/cumAdt tape writes) - // is gated on `seq_len`, so the garbage never reaches gmem or contaminates - // valid rows. - // - // Per-lane stride along the T-axis is `dt_stride_token` in both layouts - // (4D batch and 1D packed varlen); the caller bakes `head` into - // `dt_seq_base` so the inner indexing is `dt_seq_base + lane * - // dt_stride_token`. - if (lane < seq_len) { - float dt_val = toFloat(dt_ptr[dt_seq_base + (int64_t)lane * params.dt_stride_token]); - dt_val += dt_bias_val; - if (params.dt_softplus) dt_val = thresholded_softplus(dt_val); - smem.dt_proc[lane] = dt_val; - } - // cumAdt = cumsum(A * dt_proc) — warp-local Hillis-Steele shuffle. Each - // of the 4 warps runs the same reduction on identical inputs (dt_proc - // just written above) and writes the same smem.cumAdt slots. - compute_cumAdt(smem, lane, A_val); + + // Commit the cache cp.async group BEFORE the caller's `gdc_wait()` so the + // hardware actually issues the gmem→smem transfers while the wait is in + // flight (without commit, the operations sit pending and only kick off + // once the post-wait commit fires — no overlap). Placed immediately after + // the last cp.async (z); the synchronous LDGs + cumAdt scan below are not + // part of any pipeline group and run in parallel with the in-flight + // transfers. The post half issues a second group; `__pipeline_wait_prior(0)` + // there drains both. + __pipeline_commit(); + + // ── Scalar loads + cumAdt cumsum: redundant per warp. + // old_dt / old_cumAdt: load up to MAX_WINDOW lanes (cache scalars). + // dt_proc: load up to NPREDICTED lanes (new-token scalars from in_proj). + // Synchronous LDG + plain smem stores — no cp.async. Writes from 4 + // warps to the same slots are idempotent (same payloads). ── + static_assert(MAX_WINDOW <= warpSize, "MAX_WINDOW must fit in a single warp"); + if (lane < MAX_WINDOW) + { + int64_t const dt_rd_base = cache_slot * params.old_dt_stride_seq + buf_read * params.old_dt_stride_dbuf + + head * params.old_dt_stride_head; + smem.old_dt[lane] = old_dt_ptr[dt_rd_base + lane]; + + int64_t const ca_rd_base = cache_slot * params.old_cumAdt_stride_seq + buf_read * params.old_cumAdt_stride_dbuf + + head * params.old_cumAdt_stride_head; + smem.old_cumAdt[lane] = old_cumAdt_ptr[ca_rd_base + lane]; + } + // dt → softplus → smem.dt_proc. Under varlen the active lane range is + // `[0, seq_len)`; lanes `[seq_len, NPREDICTED)` are left uninitialized — + // `compute_cumAdt` will scan over them and produce garbage in the + // `cumAdt[seq_len:NPREDICTED]` tail, but every downstream consumer + // (`compute_CB_scaled_2warp` mask, output STG, dt_proc/cumAdt tape writes) + // is gated on `seq_len`, so the garbage never reaches gmem or contaminates + // valid rows. + // + // Per-lane stride along the T-axis is `dt_stride_token` in both layouts + // (4D batch and 1D packed varlen); the caller bakes `head` into + // `dt_seq_base` so the inner indexing is `dt_seq_base + lane * + // dt_stride_token`. + if (lane < seq_len) + { + float dt_val = toFloat(dt_ptr[dt_seq_base + (int64_t) lane * params.dt_stride_token]); + dt_val += dt_bias_val; + if (params.dt_softplus) + dt_val = thresholded_softplus(dt_val); + smem.dt_proc[lane] = dt_val; + } + // cumAdt = cumsum(A * dt_proc) — warp-local Hillis-Steele shuffle. Each + // of the 4 warps runs the same reduction on identical inputs (dt_proc + // just written above) and writes the same smem.cumAdt slots. + compute_cumAdt(smem, lane, A_val); } // Post-wait half. Issues cp.async for conv1d outputs (x, B, C) and drains @@ -654,59 +687,58 @@ __device__ __forceinline__ void load_pre_pdl_wait_data( // and thus can't rematerialize through). Saves ~6 registers vs. pre-computed // bases. template -__device__ __forceinline__ void load_post_pdl_wait_data(SmemT& smem, - CheckpointingSsuParams const& params, - int lane, int warp, int d_tile, int head, - int group_idx, int64_t outer, int seq_len) { - constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 - static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); - static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); - - int const d_tile_off = d_tile * D_PER_CTA; - - auto const* __restrict__ B_ptr = reinterpret_cast(params.B); - auto const* __restrict__ C_ptr = reinterpret_cast(params.C); - auto const* __restrict__ x_ptr = reinterpret_cast(params.x); - - int64_t const B_base = outer * params.B_stride_seq + (int64_t)group_idx * DSTATE; - int64_t const C_base = outer * params.C_stride_seq + (int64_t)group_idx * DSTATE; - int64_t const x_base = outer * params.x_stride_seq + head * DIM + d_tile_off; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; - // CShape: first dim shrunk to the swizzle atom's row extent so cp.async - // writes don't spill past the (also shrunk) C smem buffer. When NPREDICTED - // > ATOM_ROWS this falls back to NPREDICTED_PAD_MMA_M. - using CShape = cute::Shape, cute::Int>; - // B's smem row count = matmul-1 N-axis = NPREDICTED_PAD_MMA_N. - using BShape = cute::Shape, cute::Int>; - using XShape = cute::Shape, cute::Int>; - - // ── B: redundant on W0, W1 (both do 2-warp CB compute) ── - if (warp < 2) { - load_tile_async(smem.B, B_ptr + B_base, params.B_stride_token, lane, - seq_len); - } - // ── C: redundant on all 4 warps — chain matmul-3 reads smem.C from every - // warp, so each warp must see its own cp.async without a cross-warp sync. - // Identical payloads to same smem dest (same pattern as old_B / old_x). ── - load_tile_async(smem.C, C_ptr + C_base, params.C_stride_token, lane, seq_len); - - // ── x: W2 only (Phase-2 read, final __syncthreads makes it visible) ── - if (warp == 2) { - load_tile_async(smem.x, x_ptr + x_base, params.x_stride_token, lane, - seq_len); - } - - // Commit the conv1d cp.async group and drain BOTH groups: the cache - // group committed in `load_pre_pdl_wait_data` (pre-`gdc_wait`) and this - // conv1d group. `__pipeline_wait_prior(0)` waits for ≤0 pending groups. - // __syncwarp() provides acquire semantics across the 32 lanes of each - // warp. No cross-warp sync here; the only __syncthreads is after - // CB + replay. - __pipeline_commit(); - __pipeline_wait_prior(0); - __syncwarp(); +__device__ __forceinline__ void load_post_pdl_wait_data(SmemT& smem, CheckpointingSsuParams const& params, int lane, + int warp, int d_tile, int head, int group_idx, int64_t outer, int seq_len) +{ + constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 + static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); + static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); + + int const d_tile_off = d_tile * D_PER_CTA; + + auto const* __restrict__ B_ptr = reinterpret_cast(params.B); + auto const* __restrict__ C_ptr = reinterpret_cast(params.C); + auto const* __restrict__ x_ptr = reinterpret_cast(params.x); + + int64_t const B_base = outer * params.B_stride_seq + (int64_t) group_idx * DSTATE; + int64_t const C_base = outer * params.C_stride_seq + (int64_t) group_idx * DSTATE; + int64_t const x_base = outer * params.x_stride_seq + head * DIM + d_tile_off; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; + // CShape: first dim shrunk to the swizzle atom's row extent so cp.async + // writes don't spill past the (also shrunk) C smem buffer. When NPREDICTED + // > ATOM_ROWS this falls back to NPREDICTED_PAD_MMA_M. + using CShape = cute::Shape, cute::Int>; + // B's smem row count = matmul-1 N-axis = NPREDICTED_PAD_MMA_N. + using BShape = cute::Shape, cute::Int>; + using XShape = cute::Shape, cute::Int>; + + // ── B: redundant on W0, W1 (both do 2-warp CB compute) ── + if (warp < 2) + { + load_tile_async(smem.B, B_ptr + B_base, params.B_stride_token, lane, seq_len); + } + // ── C: redundant on all 4 warps — chain matmul-3 reads smem.C from every + // warp, so each warp must see its own cp.async without a cross-warp sync. + // Identical payloads to same smem dest (same pattern as old_B / old_x). ── + load_tile_async(smem.C, C_ptr + C_base, params.C_stride_token, lane, seq_len); + + // ── x: W2 only (Phase-2 read, final __syncthreads makes it visible) ── + if (warp == 2) + { + load_tile_async(smem.x, x_ptr + x_base, params.x_stride_token, lane, seq_len); + } + + // Commit the conv1d cp.async group and drain BOTH groups: the cache + // group committed in `load_pre_pdl_wait_data` (pre-`gdc_wait`) and this + // conv1d group. `__pipeline_wait_prior(0)` waits for ≤0 pending groups. + // __syncwarp() provides acquire semantics across the 32 lanes of each + // warp. No cross-warp sync here; the only __syncthreads is after + // CB + replay. + __pipeline_commit(); + __pipeline_wait_prior(0); + __syncwarp(); } // Single-pass load — used when `params.enable_pdl == false`. All cp.async @@ -719,105 +751,106 @@ __device__ __forceinline__ void load_post_pdl_wait_data(SmemT& smem, // the wait. When PDL is paired with an upstream conv1d, the split's // cache-load-during-wait overlap dominates; when not paired, the split is // pure overhead (gdc_wait is a no-op, but the cp.async are delayed). -template -__device__ __forceinline__ void load_data(SmemT& smem, CheckpointingSsuParams const& params, - int lane, int warp, int d_tile, int head, int group_idx, - int64_t cache_slot, int buf_read, float A_val, - float dt_bias_val, int64_t outer, int seq_len) { - constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 - static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); - static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); - - int const d_tile_off = d_tile * D_PER_CTA; - - auto const* __restrict__ B_ptr = reinterpret_cast(params.B); - auto const* __restrict__ C_ptr = reinterpret_cast(params.C); - auto const* __restrict__ x_ptr = reinterpret_cast(params.x); - auto const* __restrict__ z_ptr = reinterpret_cast(params.z); - auto const* __restrict__ old_x_ptr = reinterpret_cast(params.old_x); - auto const* __restrict__ old_B_ptr = reinterpret_cast(params.old_B); - auto const* __restrict__ old_dt_ptr = reinterpret_cast(params.old_dt); - auto const* __restrict__ old_cumAdt_ptr = reinterpret_cast(params.old_cumAdt); - auto const* __restrict__ dt_ptr = reinterpret_cast(params.dt); - - int64_t const B_base = outer * params.B_stride_seq + (int64_t)group_idx * DSTATE; - int64_t const C_base = outer * params.C_stride_seq + (int64_t)group_idx * DSTATE; - int64_t const x_base = outer * params.x_stride_seq + head * DIM + d_tile_off; - int64_t const ox_base = cache_slot * params.old_x_stride_seq + head * DIM + d_tile_off; - int64_t const oB_base = cache_slot * params.old_B_stride_seq + - buf_read * params.old_B_stride_dbuf + group_idx * DSTATE; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - using CShape = cute::Shape, cute::Int>; - using BShape = cute::Shape, cute::Int>; - using XShape = cute::Shape, cute::Int>; - using ZShape = cute::Shape, cute::Int>; - using OldBShape = cute::Shape, cute::Int>; - using OxShape = cute::Shape, cute::Int>; - - // ── State: per-CTA D-slice ([D_PER_CTA, DSTATE]) ── - { - auto const* __restrict__ state_ptr = reinterpret_cast(params.state); - int64_t const state_base = cache_slot * params.state_stride_seq + (int64_t)head * DIM * DSTATE + - (int64_t)d_tile_off * DSTATE; - if constexpr (DIM == D_PER_CTA) { - load_state_per_warp(smem, state_ptr, state_base, warp, - lane); - } else { - int const tid = warp * warpSize + lane; - load_state_cta(smem, state_ptr, state_base, tid); +template +__device__ __forceinline__ void load_data(SmemT& smem, CheckpointingSsuParams const& params, int lane, int warp, + int d_tile, int head, int group_idx, int64_t cache_slot, int buf_read, float A_val, float dt_bias_val, + int64_t outer, int seq_len) +{ + constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 + static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); + static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); + + int const d_tile_off = d_tile * D_PER_CTA; + + auto const* __restrict__ B_ptr = reinterpret_cast(params.B); + auto const* __restrict__ C_ptr = reinterpret_cast(params.C); + auto const* __restrict__ x_ptr = reinterpret_cast(params.x); + auto const* __restrict__ z_ptr = reinterpret_cast(params.z); + auto const* __restrict__ old_x_ptr = reinterpret_cast(params.old_x); + auto const* __restrict__ old_B_ptr = reinterpret_cast(params.old_B); + auto const* __restrict__ old_dt_ptr = reinterpret_cast(params.old_dt); + auto const* __restrict__ old_cumAdt_ptr = reinterpret_cast(params.old_cumAdt); + auto const* __restrict__ dt_ptr = reinterpret_cast(params.dt); + + int64_t const B_base = outer * params.B_stride_seq + (int64_t) group_idx * DSTATE; + int64_t const C_base = outer * params.C_stride_seq + (int64_t) group_idx * DSTATE; + int64_t const x_base = outer * params.x_stride_seq + head * DIM + d_tile_off; + int64_t const ox_base = cache_slot * params.old_x_stride_seq + head * DIM + d_tile_off; + int64_t const oB_base + = cache_slot * params.old_B_stride_seq + buf_read * params.old_B_stride_dbuf + group_idx * DSTATE; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + using CShape = cute::Shape, cute::Int>; + using BShape = cute::Shape, cute::Int>; + using XShape = cute::Shape, cute::Int>; + using ZShape = cute::Shape, cute::Int>; + using OldBShape = cute::Shape, cute::Int>; + using OxShape = cute::Shape, cute::Int>; + + // ── State: per-CTA D-slice ([D_PER_CTA, DSTATE]) ── + { + auto const* __restrict__ state_ptr = reinterpret_cast(params.state); + int64_t const state_base + = cache_slot * params.state_stride_seq + (int64_t) head * DIM * DSTATE + (int64_t) d_tile_off * DSTATE; + if constexpr (DIM == D_PER_CTA) + { + load_state_per_warp(smem, state_ptr, state_base, warp, lane); + } + else + { + int const tid = warp * warpSize + lane; + load_state_cta(smem, state_ptr, state_base, tid); + } + } + + if (warp < 2) + { + load_tile_async(smem.B, B_ptr + B_base, params.B_stride_token, lane, seq_len); + } + load_tile_async(smem.C, C_ptr + C_base, params.C_stride_token, lane, seq_len); + + load_tile_async(smem.old_B, old_B_ptr + oB_base, params.old_B_stride_token, lane); + load_tile_async(smem.old_x, old_x_ptr + ox_base, params.old_x_stride_token, lane); + + if (warp == 2) + { + load_tile_async(smem.x, x_ptr + x_base, params.x_stride_token, lane, seq_len); + } + if (warp == 3 && z_ptr) + { + int64_t const z_base = outer * params.z_stride_seq + head * DIM + d_tile_off; + load_tile_async(smem.z, z_ptr + z_base, params.z_stride_token, lane, seq_len); + } + + // ── Scalar loads (overlap with cp.async) + cumAdt cumsum ── + static_assert(MAX_WINDOW <= warpSize, "MAX_WINDOW must fit in a single warp"); + if (lane < MAX_WINDOW) + { + int64_t const dt_rd_base = cache_slot * params.old_dt_stride_seq + buf_read * params.old_dt_stride_dbuf + + head * params.old_dt_stride_head; + smem.old_dt[lane] = old_dt_ptr[dt_rd_base + lane]; + + int64_t const ca_rd_base = cache_slot * params.old_cumAdt_stride_seq + buf_read * params.old_cumAdt_stride_dbuf + + head * params.old_cumAdt_stride_head; + smem.old_cumAdt[lane] = old_cumAdt_ptr[ca_rd_base + lane]; } - } - - if (warp < 2) { - load_tile_async(smem.B, B_ptr + B_base, params.B_stride_token, lane, - seq_len); - } - load_tile_async(smem.C, C_ptr + C_base, params.C_stride_token, lane, seq_len); - - load_tile_async(smem.old_B, old_B_ptr + oB_base, params.old_B_stride_token, - lane); - load_tile_async(smem.old_x, old_x_ptr + ox_base, params.old_x_stride_token, - lane); - - if (warp == 2) { - load_tile_async(smem.x, x_ptr + x_base, params.x_stride_token, lane, - seq_len); - } - if (warp == 3 && z_ptr) { - int64_t const z_base = outer * params.z_stride_seq + head * DIM + d_tile_off; - load_tile_async(smem.z, z_ptr + z_base, params.z_stride_token, lane, - seq_len); - } - - // ── Scalar loads (overlap with cp.async) + cumAdt cumsum ── - static_assert(MAX_WINDOW <= warpSize, "MAX_WINDOW must fit in a single warp"); - if (lane < MAX_WINDOW) { - int64_t const dt_rd_base = cache_slot * params.old_dt_stride_seq + - buf_read * params.old_dt_stride_dbuf + - head * params.old_dt_stride_head; - smem.old_dt[lane] = old_dt_ptr[dt_rd_base + lane]; - - int64_t const ca_rd_base = cache_slot * params.old_cumAdt_stride_seq + - buf_read * params.old_cumAdt_stride_dbuf + - head * params.old_cumAdt_stride_head; - smem.old_cumAdt[lane] = old_cumAdt_ptr[ca_rd_base + lane]; - } - int64_t const dt_seq_base_local = outer * params.dt_stride_seq + head; - if (lane < seq_len) { - float dt_val = toFloat(dt_ptr[dt_seq_base_local + (int64_t)lane * params.dt_stride_token]); - dt_val += dt_bias_val; - if (params.dt_softplus) dt_val = thresholded_softplus(dt_val); - smem.dt_proc[lane] = dt_val; - } - compute_cumAdt(smem, lane, A_val); - - __pipeline_commit(); - __pipeline_wait_prior(0); - __syncwarp(); + int64_t const dt_seq_base_local = outer * params.dt_stride_seq + head; + if (lane < seq_len) + { + float dt_val = toFloat(dt_ptr[dt_seq_base_local + (int64_t) lane * params.dt_stride_token]); + dt_val += dt_bias_val; + if (params.dt_softplus) + dt_val = thresholded_softplus(dt_val); + smem.dt_proc[lane] = dt_val; + } + compute_cumAdt(smem, lane, A_val); + + __pipeline_commit(); + __pipeline_wait_prior(0); + __syncwarp(); } // (compute_cumAdt moved above load_pre_pdl_wait_data so it can be called from there) @@ -832,129 +865,132 @@ __device__ __forceinline__ void load_data(SmemT& smem, CheckpointingSsuParams co // Varlen passes the per-sequence `seq_len ≤ NPREDICTED`; rows/cols past it // get zeroed so downstream matmul-4 / chain matmul-3 see zeros there. template -__device__ __forceinline__ void compute_CB_scaled_2warp(SmemT& smem, int warp, int lane, - int seq_len) { - using namespace cute; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; - // 2-warp output split: each warp owns NPREDICTED_PAD_MMA_M / 2 cols of - // the (NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M) CB tile. Must be a - // multiple of the MMA atom's N for the partition to be atom-aligned - // (currently 8 == MMA_prop::N for NPREDICTED_PAD_MMA_M=16; if M-pad ever - // grows, this still holds as long as M-pad is a multiple of 2 * MMA_prop::N). - constexpr int N_HALF = NPREDICTED_PAD_MMA_M / 2; - static_assert(N_HALF % MMA_prop::N == 0, - "compute_CB_scaled_2warp: NPREDICTED_PAD_MMA_M / 2 must be a multiple of MMA::N"); - - // CB_scaled output tile layout (used by both warp 0 compute and warp 1 - // zero-fill when smem.B has only 8 rows). Row stride matches the buffer's - // padded width (one swizzle atom of `input_t`). - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - auto layout_cb_swz = - make_swizzled_layout_rc(); - - // ── NPREDICTED_PAD_MMA_N == 8: warp 1 has no valid B rows to read. But - // CB_scaled[:, 8:16] must still be zero so matmul-4's K-reduction sees - // zeros for j ≥ NPREDICTED. Do a simple 32-thread zero-fill and return. - if constexpr (NPREDICTED_PAD_MMA_N == 8) { - if (warp == 1) { - auto* __restrict__ cb = reinterpret_cast(smem.CB_scaled); - constexpr int COLS_TO_CLEAR = NPREDICTED_PAD_MMA_M - N_HALF; // 8 +__device__ __forceinline__ void compute_CB_scaled_2warp(SmemT& smem, int warp, int lane, int seq_len) +{ + using namespace cute; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; + // 2-warp output split: each warp owns NPREDICTED_PAD_MMA_M / 2 cols of + // the (NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M) CB tile. Must be a + // multiple of the MMA atom's N for the partition to be atom-aligned + // (currently 8 == MMA_prop::N for NPREDICTED_PAD_MMA_M=16; if M-pad ever + // grows, this still holds as long as M-pad is a multiple of 2 * MMA_prop::N). + constexpr int N_HALF = NPREDICTED_PAD_MMA_M / 2; + static_assert( + N_HALF % MMA_prop::N == 0, "compute_CB_scaled_2warp: NPREDICTED_PAD_MMA_M / 2 must be a multiple of MMA::N"); + + // CB_scaled output tile layout (used by both warp 0 compute and warp 1 + // zero-fill when smem.B has only 8 rows). Row stride matches the buffer's + // padded width (one swizzle atom of `input_t`). + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + auto layout_cb_swz = make_swizzled_layout_rc(); + + // ── NPREDICTED_PAD_MMA_N == 8: warp 1 has no valid B rows to read. But + // CB_scaled[:, 8:16] must still be zero so matmul-4's K-reduction sees + // zeros for j ≥ NPREDICTED. Do a simple 32-thread zero-fill and return. + if constexpr (NPREDICTED_PAD_MMA_N == 8) + { + if (warp == 1) + { + auto* __restrict__ cb = reinterpret_cast(smem.CB_scaled); + constexpr int COLS_TO_CLEAR = NPREDICTED_PAD_MMA_M - N_HALF; // 8 #pragma unroll - for (int i = lane; i < NPREDICTED_PAD_MMA_M * COLS_TO_CLEAR; i += warpSize) { - int const r = i / COLS_TO_CLEAR; - int const c = N_HALF + (i % COLS_TO_CLEAR); - cb[layout_cb_swz(r, c)] = MMA_prop::operand_t(0.f); - } - return; + for (int i = lane; i < NPREDICTED_PAD_MMA_M * COLS_TO_CLEAR; i += warpSize) + { + int const r = i / COLS_TO_CLEAR; + int const c = N_HALF + (i % COLS_TO_CLEAR); + cb[layout_cb_swz(r, c)] = MMA_prop::operand_t(0.f); + } + return; + } } - } - - // ── Swizzled smem views ── - // C is padded to NPREDICTED_PAD_MMA_M; B has NPREDICTED_PAD_MMA_N rows. - // Use NPREDICTED_PAD_MMA_N for smem_B so the physical layout matches the - // write layout from load_tile_async — `tile_to_shape` produces different - // outer strides for (8, 128) vs (16, 128). - // Aliased C view: physical buffer is just next_multiple_of(NPREDICTED) - // rows tall but the MMA atom needs M=16; second m-tile aliases first m-tile - // (predicated rows discarded at output store). - auto layout_C = - make_aliased_swizzled_layout_rc(); - auto layout_B = make_swizzled_layout_rc(); - Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C); - Tensor smem_B = make_tensor(make_smem_ptr(reinterpret_cast(smem.B)), layout_B); - - // ── TiledMMA: _1x_1 = 32 threads, one [16, 8] atom ── - auto tiled_mma = - make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(lane); - - // ── K-tile A operand (C): full [NPREDICTED_PAD_MMA_M, K_TILE], shared by both warps ── - constexpr int K_TILE = MMA_prop::K_BIG; - Tensor smem_C_tiled = local_tile(smem_C, make_tile(Int{}, Int{}), - make_coord(_0{}, _)); - - // ── K-tile B operand ── - // NPREDICTED_PAD_MMA_N == 16: warp 0 → N=[0,8), warp 1 → N=[8,16). - // NPREDICTED_PAD_MMA_N == 8 : only warp 0 runs (warp 1 took the early - // exit above), tile at (_0, _). - Tensor smem_B_half = - local_tile(smem_B, make_tile(Int{}, Int{}), make_coord(warp, _)); - - // ── Register fragments ── - Tensor frag_A = thr_mma.partition_fragment_A(smem_C_tiled(_, _, _0{})); - Tensor frag_B = thr_mma.partition_fragment_B(smem_B_half(_, _, _0{})); - - // ── Output accumulator: [NPREDICTED_PAD_MMA_M, N_HALF] f32 ── - auto layout_cb_half = make_layout(make_shape(Int{}, Int{})); - Tensor frag_acc = thr_mma.partition_fragment_C(make_tensor((float*)nullptr, layout_cb_half)); - clear(frag_acc); - - // ── S2R copies ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(lane); - Tensor smem_C_s2r = s2r_thr_A.partition_S(smem_C_tiled); - Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); - - auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B = s2r_B.get_slice(lane); - Tensor smem_B_s2r = s2r_thr_B.partition_S(smem_B_half); - Tensor frag_B_view = s2r_thr_B.retile_D(frag_B); - - // ── Gemm: 8 K-tiles, 1 HMMA each ── - constexpr int NUM_K_TILES = DSTATE / K_TILE; + + // ── Swizzled smem views ── + // C is padded to NPREDICTED_PAD_MMA_M; B has NPREDICTED_PAD_MMA_N rows. + // Use NPREDICTED_PAD_MMA_N for smem_B so the physical layout matches the + // write layout from load_tile_async — `tile_to_shape` produces different + // outer strides for (8, 128) vs (16, 128). + // Aliased C view: physical buffer is just next_multiple_of(NPREDICTED) + // rows tall but the MMA atom needs M=16; second m-tile aliases first m-tile + // (predicated rows discarded at output store). + auto layout_C = make_aliased_swizzled_layout_rc(); + auto layout_B = make_swizzled_layout_rc(); + Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C); + Tensor smem_B = make_tensor(make_smem_ptr(reinterpret_cast(smem.B)), layout_B); + + // ── TiledMMA: _1x_1 = 32 threads, one [16, 8] atom ── + auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(lane); + + // ── K-tile A operand (C): full [NPREDICTED_PAD_MMA_M, K_TILE], shared by both warps ── + constexpr int K_TILE = MMA_prop::K_BIG; + Tensor smem_C_tiled + = local_tile(smem_C, make_tile(Int{}, Int{}), make_coord(_0{}, _)); + + // ── K-tile B operand ── + // NPREDICTED_PAD_MMA_N == 16: warp 0 → N=[0,8), warp 1 → N=[8,16). + // NPREDICTED_PAD_MMA_N == 8 : only warp 0 runs (warp 1 took the early + // exit above), tile at (_0, _). + Tensor smem_B_half = local_tile(smem_B, make_tile(Int{}, Int{}), make_coord(warp, _)); + + // ── Register fragments ── + Tensor frag_A = thr_mma.partition_fragment_A(smem_C_tiled(_, _, _0{})); + Tensor frag_B = thr_mma.partition_fragment_B(smem_B_half(_, _, _0{})); + + // ── Output accumulator: [NPREDICTED_PAD_MMA_M, N_HALF] f32 ── + auto layout_cb_half = make_layout(make_shape(Int{}, Int{})); + Tensor frag_acc = thr_mma.partition_fragment_C(make_tensor((float*) nullptr, layout_cb_half)); + clear(frag_acc); + + // ── S2R copies ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(lane); + Tensor smem_C_s2r = s2r_thr_A.partition_S(smem_C_tiled); + Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); + + auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B = s2r_B.get_slice(lane); + Tensor smem_B_s2r = s2r_thr_B.partition_S(smem_B_half); + Tensor frag_B_view = s2r_thr_B.retile_D(frag_B); + + // ── Gemm: 8 K-tiles, 1 HMMA each ── + constexpr int NUM_K_TILES = DSTATE / K_TILE; #pragma unroll - for (int k = 0; k < NUM_K_TILES; ++k) { - cute::copy(s2r_A, smem_C_s2r(_, _, _, k), frag_A_view); - cute::copy(s2r_B, smem_B_s2r(_, _, _, k), frag_B_view); - cute::gemm(tiled_mma, frag_acc, frag_A, frag_B, frag_acc); - } - - // ── Elementwise: decay * dt_proc * causal mask, convert f32 → MMA_prop::operand_t ── - auto id_half = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_half); - - // ── Store to swizzled smem.CB_scaled ── - Tensor smem_CB = make_tensor( - make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); - // Tile into [NPREDICTED_PAD_MMA_M, N_HALF] halves; warp selects its half - Tensor smem_CB_half = local_tile(smem_CB, make_tile(Int{}, Int{}), - make_coord(_0{}, warp)); - Tensor smem_CB_part = thr_mma.partition_C(smem_CB_half); + for (int k = 0; k < NUM_K_TILES; ++k) + { + cute::copy(s2r_A, smem_C_s2r(_, _, _, k), frag_A_view); + cute::copy(s2r_B, smem_B_s2r(_, _, _, k), frag_B_view); + cute::gemm(tiled_mma, frag_acc, frag_A, frag_B, frag_acc); + } + + // ── Elementwise: decay * dt_proc * causal mask, convert f32 → MMA_prop::operand_t ── + auto id_half = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_half); + + // ── Store to swizzled smem.CB_scaled ── + Tensor smem_CB = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); + // Tile into [NPREDICTED_PAD_MMA_M, N_HALF] halves; warp selects its half + Tensor smem_CB_half + = local_tile(smem_CB, make_tile(Int{}, Int{}), make_coord(_0{}, warp)); + Tensor smem_CB_part = thr_mma.partition_C(smem_CB_half); #pragma unroll - for (int i = 0; i < size(frag_acc); ++i) { - int t = get<0>(id_part(i)); - int j = warp * N_HALF + get<1>(id_part(i)); - float val; - if (j <= t && t < seq_len && j < seq_len) { - val = frag_acc(i) * __expf(smem.cumAdt[t] - smem.cumAdt[j]) * smem.dt_proc[j]; - } else { - val = 0.f; + for (int i = 0; i < size(frag_acc); ++i) + { + int t = get<0>(id_part(i)); + int j = warp * N_HALF + get<1>(id_part(i)); + float val; + if (j <= t && t < seq_len && j < seq_len) + { + val = frag_acc(i) * __expf(smem.cumAdt[t] - smem.cumAdt[j]) * smem.dt_proc[j]; + } + else + { + val = 0.f; + } + smem_CB_part(i) = MMA_prop::operand_t(val); } - smem_CB_part(i) = MMA_prop::operand_t(val); - } } // Compute CB_old[t, i] = (C @ old_B^T)[t, i] * exp(cumAdt[t]) * dB_old(i) for @@ -975,108 +1011,108 @@ __device__ __forceinline__ void compute_CB_scaled_2warp(SmemT& smem, int warp, i // MAX_WINDOW_PAD_MMA_K == 16: warp 2 → cols [0, 8); warp 3 → cols [8, 16). // MAX_WINDOW_PAD_MMA_K == 8 : warp 2 covers all 8 cols; warp 3 returns early. template -__device__ __forceinline__ void compute_CB_old_2warp(SmemT& smem, int warp, int lane, int prev_k, - int seq_len) { - using namespace cute; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - constexpr int N_HALF = MMA_prop::N; // 8 — one m16n8 atom per warp - constexpr int NUM_N_ATOMS = MAX_WINDOW_PAD_MMA_K / N_HALF; - static_assert(MAX_WINDOW_PAD_MMA_K % N_HALF == 0, - "compute_CB_old_2warp: MAX_WINDOW_PAD_MMA_K must be a multiple of MMA::N"); - static_assert(NPREDICTED_PAD_MMA_M + MAX_WINDOW_PAD_MMA_K <= CB_ROW_STRIDE, - "CB_scaled buffer must fit both CB_new (cols [0,T_pad)) and CB_old " - "(cols [T_pad, T_pad+K_old)) within its physical row stride"); - - int const sub_warp = warp - 2; // ∈ {0, 1} - if (sub_warp >= NUM_N_ATOMS) return; - - float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - - // ── Swizzled smem views (A = C, B = old_B; same shapes as the replay path's - // C/old_B reads, so we get cache locality with no extra cp.async). ── - auto layout_C = - make_aliased_swizzled_layout_rc(); - auto layout_old_B = make_swizzled_layout_rc(); - Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C); - Tensor smem_old_B = - make_tensor(make_smem_ptr(reinterpret_cast(smem.old_B)), layout_old_B); - - // ── TiledMMA: 32 threads, single m16n8k16 atom (K-loops over DSTATE) ── - auto tiled_mma = - make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(lane); - - // ── K-tile A operand (C): full M, K-loop dim ── - constexpr int K_TILE = MMA_prop::K_BIG; - Tensor smem_C_tiled = local_tile(smem_C, make_tile(Int{}, Int{}), - make_coord(_0{}, _)); - - // ── K-tile B operand (old_B): warp picks its 8-col N-atom slice ── - Tensor smem_old_B_half = - local_tile(smem_old_B, make_tile(Int{}, Int{}), make_coord(sub_warp, _)); - - // ── Register fragments ── - Tensor frag_A = thr_mma.partition_fragment_A(smem_C_tiled(_, _, _0{})); - Tensor frag_B = thr_mma.partition_fragment_B(smem_old_B_half(_, _, _0{})); - - auto layout_cb_half = make_layout(make_shape(Int{}, Int{})); - Tensor frag_acc = thr_mma.partition_fragment_C(make_tensor((float*)nullptr, layout_cb_half)); - clear(frag_acc); - - // ── S2R copies ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(lane); - Tensor smem_C_s2r = s2r_thr_A.partition_S(smem_C_tiled); - Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); - - auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B = s2r_B.get_slice(lane); - Tensor smem_old_B_s2r = s2r_thr_B.partition_S(smem_old_B_half); - Tensor frag_B_view = s2r_thr_B.retile_D(frag_B); - - // ── GEMM: DSTATE / K_BIG = 8 K-tiles ── - constexpr int NUM_K_TILES = DSTATE / K_TILE; +__device__ __forceinline__ void compute_CB_old_2warp(SmemT& smem, int warp, int lane, int prev_k, int seq_len) +{ + using namespace cute; + + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; + constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; + constexpr int N_HALF = MMA_prop::N; // 8 — one m16n8 atom per warp + constexpr int NUM_N_ATOMS = MAX_WINDOW_PAD_MMA_K / N_HALF; + static_assert( + MAX_WINDOW_PAD_MMA_K % N_HALF == 0, "compute_CB_old_2warp: MAX_WINDOW_PAD_MMA_K must be a multiple of MMA::N"); + static_assert(NPREDICTED_PAD_MMA_M + MAX_WINDOW_PAD_MMA_K <= CB_ROW_STRIDE, + "CB_scaled buffer must fit both CB_new (cols [0,T_pad)) and CB_old " + "(cols [T_pad, T_pad+K_old)) within its physical row stride"); + + int const sub_warp = warp - 2; // ∈ {0, 1} + if (sub_warp >= NUM_N_ATOMS) + return; + + float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; + + // ── Swizzled smem views (A = C, B = old_B; same shapes as the replay path's + // C/old_B reads, so we get cache locality with no extra cp.async). ── + auto layout_C = make_aliased_swizzled_layout_rc(); + auto layout_old_B = make_swizzled_layout_rc(); + Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C); + Tensor smem_old_B = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_B)), layout_old_B); + + // ── TiledMMA: 32 threads, single m16n8k16 atom (K-loops over DSTATE) ── + auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); + auto thr_mma = tiled_mma.get_slice(lane); + + // ── K-tile A operand (C): full M, K-loop dim ── + constexpr int K_TILE = MMA_prop::K_BIG; + Tensor smem_C_tiled + = local_tile(smem_C, make_tile(Int{}, Int{}), make_coord(_0{}, _)); + + // ── K-tile B operand (old_B): warp picks its 8-col N-atom slice ── + Tensor smem_old_B_half = local_tile(smem_old_B, make_tile(Int{}, Int{}), make_coord(sub_warp, _)); + + // ── Register fragments ── + Tensor frag_A = thr_mma.partition_fragment_A(smem_C_tiled(_, _, _0{})); + Tensor frag_B = thr_mma.partition_fragment_B(smem_old_B_half(_, _, _0{})); + + auto layout_cb_half = make_layout(make_shape(Int{}, Int{})); + Tensor frag_acc = thr_mma.partition_fragment_C(make_tensor((float*) nullptr, layout_cb_half)); + clear(frag_acc); + + // ── S2R copies ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(lane); + Tensor smem_C_s2r = s2r_thr_A.partition_S(smem_C_tiled); + Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); + + auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); + auto s2r_thr_B = s2r_B.get_slice(lane); + Tensor smem_old_B_s2r = s2r_thr_B.partition_S(smem_old_B_half); + Tensor frag_B_view = s2r_thr_B.retile_D(frag_B); + + // ── GEMM: DSTATE / K_BIG = 8 K-tiles ── + constexpr int NUM_K_TILES = DSTATE / K_TILE; #pragma unroll - for (int k = 0; k < NUM_K_TILES; ++k) { - cute::copy(s2r_A, smem_C_s2r(_, _, _, k), frag_A_view); - cute::copy(s2r_B, smem_old_B_s2r(_, _, _, k), frag_B_view); - cute::gemm(tiled_mma, frag_acc, frag_A, frag_B, frag_acc); - } - - // ── Identity coords for elementwise / store ── - auto id_half = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_half); - - // ── Store to swizzled smem.CB_scaled at the CB_old region (cols [T_pad, T_pad+K_old)). - // Use the full physical (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE) padded swizzle view. - // Byte-compatible with compute_CB_scaled_2warp's (T_pad, T_pad, CB_ROW_STRIDE) - // padded view: both produce inner offset r*CB_ROW_STRIDE + c, same Swizzle. ── - auto layout_cb_full = make_swizzled_layout_rc(); - Tensor smem_CB = make_tensor( - make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); - // (16, CB_ROW_STRIDE) tiled by (16, N_HALF) → (1, CB_ROW_STRIDE/N_HALF) tiles. - // Coord (0, T_pad/N_HALF + sub_warp) lands inside the CB_old region. - constexpr int CB_OLD_TILE_BASE = NPREDICTED_PAD_MMA_M / N_HALF; - Tensor smem_CB_half = local_tile(smem_CB, make_tile(Int{}, Int{}), - make_coord(_0{}, CB_OLD_TILE_BASE + sub_warp)); - Tensor smem_CB_part = thr_mma.partition_C(smem_CB_half); + for (int k = 0; k < NUM_K_TILES; ++k) + { + cute::copy(s2r_A, smem_C_s2r(_, _, _, k), frag_A_view); + cute::copy(s2r_B, smem_old_B_s2r(_, _, _, k), frag_B_view); + cute::gemm(tiled_mma, frag_acc, frag_A, frag_B, frag_acc); + } + + // ── Identity coords for elementwise / store ── + auto id_half = make_identity_tensor(make_shape(Int{}, Int{})); + auto id_part = thr_mma.partition_C(id_half); + + // ── Store to swizzled smem.CB_scaled at the CB_old region (cols [T_pad, T_pad+K_old)). + // Use the full physical (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE) padded swizzle view. + // Byte-compatible with compute_CB_scaled_2warp's (T_pad, T_pad, CB_ROW_STRIDE) + // padded view: both produce inner offset r*CB_ROW_STRIDE + c, same Swizzle. ── + auto layout_cb_full = make_swizzled_layout_rc(); + Tensor smem_CB = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); + // (16, CB_ROW_STRIDE) tiled by (16, N_HALF) → (1, CB_ROW_STRIDE/N_HALF) tiles. + // Coord (0, T_pad/N_HALF + sub_warp) lands inside the CB_old region. + constexpr int CB_OLD_TILE_BASE = NPREDICTED_PAD_MMA_M / N_HALF; + Tensor smem_CB_half = local_tile( + smem_CB, make_tile(Int{}, Int{}), make_coord(_0{}, CB_OLD_TILE_BASE + sub_warp)); + Tensor smem_CB_part = thr_mma.partition_C(smem_CB_half); #pragma unroll - for (int i = 0; i < size(frag_acc); ++i) { - int t = get<0>(id_part(i)); - int j = sub_warp * N_HALF + get<1>(id_part(i)); - float val; - if (j < prev_k && t < seq_len) { - val = frag_acc(i) * __expf(smem.cumAdt[t] + total_old_cumAdt - smem.old_cumAdt[j]) * - smem.old_dt[j]; - } else { - val = 0.f; + for (int i = 0; i < size(frag_acc); ++i) + { + int t = get<0>(id_part(i)); + int j = sub_warp * N_HALF + get<1>(id_part(i)); + float val; + if (j < prev_k && t < seq_len) + { + val = frag_acc(i) * __expf(smem.cumAdt[t] + total_old_cumAdt - smem.old_cumAdt[j]) * smem.old_dt[j]; + } + else + { + val = 0.f; + } + smem_CB_part(i) = MMA_prop::operand_t(val); } - smem_CB_part(i) = MMA_prop::operand_t(val); - } } // ============================================================================= @@ -1095,19 +1131,19 @@ __device__ __forceinline__ void compute_CB_old_2warp(SmemT& smem, int warp, int // m16n8k8 B frag (2 elts): 0→K_base, 1→K_base+1 // ============================================================================= template -__device__ __forceinline__ void precompute_dB_coeff(float coeff[DB_COEFFS_PER_LANE], - SmemT const& smem, float total_cumAdt, - int prev_k, int lane) { - static_assert(DB_COEFFS_PER_LANE == 2 || DB_COEFFS_PER_LANE == 4, - "DB_COEFFS_PER_LANE must be 2 (k8) or 4 (k16)"); - int const K_base = (lane % 4) * 2; +__device__ __forceinline__ void precompute_dB_coeff( + float coeff[DB_COEFFS_PER_LANE], SmemT const& smem, float total_cumAdt, int prev_k, int lane) +{ + static_assert(DB_COEFFS_PER_LANE == 2 || DB_COEFFS_PER_LANE == 4, "DB_COEFFS_PER_LANE must be 2 (k8) or 4 (k16)"); + int const K_base = (lane % 4) * 2; #pragma unroll - for (int i = 0; i < DB_COEFFS_PER_LANE; ++i) { - // m16n8k_ V-index → K-offset: (V & 1) is the col-pair offset; (V & 2) ? 8 : 0 - // covers the second K-tile inside the K_BIG (k16) atom. - int const k = K_base + (i & 1) + ((i & 2) << 2); - coeff[i] = (k < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[k]) * smem.old_dt[k] : 0.f; - } + for (int i = 0; i < DB_COEFFS_PER_LANE; ++i) + { + // m16n8k_ V-index → K-offset: (V & 1) is the col-pair offset; (V & 2) ? 8 : 0 + // covers the second K-tile inside the K_BIG (k16) atom. + int const k = K_base + (i & 1) + ((i & 2) << 2); + coeff[i] = (k < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[k]) * smem.old_dt[k] : 0.f; + } } // Apply precomputed dB coefficients to frag_B in-place. @@ -1116,15 +1152,16 @@ __device__ __forceinline__ void precompute_dB_coeff(float coeff[DB_COEFFS_PER_LA // coeff[i] = 0 encodes both causal mask and zero-fill for k >= prev_k. // ============================================================================= template -__device__ __forceinline__ void compute_dB_scaling(FragB& frag_B, - float const coeff[DB_COEFFS_PER_LANE]) { - using namespace cute; - static_assert(size(FragB{}) == DB_COEFFS_PER_LANE, "frag_B size must match DB_COEFFS_PER_LANE"); - using frag_t = typename FragB::value_type; +__device__ __forceinline__ void compute_dB_scaling(FragB& frag_B, float const coeff[DB_COEFFS_PER_LANE]) +{ + using namespace cute; + static_assert(size(FragB{}) == DB_COEFFS_PER_LANE, "frag_B size must match DB_COEFFS_PER_LANE"); + using frag_t = typename FragB::value_type; #pragma unroll - for (int i = 0; i < DB_COEFFS_PER_LANE; ++i) { - frag_B(i) = frag_t(toFloat(frag_B(i)) * coeff[i]); - } + for (int i = 0; i < DB_COEFFS_PER_LANE; ++i) + { + frag_B(i) = frag_t(toFloat(frag_B(i)) * coeff[i]); + } } // Scale frag_A by dB coefficients ONCE before the N-pass loop, replacing @@ -1144,41 +1181,45 @@ __device__ __forceinline__ void compute_dB_scaling(FragB& frag_B, // {4,6}→K_base+8, {5,7}→K_base+9 (4 unique K) // ============================================================================= template -__device__ __forceinline__ void apply_dA_coeff(FragA& frag_A, SmemT const& smem, float total_cumAdt, - int prev_k, int lane) { - using namespace cute; - constexpr int FRAG_A_SIZE = size(FragA{}); - static_assert((MAX_WINDOW_PAD_MMA_K == 16 && FRAG_A_SIZE == 8) || - (MAX_WINDOW_PAD_MMA_K == 8 && FRAG_A_SIZE == 4), - "apply_dA_coeff: unsupported MMA K / frag_A size combination"); - using frag_t = typename FragA::value_type; - - int const K_base = (lane % 4) * 2; - - if constexpr (MAX_WINDOW_PAD_MMA_K == 8) { - float const c0 = (K_base < prev_k) - ? __expf(total_cumAdt - smem.old_cumAdt[K_base]) * smem.old_dt[K_base] - : 0.f; - float const c1 = (K_base + 1 < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[K_base + 1]) * - smem.old_dt[K_base + 1] - : 0.f; +__device__ __forceinline__ void apply_dA_coeff( + FragA& frag_A, SmemT const& smem, float total_cumAdt, int prev_k, int lane) +{ + using namespace cute; + constexpr int FRAG_A_SIZE = size(FragA{}); + static_assert((MAX_WINDOW_PAD_MMA_K == 16 && FRAG_A_SIZE == 8) || (MAX_WINDOW_PAD_MMA_K == 8 && FRAG_A_SIZE == 4), + "apply_dA_coeff: unsupported MMA K / frag_A size combination"); + using frag_t = typename FragA::value_type; + + int const K_base = (lane % 4) * 2; + + if constexpr (MAX_WINDOW_PAD_MMA_K == 8) + { + float const c0 = (K_base < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[K_base]) * smem.old_dt[K_base] : 0.f; + float const c1 = (K_base + 1 < prev_k) + ? __expf(total_cumAdt - smem.old_cumAdt[K_base + 1]) * smem.old_dt[K_base + 1] + : 0.f; #pragma unroll - for (int i = 0; i < 4; ++i) { - frag_A(i) = frag_t(toFloat(frag_A(i)) * ((i & 1) ? c1 : c0)); + for (int i = 0; i < 4; ++i) + { + frag_A(i) = frag_t(toFloat(frag_A(i)) * ((i & 1) ? c1 : c0)); + } } - } else { - float c[4]; + else + { + float c[4]; #pragma unroll - for (int j = 0; j < 4; ++j) { - int const k = K_base + (j & 1) + ((j & 2) ? 8 : 0); - c[j] = (k < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[k]) * smem.old_dt[k] : 0.f; - } + for (int j = 0; j < 4; ++j) + { + int const k = K_base + (j & 1) + ((j & 2) ? 8 : 0); + c[j] = (k < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[k]) * smem.old_dt[k] : 0.f; + } #pragma unroll - for (int i = 0; i < 8; ++i) { - int const ci = (i & 1) | ((i & 4) >> 1); - frag_A(i) = frag_t(toFloat(frag_A(i)) * c[ci]); + for (int i = 0; i < 8; ++i) + { + int const ci = (i & 1) | ((i & 4) >> 1); + frag_A(i) = frag_t(toFloat(frag_A(i)) * c[ci]); + } } - } } // ── CuTe mma.sync output sub-functions ────────────────────────────────────── @@ -1191,14 +1232,17 @@ __device__ __forceinline__ void apply_dA_coeff(FragA& frag_A, SmemT const& smem, // dispatches to the native packed cvt for the destination type (e.g. // cvt.rn.bf16x2.f32 for bf16). template -__device__ __forceinline__ void convert_frag(Frag& frag) { - if constexpr (!std::is_same_v) { +__device__ __forceinline__ void convert_frag(Frag& frag) +{ + if constexpr (!std::is_same_v) + { #pragma unroll - for (int i = 0; i < cute::size(frag); i += 2) { - float2 const vals = toFloat2(reinterpret_cast(&frag(i))); - *reinterpret_cast*>(&frag(i)) = pack_float2(vals); + for (int i = 0; i < cute::size(frag); i += 2) + { + float2 const vals = toFloat2(reinterpret_cast(&frag(i))); + *reinterpret_cast*>(&frag(i)) = pack_float2(vals); + } } - } } // State → MMA B operand: dtype-aware TiledCopy. @@ -1206,14 +1250,18 @@ __device__ __forceinline__ void convert_frag(Frag& frag) { // 4-byte smem: scalar UniversalCopy; pairs are converted to // bf16 in registers by `convert_frag` after the load. template -__device__ __forceinline__ auto make_state_b_s2r(TiledMma const& tm) { - using namespace cute; - if constexpr (sizeof(state_t) == 2) { - return make_tiled_copy_B(Copy_Atom{}, tm); - } else { - static_assert(sizeof(state_t) == 4, "wide state path expects 4-byte smem"); - return make_tiled_copy_B(Copy_Atom, state_t>{}, tm); - } +__device__ __forceinline__ auto make_state_b_s2r(TiledMma const& tm) +{ + using namespace cute; + if constexpr (sizeof(state_t) == 2) + { + return make_tiled_copy_B(Copy_Atom{}, tm); + } + else + { + static_assert(sizeof(state_t) == 4, "wide state path expects 4-byte smem"); + return make_tiled_copy_B(Copy_Atom, state_t>{}, tm); + } } // Src → dst fragment conversion — a strict superset of the in-place overload @@ -1225,50 +1273,56 @@ __device__ __forceinline__ auto make_state_b_s2r(TiledMma const& tm) { // Works in-place when `src` aliases `dst`. // (3) Different width (e.g. f32 → bf16): paired element load + pack_float2. template -__device__ __forceinline__ void convert_frag(SrcFrag const& src, DstFrag& dst) { - using namespace cute; - if constexpr (std::is_same_v) { +__device__ __forceinline__ void convert_frag(SrcFrag const& src, DstFrag& dst) +{ + using namespace cute; + if constexpr (std::is_same_v) + { #pragma unroll - for (int i = 0; i < size(src); i += 2) { - *reinterpret_cast*>(&dst(i)) = *reinterpret_cast const*>(&src(i)); + for (int i = 0; i < size(src); i += 2) + { + *reinterpret_cast*>(&dst(i)) = *reinterpret_cast const*>(&src(i)); + } } - } else if constexpr (sizeof(src_t) == sizeof(dst_t)) { + else if constexpr (sizeof(src_t) == sizeof(dst_t)) + { #pragma unroll - for (int i = 0; i < size(src); i += 2) { - float2 const vals = toFloat2(reinterpret_cast(&src(i))); - *reinterpret_cast*>(&dst(i)) = pack_float2(vals); + for (int i = 0; i < size(src); i += 2) + { + float2 const vals = toFloat2(reinterpret_cast(&src(i))); + *reinterpret_cast*>(&dst(i)) = pack_float2(vals); + } } - } else { - static_assert(sizeof(dst_t) == 2, "only narrowing to 2-byte dst supported"); + else + { + static_assert(sizeof(dst_t) == 2, "only narrowing to 2-byte dst supported"); #pragma unroll - for (int i = 0; i < size(src); i += 2) { - *reinterpret_cast*>(&dst(i)) = - pack_float2(make_float2(src(i), src(i + 1))); + for (int i = 0; i < size(src); i += 2) + { + *reinterpret_cast*>(&dst(i)) = pack_float2(make_float2(src(i), src(i + 1))); + } } - } } // 2b. frag_y += CB_scaled @ x (matmul 4, single K-tile) // CB_scaled A operand loaded from swizzled smem via LDSM (precomputed by warps 0,1). // x B operand loaded from smem via ldmatrix.trans. -template -__device__ __forceinline__ void add_cb_x(FragY& frag_y, FragCB const& frag_CB, - SmemXTrans const& smem_x_trans, - S2RBTrans const& s2r_B_trans, - S2RThrBTrans const& s2r_thr_B_trans, ThrMma const& thr_mma, - TiledMma const& tiled_mma, int n) { - using namespace cute; - Tensor smem_x_trans_ntile = local_tile( - smem_x_trans, make_tile(Int{}, Int{}), make_coord(n, _0{})); - auto smem_x_trans_s2r = s2r_thr_B_trans.partition_S(smem_x_trans_ntile); - auto frag_B_x = thr_mma.partition_fragment_B( - make_tensor((MmaT*)0x0, make_shape(Int{}, Int{}))); - auto frag_B_x_view = s2r_thr_B_trans.retile_D(frag_B_x); - - cute::copy(s2r_B_trans, smem_x_trans_s2r, frag_B_x_view); - cute::gemm(tiled_mma, frag_y, frag_CB, frag_B_x, frag_y); +template +__device__ __forceinline__ void add_cb_x(FragY& frag_y, FragCB const& frag_CB, SmemXTrans const& smem_x_trans, + S2RBTrans const& s2r_B_trans, S2RThrBTrans const& s2r_thr_B_trans, ThrMma const& thr_mma, TiledMma const& tiled_mma, + int n) +{ + using namespace cute; + Tensor smem_x_trans_ntile + = local_tile(smem_x_trans, make_tile(Int{}, Int{}), make_coord(n, _0{})); + auto smem_x_trans_s2r = s2r_thr_B_trans.partition_S(smem_x_trans_ntile); + auto frag_B_x = thr_mma.partition_fragment_B( + make_tensor((MmaT*) 0x0, make_shape(Int{}, Int{}))); + auto frag_B_x_view = s2r_thr_B_trans.retile_D(frag_B_x); + + cute::copy(s2r_B_trans, smem_x_trans_s2r, frag_B_x_view); + cute::gemm(tiled_mma, frag_y, frag_CB, frag_B_x, frag_y); } // 2c. frag_y += CB_old @ old_x (matmul-4 over old tokens; sibling of add_cb_x). @@ -1278,64 +1332,63 @@ __device__ __forceinline__ void add_cb_x(FragY& frag_y, FragCB const& frag_CB, // matching m16n8k_OLD atom (K_BIG=16 or K_SMALL=8). frag_y partitioned by a // different (K_BIG) tiled_mma is layout-compatible — the m16n8 C-frag shape is // the same regardless of K. -template +template __device__ __forceinline__ void add_cb_old_x(FragY& frag_y, FragCBOld const& frag_CB_old, - SmemOldXTrans const& smem_old_x_trans, - S2RBTransOld const& s2r_B_trans_old, - S2RThrBTransOld const& s2r_thr_B_trans_old, - ThrMmaOld const& thr_mma_old, - TiledMmaOld const& tiled_mma_old, int n) { - using namespace cute; - Tensor smem_old_x_ntile = local_tile( - smem_old_x_trans, make_tile(Int{}, Int{}), make_coord(n, _0{})); - auto smem_old_x_s2r = s2r_thr_B_trans_old.partition_S(smem_old_x_ntile); - auto frag_B_old_x = thr_mma_old.partition_fragment_B( - make_tensor((MmaT*)0x0, make_shape(Int{}, Int{}))); - auto frag_B_old_x_view = s2r_thr_B_trans_old.retile_D(frag_B_old_x); - - cute::copy(s2r_B_trans_old, smem_old_x_s2r, frag_B_old_x_view); - cute::gemm(tiled_mma_old, frag_y, frag_CB_old, frag_B_old_x, frag_y); + SmemOldXTrans const& smem_old_x_trans, S2RBTransOld const& s2r_B_trans_old, + S2RThrBTransOld const& s2r_thr_B_trans_old, ThrMmaOld const& thr_mma_old, TiledMmaOld const& tiled_mma_old, int n) +{ + using namespace cute; + Tensor smem_old_x_ntile + = local_tile(smem_old_x_trans, make_tile(Int{}, Int{}), make_coord(n, _0{})); + auto smem_old_x_s2r = s2r_thr_B_trans_old.partition_S(smem_old_x_ntile); + auto frag_B_old_x = thr_mma_old.partition_fragment_B( + make_tensor((MmaT*) 0x0, make_shape(Int{}, Int{}))); + auto frag_B_old_x_view = s2r_thr_B_trans_old.retile_D(frag_B_old_x); + + cute::copy(s2r_B_trans_old, smem_old_x_s2r, frag_B_old_x_view); + cute::gemm(tiled_mma_old, frag_y, frag_CB_old, frag_B_old_x, frag_y); } // 3b. frag_y += D * x[t, d] (per-thread skip connection via partition_C) -template -__device__ __forceinline__ void add_D_skip(FragY& frag_y, SmemX const& smem_x, - ThrMma const& thr_mma, float D_val, int n) { - using namespace cute; - if (D_val == 0.f) return; - Tensor smem_x_tile = local_tile(smem_x, make_tile(Int{}, Int{}), - make_coord(_0{}, n)); - Tensor x_part = thr_mma.partition_C(smem_x_tile); - // Load pairs of consecutive bf16 elements and convert via paired toFloat2. - // m16n8k16 partition_C places consecutive N-column pairs adjacent in smem. - static_assert(sizeof(input_t) == 2, "vectorized D_skip requires 2-byte input_t"); +template +__device__ __forceinline__ void add_D_skip( + FragY& frag_y, SmemX const& smem_x, ThrMma const& thr_mma, float D_val, int n) +{ + using namespace cute; + if (D_val == 0.f) + return; + Tensor smem_x_tile = local_tile(smem_x, make_tile(Int{}, Int{}), make_coord(_0{}, n)); + Tensor x_part = thr_mma.partition_C(smem_x_tile); + // Load pairs of consecutive bf16 elements and convert via paired toFloat2. + // m16n8k16 partition_C places consecutive N-column pairs adjacent in smem. + static_assert(sizeof(input_t) == 2, "vectorized D_skip requires 2-byte input_t"); #pragma unroll - for (int i = 0; i < size(frag_y); i += 2) { - float2 vals = toFloat2(reinterpret_cast(&x_part(i))); - frag_y(i) += D_val * vals.x; - frag_y(i + 1) += D_val * vals.y; - } + for (int i = 0; i < size(frag_y); i += 2) + { + float2 vals = toFloat2(reinterpret_cast(&x_part(i))); + frag_y(i) += D_val * vals.x; + frag_y(i + 1) += D_val * vals.y; + } } // 4b. frag_y *= z * sigmoid(z) (z-gating via partition_C) -template -__device__ __forceinline__ void compute_z_gating(FragY& frag_y, SmemZ const& smem_z, - ThrMma const& thr_mma, void const* z_ptr, int n) { - using namespace cute; - if (!z_ptr) return; - Tensor smem_z_tile = local_tile(smem_z, make_tile(Int{}, Int{}), - make_coord(_0{}, n)); - Tensor z_part = thr_mma.partition_C(smem_z_tile); +template +__device__ __forceinline__ void compute_z_gating( + FragY& frag_y, SmemZ const& smem_z, ThrMma const& thr_mma, void const* z_ptr, int n) +{ + using namespace cute; + if (!z_ptr) + return; + Tensor smem_z_tile = local_tile(smem_z, make_tile(Int{}, Int{}), make_coord(_0{}, n)); + Tensor z_part = thr_mma.partition_C(smem_z_tile); #pragma unroll - for (int i = 0; i < size(frag_y); i += 2) { - float2 const z = toFloat2(reinterpret_cast(&z_part(i))); - frag_y(i) *= z.x * __fdividef(1.f, (1.f + __expf(-z.x))); - frag_y(i + 1) *= z.y * __fdividef(1.f, (1.f + __expf(-z.y))); - } + for (int i = 0; i < size(frag_y); i += 2) + { + float2 const z = toFloat2(reinterpret_cast(&z_part(i))); + frag_y(i) *= z.x * __fdividef(1.f, (1.f + __expf(-z.x))); + frag_y(i + 1) *= z.y * __fdividef(1.f, (1.f + __expf(-z.y))); + } } // ============================================================================= @@ -1352,117 +1405,128 @@ __device__ __forceinline__ void compute_z_gating(FragY& frag_y, SmemZ const& sme // Used by matmul 3 (init_out += C @ state^T): A = C (shared), B = state. // NumNTiles = sizeof...(FragY) = D_PER_CTA / N_TILE (1 for D_SPLIT=2, 2 for // D_SPLIT=1). -template -__device__ __forceinline__ void pipelined_kloop_gemm(TiledMma const& tiled_mma, - ThrMma const& thr_mma, int tid, - SmemAKtiled const& smem_A_ktiled, - SmemB const& smem_B, FragY&... frag_y) { - using namespace cute; - constexpr int NumNTiles = sizeof...(FragY); - static_assert(NumStages >= 2, "NumStages must be >= 2 for pipelining"); - static_assert(NumKTiles >= NumStages - 1, "NumKTiles must be >= NumStages - 1 for full prologue"); - static_assert(NumNTiles >= 1, "NumNTiles must be >= 1"); - - constexpr int N_TILE = cute::tile_size<1>(TiledMma{}); - constexpr int K_TILE = cute::tile_size<2>(TiledMma{}); - - // ── S2R copies ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - auto s2r_B = make_state_b_s2r(tiled_mma); - auto s2r_thr_B = s2r_B.get_slice(tid); - - // ── Tile B by (N, K): shape (N_TILE, K_TILE, N_OUTER, NumKTiles) ── - auto smem_B_tiled = local_tile(smem_B, make_tile(Int{}, Int{}), make_coord(_, _)); - - // ── Partitioned smem (A shared, B per-N-tile) ── - auto smem_A_s2r = s2r_thr_A.partition_S(smem_A_ktiled); - auto sample_smem_B_n = smem_B_tiled(_, _, _0{}, _); - using SmemBS2RType = decltype(s2r_thr_B.partition_S(sample_smem_B_n)); - SmemBS2RType smem_B_s2r[NumNTiles]; - CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) { - smem_B_s2r[n] = s2r_thr_B.partition_S(smem_B_tiled(_, _, n, _)); - } - - // ── Fragment / view types ── - using FragA = decltype(thr_mma.partition_fragment_A(smem_A_ktiled(_, _, _0{}))); - using FragB = decltype(thr_mma.partition_fragment_B(sample_smem_B_n(_, _, _0{}))); - using b_view_t = std::conditional_t; - using FragBStg = decltype(make_fragment_like(std::declval())); - using FragAView = decltype(s2r_thr_A.retile_D(std::declval())); - using FragBStgView = decltype(s2r_thr_B.retile_D(std::declval())); - - // ── Multi-stage register fragments ── - // Storage type matches the MMA fragment for A; for B the staging buffer is - // BTypeIn-typed (when narrowing) or MmaT-typed (when widths match — the two - // alias the same registers and `convert_frag` collapses to a bit-copy / - // in-place reinterpret). - FragA frag_A[NumStages]; - FragB frag_B[NumNTiles][NumStages]; - FragBStg frag_B_stg[NumNTiles][NumStages]; - FragAView frag_A_view[NumStages]; - FragBStgView frag_B_stg_view[NumNTiles][NumStages]; - CUTE_UNROLL - for (int s = 0; s < NumStages; ++s) { - frag_A_view[s] = s2r_thr_A.retile_D(frag_A[s]); +template +__device__ __forceinline__ void pipelined_kloop_gemm(TiledMma const& tiled_mma, ThrMma const& thr_mma, int tid, + SmemAKtiled const& smem_A_ktiled, SmemB const& smem_B, FragY&... frag_y) +{ + using namespace cute; + constexpr int NumNTiles = sizeof...(FragY); + static_assert(NumStages >= 2, "NumStages must be >= 2 for pipelining"); + static_assert(NumKTiles >= NumStages - 1, "NumKTiles must be >= NumStages - 1 for full prologue"); + static_assert(NumNTiles >= 1, "NumNTiles must be >= 1"); + + constexpr int N_TILE = cute::tile_size<1>(TiledMma{}); + constexpr int K_TILE = cute::tile_size<2>(TiledMma{}); + + // ── S2R copies ── + auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); + auto s2r_thr_A = s2r_A.get_slice(tid); + auto s2r_B = make_state_b_s2r(tiled_mma); + auto s2r_thr_B = s2r_B.get_slice(tid); + + // ── Tile B by (N, K): shape (N_TILE, K_TILE, N_OUTER, NumKTiles) ── + auto smem_B_tiled = local_tile(smem_B, make_tile(Int{}, Int{}), make_coord(_, _)); + + // ── Partitioned smem (A shared, B per-N-tile) ── + auto smem_A_s2r = s2r_thr_A.partition_S(smem_A_ktiled); + auto sample_smem_B_n = smem_B_tiled(_, _, _0{}, _); + using SmemBS2RType = decltype(s2r_thr_B.partition_S(sample_smem_B_n)); + SmemBS2RType smem_B_s2r[NumNTiles]; CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) { - frag_B_stg_view[n][s] = s2r_thr_B.retile_D(frag_B_stg[n][s]); + for (int n = 0; n < NumNTiles; ++n) + { + smem_B_s2r[n] = s2r_thr_B.partition_S(smem_B_tiled(_, _, n, _)); } - } - - // Pack frag_y into a pointer array for indexed access (replay kernel pattern). - using FragY0 = std::tuple_element_t<0, std::tuple>; - static_assert((std::is_same_v && ...), - "all FragY parameters must be the same type"); - FragY0* frag_y_p[NumNTiles] = {(&frag_y)...}; - // ── Per-stage operations (slot is constant after #pragma unroll) ── - auto load_one = [&](int k_src, int slot) { - cute::copy(s2r_A, smem_A_s2r(_, _, _, k_src), frag_A_view[slot]); + // ── Fragment / view types ── + using FragA = decltype(thr_mma.partition_fragment_A(smem_A_ktiled(_, _, _0{}))); + using FragB = decltype(thr_mma.partition_fragment_B(sample_smem_B_n(_, _, _0{}))); + using b_view_t = std::conditional_t; + using FragBStg = decltype(make_fragment_like(std::declval())); + using FragAView = decltype(s2r_thr_A.retile_D(std::declval())); + using FragBStgView = decltype(s2r_thr_B.retile_D(std::declval())); + + // ── Multi-stage register fragments ── + // Storage type matches the MMA fragment for A; for B the staging buffer is + // BTypeIn-typed (when narrowing) or MmaT-typed (when widths match — the two + // alias the same registers and `convert_frag` collapses to a bit-copy / + // in-place reinterpret). + FragA frag_A[NumStages]; + FragB frag_B[NumNTiles][NumStages]; + FragBStg frag_B_stg[NumNTiles][NumStages]; + FragAView frag_A_view[NumStages]; + FragBStgView frag_B_stg_view[NumNTiles][NumStages]; CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) { - cute::copy(s2r_B, smem_B_s2r[n](_, _, _, k_src), frag_B_stg_view[n][slot]); + for (int s = 0; s < NumStages; ++s) + { + frag_A_view[s] = s2r_thr_A.retile_D(frag_A[s]); + CUTE_UNROLL + for (int n = 0; n < NumNTiles; ++n) + { + frag_B_stg_view[n][s] = s2r_thr_B.retile_D(frag_B_stg[n][s]); + } } - }; - auto convert_one = [&](int slot) { - convert_frag(frag_A[slot]); + + // Pack frag_y into a pointer array for indexed access (replay kernel pattern). + using FragY0 = std::tuple_element_t<0, std::tuple>; + static_assert((std::is_same_v && ...), "all FragY parameters must be the same type"); + FragY0* frag_y_p[NumNTiles] = {(&frag_y)...}; + + // ── Per-stage operations (slot is constant after #pragma unroll) ── + auto load_one = [&](int k_src, int slot) + { + cute::copy(s2r_A, smem_A_s2r(_, _, _, k_src), frag_A_view[slot]); + CUTE_UNROLL + for (int n = 0; n < NumNTiles; ++n) + { + cute::copy(s2r_B, smem_B_s2r[n](_, _, _, k_src), frag_B_stg_view[n][slot]); + } + }; + auto convert_one = [&](int slot) + { + convert_frag(frag_A[slot]); + CUTE_UNROLL + for (int n = 0; n < NumNTiles; ++n) + { + convert_frag(frag_B_stg[n][slot], frag_B[n][slot]); + } + }; + auto compute_one = [&](int slot) + { + CUTE_UNROLL + for (int n = 0; n < NumNTiles; ++n) + { + cute::gemm(tiled_mma, *frag_y_p[n], frag_A[slot], frag_B[n][slot], *frag_y_p[n]); + } + }; + + // ── Clear accumulators ── CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) { - convert_frag(frag_B_stg[n][slot], frag_B[n][slot]); - } - }; - auto compute_one = [&](int slot) { + for (int n = 0; n < NumNTiles; ++n) + clear(*frag_y_p[n]); + + // ── Prologue: load + convert stages 0..NumStages-2 ── CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) { - cute::gemm(tiled_mma, *frag_y_p[n], frag_A[slot], frag_B[n][slot], *frag_y_p[n]); + for (int s = 0; s < NumStages - 1; ++s) + { + load_one(s, s); + convert_one(s); } - }; - - // ── Clear accumulators ── - CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) clear(*frag_y_p[n]); - // ── Prologue: load + convert stages 0..NumStages-2 ── - CUTE_UNROLL - for (int s = 0; s < NumStages - 1; ++s) { - load_one(s, s); - convert_one(s); - } - - // ── Main K-loop: load slot (k+NumStages-1) % NumStages, compute slot k % NumStages ── + // ── Main K-loop: load slot (k+NumStages-1) % NumStages, compute slot k % NumStages ── #pragma unroll - for (int k = 0; k < NumKTiles; ++k) { - int const k_load = k + NumStages - 1; - int const slot_load = k_load % NumStages; - int const slot_compute = k % NumStages; - if (k_load < NumKTiles) load_one(k_load, slot_load); - compute_one(slot_compute); - if (k_load < NumKTiles) convert_one(slot_load); - } + for (int k = 0; k < NumKTiles; ++k) + { + int const k_load = k + NumStages - 1; + int const slot_load = k_load % NumStages; + int const slot_compute = k % NumStages; + if (k_load < NumKTiles) + load_one(k_load, slot_load); + compute_one(slot_compute); + if (k_load < NumKTiles) + convert_one(slot_load); + } } // ── Matmul 3: init_out = C @ state^T ──────────────────────────────────────── @@ -1474,39 +1538,38 @@ __device__ __forceinline__ void pipelined_kloop_gemm(TiledMma const& tiled_mma, // reinterpret-cast to MMA_prop::operand_t so the 16-bit LDSM atom matches the view // (actual element type recovered inside `convert_frag`); ≥4-byte smem keeps // the native dtype and uses scalar UniversalCopy + register conversion. -template -__device__ __forceinline__ void add_init_out(SmemT const& smem, TiledMma const& tiled_mma, - ThrMma const& thr_mma, int tid, FragY&... frag_y) { - using namespace cute; - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int K_TILE = cute::tile_size<2>(TiledMma{}); - constexpr int NUM_K_TILES = DSTATE / K_TILE; - // Smem source dtype for matmul-3 (generic kernel only; 8-bit state goes - // through the dedicated `checkpointing_ssu_kernel_8bit` path): - // - sizeof(state_t) == 2 (fp16/bf16): LDSM the native 16-bit, view as bf16. - // - sizeof(state_t) == 4 (fp32): scalar UniversalCopy + on-the-fly convert. - static_assert(sizeof(state_t) != 1, - "add_init_out is the 2/4-byte path; 1-byte state goes through " - "compute_output_8bit"); - constexpr bool is_2byte_smem = (sizeof(state_t) == 2); - using state_view_t = std::conditional_t; - using BTypeIn = state_t; - - auto layout_C_swz = - make_aliased_swizzled_layout_rc(); - Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), - layout_C_swz); - Tensor smem_C_ktiled = local_tile(smem_C, make_tile(Int{}, Int{}), - make_coord(_0{}, _)); - - // Swizzle layout matches the dtype of the buffer being viewed. - auto const layout_state_swz = make_swizzled_layout_rc(); - state_view_t const* smem_state_ptr = reinterpret_cast(smem.state); - Tensor smem_state = make_tensor(make_smem_ptr(smem_state_ptr), layout_state_swz); - - pipelined_kloop_gemm<3, NUM_K_TILES, input_t, BTypeIn, MMA_prop::operand_t>( - tiled_mma, thr_mma, tid, smem_C_ktiled, smem_state, frag_y...); +template +__device__ __forceinline__ void add_init_out( + SmemT const& smem, TiledMma const& tiled_mma, ThrMma const& thr_mma, int tid, FragY&... frag_y) +{ + using namespace cute; + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + constexpr int K_TILE = cute::tile_size<2>(TiledMma{}); + constexpr int NUM_K_TILES = DSTATE / K_TILE; + // Smem source dtype for matmul-3 (generic kernel only; 8-bit state goes + // through the dedicated `checkpointing_ssu_kernel_8bit` path): + // - sizeof(state_t) == 2 (fp16/bf16): LDSM the native 16-bit, view as bf16. + // - sizeof(state_t) == 4 (fp32): scalar UniversalCopy + on-the-fly convert. + static_assert(sizeof(state_t) != 1, + "add_init_out is the 2/4-byte path; 1-byte state goes through " + "compute_output_8bit"); + constexpr bool is_2byte_smem = (sizeof(state_t) == 2); + using state_view_t = std::conditional_t; + using BTypeIn = state_t; + + auto layout_C_swz = make_aliased_swizzled_layout_rc(); + Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C_swz); + Tensor smem_C_ktiled + = local_tile(smem_C, make_tile(Int{}, Int{}), make_coord(_0{}, _)); + + // Swizzle layout matches the dtype of the buffer being viewed. + auto const layout_state_swz = make_swizzled_layout_rc(); + state_view_t const* smem_state_ptr = reinterpret_cast(smem.state); + Tensor smem_state = make_tensor(make_smem_ptr(smem_state_ptr), layout_state_swz); + + pipelined_kloop_gemm<3, NUM_K_TILES, input_t, BTypeIn, MMA_prop::operand_t>( + tiled_mma, thr_mma, tid, smem_C_ktiled, smem_state, frag_y...); } // store_state: vectorized smem → gmem state writeback (128 threads). @@ -1516,31 +1579,29 @@ __device__ __forceinline__ void add_init_out(SmemT const& smem, TiledMma const& // epilogue. smem and gmem hold the same dtype now (no on-egress // conversion) so this is always a direct 128-bit copy. template -__device__ __forceinline__ void store_state(SmemT& smem, CheckpointingSsuParams const& params, - int warp, int lane, int d_tile, int head, - int64_t cache_slot) { - using namespace cute; - int const flat_tid = warp * warpSize + lane; - auto* __restrict__ state_w = reinterpret_cast(params.state); - // gmem dest = head's full state base + d_tile's row slice. - int64_t const state_base = cache_slot * params.state_stride_seq + (int64_t)head * DIM * DSTATE + - (int64_t)d_tile * D_PER_CTA * DSTATE; - - // ── Per-CTA smem swizzle layout [D_PER_CTA, DSTATE]. ── - auto layout_smem_swz = make_swizzled_layout_rc(); - state_t const* smem_state_base = reinterpret_cast(smem.state); - - Tensor sState = make_tensor(make_smem_ptr(smem_state_base), layout_smem_swz); - Tensor gState = make_tensor(make_gmem_ptr(state_w + state_base), - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - // Each store is 16 bytes — adjust val cols to the dtype. - constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); - auto s2g = - make_tiled_copy(Copy_Atom, state_t>{}, - Layout, Stride<_8, _1>>{}, Layout>>{}); - auto thr = s2g.get_slice(flat_tid); - copy(s2g, thr.partition_S(sState), thr.partition_D(gState)); +__device__ __forceinline__ void store_state( + SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int d_tile, int head, int64_t cache_slot) +{ + using namespace cute; + int const flat_tid = warp * warpSize + lane; + auto* __restrict__ state_w = reinterpret_cast(params.state); + // gmem dest = head's full state base + d_tile's row slice. + int64_t const state_base + = cache_slot * params.state_stride_seq + (int64_t) head * DIM * DSTATE + (int64_t) d_tile * D_PER_CTA * DSTATE; + + // ── Per-CTA smem swizzle layout [D_PER_CTA, DSTATE]. ── + auto layout_smem_swz = make_swizzled_layout_rc(); + state_t const* smem_state_base = reinterpret_cast(smem.state); + + Tensor sState = make_tensor(make_smem_ptr(smem_state_base), layout_smem_swz); + Tensor gState = make_tensor(make_gmem_ptr(state_w + state_base), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{}))); + // Each store is 16 bytes — adjust val cols to the dtype. + constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); + auto s2g = make_tiled_copy(Copy_Atom, state_t>{}, Layout, Stride<_8, _1>>{}, + Layout>>{}); + auto thr = s2g.get_slice(flat_tid); + copy(s2g, thr.partition_S(sState), thr.partition_D(gState)); } // ── Store functions (called from kernel after compute_y + sync) ── @@ -1548,54 +1609,53 @@ __device__ __forceinline__ void store_state(SmemT& smem, CheckpointingSsuParams // the state-writeback hoist.) template -__device__ __forceinline__ void store_old_x(SmemT& smem, CheckpointingSsuParams const& params, - int warp, int lane, int d_tile, int head, - int64_t cache_slot, int write_offset, int seq_len) { - using namespace cute; - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - int const flat_tid = warp * warpSize + lane; - - auto* __restrict__ old_x_w = reinterpret_cast(params.old_x); - // gmem dest = head's full slot + d_tile's D-slice offset, shifted by - // `write_offset` along the T-axis (must_checkpoint ? 0 : prev_k). - int64_t const ox_w_base = cache_slot * params.old_x_stride_seq + - (int64_t)write_offset * params.old_x_stride_token + head * DIM + - (int64_t)d_tile * D_PER_CTA; - - // Smem and gmem are both viewed at the full atom-padded width D_SMEM_COLS. - // The wide thread layout (16 row × 8 col × 1×8 val = 16 rows × 64 cols/pass - // for bf16) covers one full atom width per thread-row, which is the - // swizzle's bank-conflict-free contract on the LDS side (load-from-smem). - // A narrow layout would (a) waste 64 threads (warps 2, 3 idle) and - // (b) cause LDS bank conflicts on the smem-read side (observed as - // 4-way LDS conflict in d_split=2 ncu). Cols ≥ D_PER_CTA are predicated - // off via copy_if so STG never fires for them — no OOB write into the - // next d_tile / next head's gmem region. - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - auto layout_x_swz = make_swizzled_layout_rc(); - Tensor sX = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_swz); - Tensor gX = make_tensor(make_gmem_ptr(old_x_w + ox_w_base), - make_layout(make_shape(Int{}, Int{}), - make_stride(params.old_x_stride_token, Int<1>{}))); - - using ThrLayoutX = Layout, Stride<_8, _1>>; - auto s2g = make_tiled_copy(Copy_Atom, input_t>{}, ThrLayoutX{}, - Layout>{}); - auto thr_s2g = s2g.get_slice(flat_tid); - - auto tSsX = thr_s2g.partition_S(sX); - auto tSgX = thr_s2g.partition_D(gX); - - // Per-(row, col) predicate: skip rows ≥ NPREDICTED (m-padding) and cols ≥ - // D_PER_CTA (atom-padding past the d_tile's data). - auto cX = make_identity_tensor(make_shape(Int{}, Int{})); - auto tScX = thr_s2g.partition_D(cX); - auto pred = make_tensor(shape(tScX)); - CUTE_UNROLL - for (int i = 0; i < size(pred); ++i) { - pred(i) = (get<0>(tScX(i)) < seq_len) && (get<1>(tScX(i)) < D_PER_CTA); - } - copy_if(s2g, pred, tSsX, tSgX); +__device__ __forceinline__ void store_old_x(SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, + int d_tile, int head, int64_t cache_slot, int write_offset, int seq_len) +{ + using namespace cute; + constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; + int const flat_tid = warp * warpSize + lane; + + auto* __restrict__ old_x_w = reinterpret_cast(params.old_x); + // gmem dest = head's full slot + d_tile's D-slice offset, shifted by + // `write_offset` along the T-axis (must_checkpoint ? 0 : prev_k). + int64_t const ox_w_base = cache_slot * params.old_x_stride_seq + (int64_t) write_offset * params.old_x_stride_token + + head * DIM + (int64_t) d_tile * D_PER_CTA; + + // Smem and gmem are both viewed at the full atom-padded width D_SMEM_COLS. + // The wide thread layout (16 row × 8 col × 1×8 val = 16 rows × 64 cols/pass + // for bf16) covers one full atom width per thread-row, which is the + // swizzle's bank-conflict-free contract on the LDS side (load-from-smem). + // A narrow layout would (a) waste 64 threads (warps 2, 3 idle) and + // (b) cause LDS bank conflicts on the smem-read side (observed as + // 4-way LDS conflict in d_split=2 ncu). Cols ≥ D_PER_CTA are predicated + // off via copy_if so STG never fires for them — no OOB write into the + // next d_tile / next head's gmem region. + constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; + auto layout_x_swz = make_swizzled_layout_rc(); + Tensor sX = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_swz); + Tensor gX = make_tensor(make_gmem_ptr(old_x_w + ox_w_base), + make_layout(make_shape(Int{}, Int{}), + make_stride(params.old_x_stride_token, Int<1>{}))); + + using ThrLayoutX = Layout, Stride<_8, _1>>; + auto s2g = make_tiled_copy(Copy_Atom, input_t>{}, ThrLayoutX{}, Layout>{}); + auto thr_s2g = s2g.get_slice(flat_tid); + + auto tSsX = thr_s2g.partition_S(sX); + auto tSgX = thr_s2g.partition_D(gX); + + // Per-(row, col) predicate: skip rows ≥ NPREDICTED (m-padding) and cols ≥ + // D_PER_CTA (atom-padding past the d_tile's data). + auto cX = make_identity_tensor(make_shape(Int{}, Int{})); + auto tScX = thr_s2g.partition_D(cX); + auto pred = make_tensor(shape(tScX)); + CUTE_UNROLL + for (int i = 0; i < size(pred); ++i) + { + pred(i) = (get<0>(tScX(i)) < seq_len) && (get<1>(tScX(i)) < D_PER_CTA); + } + copy_if(s2g, pred, tSsX, tSgX); } // store_old_B runs on W0, W1 only (64 threads). Caller must gate @@ -1615,57 +1675,59 @@ __device__ __forceinline__ void store_old_x(SmemT& smem, CheckpointingSsuParams // =8: iters (1, 2) = 2 tiles, each thread owns 1 row. The per-element // predicate works for both. template -__device__ __forceinline__ void store_old_B(SmemT& smem, CheckpointingSsuParams const& params, - int warp, int lane, int head, int group_idx, - int64_t cache_slot, int buf_write, int write_offset, - int seq_len) { - using namespace cute; - if (head % HEADS_PER_GROUP != 0) return; - constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; // matches smem.B row count - // Called only from warps 0, 1 — flat_tid ∈ [0, 64). - int const flat_tid = warp * warpSize + lane; - - auto* __restrict__ old_B_w = reinterpret_cast(params.old_B); - int64_t const oB_base = cache_slot * params.old_B_stride_seq + - buf_write * params.old_B_stride_dbuf + - (int64_t)write_offset * params.old_B_stride_token + group_idx * DSTATE; - - auto layout_B_swz = make_swizzled_layout_rc(); - Tensor sB = make_tensor(make_smem_ptr(reinterpret_cast(smem.B)), layout_B_swz); - Tensor gB = make_tensor(make_gmem_ptr(old_B_w + oB_base), - make_layout(make_shape(Int{}, Int{}), - make_stride(params.old_B_stride_token, Int<1>{}))); - - // 64 threads, (8, 8) × (1, 8) = atom-aligned per-tile (8, 64). - auto s2g = make_tiled_copy(Copy_Atom, input_t>{}, - Layout, Stride<_8, _1>>{}, Layout>{}); - auto thr_s2g = s2g.get_slice(flat_tid); - auto tSsB = thr_s2g.partition_S(sB); - auto tSgB = thr_s2g.partition_D(gB); - - // Fast path: no smem-side row padding AND no varlen-side truncation. - // The runtime `seq_len == NPREDICTED` is a constexpr-foldable compare in - // the non-varlen path (kernel prologue assigns `seq_len = NPREDICTED`), - // so it eliminates at -O3. In varlen with `seq_len == NPREDICTED` it's - // a runtime check that picks the cheaper unpredicated STG. - if constexpr (NPREDICTED == NPREDICTED_PAD_MMA_N) { - if (seq_len == NPREDICTED) { - copy(s2g, tSsB, tSgB); - return; +__device__ __forceinline__ void store_old_B(SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, + int head, int group_idx, int64_t cache_slot, int buf_write, int write_offset, int seq_len) +{ + using namespace cute; + if (head % HEADS_PER_GROUP != 0) + return; + constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; // matches smem.B row count + // Called only from warps 0, 1 — flat_tid ∈ [0, 64). + int const flat_tid = warp * warpSize + lane; + + auto* __restrict__ old_B_w = reinterpret_cast(params.old_B); + int64_t const oB_base = cache_slot * params.old_B_stride_seq + buf_write * params.old_B_stride_dbuf + + (int64_t) write_offset * params.old_B_stride_token + group_idx * DSTATE; + + auto layout_B_swz = make_swizzled_layout_rc(); + Tensor sB = make_tensor(make_smem_ptr(reinterpret_cast(smem.B)), layout_B_swz); + Tensor gB = make_tensor(make_gmem_ptr(old_B_w + oB_base), + make_layout( + make_shape(Int{}, Int{}), make_stride(params.old_B_stride_token, Int<1>{}))); + + // 64 threads, (8, 8) × (1, 8) = atom-aligned per-tile (8, 64). + auto s2g = make_tiled_copy(Copy_Atom, input_t>{}, Layout, Stride<_8, _1>>{}, + Layout>{}); + auto thr_s2g = s2g.get_slice(flat_tid); + auto tSsB = thr_s2g.partition_S(sB); + auto tSgB = thr_s2g.partition_D(gB); + + // Fast path: no smem-side row padding AND no varlen-side truncation. + // The runtime `seq_len == NPREDICTED` is a constexpr-foldable compare in + // the non-varlen path (kernel prologue assigns `seq_len = NPREDICTED`), + // so it eliminates at -O3. In varlen with `seq_len == NPREDICTED` it's + // a runtime check that picks the cheaper unpredicated STG. + if constexpr (NPREDICTED == NPREDICTED_PAD_MMA_N) + { + if (seq_len == NPREDICTED) + { + copy(s2g, tSsB, tSgB); + return; + } + } + // Predicated: either smem rows > NPREDICTED (m-padding) OR varlen with + // seq_len < NPREDICTED. Mask each iter against `seq_len`. + auto cB = make_identity_tensor(make_shape(Int{}, Int{})); + auto tScB = thr_s2g.partition_D(cB); + auto pred = make_tensor(shape(tScB)); + CUTE_UNROLL + for (int i = 0; i < size(pred); ++i) + { + pred(i) = get<0>(tScB(i)) < seq_len; } - } - // Predicated: either smem rows > NPREDICTED (m-padding) OR varlen with - // seq_len < NPREDICTED. Mask each iter against `seq_len`. - auto cB = make_identity_tensor(make_shape(Int{}, Int{})); - auto tScB = thr_s2g.partition_D(cB); - auto pred = make_tensor(shape(tScB)); - CUTE_UNROLL - for (int i = 0; i < size(pred); ++i) { - pred(i) = get<0>(tScB(i)) < seq_len; - } - copy_if(s2g, pred, tSsB, tSgB); + copy_if(s2g, pred, tSsB, tSgB); } -} // namespace flashinfer::mamba::checkpointing +} // namespace flashinfer::mamba::checkpointing -#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_COMMON_CUH_ +#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_COMMON_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh index 754217d204da..c338e2a9eb90 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh @@ -22,7 +22,8 @@ #include "kernel_checkpointing_ssu.cuh" #include "kernel_checkpointing_ssu_8bit.cuh" -namespace flashinfer::mamba::checkpointing { +namespace flashinfer::mamba::checkpointing +{ // ── Dispatcher ───────────────────────────────────────────────────────────── // `D_SPLIT` splits each head's DIM axis across `D_SPLIT` CTAs. @@ -30,116 +31,123 @@ namespace flashinfer::mamba::checkpointing { // `launchCheckpointingSsuImpl` is the per-(D_SPLIT, VARLEN) specialization; // `launchCheckpointingSsu` (below) is the runtime dispatcher. template -void launchCheckpointingSsuImpl(CheckpointingSsuParams& params, cudaStream_t stream) { - constexpr int NUM_WARPS = 4; - - FLASHINFER_CHECK(params.nheads % params.ngroups == 0, "nheads (", params.nheads, - ") must be divisible by ngroups (", params.ngroups, ")"); - - // cp.async.ca with .L2::128B requires 16B-aligned pointers (128-bit / sizeof element). - // The .L2::128B hint further requires the base address to be 128B-aligned for full - // cache line utilization, but the hardware only faults on < 16B alignment. - // All cp.async-loaded operands need 16B alignment; output is also vectorized - // (Pair stores partitioned by m16n8k16 partition_C — base must be at - // least 16B-aligned for the stride math to keep per-thread stores aligned). - FLASHINFER_CHECK_ALIGNMENT(params.B, 16); - FLASHINFER_CHECK_ALIGNMENT(params.C, 16); - FLASHINFER_CHECK_ALIGNMENT(params.x, 16); - FLASHINFER_CHECK_ALIGNMENT(params.state, 16); - FLASHINFER_CHECK_ALIGNMENT(params.old_x, 16); - FLASHINFER_CHECK_ALIGNMENT(params.old_B, 16); - FLASHINFER_CHECK_ALIGNMENT(params.output, 16); - if (params.z != nullptr) { - FLASHINFER_CHECK_ALIGNMENT(params.z, 16); - } - - // Per-CTA D = DIM / D_SPLIT. Smem footprint shrinks for D-owned - // buffers (state, x, z, old_x); non-D buffers (B, C, old_B, scalars) unchanged. - constexpr int D_PER_CTA = DIM / D_SPLIT; - - // HEADS_PER_GROUP is JIT-stamped via the customize_config jinja, so only - // one (nheads / ngroups) specialization gets baked into this .so. The - // wrapper has already validated `nheads / ngroups == HEADS_PER_GROUP` - // before reaching us — the kernel cross-checks with an assert below. - FLASHINFER_CHECK(params.nheads / params.ngroups == HEADS_PER_GROUP, - "nheads/ngroups (=", params.nheads / params.ngroups, - ") must match JIT HEADS_PER_GROUP=", HEADS_PER_GROUP); - // PDL launch attribute. ENABLE_PDL is JIT-stamped (see - // checkpointing_ssu_customize_config.jinja); the kernel's body has its - // PDL PTX gated on the same constexpr via `if constexpr (ENABLE_PDL)`, so - // the .so contains exactly one load path. When ENABLE_PDL is false the - // attribute is set to 0 (effectively no PDL) — cudaLaunchKernelEx is - // used either way per FlashInfer convention (see norm.cuh:135). - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = ENABLE_PDL ? 1 : 0; - - auto launch_kernel = [&]() { - if constexpr (sizeof(state_t) == 1) { - // int8 chain rewrite — uses checkpointing_ssu_kernel_8bit + - // CheckpointingSsuStorage8bit. Only D_SPLIT == 1 is valid (the wrapper - // asserts this); D_SPLIT == 2 still gets template-instantiated by the - // public dispatcher's switch but is unreachable at runtime — gate the - // body with `if constexpr (D_SPLIT == 1)` so that path doesn't launch. - if constexpr (D_SPLIT == 1) { - auto func = - checkpointing_ssu_kernel_8bit; - constexpr size_t smem_size = - sizeof(CheckpointingSsuStorage8bit); - - if constexpr (smem_size > 0) { - FLASHINFER_CUDA_CHECK( - cudaFuncSetAttribute(func, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - } - - cudaLaunchConfig_t config; - config.gridDim = dim3(D_SPLIT, params.batch, params.nheads); - config.blockDim = dim3(warpSize, NUM_WARPS); - config.dynamicSmemBytes = smem_size; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = 1; - FLASHINFER_CUDA_CHECK(cudaLaunchKernelEx(&config, func, params)); - } else { - FLASHINFER_CHECK(false, - "checkpointing_ssu_kernel_8bit: unsupported D_SPLIT != 1 for 8-bit " - "state_t (got D_SPLIT=", - D_SPLIT, ")"); - } - } else { - // Generic kernel: bf16 / fp16 / fp32 state, supports D_SPLIT ∈ {1, 2}. - auto func = - checkpointing_ssu_kernel; - - constexpr size_t smem_size = sizeof( - CheckpointingSsuStorage); - - if constexpr (smem_size > 0) { - FLASHINFER_CUDA_CHECK( - cudaFuncSetAttribute(func, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - } - - // Grid is (D_SPLIT, batch, nheads). D-tile is the fastest axis so the - // `D_SPLIT` CTAs of the same head land on adjacent SMs and share L2 - // lines for the redundantly-loaded inputs (C, B, dt, ...). - cudaLaunchConfig_t config; - config.gridDim = dim3(D_SPLIT, params.batch, params.nheads); - config.blockDim = dim3(warpSize, NUM_WARPS); - config.dynamicSmemBytes = smem_size; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = 1; - FLASHINFER_CUDA_CHECK(cudaLaunchKernelEx(&config, func, params)); + typename stateIndex_t, typename state_scale_t, int D_SPLIT, bool VARLEN> +void launchCheckpointingSsuImpl(CheckpointingSsuParams& params, cudaStream_t stream) +{ + constexpr int NUM_WARPS = 4; + + FLASHINFER_CHECK(params.nheads % params.ngroups == 0, "nheads (", params.nheads, ") must be divisible by ngroups (", + params.ngroups, ")"); + + // cp.async.ca with .L2::128B requires 16B-aligned pointers (128-bit / sizeof element). + // The .L2::128B hint further requires the base address to be 128B-aligned for full + // cache line utilization, but the hardware only faults on < 16B alignment. + // All cp.async-loaded operands need 16B alignment; output is also vectorized + // (Pair stores partitioned by m16n8k16 partition_C — base must be at + // least 16B-aligned for the stride math to keep per-thread stores aligned). + FLASHINFER_CHECK_ALIGNMENT(params.B, 16); + FLASHINFER_CHECK_ALIGNMENT(params.C, 16); + FLASHINFER_CHECK_ALIGNMENT(params.x, 16); + FLASHINFER_CHECK_ALIGNMENT(params.state, 16); + FLASHINFER_CHECK_ALIGNMENT(params.old_x, 16); + FLASHINFER_CHECK_ALIGNMENT(params.old_B, 16); + FLASHINFER_CHECK_ALIGNMENT(params.output, 16); + if (params.z != nullptr) + { + FLASHINFER_CHECK_ALIGNMENT(params.z, 16); } - }; - launch_kernel(); + // Per-CTA D = DIM / D_SPLIT. Smem footprint shrinks for D-owned + // buffers (state, x, z, old_x); non-D buffers (B, C, old_B, scalars) unchanged. + constexpr int D_PER_CTA = DIM / D_SPLIT; + + // HEADS_PER_GROUP is JIT-stamped via the customize_config jinja, so only + // one (nheads / ngroups) specialization gets baked into this .so. The + // wrapper has already validated `nheads / ngroups == HEADS_PER_GROUP` + // before reaching us — the kernel cross-checks with an assert below. + FLASHINFER_CHECK(params.nheads / params.ngroups == HEADS_PER_GROUP, + "nheads/ngroups (=", params.nheads / params.ngroups, ") must match JIT HEADS_PER_GROUP=", HEADS_PER_GROUP); + // PDL launch attribute. ENABLE_PDL is JIT-stamped (see + // checkpointing_ssu_customize_config.jinja); the kernel's body has its + // PDL PTX gated on the same constexpr via `if constexpr (ENABLE_PDL)`, so + // the .so contains exactly one load path. When ENABLE_PDL is false the + // attribute is set to 0 (effectively no PDL) — cudaLaunchKernelEx is + // used either way per FlashInfer convention (see norm.cuh:135). + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = ENABLE_PDL ? 1 : 0; + + auto launch_kernel = [&]() + { + if constexpr (sizeof(state_t) == 1) + { + // int8 chain rewrite — uses checkpointing_ssu_kernel_8bit + + // CheckpointingSsuStorage8bit. Only D_SPLIT == 1 is valid (the wrapper + // asserts this); D_SPLIT == 2 still gets template-instantiated by the + // public dispatcher's switch but is unreachable at runtime — gate the + // body with `if constexpr (D_SPLIT == 1)` so that path doesn't launch. + if constexpr (D_SPLIT == 1) + { + auto func = checkpointing_ssu_kernel_8bit; + constexpr size_t smem_size + = sizeof(CheckpointingSsuStorage8bit); + + if constexpr (smem_size > 0) + { + FLASHINFER_CUDA_CHECK( + cudaFuncSetAttribute(func, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + + cudaLaunchConfig_t config; + config.gridDim = dim3(D_SPLIT, params.batch, params.nheads); + config.blockDim = dim3(warpSize, NUM_WARPS); + config.dynamicSmemBytes = smem_size; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = 1; + FLASHINFER_CUDA_CHECK(cudaLaunchKernelEx(&config, func, params)); + } + else + { + FLASHINFER_CHECK(false, + "checkpointing_ssu_kernel_8bit: unsupported D_SPLIT != 1 for 8-bit " + "state_t (got D_SPLIT=", + D_SPLIT, ")"); + } + } + else + { + // Generic kernel: bf16 / fp16 / fp32 state, supports D_SPLIT ∈ {1, 2}. + auto func + = checkpointing_ssu_kernel; + + constexpr size_t smem_size + = sizeof(CheckpointingSsuStorage); + + if constexpr (smem_size > 0) + { + FLASHINFER_CUDA_CHECK( + cudaFuncSetAttribute(func, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + + // Grid is (D_SPLIT, batch, nheads). D-tile is the fastest axis so the + // `D_SPLIT` CTAs of the same head land on adjacent SMs and share L2 + // lines for the redundantly-loaded inputs (C, B, dt, ...). + cudaLaunchConfig_t config; + config.gridDim = dim3(D_SPLIT, params.batch, params.nheads); + config.blockDim = dim3(warpSize, NUM_WARPS); + config.dynamicSmemBytes = smem_size; + config.stream = stream; + config.attrs = attrs; + config.numAttrs = 1; + FLASHINFER_CUDA_CHECK(cudaLaunchKernelEx(&config, func, params)); + } + }; + + launch_kernel(); } // Public dispatcher: routes on `params.d_split` ({1, 2}) and varlen @@ -148,34 +156,37 @@ void launchCheckpointingSsuImpl(CheckpointingSsuParams& params, cudaStream_t str // only via `d_split` today, so the same compiled `.so` will hold all four // specializations after this commit. template -void launchCheckpointingSsu(CheckpointingSsuParams& params, cudaStream_t stream) { - bool const is_varlen = (params.cu_seqlens != nullptr); - auto launch = [&]() { - launchCheckpointingSsuImpl(params, stream); - }; - auto launch_d_split = [&]() { - if (is_varlen) { - launch.template operator()(); - } else { - launch.template operator()(); - } - }; - switch (params.d_split) { - case 1: - launch_d_split.template operator()<1>(); - break; - case 2: - launch_d_split.template operator()<2>(); - break; + typename stateIndex_t, typename state_scale_t> +void launchCheckpointingSsu(CheckpointingSsuParams& params, cudaStream_t stream) +{ + bool const is_varlen = (params.cu_seqlens != nullptr); + auto launch = [&]() + { + launchCheckpointingSsuImpl(params, stream); + }; + auto launch_d_split = [&]() + { + if (is_varlen) + { + launch.template operator()(); + } + else + { + launch.template operator()(); + } + }; + switch (params.d_split) + { + case 1: launch_d_split.template operator()<1>(); break; + case 2: launch_d_split.template operator()<2>(); break; default: - FLASHINFER_CHECK(false, "Unsupported d_split: ", params.d_split, - ". Allowed values: {1, 2}. d_split=4 needs " - "warp-count restructure."); - } + FLASHINFER_CHECK(false, "Unsupported d_split: ", params.d_split, + ". Allowed values: {1, 2}. d_split=4 needs " + "warp-count restructure."); + } } -} // namespace flashinfer::mamba::checkpointing +} // namespace flashinfer::mamba::checkpointing -#endif // FLASHINFER_MAMBA_LAUNCH_CHECKPOINTING_SSU_CUH_ +#endif // FLASHINFER_MAMBA_LAUNCH_CHECKPOINTING_SSU_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh index bb9accba44d9..8c6d8e93dd5b 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh @@ -22,40 +22,51 @@ #include "conversion.cuh" -namespace flashinfer::mamba::mtp { +namespace flashinfer::mamba::mtp +{ // Round up to next power of 2 (compile-time). -constexpr int nextPow2(int v) { - v--; - v |= v >> 1; - v |= v >> 2; - v |= v >> 4; - v |= v >> 8; - v |= v >> 16; - return v + 1; +constexpr int nextPow2(int v) +{ + v--; + v |= v >> 1; + v |= v >> 2; + v |= v >> 4; + v |= v >> 8; + v |= v >> 16; + return v + 1; } using barrier_t = cuda::barrier; -enum class WarpRole { kCompute, kTMALoad, kEpilogue }; - -__device__ __forceinline__ WarpRole get_warp_role(int warp) { - if (warp < 12) return WarpRole::kCompute; - if (warp < 15) return WarpRole::kTMALoad; - return WarpRole::kEpilogue; +enum class WarpRole +{ + kCompute, + kTMALoad, + kEpilogue +}; + +__device__ __forceinline__ WarpRole get_warp_role(int warp) +{ + if (warp < 12) + return WarpRole::kCompute; + if (warp < 15) + return WarpRole::kTMALoad; + return WarpRole::kEpilogue; } // XOR-based bank-conflict-free swizzle for horizontal state traversal. // Operates on flat byte addresses: XORs the bank index with the row (cycle) index. // cycle_length = row stride in bytes, bank_size = sizeof(uint32_t). template -__device__ __forceinline__ int xor_swizzle(int address) { - int const cycle = address / cycle_length; - int const delta = address % cycle_length; - int const bank_idx = delta / bank_size; - int const intra_bank = delta % bank_size; - int const new_bank_idx = bank_idx ^ cycle; - return cycle * cycle_length + new_bank_idx * bank_size + intra_bank; +__device__ __forceinline__ int xor_swizzle(int address) +{ + int const cycle = address / cycle_length; + int const delta = address % cycle_length; + int const bank_idx = delta / bank_size; + int const intra_bank = delta % bank_size; + int const new_bank_idx = bank_idx ^ cycle; + return cycle * cycle_length + new_bank_idx * bank_size + intra_bank; } // ── Parity-based barrier helpers (tight spin, no NANOSLEEP) ───────────────── @@ -65,22 +76,24 @@ __device__ __forceinline__ int xor_swizzle(int address) { // mbarrier.try_wait.parity instruction does a tight spin instead. // See CUDA Programming Guide §4.9.3 "Explicit Phase Tracking". -__device__ __forceinline__ void arrive_and_wait_parity(barrier_t& bar, uint32_t& parity) { - uint32_t const smem_addr = - static_cast(__cvta_generic_to_shared(cuda::device::barrier_native_handle(bar))); - asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(smem_addr) : "memory"); - uint32_t ready = 0; - while (!ready) { - asm volatile( - "{\n" - ".reg .pred p;\n" - "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n" - "selp.b32 %0, 1, 0, p;\n" - "}\n" - : "=r"(ready) - : "r"(smem_addr), "r"(parity)); - } - parity ^= 1; +__device__ __forceinline__ void arrive_and_wait_parity(barrier_t& bar, uint32_t& parity) +{ + uint32_t const smem_addr + = static_cast(__cvta_generic_to_shared(cuda::device::barrier_native_handle(bar))); + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(smem_addr) : "memory"); + uint32_t ready = 0; + while (!ready) + { + asm volatile( + "{\n" + ".reg .pred p;\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n" + "selp.b32 %0, 1, 0, p;\n" + "}\n" + : "=r"(ready) + : "r"(smem_addr), "r"(parity)); + } + parity ^= 1; } // ── SM100 f32x2 packed SIMD helpers ────────────────────────────────────────── @@ -91,27 +104,28 @@ __device__ __forceinline__ void arrive_and_wait_parity(barrier_t& bar, uint32_t& // On older architectures the fallback is two scalar ops — zero overhead. // See: https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/simd_sm100.hpp -__device__ __forceinline__ void mul_f32x2(float2& c, float2 const& a, float2 const& b) { +__device__ __forceinline__ void mul_f32x2(float2& c, float2 const& a, float2 const& b) +{ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 - asm("mul.f32x2 %0, %1, %2;\n" - : "=l"(reinterpret_cast(c)) - : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); + asm("mul.f32x2 %0, %1, %2;\n" + : "=l"(reinterpret_cast(c)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); #else - c.x = a.x * b.x; - c.y = a.y * b.y; + c.x = a.x * b.x; + c.y = a.y * b.y; #endif } -__device__ __forceinline__ void fma_f32x2(float2& d, float2 const& a, float2 const& b, - float2 const& c) { +__device__ __forceinline__ void fma_f32x2(float2& d, float2 const& a, float2 const& b, float2 const& c) +{ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 - asm("fma.rn.f32x2 %0, %1, %2, %3;\n" - : "=l"(reinterpret_cast(d)) - : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), + asm("fma.rn.f32x2 %0, %1, %2, %3;\n" + : "=l"(reinterpret_cast(d)) + : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), "l"(reinterpret_cast(c))); #else - d.x = a.x * b.x + c.x; - d.y = a.y * b.y + c.y; + d.x = a.x * b.x + c.x; + d.y = a.y * b.y + c.y; #endif } @@ -123,22 +137,24 @@ __device__ __forceinline__ void fma_f32x2(float2& d, float2 const& a, float2 con // ============================================================================= template -__device__ __forceinline__ void convertAndStoreSRHorizontal(state_t& out0, state_t& out1, float s0, - float s1, int64_t rand_seed, - int state_ptr_offset, int dd, int col0, - int e, uint32_t (&rand_ints)[4]) { - using namespace conversion; - if constexpr (PHILOX_ROUNDS > 0) { - if (e % 4 == 0) - philox_randint4x(rand_seed, state_ptr_offset + dd * DSTATE + col0 + e, - rand_ints[0], rand_ints[1], rand_ints[2], rand_ints[3]); - uint32_t packed = cvt_rs_f16x2_f32(s0, s1, rand_ints[e / 2 % 2]); - out0 = __ushort_as_half(static_cast(packed & 0xFFFFu)); - out1 = __ushort_as_half(static_cast(packed >> 16)); - } else { - convertAndStore(&out0, s0); - convertAndStore(&out1, s1); - } +__device__ __forceinline__ void convertAndStoreSRHorizontal(state_t& out0, state_t& out1, float s0, float s1, + int64_t rand_seed, int state_ptr_offset, int dd, int col0, int e, uint32_t (&rand_ints)[4]) +{ + using namespace conversion; + if constexpr (PHILOX_ROUNDS > 0) + { + if (e % 4 == 0) + philox_randint4x(rand_seed, state_ptr_offset + dd * DSTATE + col0 + e, rand_ints[0], + rand_ints[1], rand_ints[2], rand_ints[3]); + uint32_t packed = cvt_rs_f16x2_f32(s0, s1, rand_ints[e / 2 % 2]); + out0 = __ushort_as_half(static_cast(packed & 0xFFFFu)); + out1 = __ushort_as_half(static_cast(packed >> 16)); + } + else + { + convertAndStore(&out0, s0); + convertAndStore(&out1, s1); + } } -} // namespace flashinfer::mamba::mtp +} // namespace flashinfer::mamba::mtp diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh index 787a6de6d656..334cd3c5f2ee 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh @@ -38,399 +38,496 @@ #endif #ifndef NDEBUG -#define FLASHINFER_CUDA_CALL(func, ...) \ - { \ - cudaError_t e = (func); \ - if (e != cudaSuccess) { \ - std::cerr << "CUDA Error: " << cudaGetErrorString(e) << " (" << e << ") " << __FILE__ \ - << ": line " << __LINE__ << " at function " << STR(func) << std::endl; \ - return e; \ - } \ - } +#define FLASHINFER_CUDA_CALL(func, ...) \ + { \ + cudaError_t e = (func); \ + if (e != cudaSuccess) \ + { \ + std::cerr << "CUDA Error: " << cudaGetErrorString(e) << " (" << e << ") " << __FILE__ << ": line " \ + << __LINE__ << " at function " << STR(func) << std::endl; \ + return e; \ + } \ + } #else -#define FLASHINFER_CUDA_CALL(func, ...) \ - { \ - cudaError_t e = (func); \ - if (e != cudaSuccess) { \ - return e; \ - } \ - } +#define FLASHINFER_CUDA_CALL(func, ...) \ + { \ + cudaError_t e = (func); \ + if (e != cudaSuccess) \ + { \ + return e; \ + } \ + } #endif -#define FLASHINFER_CUDA_CHECK(func) \ - do { \ - cudaError_t e = (func); \ - FLASHINFER_CHECK(e == cudaSuccess, "CUDA Error: ", cudaGetErrorString(e), " (", int(e), \ - ") at ", __FILE__, ":", __LINE__, " in ", STR(func)); \ - } while (0) +#define FLASHINFER_CUDA_CHECK(func) \ + do \ + { \ + cudaError_t e = (func); \ + FLASHINFER_CHECK(e == cudaSuccess, "CUDA Error: ", cudaGetErrorString(e), " (", int(e), ") at ", __FILE__, \ + ":", __LINE__, " in ", STR(func)); \ + } while (0) -#define FLASHINFER_CHECK_ALIGNMENT(ptr, size_bytes) \ - FLASHINFER_CHECK(reinterpret_cast(ptr) % (size_bytes) == 0, #ptr, \ - " must be aligned to ", (size_bytes), " bytes, got address ", (uintptr_t)(ptr)) +#define FLASHINFER_CHECK_ALIGNMENT(ptr, size_bytes) \ + FLASHINFER_CHECK(reinterpret_cast(ptr) % (size_bytes) == 0, #ptr, " must be aligned to ", (size_bytes), \ + " bytes, got address ", (uintptr_t) (ptr)) #define FLASHINFER_CHECK_TMA_ALIGNED(ptr) FLASHINFER_CHECK_ALIGNMENT(ptr, 128) -#define DISPATCH_USE_FP16_QK_REDUCTION(use_fp16_qk_reduction, USE_FP16_QK_REDUCTION, ...) \ - if (use_fp16_qk_reduction) { \ - FLASHINFER_ERROR("FP16_QK_REDUCTION disabled at compile time"); \ - } else { \ - constexpr bool USE_FP16_QK_REDUCTION = false; \ - __VA_ARGS__ \ - } - -#define DISPATCH_NUM_MMA_Q(num_mma_q, NUM_MMA_Q, ...) \ - if (num_mma_q == 1) { \ - constexpr size_t NUM_MMA_Q = 1; \ - __VA_ARGS__ \ - } else if (num_mma_q == 2) { \ - constexpr size_t NUM_MMA_Q = 2; \ - __VA_ARGS__ \ - } else { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported num_mma_q: " << num_mma_q; \ - FLASHINFER_ERROR(err_msg.str()); \ - } - -#define DISPATCH_NUM_MMA_KV(max_mma_kv, NUM_MMA_KV, ...) \ - if (max_mma_kv >= 8) { \ - constexpr size_t NUM_MMA_KV = 8; \ - __VA_ARGS__ \ - } else if (max_mma_kv >= 4) { \ - constexpr size_t NUM_MMA_KV = 4; \ - __VA_ARGS__ \ - } else if (max_mma_kv >= 2) { \ - constexpr size_t NUM_MMA_KV = 2; \ - __VA_ARGS__ \ - } else if (max_mma_kv >= 1) { \ - constexpr size_t NUM_MMA_KV = 1; \ - __VA_ARGS__ \ - } else { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported max_mma_kv: " << max_mma_kv; \ - FLASHINFER_ERROR(err_msg.str()); \ - } - -#define DISPATCH_CTA_TILE_Q(cta_tile_q, CTA_TILE_Q, ...) \ - switch (cta_tile_q) { \ - case 128: { \ - constexpr uint32_t CTA_TILE_Q = 128; \ - __VA_ARGS__ \ - break; \ - } \ - case 64: { \ - constexpr uint32_t CTA_TILE_Q = 64; \ - __VA_ARGS__ \ - break; \ - } \ - case 16: { \ - constexpr uint32_t CTA_TILE_Q = 16; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported cta_tile_q: " << cta_tile_q; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_GQA_GROUP_SIZE(group_size, GROUP_SIZE, ...) \ - if (group_size == 1) { \ - constexpr size_t GROUP_SIZE = 1; \ - __VA_ARGS__ \ - } else if (group_size == 2) { \ - constexpr size_t GROUP_SIZE = 2; \ - __VA_ARGS__ \ - } else if (group_size == 3) { \ - constexpr size_t GROUP_SIZE = 3; \ - __VA_ARGS__ \ - } else if (group_size == 4) { \ - constexpr size_t GROUP_SIZE = 4; \ - __VA_ARGS__ \ - } else if (group_size == 8) { \ - constexpr size_t GROUP_SIZE = 8; \ - __VA_ARGS__ \ - } else { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported group_size: " << group_size; \ - FLASHINFER_ERROR(err_msg.str()); \ - } - -#define DISPATCH_MASK_MODE(mask_mode, MASK_MODE, ...) \ - switch (mask_mode) { \ - case MaskMode::kNone: { \ - constexpr MaskMode MASK_MODE = MaskMode::kNone; \ - __VA_ARGS__ \ - break; \ - } \ - case MaskMode::kCausal: { \ - constexpr MaskMode MASK_MODE = MaskMode::kCausal; \ - __VA_ARGS__ \ - break; \ - } \ - case MaskMode::kCustom: { \ - constexpr MaskMode MASK_MODE = MaskMode::kCustom; \ - __VA_ARGS__ \ - break; \ - } \ - case MaskMode::kMultiItemScoring: { \ - constexpr MaskMode MASK_MODE = MaskMode::kMultiItemScoring; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported mask_mode: " << int(mask_mode); \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } +#define DISPATCH_USE_FP16_QK_REDUCTION(use_fp16_qk_reduction, USE_FP16_QK_REDUCTION, ...) \ + if (use_fp16_qk_reduction) \ + { \ + FLASHINFER_ERROR("FP16_QK_REDUCTION disabled at compile time"); \ + } \ + else \ + { \ + constexpr bool USE_FP16_QK_REDUCTION = false; \ + __VA_ARGS__ \ + } + +#define DISPATCH_NUM_MMA_Q(num_mma_q, NUM_MMA_Q, ...) \ + if (num_mma_q == 1) \ + { \ + constexpr size_t NUM_MMA_Q = 1; \ + __VA_ARGS__ \ + } \ + else if (num_mma_q == 2) \ + { \ + constexpr size_t NUM_MMA_Q = 2; \ + __VA_ARGS__ \ + } \ + else \ + { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported num_mma_q: " << num_mma_q; \ + FLASHINFER_ERROR(err_msg.str()); \ + } + +#define DISPATCH_NUM_MMA_KV(max_mma_kv, NUM_MMA_KV, ...) \ + if (max_mma_kv >= 8) \ + { \ + constexpr size_t NUM_MMA_KV = 8; \ + __VA_ARGS__ \ + } \ + else if (max_mma_kv >= 4) \ + { \ + constexpr size_t NUM_MMA_KV = 4; \ + __VA_ARGS__ \ + } \ + else if (max_mma_kv >= 2) \ + { \ + constexpr size_t NUM_MMA_KV = 2; \ + __VA_ARGS__ \ + } \ + else if (max_mma_kv >= 1) \ + { \ + constexpr size_t NUM_MMA_KV = 1; \ + __VA_ARGS__ \ + } \ + else \ + { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported max_mma_kv: " << max_mma_kv; \ + FLASHINFER_ERROR(err_msg.str()); \ + } + +#define DISPATCH_CTA_TILE_Q(cta_tile_q, CTA_TILE_Q, ...) \ + switch (cta_tile_q) \ + { \ + case 128: \ + { \ + constexpr uint32_t CTA_TILE_Q = 128; \ + __VA_ARGS__ \ + break; \ + } \ + case 64: \ + { \ + constexpr uint32_t CTA_TILE_Q = 64; \ + __VA_ARGS__ \ + break; \ + } \ + case 16: \ + { \ + constexpr uint32_t CTA_TILE_Q = 16; \ + __VA_ARGS__ \ + break; \ + } \ + default: \ + { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported cta_tile_q: " << cta_tile_q; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_GQA_GROUP_SIZE(group_size, GROUP_SIZE, ...) \ + if (group_size == 1) \ + { \ + constexpr size_t GROUP_SIZE = 1; \ + __VA_ARGS__ \ + } \ + else if (group_size == 2) \ + { \ + constexpr size_t GROUP_SIZE = 2; \ + __VA_ARGS__ \ + } \ + else if (group_size == 3) \ + { \ + constexpr size_t GROUP_SIZE = 3; \ + __VA_ARGS__ \ + } \ + else if (group_size == 4) \ + { \ + constexpr size_t GROUP_SIZE = 4; \ + __VA_ARGS__ \ + } \ + else if (group_size == 8) \ + { \ + constexpr size_t GROUP_SIZE = 8; \ + __VA_ARGS__ \ + } \ + else \ + { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported group_size: " << group_size; \ + FLASHINFER_ERROR(err_msg.str()); \ + } + +#define DISPATCH_MASK_MODE(mask_mode, MASK_MODE, ...) \ + switch (mask_mode) \ + { \ + case MaskMode::kNone: \ + { \ + constexpr MaskMode MASK_MODE = MaskMode::kNone; \ + __VA_ARGS__ \ + break; \ + } \ + case MaskMode::kCausal: \ + { \ + constexpr MaskMode MASK_MODE = MaskMode::kCausal; \ + __VA_ARGS__ \ + break; \ + } \ + case MaskMode::kCustom: \ + { \ + constexpr MaskMode MASK_MODE = MaskMode::kCustom; \ + __VA_ARGS__ \ + break; \ + } \ + case MaskMode::kMultiItemScoring: \ + { \ + constexpr MaskMode MASK_MODE = MaskMode::kMultiItemScoring; \ + __VA_ARGS__ \ + break; \ + } \ + default: \ + { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported mask_mode: " << int(mask_mode); \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } // convert head_dim to compile-time constant -#define DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, ...) \ - switch (head_dim) { \ - case 64: { \ - constexpr size_t HEAD_DIM = 64; \ - __VA_ARGS__ \ - break; \ - } \ - case 128: { \ - constexpr size_t HEAD_DIM = 128; \ - __VA_ARGS__ \ - break; \ - } \ - case 256: { \ - constexpr size_t HEAD_DIM = 256; \ - __VA_ARGS__ \ - break; \ - } \ - case 512: { \ - constexpr size_t HEAD_DIM = 512; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported head_dim: " << head_dim; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } +#define DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, ...) \ + switch (head_dim) \ + { \ + case 64: \ + { \ + constexpr size_t HEAD_DIM = 64; \ + __VA_ARGS__ \ + break; \ + } \ + case 128: \ + { \ + constexpr size_t HEAD_DIM = 128; \ + __VA_ARGS__ \ + break; \ + } \ + case 256: \ + { \ + constexpr size_t HEAD_DIM = 256; \ + __VA_ARGS__ \ + break; \ + } \ + case 512: \ + { \ + constexpr size_t HEAD_DIM = 512; \ + __VA_ARGS__ \ + break; \ + } \ + default: \ + { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported head_dim: " << head_dim; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } // convert interleave to compile-time constant -#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \ - if (interleave) { \ - constexpr bool INTERLEAVE = true; \ - __VA_ARGS__ \ - } else { \ - constexpr bool INTERLEAVE = false; \ - __VA_ARGS__ \ - } - -#define DISPATCH_ROPE_DIM(rope_dim, ROPE_DIM, ...) \ - switch (rope_dim) { \ - case 16: { \ - constexpr uint32_t ROPE_DIM = 16; \ - __VA_ARGS__ \ - break; \ - } \ - case 32: { \ - constexpr uint32_t ROPE_DIM = 32; \ - __VA_ARGS__ \ - break; \ - } \ - case 64: { \ - constexpr uint32_t ROPE_DIM = 64; \ - __VA_ARGS__ \ - break; \ - } \ - case 128: { \ - constexpr uint32_t ROPE_DIM = 128; \ - __VA_ARGS__ \ - break; \ - } \ - case 256: { \ - constexpr uint32_t ROPE_DIM = 256; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported ROPE_DIM: " << rope_dim; \ - err_msg << ". Supported values: 16, 32, 64, 128, 256"; \ - err_msg << " in DISPATCH_ROPE_DIM"; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_POS_ENCODING_MODE(pos_encoding_mode, POS_ENCODING_MODE, ...) \ - switch (pos_encoding_mode) { \ - case PosEncodingMode::kNone: { \ - constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kNone; \ - __VA_ARGS__ \ - break; \ - } \ - case PosEncodingMode::kRoPELlama: { \ - constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kRoPELlama; \ - __VA_ARGS__ \ - break; \ - } \ - case PosEncodingMode::kALiBi: { \ - constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kALiBi; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported pos_encoding_mode: " << int(pos_encoding_mode); \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_ALIGNED_VEC_SIZE(aligned_vec_size, ALIGNED_VEC_SIZE, ...) \ - switch (aligned_vec_size) { \ - case 16: { \ - constexpr size_t ALIGNED_VEC_SIZE = 16; \ - __VA_ARGS__ \ - break; \ - } \ - case 8: { \ - constexpr size_t ALIGNED_VEC_SIZE = 8; \ - __VA_ARGS__ \ - break; \ - } \ - case 4: { \ - constexpr size_t ALIGNED_VEC_SIZE = 4; \ - __VA_ARGS__ \ - break; \ - } \ - case 2: { \ - constexpr size_t ALIGNED_VEC_SIZE = 2; \ - __VA_ARGS__ \ - break; \ - } \ - case 1: { \ - constexpr size_t ALIGNED_VEC_SIZE = 1; \ - __VA_ARGS__ \ - break; \ - } \ - default: { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported aligned_vec_size: " << aligned_vec_size; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, ...) \ - if (compute_capacity.first >= 8) { \ - constexpr uint32_t NUM_STAGES_SMEM = 2; \ - __VA_ARGS__ \ - } else { \ - constexpr uint32_t NUM_STAGES_SMEM = 1; \ - __VA_ARGS__ \ - } - -namespace flashinfer { +#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \ + if (interleave) \ + { \ + constexpr bool INTERLEAVE = true; \ + __VA_ARGS__ \ + } \ + else \ + { \ + constexpr bool INTERLEAVE = false; \ + __VA_ARGS__ \ + } + +#define DISPATCH_ROPE_DIM(rope_dim, ROPE_DIM, ...) \ + switch (rope_dim) \ + { \ + case 16: \ + { \ + constexpr uint32_t ROPE_DIM = 16; \ + __VA_ARGS__ \ + break; \ + } \ + case 32: \ + { \ + constexpr uint32_t ROPE_DIM = 32; \ + __VA_ARGS__ \ + break; \ + } \ + case 64: \ + { \ + constexpr uint32_t ROPE_DIM = 64; \ + __VA_ARGS__ \ + break; \ + } \ + case 128: \ + { \ + constexpr uint32_t ROPE_DIM = 128; \ + __VA_ARGS__ \ + break; \ + } \ + case 256: \ + { \ + constexpr uint32_t ROPE_DIM = 256; \ + __VA_ARGS__ \ + break; \ + } \ + default: \ + { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported ROPE_DIM: " << rope_dim; \ + err_msg << ". Supported values: 16, 32, 64, 128, 256"; \ + err_msg << " in DISPATCH_ROPE_DIM"; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_POS_ENCODING_MODE(pos_encoding_mode, POS_ENCODING_MODE, ...) \ + switch (pos_encoding_mode) \ + { \ + case PosEncodingMode::kNone: \ + { \ + constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kNone; \ + __VA_ARGS__ \ + break; \ + } \ + case PosEncodingMode::kRoPELlama: \ + { \ + constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kRoPELlama; \ + __VA_ARGS__ \ + break; \ + } \ + case PosEncodingMode::kALiBi: \ + { \ + constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kALiBi; \ + __VA_ARGS__ \ + break; \ + } \ + default: \ + { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported pos_encoding_mode: " << int(pos_encoding_mode); \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_ALIGNED_VEC_SIZE(aligned_vec_size, ALIGNED_VEC_SIZE, ...) \ + switch (aligned_vec_size) \ + { \ + case 16: \ + { \ + constexpr size_t ALIGNED_VEC_SIZE = 16; \ + __VA_ARGS__ \ + break; \ + } \ + case 8: \ + { \ + constexpr size_t ALIGNED_VEC_SIZE = 8; \ + __VA_ARGS__ \ + break; \ + } \ + case 4: \ + { \ + constexpr size_t ALIGNED_VEC_SIZE = 4; \ + __VA_ARGS__ \ + break; \ + } \ + case 2: \ + { \ + constexpr size_t ALIGNED_VEC_SIZE = 2; \ + __VA_ARGS__ \ + break; \ + } \ + case 1: \ + { \ + constexpr size_t ALIGNED_VEC_SIZE = 1; \ + __VA_ARGS__ \ + break; \ + } \ + default: \ + { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported aligned_vec_size: " << aligned_vec_size; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, ...) \ + if (compute_capacity.first >= 8) \ + { \ + constexpr uint32_t NUM_STAGES_SMEM = 2; \ + __VA_ARGS__ \ + } \ + else \ + { \ + constexpr uint32_t NUM_STAGES_SMEM = 1; \ + __VA_ARGS__ \ + } + +namespace flashinfer +{ template -__forceinline__ __device__ __host__ constexpr T1 ceil_div(const T1 x, const T2 y) noexcept { - return (x + y - 1) / y; +__forceinline__ __device__ __host__ constexpr T1 ceil_div(const T1 x, const T2 y) noexcept +{ + return (x + y - 1) / y; } template -__forceinline__ __device__ __host__ constexpr T1 round_up(const T1 x, const T2 y) noexcept { - return ceil_div(x, y) * y; +__forceinline__ __device__ __host__ constexpr T1 round_up(const T1 x, const T2 y) noexcept +{ + return ceil_div(x, y) * y; } template -__forceinline__ __device__ __host__ constexpr T1 round_down(const T1 x, const T2 y) noexcept { - return (x / y) * y; +__forceinline__ __device__ __host__ constexpr T1 round_down(const T1 x, const T2 y) noexcept +{ + return (x / y) * y; } -inline std::pair GetCudaComputeCapability() { - int device_id = 0; - cudaGetDevice(&device_id); - int major = 0, minor = 0; - cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_id); - cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device_id); - return std::make_pair(major, minor); +inline std::pair GetCudaComputeCapability() +{ + int device_id = 0; + cudaGetDevice(&device_id); + int major = 0, minor = 0; + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_id); + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device_id); + return std::make_pair(major, minor); } // This function is thread-safe and cached the sm_count. // But it will only check the current CUDA device, thus assuming each process handles single GPU. -inline int GetCudaMultiProcessorCount() { - static std::atomic sm_count{0}; - int cached = sm_count.load(std::memory_order_relaxed); - if (cached == 0) { - int device_id; - cudaGetDevice(&device_id); - cudaDeviceProp device_prop; - cudaGetDeviceProperties(&device_prop, device_id); - cached = device_prop.multiProcessorCount; - sm_count.store(cached, std::memory_order_relaxed); - } - return cached; +inline int GetCudaMultiProcessorCount() +{ + static std::atomic sm_count{0}; + int cached = sm_count.load(std::memory_order_relaxed); + if (cached == 0) + { + int device_id; + cudaGetDevice(&device_id); + cudaDeviceProp device_prop; + cudaGetDeviceProperties(&device_prop, device_id); + cached = device_prop.multiProcessorCount; + sm_count.store(cached, std::memory_order_relaxed); + } + return cached; } template -inline void DebugPrintCUDAArray(T* device_ptr, size_t size, std::string prefix = "") { - std::vector host_array(size); - std::cout << prefix; - cudaMemcpy(host_array.data(), device_ptr, size * sizeof(T), cudaMemcpyDeviceToHost); - for (size_t i = 0; i < size; ++i) { - std::cout << host_array[i] << " "; - } - std::cout << std::endl; +inline void DebugPrintCUDAArray(T* device_ptr, size_t size, std::string prefix = "") +{ + std::vector host_array(size); + std::cout << prefix; + cudaMemcpy(host_array.data(), device_ptr, size * sizeof(T), cudaMemcpyDeviceToHost); + for (size_t i = 0; i < size; ++i) + { + std::cout << host_array[i] << " "; + } + std::cout << std::endl; } -inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_dim) { - if (avg_packed_qo_len > 64 && head_dim < 256) { - return 128; - } else { - auto compute_capacity = GetCudaComputeCapability(); - if (compute_capacity.first >= 8) { - // Ampere or newer - if (avg_packed_qo_len > 16) { - // avg_packed_qo_len <= 64 - return 64; - } else { - // avg_packed_qo_len <= 16 - return 16; - } - } else { - // NOTE(Zihao): not enough shared memory on Turing for 1x4 warp layout - return 64; +inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_dim) +{ + if (avg_packed_qo_len > 64 && head_dim < 256) + { + return 128; + } + else + { + auto compute_capacity = GetCudaComputeCapability(); + if (compute_capacity.first >= 8) + { + // Ampere or newer + if (avg_packed_qo_len > 16) + { + // avg_packed_qo_len <= 64 + return 64; + } + else + { + // avg_packed_qo_len <= 16 + return 16; + } + } + else + { + // NOTE(Zihao): not enough shared memory on Turing for 1x4 warp layout + return 64; + } } - } } -inline int UpPowerOfTwo(int x) { - // Returns the smallest power of two greater than or equal to x - if (x <= 0) return 1; - --x; - x |= x >> 1; - x |= x >> 2; - x |= x >> 4; - x |= x >> 8; - x |= x >> 16; - return x + 1; +inline int UpPowerOfTwo(int x) +{ + // Returns the smallest power of two greater than or equal to x + if (x <= 0) + return 1; + --x; + x |= x >> 1; + x |= x >> 2; + x |= x >> 4; + x |= x >> 8; + x |= x >> 16; + return x + 1; } -#define LOOP_SPLIT_MASK(iter, COND1, COND2, ...) \ - { \ - _Pragma("unroll 1") for (; (COND1); (iter) -= 1) { \ - constexpr bool WITH_MASK = true; \ - __VA_ARGS__ \ - } \ - _Pragma("unroll 1") for (; (COND2); (iter) -= 1) { \ - constexpr bool WITH_MASK = false; \ - __VA_ARGS__ \ - } \ - } +#define LOOP_SPLIT_MASK(iter, COND1, COND2, ...) \ + { \ + _Pragma("unroll 1") for (; (COND1); (iter) -= 1) \ + { \ + constexpr bool WITH_MASK = true; \ + __VA_ARGS__ \ + } \ + _Pragma("unroll 1") for (; (COND2); (iter) -= 1) \ + { \ + constexpr bool WITH_MASK = false; \ + __VA_ARGS__ \ + } \ + } /*! * \brief Return x - y if x > y, otherwise return 0. */ -__device__ __forceinline__ uint32_t sub_if_greater_or_zero(uint32_t x, uint32_t y) { - return (x > y) ? x - y : 0U; +__device__ __forceinline__ uint32_t sub_if_greater_or_zero(uint32_t x, uint32_t y) +{ + return (x > y) ? x - y : 0U; } // ======================= PTX Memory Utility Functions ======================= @@ -440,100 +537,112 @@ __device__ __forceinline__ uint32_t sub_if_greater_or_zero(uint32_t x, uint32_t /*! * \brief Get the lane ID within a warp (0-31) */ -__forceinline__ __device__ int get_lane_id() { - int lane_id; - asm("mov.u32 %0, %%laneid;" : "=r"(lane_id)); - return lane_id; +__forceinline__ __device__ int get_lane_id() +{ + int lane_id; + asm("mov.u32 %0, %%laneid;" : "=r"(lane_id)); + return lane_id; } /*! * \brief Non-atomic global load for short (2 bytes) with cache streaming hint */ -__forceinline__ __device__ short ld_na_global_s16(const short* addr) { - short val; - asm volatile("ld.global.cs.b16 %0, [%1];" : "=h"(val) : "l"(addr)); - return val; +__forceinline__ __device__ short ld_na_global_s16(short const* addr) +{ + short val; + asm volatile("ld.global.cs.b16 %0, [%1];" : "=h"(val) : "l"(addr)); + return val; } /*! * \brief Non-atomic global store for short (2 bytes) with cache streaming hint */ -__forceinline__ __device__ void st_na_global_s16(short* addr, short val) { - asm volatile("st.global.cs.b16 [%0], %1;" ::"l"(addr), "h"(val)); +__forceinline__ __device__ void st_na_global_s16(short* addr, short val) +{ + asm volatile("st.global.cs.b16 [%0], %1;" ::"l"(addr), "h"(val)); } /*! * \brief Non-atomic global load for int (4 bytes) with cache streaming hint */ -__forceinline__ __device__ int ld_na_global_v1(const int* addr) { - int val; - asm volatile("ld.global.cs.b32 %0, [%1];" : "=r"(val) : "l"(addr)); - return val; +__forceinline__ __device__ int ld_na_global_v1(int const* addr) +{ + int val; + asm volatile("ld.global.cs.b32 %0, [%1];" : "=r"(val) : "l"(addr)); + return val; } /*! * \brief Non-atomic global load for int2 (8 bytes) with cache streaming hint */ -__forceinline__ __device__ int2 ld_na_global_v2(const int2* addr) { - int2 val; - asm volatile("ld.global.cs.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr)); - return val; +__forceinline__ __device__ int2 ld_na_global_v2(int2 const* addr) +{ + int2 val; + asm volatile("ld.global.cs.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr)); + return val; } /*! * \brief Non-atomic global store for int (4 bytes) with cache streaming hint */ -__forceinline__ __device__ void st_na_global_v1(int* addr, int val) { - asm volatile("st.global.cs.b32 [%0], %1;" ::"l"(addr), "r"(val)); +__forceinline__ __device__ void st_na_global_v1(int* addr, int val) +{ + asm volatile("st.global.cs.b32 [%0], %1;" ::"l"(addr), "r"(val)); } /*! * \brief Non-atomic global store for int2 (8 bytes) with cache streaming hint */ -__forceinline__ __device__ void st_na_global_v2(int2* addr, int2 val) { - asm volatile("st.global.cs.v2.b32 [%0], {%1, %2};" ::"l"(addr), "r"(val.x), "r"(val.y)); +__forceinline__ __device__ void st_na_global_v2(int2* addr, int2 val) +{ + asm volatile("st.global.cs.v2.b32 [%0], {%1, %2};" ::"l"(addr), "r"(val.x), "r"(val.y)); } /*! * \brief Prefetch data to L2 cache */ template -__forceinline__ __device__ void prefetch_L2(const T* addr) { - asm volatile("prefetch.global.L2 [%0];" ::"l"(addr)); +__forceinline__ __device__ void prefetch_L2(T const* addr) +{ + asm volatile("prefetch.global.L2 [%0];" ::"l"(addr)); } -__device__ __forceinline__ void swap(uint32_t& a, uint32_t& b) { - uint32_t tmp = a; - a = b; - b = tmp; +__device__ __forceinline__ void swap(uint32_t& a, uint32_t& b) +{ + uint32_t tmp = a; + a = b; + b = tmp; } -__device__ __forceinline__ uint32_t dim2_offset(const uint32_t& dim_a, const uint32_t& idx_b, - const uint32_t& idx_a) { - return idx_b * dim_a + idx_a; +__device__ __forceinline__ uint32_t dim2_offset(uint32_t const& dim_a, uint32_t const& idx_b, uint32_t const& idx_a) +{ + return idx_b * dim_a + idx_a; } -__device__ __forceinline__ uint32_t dim3_offset(const uint32_t& dim_b, const uint32_t& dim_a, - const uint32_t& idx_c, const uint32_t& idx_b, - const uint32_t& idx_a) { - return (idx_c * dim_b + idx_b) * dim_a + idx_a; +__device__ __forceinline__ uint32_t dim3_offset( + uint32_t const& dim_b, uint32_t const& dim_a, uint32_t const& idx_c, uint32_t const& idx_b, uint32_t const& idx_a) +{ + return (idx_c * dim_b + idx_b) * dim_a + idx_a; } -__device__ __forceinline__ uint32_t dim4_offset(const uint32_t& dim_c, const uint32_t& dim_b, - const uint32_t& dim_a, const uint32_t& idx_d, - const uint32_t& idx_c, const uint32_t& idx_b, - const uint32_t& idx_a) { - return ((idx_d * dim_c + idx_c) * dim_b + idx_b) * dim_a + idx_a; +__device__ __forceinline__ uint32_t dim4_offset(uint32_t const& dim_c, uint32_t const& dim_b, uint32_t const& dim_a, + uint32_t const& idx_d, uint32_t const& idx_c, uint32_t const& idx_b, uint32_t const& idx_a) +{ + return ((idx_d * dim_c + idx_c) * dim_b + idx_b) * dim_a + idx_a; } -#define DEFINE_HAS_MEMBER(member) \ - template \ - struct has_##member : std::false_type {}; \ - template \ - struct has_##member().member)>> : std::true_type {}; \ - template \ - inline constexpr bool has_##member##_v = has_##member::value; - -} // namespace flashinfer - -#endif // FLASHINFER_UTILS_CUH_ +#define DEFINE_HAS_MEMBER(member) \ + template \ + struct has_##member : std::false_type \ + { \ + }; \ + template \ + struct has_##member().member)>> : std::true_type \ + { \ + }; \ + template \ + inline constexpr bool has_##member##_v = has_##member::value; + +} // namespace flashinfer + +#endif // FLASHINFER_UTILS_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh index 25c3b6fc60d4..54ebaebb6e7d 100644 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh +++ b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh @@ -27,7 +27,8 @@ #include -namespace flashinfer { +namespace flashinfer +{ #if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 900)) #define FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED @@ -35,176 +36,230 @@ namespace flashinfer { #define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__ -__device__ __forceinline__ void st_global_release(int4 const& val, int4* addr) { - asm volatile("st.release.global.sys.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), - "r"(val.z), "r"(val.w), "l"(addr)); +__device__ __forceinline__ void st_global_release(int4 const& val, int4* addr) +{ + asm volatile("st.release.global.sys.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), "r"(val.z), + "r"(val.w), "l"(addr)); } -__device__ __forceinline__ int4 ld_global_acquire(int4* addr) { - int4 val; - asm volatile("ld.acquire.global.sys.v4.b32 {%0, %1, %2, %3}, [%4];" - : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) - : "l"(addr)); - return val; +__device__ __forceinline__ int4 ld_global_acquire(int4* addr) +{ + int4 val; + asm volatile("ld.acquire.global.sys.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(addr)); + return val; } -__device__ __forceinline__ void st_global_volatile(int4 const& val, int4* addr) { - asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), - "r"(val.z), "r"(val.w), "l"(addr)); +__device__ __forceinline__ void st_global_volatile(int4 const& val, int4* addr) +{ + asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), + "l"(addr)); } -__device__ __forceinline__ int4 ld_global_volatile(int4* addr) { - int4 val; - asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" - : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) - : "l"(addr)); - return val; +__device__ __forceinline__ int4 ld_global_volatile(int4* addr) +{ + int4 val; + asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(addr)); + return val; } -#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 < 120200) && \ - (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)) +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 < 120200) \ + && (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)) // CUDA version < 12.2 and GPU architecture < 80 -FLASHINFER_INLINE __nv_bfloat162 make_bfloat162(const __nv_bfloat16 x, const __nv_bfloat16 y) { - __nv_bfloat162 t; - t.x = x; - t.y = y; - return t; +FLASHINFER_INLINE __nv_bfloat162 make_bfloat162(const __nv_bfloat16 x, const __nv_bfloat16 y) +{ + __nv_bfloat162 t; + t.x = x; + t.y = y; + return t; } -FLASHINFER_INLINE __nv_bfloat16 __hmul(const __nv_bfloat16 a, const __nv_bfloat16 b) { - __nv_bfloat16 val; - const float fa = __bfloat162float(a); - const float fb = __bfloat162float(b); - // avoid ftz in device code - val = __float2bfloat16(__fmaf_ieee_rn(fa, fb, -0.0f)); - return val; +FLASHINFER_INLINE __nv_bfloat16 __hmul(const __nv_bfloat16 a, const __nv_bfloat16 b) +{ + __nv_bfloat16 val; + float const fa = __bfloat162float(a); + float const fb = __bfloat162float(b); + // avoid ftz in device code + val = __float2bfloat16(__fmaf_ieee_rn(fa, fb, -0.0f)); + return val; } -FLASHINFER_INLINE __nv_bfloat162 __hmul2(const __nv_bfloat162 a, const __nv_bfloat162 b) { - __nv_bfloat162 val; - val.x = __hmul(a.x, b.x); - val.y = __hmul(a.y, b.y); - return val; +FLASHINFER_INLINE __nv_bfloat162 __hmul2(const __nv_bfloat162 a, const __nv_bfloat162 b) +{ + __nv_bfloat162 val; + val.x = __hmul(a.x, b.x); + val.y = __hmul(a.y, b.y); + return val; } -FLASHINFER_INLINE __nv_bfloat162 __floats2bfloat162_rn(const float a, const float b) { - __nv_bfloat162 val; - val = __nv_bfloat162(__float2bfloat16_rn(a), __float2bfloat16_rn(b)); - return val; +FLASHINFER_INLINE __nv_bfloat162 __floats2bfloat162_rn(float const a, float const b) +{ + __nv_bfloat162 val; + val = __nv_bfloat162(__float2bfloat16_rn(a), __float2bfloat16_rn(b)); + return val; } -FLASHINFER_INLINE __nv_bfloat162 __float22bfloat162_rn(const float2 a) { - __nv_bfloat162 val = __floats2bfloat162_rn(a.x, a.y); - return val; +FLASHINFER_INLINE __nv_bfloat162 __float22bfloat162_rn(const float2 a) +{ + __nv_bfloat162 val = __floats2bfloat162_rn(a.x, a.y); + return val; } -FLASHINFER_INLINE float2 __bfloat1622float2(const __nv_bfloat162 a) { - float hi_float; - float lo_float; - lo_float = __internal_bfloat162float(((__nv_bfloat162_raw)a).x); - hi_float = __internal_bfloat162float(((__nv_bfloat162_raw)a).y); - return make_float2(lo_float, hi_float); + +FLASHINFER_INLINE float2 __bfloat1622float2(const __nv_bfloat162 a) +{ + float hi_float; + float lo_float; + lo_float = __internal_bfloat162float(((__nv_bfloat162_raw) a).x); + hi_float = __internal_bfloat162float(((__nv_bfloat162_raw) a).y); + return make_float2(lo_float, hi_float); } #endif /******************* vec_t type cast *******************/ template -struct vec_cast { - template - FLASHINFER_INLINE static void cast(dst_t* dst, const src_t* src) { +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(dst_t* dst, src_t const* src) + { #pragma unroll - for (size_t i = 0; i < vec_size; ++i) { - dst[i] = (dst_t)src[i]; + for (size_t i = 0; i < vec_size; ++i) + { + dst[i] = (dst_t) src[i]; + } } - } }; template <> -struct vec_cast<__nv_fp8_e4m3, float> { - template - FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, const float* src) { - if constexpr (vec_size == 1) { - dst[0] = __nv_fp8_e4m3(src[0]); - } else { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - ((__nv_fp8x2_storage_t*)dst)[i] = - __nv_cvt_float2_to_fp8x2(((float2*)src)[i], __NV_SATFINITE, __NV_E4M3); - } - } - } +struct vec_cast<__nv_fp8_e4m3, float> +{ + template + FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, float const* src) + { + if constexpr (vec_size == 1) + { + dst[0] = __nv_fp8_e4m3(src[0]); + } + else + { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + ((__nv_fp8x2_storage_t*) dst)[i] + = __nv_cvt_float2_to_fp8x2(((float2*) src)[i], __NV_SATFINITE, __NV_E4M3); + } + } + } }; template <> -struct vec_cast<__nv_fp8_e5m2, float> { - template - FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, const float* src) { - if constexpr (vec_size == 1) { - dst[0] = __nv_fp8_e5m2(src[0]); - } else { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - ((__nv_fp8x2_storage_t*)dst)[i] = - __nv_cvt_float2_to_fp8x2(((float2*)src)[i], __NV_SATFINITE, __NV_E5M2); - } - } - } +struct vec_cast<__nv_fp8_e5m2, float> +{ + template + FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, float const* src) + { + if constexpr (vec_size == 1) + { + dst[0] = __nv_fp8_e5m2(src[0]); + } + else + { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + ((__nv_fp8x2_storage_t*) dst)[i] + = __nv_cvt_float2_to_fp8x2(((float2*) src)[i], __NV_SATFINITE, __NV_E5M2); + } + } + } }; template <> -struct vec_cast { - template - FLASHINFER_INLINE static void cast(float* dst, const half* src) { - if constexpr (vec_size == 1) { - dst[0] = (float)src[0]; - } else { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - ((float2*)dst)[i] = __half22float2(((half2*)src)[i]); - } - } - } +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(float* dst, half const* src) + { + if constexpr (vec_size == 1) + { + dst[0] = (float) src[0]; + } + else + { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + ((float2*) dst)[i] = __half22float2(((half2*) src)[i]); + } + } + } }; template <> -struct vec_cast { - template - FLASHINFER_INLINE static void cast(half* dst, const float* src) { - if constexpr (vec_size == 1) { - dst[0] = __float2half(src[0]); - } else { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - ((half2*)dst)[i] = __float22half2_rn(((float2*)src)[i]); - } - } - } +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(half* dst, float const* src) + { + if constexpr (vec_size == 1) + { + dst[0] = __float2half(src[0]); + } + else + { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + ((half2*) dst)[i] = __float22half2_rn(((float2*) src)[i]); + } + } + } }; template -constexpr FLASHINFER_INLINE int get_exponent_bits() { - if constexpr (std::is_same_v) { - return 4; - } else if constexpr (std::is_same_v) { - return 5; - } else if constexpr (std::is_same_v) { - return 5; - } else if constexpr (std::is_same_v) { - return 8; - } +constexpr FLASHINFER_INLINE int get_exponent_bits() +{ + if constexpr (std::is_same_v) + { + return 4; + } + else if constexpr (std::is_same_v) + { + return 5; + } + else if constexpr (std::is_same_v) + { + return 5; + } + else if constexpr (std::is_same_v) + { + return 8; + } } template -constexpr FLASHINFER_INLINE int get_mantissa_bits() { - if constexpr (std::is_same_v) { - return 3; - } else if constexpr (std::is_same_v) { - return 2; - } else if constexpr (std::is_same_v) { - return 11; - } else if constexpr (std::is_same_v) { - return 7; - } +constexpr FLASHINFER_INLINE int get_mantissa_bits() +{ + if constexpr (std::is_same_v) + { + return 3; + } + else if constexpr (std::is_same_v) + { + return 2; + } + else if constexpr (std::is_same_v) + { + return 11; + } + else if constexpr (std::is_same_v) + { + return 7; + } } /*! @@ -216,202 +271,259 @@ constexpr FLASHINFER_INLINE int get_mantissa_bits() { * https://github.com/vllm-project/vllm/blob/6dffa4b0a6120159ef2fe44d695a46817aff65bc/csrc/quantization/fp8/fp8_marlin.cu#L120 */ template -__device__ void fast_dequant_f8f16x4(uint32_t* input, uint2* output) { - uint32_t q = *input; - if constexpr (std::is_same_v && std::is_same_v) { - output->x = __byte_perm(0U, q, 0x5140); - output->y = __byte_perm(0U, q, 0x7362); - } else { - constexpr int FP8_EXPONENT = get_exponent_bits(); - constexpr int FP8_MANTISSA = get_mantissa_bits(); - constexpr int FP16_EXPONENT = get_exponent_bits(); - - constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT; - // Calculate MASK for extracting mantissa and exponent - constexpr int MASK1 = 0x80000000; - constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA); - constexpr int MASK3 = MASK2 & 0x7fffffff; - constexpr int MASK = MASK3 | (MASK3 >> 16); - q = __byte_perm(q, q, 0x1302); - - // Extract and shift FP8 values to FP16 format - uint32_t Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - uint32_t Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT); - - constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); - // Construct and apply exponent bias - if constexpr (std::is_same_v) { - const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); - - // Convert to half2 and apply bias - *(half2*)&(output->x) = __hmul2(*reinterpret_cast(&Out1), bias_reg); - *(half2*)&(output->y) = __hmul2(*reinterpret_cast(&Out2), bias_reg); - } else { - constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; - const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); - // Convert to bfloat162 and apply bias - *(nv_bfloat162*)&(output->x) = - __hmul2(*reinterpret_cast(&Out1), bias_reg); - *(nv_bfloat162*)&(output->y) = - __hmul2(*reinterpret_cast(&Out2), bias_reg); - } - } +__device__ void fast_dequant_f8f16x4(uint32_t* input, uint2* output) +{ + uint32_t q = *input; + if constexpr (std::is_same_v && std::is_same_v) + { + output->x = __byte_perm(0U, q, 0x5140); + output->y = __byte_perm(0U, q, 0x7362); + } + else + { + constexpr int FP8_EXPONENT = get_exponent_bits(); + constexpr int FP8_MANTISSA = get_mantissa_bits(); + constexpr int FP16_EXPONENT = get_exponent_bits(); + + constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT; + // Calculate MASK for extracting mantissa and exponent + constexpr int MASK1 = 0x80000000; + constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA); + constexpr int MASK3 = MASK2 & 0x7fffffff; + constexpr int MASK = MASK3 | (MASK3 >> 16); + q = __byte_perm(q, q, 0x1302); + + // Extract and shift FP8 values to FP16 format + uint32_t Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + uint32_t Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT); + + constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); + // Construct and apply exponent bias + if constexpr (std::is_same_v) + { + const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); + + // Convert to half2 and apply bias + *(half2*) &(output->x) = __hmul2(*reinterpret_cast(&Out1), bias_reg); + *(half2*) &(output->y) = __hmul2(*reinterpret_cast(&Out2), bias_reg); + } + else + { + constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; + const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); + // Convert to bfloat162 and apply bias + *(nv_bfloat162*) &(output->x) = __hmul2(*reinterpret_cast(&Out1), bias_reg); + *(nv_bfloat162*) &(output->y) = __hmul2(*reinterpret_cast(&Out2), bias_reg); + } + } } template <> -struct vec_cast { - template - FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp8_e4m3* src) { - if constexpr (vec_size == 1) { - dst[0] = nv_bfloat16(src[0]); - } else if constexpr (vec_size == 2) { - dst[0] = nv_bfloat16(src[0]); - dst[1] = nv_bfloat16(src[1]); - } else { - static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); -#pragma unroll - for (uint32_t i = 0; i < vec_size / 4; ++i) { - fast_dequant_f8f16x4<__nv_fp8_e4m3, nv_bfloat16>((uint32_t*)&src[i * 4], - (uint2*)&dst[i * 4]); - } - } - } +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, __nv_fp8_e4m3 const* src) + { + if constexpr (vec_size == 1) + { + dst[0] = nv_bfloat16(src[0]); + } + else if constexpr (vec_size == 2) + { + dst[0] = nv_bfloat16(src[0]); + dst[1] = nv_bfloat16(src[1]); + } + else + { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) + { + fast_dequant_f8f16x4<__nv_fp8_e4m3, nv_bfloat16>((uint32_t*) &src[i * 4], (uint2*) &dst[i * 4]); + } + } + } }; template <> -struct vec_cast { - template - FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp8_e5m2* src) { - if constexpr (vec_size == 1) { - dst[0] = nv_bfloat16(src[0]); - } else if constexpr (vec_size == 2) { - dst[0] = nv_bfloat16(src[0]); - dst[1] = nv_bfloat16(src[1]); - } else { - static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); -#pragma unroll - for (uint32_t i = 0; i < vec_size / 4; ++i) { - fast_dequant_f8f16x4<__nv_fp8_e5m2, nv_bfloat16>((uint32_t*)&src[i * 4], - (uint2*)&dst[i * 4]); - } - } - } +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, __nv_fp8_e5m2 const* src) + { + if constexpr (vec_size == 1) + { + dst[0] = nv_bfloat16(src[0]); + } + else if constexpr (vec_size == 2) + { + dst[0] = nv_bfloat16(src[0]); + dst[1] = nv_bfloat16(src[1]); + } + else + { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) + { + fast_dequant_f8f16x4<__nv_fp8_e5m2, nv_bfloat16>((uint32_t*) &src[i * 4], (uint2*) &dst[i * 4]); + } + } + } }; template <> -struct vec_cast<__nv_fp8_e4m3, half> { - template - FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, const half* src) { +struct vec_cast<__nv_fp8_e4m3, half> +{ + template + FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, half const* src) + { #ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - if constexpr (vec_size == 1) { - dst[0] = __nv_fp8_e4m3(src[0]); - } else { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - uint16_t y; - uint32_t x = *(uint32_t*)&src[i * 2]; - asm volatile("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(y) : "r"(x)); - *(uint16_t*)&dst[i * 2] = y; - } - } + if constexpr (vec_size == 1) + { + dst[0] = __nv_fp8_e4m3(src[0]); + } + else + { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + uint16_t y; + uint32_t x = *(uint32_t*) &src[i * 2]; + asm volatile("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(y) : "r"(x)); + *(uint16_t*) &dst[i * 2] = y; + } + } #else #pragma unroll - for (size_t i = 0; i < vec_size; ++i) { - dst[i] = __nv_fp8_e4m3(src[i]); + for (size_t i = 0; i < vec_size; ++i) + { + dst[i] = __nv_fp8_e4m3(src[i]); + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED } -#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - } }; template <> -struct vec_cast<__nv_fp8_e5m2, half> { - template - FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, const half* src) { +struct vec_cast<__nv_fp8_e5m2, half> +{ + template + FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, half const* src) + { #ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - if constexpr (vec_size == 1) { - dst[0] = __nv_fp8_e5m2(src[0]); - } else { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - uint16_t y; - uint32_t x = *(uint32_t*)&src[i * 2]; - asm volatile("cvt.rn.satfinite.e5m2x2.f16x2 %0, %1;" : "=h"(y) : "r"(x)); - *(uint16_t*)&dst[i * 2] = y; - } - } + if constexpr (vec_size == 1) + { + dst[0] = __nv_fp8_e5m2(src[0]); + } + else + { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + uint16_t y; + uint32_t x = *(uint32_t*) &src[i * 2]; + asm volatile("cvt.rn.satfinite.e5m2x2.f16x2 %0, %1;" : "=h"(y) : "r"(x)); + *(uint16_t*) &dst[i * 2] = y; + } + } #else #pragma unroll - for (size_t i = 0; i < vec_size; ++i) { - dst[i] = __nv_fp8_e5m2(src[i]); + for (size_t i = 0; i < vec_size; ++i) + { + dst[i] = __nv_fp8_e5m2(src[i]); + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED } -#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - } }; template <> -struct vec_cast { - template - FLASHINFER_INLINE static void cast(half* dst, const __nv_fp8_e4m3* src) { +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(half* dst, __nv_fp8_e4m3 const* src) + { #ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - if constexpr (vec_size == 1) { - dst[0] = half(src[0]); - } else { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - uint32_t y; - uint16_t x = *(uint16_t*)&src[i * 2]; - asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(y) : "h"(x)); - *(uint32_t*)&dst[i * 2] = y; - } - } + if constexpr (vec_size == 1) + { + dst[0] = half(src[0]); + } + else + { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + uint32_t y; + uint16_t x = *(uint16_t*) &src[i * 2]; + asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(y) : "h"(x)); + *(uint32_t*) &dst[i * 2] = y; + } + } #else - if constexpr (vec_size == 1) { - dst[0] = half(src[0]); - } else if constexpr (vec_size == 2) { - dst[0] = half(src[0]); - dst[1] = half(src[1]); - } else { - static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); -#pragma unroll - for (uint32_t i = 0; i < vec_size / 4; ++i) { - fast_dequant_f8f16x4<__nv_fp8_e4m3, half>((uint32_t*)&src[i * 4], (uint2*)&dst[i * 4]); - } - } -#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - } + if constexpr (vec_size == 1) + { + dst[0] = half(src[0]); + } + else if constexpr (vec_size == 2) + { + dst[0] = half(src[0]); + dst[1] = half(src[1]); + } + else + { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) + { + fast_dequant_f8f16x4<__nv_fp8_e4m3, half>((uint32_t*) &src[i * 4], (uint2*) &dst[i * 4]); + } + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + } }; template <> -struct vec_cast { - template - FLASHINFER_INLINE static void cast(half* dst, const __nv_fp8_e5m2* src) { +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(half* dst, __nv_fp8_e5m2 const* src) + { #ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - if constexpr (vec_size == 1) { - dst[0] = half(src[0]); - } else { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - uint32_t y; - uint16_t x = *(uint16_t*)&src[i * 2]; - asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;" : "=r"(y) : "h"(x)); - *(uint32_t*)&dst[i * 2] = y; - } - } + if constexpr (vec_size == 1) + { + dst[0] = half(src[0]); + } + else + { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + uint32_t y; + uint16_t x = *(uint16_t*) &src[i * 2]; + asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;" : "=r"(y) : "h"(x)); + *(uint32_t*) &dst[i * 2] = y; + } + } #else - if constexpr (vec_size == 1) { - dst[0] = half(src[0]); - } else if constexpr (vec_size == 2) { - dst[0] = half(src[0]); - dst[1] = half(src[1]); - } else { - static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); -#pragma unroll - for (uint32_t i = 0; i < vec_size / 4; ++i) { - fast_dequant_f8f16x4<__nv_fp8_e5m2, half>((uint32_t*)&src[i * 4], (uint2*)&dst[i * 4]); - } - } -#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - } + if constexpr (vec_size == 1) + { + dst[0] = half(src[0]); + } + else if constexpr (vec_size == 2) + { + dst[0] = half(src[0]); + dst[1] = half(src[1]); + } + else + { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) + { + fast_dequant_f8f16x4<__nv_fp8_e5m2, half>((uint32_t*) &src[i * 4], (uint2*) &dst[i * 4]); + } + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + } }; #if defined(FLASHINFER_ENABLE_FP4_E2M1) && CUDA_VERSION >= 12080 @@ -421,1774 +533,2669 @@ struct vec_cast { // src[2] holds x2,x3 src[3] is padding ... etc. // Each valid byte encodes 2 fp4 values -> 2 fp16 via cvt.rn.f16x2.e2m1x2. template <> -struct vec_cast { - template - FLASHINFER_INLINE static void cast(half* dst, const __nv_fp4x2_e2m1* src) { - static_assert(vec_size % 2 == 0, "vec_size must be even for fp4x2 dequantization"); +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(half* dst, __nv_fp4x2_e2m1 const* src) + { + static_assert(vec_size % 2 == 0, "vec_size must be even for fp4x2 dequantization"); #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) #pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - uint32_t y; - // Valid fp4x2 bytes are at even positions (stride 2); odd positions are padding. - uint32_t b = reinterpret_cast(src)[i * 2]; - asm volatile( - "{\n" - ".reg .b8 fp4_byte;\n" - "mov.b32 {fp4_byte, _, _, _}, %1;\n" - "cvt.rn.f16x2.e2m1x2 %0, fp4_byte;\n" - "}" - : "=r"(y) - : "r"(b)); - reinterpret_cast(dst)[i] = y; - } + for (size_t i = 0; i < vec_size / 2; ++i) + { + uint32_t y; + // Valid fp4x2 bytes are at even positions (stride 2); odd positions are padding. + uint32_t b = reinterpret_cast(src)[i * 2]; + asm volatile( + "{\n" + ".reg .b8 fp4_byte;\n" + "mov.b32 {fp4_byte, _, _, _}, %1;\n" + "cvt.rn.f16x2.e2m1x2 %0, fp4_byte;\n" + "}" + : "=r"(y) + : "r"(b)); + reinterpret_cast(dst)[i] = y; + } #else - // Software LUT fallback for arch < SM100. - // e2m1 encoding: bit[3]=sign, bit[2:0]=magnitude index in {0,0.5,1,1.5,2,3,4,6}. - // Each packed byte holds two fp4 values: bits[3:0]=first, bits[7:4]=second. - constexpr uint16_t lut[16] = { - 0x0000, // +0.0 - 0x3800, // +0.5 - 0x3C00, // +1.0 - 0x3E00, // +1.5 - 0x4000, // +2.0 - 0x4200, // +3.0 - 0x4400, // +4.0 - 0x4600, // +6.0 - 0x8000, // -0.0 - 0xB800, // -0.5 - 0xBC00, // -1.0 - 0xBE00, // -1.5 - 0xC000, // -2.0 - 0xC200, // -3.0 - 0xC400, // -4.0 - 0xC600, // -6.0 - }; -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - uint8_t b = reinterpret_cast(src)[i * 2]; - reinterpret_cast(dst)[i * 2 + 0] = lut[b & 0x0F]; - reinterpret_cast(dst)[i * 2 + 1] = lut[(b >> 4) & 0x0F]; - } + // Software LUT fallback for arch < SM100. + // e2m1 encoding: bit[3]=sign, bit[2:0]=magnitude index in {0,0.5,1,1.5,2,3,4,6}. + // Each packed byte holds two fp4 values: bits[3:0]=first, bits[7:4]=second. + constexpr uint16_t lut[16] = { + 0x0000, // +0.0 + 0x3800, // +0.5 + 0x3C00, // +1.0 + 0x3E00, // +1.5 + 0x4000, // +2.0 + 0x4200, // +3.0 + 0x4400, // +4.0 + 0x4600, // +6.0 + 0x8000, // -0.0 + 0xB800, // -0.5 + 0xBC00, // -1.0 + 0xBE00, // -1.5 + 0xC000, // -2.0 + 0xC200, // -3.0 + 0xC400, // -4.0 + 0xC600, // -6.0 + }; +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + uint8_t b = reinterpret_cast(src)[i * 2]; + reinterpret_cast(dst)[i * 2 + 0] = lut[b & 0x0F]; + reinterpret_cast(dst)[i * 2 + 1] = lut[(b >> 4) & 0x0F]; + } #endif - } + } }; + template <> -struct vec_cast { - template - FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp4x2_e2m1* src) { - static_assert(vec_size % 2 == 0, "vec_size must be even for fp4x2 dequantization"); +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, __nv_fp4x2_e2m1 const* src) + { + static_assert(vec_size % 2 == 0, "vec_size must be even for fp4x2 dequantization"); #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) #pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - uint32_t y; - // Valid fp4x2 bytes are at even positions (stride 2); odd positions are padding. - uint32_t b = reinterpret_cast(src)[i * 2]; -#if (defined __CUDACC_VER_MAJOR__) && (defined __CUDACC_VER_MINOR__) && \ - ((__CUDACC_VER_MAJOR__ > 13) || ((__CUDACC_VER_MAJOR__ == 13) && (__CUDACC_VER_MINOR__ >= 2))) - // cvt.rn.bf16x2.e2m1x2 requires CUDA Toolkit >= 13.2 - asm volatile( - "{\n" - ".reg .b8 fp4_byte;\n" - "mov.b32 {fp4_byte, _, _, _}, %1;\n" - "cvt.rn.bf16x2.e2m1x2 %0, fp4_byte;\n" - "}" - : "=r"(y) - : "r"(b)); + for (size_t i = 0; i < vec_size / 2; ++i) + { + uint32_t y; + // Valid fp4x2 bytes are at even positions (stride 2); odd positions are padding. + uint32_t b = reinterpret_cast(src)[i * 2]; +#if (defined __CUDACC_VER_MAJOR__) && (defined __CUDACC_VER_MINOR__) \ + && ((__CUDACC_VER_MAJOR__ > 13) || ((__CUDACC_VER_MAJOR__ == 13) && (__CUDACC_VER_MINOR__ >= 2))) + // cvt.rn.bf16x2.e2m1x2 requires CUDA Toolkit >= 13.2 + asm volatile( + "{\n" + ".reg .b8 fp4_byte;\n" + "mov.b32 {fp4_byte, _, _, _}, %1;\n" + "cvt.rn.bf16x2.e2m1x2 %0, fp4_byte;\n" + "}" + : "=r"(y) + : "r"(b)); #else - // Fallback: convert e2m1 -> fp16 -> bf16 when cvt.rn.bf16x2.e2m1x2 is unavailable - uint32_t fp16x2; - asm volatile( - "{\n" - ".reg .b8 fp4_byte;\n" - "mov.b32 {fp4_byte, _, _, _}, %1;\n" - "cvt.rn.f16x2.e2m1x2 %0, fp4_byte;\n" - "}" - : "=r"(fp16x2) - : "r"(b)); - __half2 h2 = reinterpret_cast<__half2&>(fp16x2); - __nv_bfloat162 bf16x2 = __float22bfloat162_rn(__half22float2(h2)); - y = reinterpret_cast(bf16x2); + // Fallback: convert e2m1 -> fp16 -> bf16 when cvt.rn.bf16x2.e2m1x2 is unavailable + uint32_t fp16x2; + asm volatile( + "{\n" + ".reg .b8 fp4_byte;\n" + "mov.b32 {fp4_byte, _, _, _}, %1;\n" + "cvt.rn.f16x2.e2m1x2 %0, fp4_byte;\n" + "}" + : "=r"(fp16x2) + : "r"(b)); + __half2 h2 = reinterpret_cast<__half2&>(fp16x2); + __nv_bfloat162 bf16x2 = __float22bfloat162_rn(__half22float2(h2)); + y = reinterpret_cast(bf16x2); #endif - reinterpret_cast(dst)[i] = y; - } + reinterpret_cast(dst)[i] = y; + } #else - // Software LUT fallback for arch < SM100. - // e2m1 encoding: bit[3]=sign, bit[2:0]=magnitude index in {0,0.5,1,1.5,2,3,4,6}. - // Each packed byte holds two fp4 values: bits[3:0]=first, bits[7:4]=second. - constexpr uint16_t lut[16] = { - 0x0000, // +0.0 - 0x3F00, // +0.5 - 0x3F80, // +1.0 - 0x3FC0, // +1.5 - 0x4000, // +2.0 - 0x4040, // +3.0 - 0x4080, // +4.0 - 0x40C0, // +6.0 - 0x8000, // -0.0 - 0xBF00, // -0.5 - 0xBF80, // -1.0 - 0xBFC0, // -1.5 - 0xC000, // -2.0 - 0xC040, // -3.0 - 0xC080, // -4.0 - 0xC0C0, // -6.0 - }; -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - uint8_t b = reinterpret_cast(src)[i * 2]; - reinterpret_cast(dst)[i * 2 + 0] = lut[b & 0x0F]; - reinterpret_cast(dst)[i * 2 + 1] = lut[(b >> 4) & 0x0F]; - } + // Software LUT fallback for arch < SM100. + // e2m1 encoding: bit[3]=sign, bit[2:0]=magnitude index in {0,0.5,1,1.5,2,3,4,6}. + // Each packed byte holds two fp4 values: bits[3:0]=first, bits[7:4]=second. + constexpr uint16_t lut[16] = { + 0x0000, // +0.0 + 0x3F00, // +0.5 + 0x3F80, // +1.0 + 0x3FC0, // +1.5 + 0x4000, // +2.0 + 0x4040, // +3.0 + 0x4080, // +4.0 + 0x40C0, // +6.0 + 0x8000, // -0.0 + 0xBF00, // -0.5 + 0xBF80, // -1.0 + 0xBFC0, // -1.5 + 0xC000, // -2.0 + 0xC040, // -3.0 + 0xC080, // -4.0 + 0xC0C0, // -6.0 + }; +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + uint8_t b = reinterpret_cast(src)[i * 2]; + reinterpret_cast(dst)[i * 2 + 0] = lut[b & 0x0F]; + reinterpret_cast(dst)[i * 2 + 1] = lut[(b >> 4) & 0x0F]; + } #endif - } + } }; -#endif // FLASHINFER_ENABLE_FP4_E2M1 && CUDA_VERSION >= 12080 +#endif // FLASHINFER_ENABLE_FP4_E2M1 && CUDA_VERSION >= 12080 template <> -struct vec_cast { - template - FLASHINFER_INLINE static void cast(float* dst, const nv_bfloat16* src) { - if constexpr (vec_size == 1) { - dst[0] = (float)src[0]; - } else { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - ((float2*)dst)[i] = __bfloat1622float2(((nv_bfloat162*)src)[i]); - } - } - } +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(float* dst, nv_bfloat16 const* src) + { + if constexpr (vec_size == 1) + { + dst[0] = (float) src[0]; + } + else + { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + ((float2*) dst)[i] = __bfloat1622float2(((nv_bfloat162*) src)[i]); + } + } + } }; template <> -struct vec_cast { - template - FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const float* src) { - if constexpr (vec_size == 1) { - dst[0] = nv_bfloat16(src[0]); - } else { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) { - ((nv_bfloat162*)dst)[i] = __float22bfloat162_rn(((float2*)src)[i]); - } - } - } +struct vec_cast +{ + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, float const* src) + { + if constexpr (vec_size == 1) + { + dst[0] = nv_bfloat16(src[0]); + } + else + { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) + { + ((nv_bfloat162*) dst)[i] = __float22bfloat162_rn(((float2*) src)[i]); + } + } + } }; template -struct vec_t { - FLASHINFER_INLINE float_t& operator[](size_t i); - FLASHINFER_INLINE const float_t& operator[](size_t i) const; - FLASHINFER_INLINE void fill(float_t val); - FLASHINFER_INLINE void load(const float_t* ptr); - FLASHINFER_INLINE void store(float_t* ptr) const; - FLASHINFER_INLINE void load_global_acquire(float* addr); - FLASHINFER_INLINE void store_global_release(float* addr) const; - FLASHINFER_INLINE void load_global_volatile(float* addr); - FLASHINFER_INLINE void store_global_volatile(float* addr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src); - template - FLASHINFER_INLINE void cast_load(const T* ptr); - template - FLASHINFER_INLINE void cast_store(T* ptr) const; - FLASHINFER_INLINE static void memcpy(float_t* dst, const float_t* src); - FLASHINFER_INLINE float_t* ptr(); +struct vec_t +{ + FLASHINFER_INLINE float_t& operator[](size_t i); + FLASHINFER_INLINE float_t const& operator[](size_t i) const; + FLASHINFER_INLINE void fill(float_t val); + FLASHINFER_INLINE void load(float_t const* ptr); + FLASHINFER_INLINE void store(float_t* ptr) const; + FLASHINFER_INLINE void load_global_acquire(float* addr); + FLASHINFER_INLINE void store_global_release(float* addr) const; + FLASHINFER_INLINE void load_global_volatile(float* addr); + FLASHINFER_INLINE void store_global_volatile(float* addr) const; + template + FLASHINFER_INLINE void cast_from(vec_t const& src); + template + FLASHINFER_INLINE void cast_load(T const* ptr); + template + FLASHINFER_INLINE void cast_store(T* ptr) const; + FLASHINFER_INLINE static void memcpy(float_t* dst, float_t const* src); + FLASHINFER_INLINE float_t* ptr(); }; template -FLASHINFER_INLINE void cast_from_impl(vec_t& dst, - const vec_t& src) { - vec_cast::cast( - dst.ptr(), const_cast*>(&src)->ptr()); +FLASHINFER_INLINE void cast_from_impl(vec_t& dst, vec_t const& src) +{ + vec_cast::cast( + dst.ptr(), const_cast*>(&src)->ptr()); } template -FLASHINFER_INLINE void cast_load_impl(vec_t& dst, - const src_float_t* src_ptr) { - if constexpr (std::is_same_v) { - dst.load(src_ptr); - } else { - vec_t tmp; - tmp.load(src_ptr); - dst.cast_from(tmp); - } +FLASHINFER_INLINE void cast_load_impl(vec_t& dst, src_float_t const* src_ptr) +{ + if constexpr (std::is_same_v) + { + dst.load(src_ptr); + } + else + { + vec_t tmp; + tmp.load(src_ptr); + dst.cast_from(tmp); + } } template -FLASHINFER_INLINE void cast_store_impl(tgt_float_t* dst_ptr, - const vec_t& src) { - if constexpr (std::is_same_v) { - src.store(dst_ptr); - } else { - vec_t tmp; - tmp.cast_from(src); - tmp.store(dst_ptr); - } +FLASHINFER_INLINE void cast_store_impl(tgt_float_t* dst_ptr, vec_t const& src) +{ + if constexpr (std::is_same_v) + { + src.store(dst_ptr); + } + else + { + vec_t tmp; + tmp.cast_from(src); + tmp.store(dst_ptr); + } } /******************* vec_t<__nv_fp8_e4m3> *******************/ // __nv_fp8_e4m3 x 1 template <> -struct vec_t<__nv_fp8_e4m3, 1> { - __nv_fp8_e4m3 data; - - FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } - FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { - return ((const __nv_fp8_e4m3*)(&data))[i]; - } - FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); - FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); -}; - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::fill(__nv_fp8_e4m3 val) { data = val; } - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::load(const __nv_fp8_e4m3* ptr) { data = *ptr; } +struct vec_t<__nv_fp8_e4m3, 1> +{ + __nv_fp8_e4m3 data; -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::store(__nv_fp8_e4m3* ptr) const { *ptr = data; } - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::memcpy(__nv_fp8_e4m3* dst, - const __nv_fp8_e4m3* src) { - *dst = *src; -} + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) + { + return ((__nv_fp8_e4m3*) (&data))[i]; + } -// __nv_fp8_e4m3 x 2 -template <> -struct vec_t<__nv_fp8_e4m3, 2> { - __nv_fp8x2_e4m3 data; - - FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } - FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { - return ((const __nv_fp8_e4m3*)(&data))[i]; - } - FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); - FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); -}; + FLASHINFER_INLINE __nv_fp8_e4m3 const& operator[](size_t i) const + { + return ((__nv_fp8_e4m3 const*) (&data))[i]; + } -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::fill(__nv_fp8_e4m3 val) { - data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x); -} + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() + { + return reinterpret_cast<__nv_fp8_e4m3*>(&data); + } -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::load(const __nv_fp8_e4m3* ptr) { - data = *((__nv_fp8x2_e4m3*)ptr); -} + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(__nv_fp8_e4m3 const* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::store(__nv_fp8_e4m3* ptr) const { - *((__nv_fp8x2_e4m3*)ptr) = data; -} + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::memcpy(__nv_fp8_e4m3* dst, - const __nv_fp8_e4m3* src) { - *((__nv_fp8x2_e4m3*)dst) = *((__nv_fp8x2_e4m3*)src); -} + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } -// __nv_fp8_e4m3 x 4 + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } -template <> -struct vec_t<__nv_fp8_e4m3, 4> { - __nv_fp8x4_e4m3 data; - - FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } - FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { - return ((const __nv_fp8_e4m3*)(&data))[i]; - } - FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); - FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src); }; -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::fill(__nv_fp8_e4m3 val) { - data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::fill(__nv_fp8_e4m3 val) +{ + data = val; } -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::load(const __nv_fp8_e4m3* ptr) { - data = *((__nv_fp8x4_e4m3*)ptr); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::load(__nv_fp8_e4m3 const* ptr) +{ + data = *ptr; } -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::store(__nv_fp8_e4m3* ptr) const { - *((__nv_fp8x4_e4m3*)ptr) = data; +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::store(__nv_fp8_e4m3* ptr) const +{ + *ptr = data; } -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::memcpy(__nv_fp8_e4m3* dst, - const __nv_fp8_e4m3* src) { - *((__nv_fp8x4_e4m3*)dst) = *((__nv_fp8x4_e4m3*)src); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src) +{ + *dst = *src; } -// __nv_fp8_e4m3 x 8 - +// __nv_fp8_e4m3 x 2 template <> -struct vec_t<__nv_fp8_e4m3, 8> { - uint2 data; - - FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } - FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { - return ((const __nv_fp8_e4m3*)(&data))[i]; - } - FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); - FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); -}; - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::fill(__nv_fp8_e4m3 val) { - ((__nv_fp8x4_e4m3*)(&data.x))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e4m3*)(&data.y))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::load(const __nv_fp8_e4m3* ptr) { - data = *((uint2*)ptr); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::store(__nv_fp8_e4m3* ptr) const { - *((uint2*)ptr) = data; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::memcpy(__nv_fp8_e4m3* dst, - const __nv_fp8_e4m3* src) { - *((uint2*)dst) = *((uint2*)src); -} +struct vec_t<__nv_fp8_e4m3, 2> +{ + __nv_fp8x2_e4m3 data; -// __nv_fp8_e4m3 x 16 or more -template -struct vec_t<__nv_fp8_e4m3, vec_size> { - static_assert(vec_size % 16 == 0, "Invalid vector size"); - int4 data[vec_size / 16]; - - FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)data)[i]; } - FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { - return ((const __nv_fp8_e4m3*)data)[i]; - } - FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - ((__nv_fp8x4_e4m3*)(&(data[i].x)))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e4m3*)(&(data[i].y)))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e4m3*)(&(data[i].z)))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e4m3*)(&(data[i].w)))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - } - } - FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - data[i] = ((int4*)ptr)[i]; - } - } - FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - ((int4*)ptr)[i] = data[i]; - } - } - FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e4m3* addr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - *((int4*)(data + i)) = ld_global_acquire((int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void store_global_release(__nv_fp8_e4m3* addr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - st_global_release(data[i], (int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e4m3* addr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - data[i] = ld_global_volatile((int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e4m3* addr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - st_global_volatile(data[i], (int4*)(addr + i * 16)); - } - } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - ((int4*)dst)[i] = ((int4*)src)[i]; - } - } -}; + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) + { + return ((__nv_fp8_e4m3*) (&data))[i]; + } -/******************* vec_t<__nv_fp8_e5m2> *******************/ + FLASHINFER_INLINE __nv_fp8_e4m3 const& operator[](size_t i) const + { + return ((__nv_fp8_e4m3 const*) (&data))[i]; + } -// __nv_fp8_e5m2 x 1 -template <> -struct vec_t<__nv_fp8_e5m2, 1> { - __nv_fp8_e5m2 data; - - FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } - FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { - return ((const __nv_fp8_e5m2*)(&data))[i]; - } - FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); - FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); -}; + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() + { + return reinterpret_cast<__nv_fp8_e4m3*>(&data); + } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::fill(__nv_fp8_e5m2 val) { data = val; } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(__nv_fp8_e4m3 const* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::load(const __nv_fp8_e5m2* ptr) { data = *ptr; } + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::store(__nv_fp8_e5m2* ptr) const { *ptr = data; } + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::memcpy(__nv_fp8_e5m2* dst, - const __nv_fp8_e5m2* src) { - *dst = *src; -} + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } -// __nv_fp8_e5m2 x 2 -template <> -struct vec_t<__nv_fp8_e5m2, 2> { - __nv_fp8x2_e5m2 data; - - FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } - FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { - return ((const __nv_fp8_e5m2*)(&data))[i]; - } - FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); - FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src); }; -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::fill(__nv_fp8_e5m2 val) { - data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::fill(__nv_fp8_e4m3 val) +{ + data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x); } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::load(const __nv_fp8_e5m2* ptr) { - data = *((__nv_fp8x2_e5m2*)ptr); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::load(__nv_fp8_e4m3 const* ptr) +{ + data = *((__nv_fp8x2_e4m3*) ptr); } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::store(__nv_fp8_e5m2* ptr) const { - *((__nv_fp8x2_e5m2*)ptr) = data; +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::store(__nv_fp8_e4m3* ptr) const +{ + *((__nv_fp8x2_e4m3*) ptr) = data; } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::memcpy(__nv_fp8_e5m2* dst, - const __nv_fp8_e5m2* src) { - *((__nv_fp8x2_e5m2*)dst) = *((__nv_fp8x2_e5m2*)src); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src) +{ + *((__nv_fp8x2_e4m3*) dst) = *((__nv_fp8x2_e4m3*) src); } -// __nv_fp8_e5m2 x 4 +// __nv_fp8_e4m3 x 4 template <> -struct vec_t<__nv_fp8_e5m2, 4> { - __nv_fp8x4_e5m2 data; - - FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } - FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { - return ((const __nv_fp8_e5m2*)(&data))[i]; - } - FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); - FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); -}; +struct vec_t<__nv_fp8_e4m3, 4> +{ + __nv_fp8x4_e4m3 data; -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::fill(__nv_fp8_e5m2 val) { - data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); -} + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) + { + return ((__nv_fp8_e4m3*) (&data))[i]; + } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::load(const __nv_fp8_e5m2* ptr) { - data = *((__nv_fp8x4_e5m2*)ptr); -} + FLASHINFER_INLINE __nv_fp8_e4m3 const& operator[](size_t i) const + { + return ((__nv_fp8_e4m3 const*) (&data))[i]; + } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::store(__nv_fp8_e5m2* ptr) const { - *((__nv_fp8x4_e5m2*)ptr) = data; -} + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() + { + return reinterpret_cast<__nv_fp8_e4m3*>(&data); + } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::memcpy(__nv_fp8_e5m2* dst, - const __nv_fp8_e5m2* src) { - *((__nv_fp8x4_e5m2*)dst) = *((__nv_fp8x4_e5m2*)src); -} + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(__nv_fp8_e4m3 const* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; -// __nv_fp8_e5m2 x 8 + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } -template <> -struct vec_t<__nv_fp8_e5m2, 8> { - uint2 data; - - FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } - FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { - return ((const __nv_fp8_e5m2*)(&data))[i]; - } - FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); - FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src); }; -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::fill(__nv_fp8_e5m2 val) { - ((__nv_fp8x4_e5m2*)(&data.x))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e5m2*)(&data.y))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::fill(__nv_fp8_e4m3 val) +{ + data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) + | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::load(const __nv_fp8_e5m2* ptr) { - data = *((uint2*)ptr); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::load(__nv_fp8_e4m3 const* ptr) +{ + data = *((__nv_fp8x4_e4m3*) ptr); } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::store(__nv_fp8_e5m2* ptr) const { - *((uint2*)ptr) = data; +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::store(__nv_fp8_e4m3* ptr) const +{ + *((__nv_fp8x4_e4m3*) ptr) = data; } -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::memcpy(__nv_fp8_e5m2* dst, - const __nv_fp8_e5m2* src) { - *((uint2*)dst) = *((uint2*)src); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src) +{ + *((__nv_fp8x4_e4m3*) dst) = *((__nv_fp8x4_e4m3*) src); } -// __nv_fp8_e5m2 x 16 or more - -template -struct vec_t<__nv_fp8_e5m2, vec_size> { - static_assert(vec_size % 16 == 0, "Invalid vector size"); - int4 data[vec_size / 16]; - - FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)data)[i]; } - FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { - return ((const __nv_fp8_e5m2*)data)[i]; - } - FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - ((__nv_fp8x4_e5m2*)(&(data[i].x)))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e5m2*)(&(data[i].y)))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e5m2*)(&(data[i].z)))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e5m2*)(&(data[i].w)))->__x = - (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | - (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - } - } - FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - data[i] = ((int4*)ptr)[i]; - } - } - FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - ((int4*)ptr)[i] = data[i]; - } - } - FLASHINFER_INLINE void store_global_release(__nv_fp8_e5m2* addr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - st_global_release(data[i], (int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e5m2* addr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - data[i] = ld_global_acquire((int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e5m2* addr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - st_global_volatile(data[i], (int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e5m2* addr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - data[i] = ld_global_volatile((int4*)(addr + i * 16)); - } - } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - ((int4*)dst)[i] = ((int4*)src)[i]; - } - } -}; - -#if defined(FLASHINFER_ENABLE_FP4_E2M1) && CUDA_VERSION >= 12080 -/******************* vec_t<__nv_fp4_e2m1> *******************/ - -// __nv_fp4_e2m1 x 2 -template <> -struct vec_t<__nv_fp4_e2m1, 2> { - uint8_t data; - // index access is not supported for sub-byte data type - FLASHINFER_INLINE __nv_fp4_e2m1* ptr() { return reinterpret_cast<__nv_fp4_e2m1*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) { - data = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); - } - FLASHINFER_INLINE void load(const __nv_fp4_e2m1* ptr) { data = *((uint8_t*)ptr); } - FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const { *((uint8_t*)ptr) = data; } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, const __nv_fp4_e2m1* src) { - *((uint8_t*)dst) = *((uint8_t*)src); - } -}; - -// __nv_fp4_e2m1 x 4 -template <> -struct vec_t<__nv_fp4_e2m1, 4> { - uint16_t data; - - FLASHINFER_INLINE __nv_fp4_e2m1* ptr() { return reinterpret_cast<__nv_fp4_e2m1*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) { - __nv_fp4x2_storage_t val8 = - (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); - data = (uint16_t(val8) << 8) | uint16_t(val8); - } - FLASHINFER_INLINE void load(const __nv_fp4_e2m1* ptr) { data = *((uint16_t*)ptr); } - FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const { *((uint16_t*)ptr) = data; } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, const __nv_fp4_e2m1* src) { - *((uint16_t*)dst) = *((uint16_t*)src); - } -}; - -// __nv_fp4_e2m1 x 8 -template <> -struct vec_t<__nv_fp4_e2m1, 8> { - uint32_t data; - - FLASHINFER_INLINE __nv_fp4_e2m1* ptr() { return reinterpret_cast<__nv_fp4_e2m1*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) { - __nv_fp4x2_storage_t val8 = - (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); - uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); - data = (uint32_t(val16) << 16) | uint32_t(val16); - } - FLASHINFER_INLINE void load(const __nv_fp4_e2m1* ptr) { data = *((uint32_t*)ptr); } - FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const { *((uint32_t*)ptr) = data; } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, const __nv_fp4_e2m1* src) { - *((uint32_t*)dst) = *((uint32_t*)src); - } -}; +// __nv_fp8_e4m3 x 8 template <> -struct vec_t<__nv_fp4_e2m1, 16> { - uint2 data; - FLASHINFER_INLINE __nv_fp4_e2m1* ptr() { return reinterpret_cast<__nv_fp4_e2m1*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) { - __nv_fp4x2_storage_t val8 = - (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); - uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); - uint32_t val32 = (uint32_t(val16) << 16) | uint32_t(val16); - data.x = val32; - data.y = val32; - } - FLASHINFER_INLINE void load(const __nv_fp4_e2m1* ptr) { data = *((uint2*)ptr); } - FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const { *((uint2*)ptr) = data; } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, const __nv_fp4_e2m1* src) { - *((uint2*)dst) = *((uint2*)src); - } -}; - -// __nv_fp4_e2m1 x 32 or more -template -struct vec_t<__nv_fp4_e2m1, vec_size> { - static_assert(vec_size % 32 == 0, "Invalid vector size"); - int4 data[vec_size / 32]; - - FLASHINFER_INLINE __nv_fp4_e2m1* ptr() { return reinterpret_cast<__nv_fp4_e2m1*>(&data); } - FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) { - __nv_fp4x2_storage_t val8 = - (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); - uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); - uint32_t val32 = (uint32_t(val16) << 16) | uint32_t(val16); -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) { - data[i].x = val32; - data[i].y = val32; - data[i].z = val32; - data[i].w = val32; - } - } - FLASHINFER_INLINE void load(const __nv_fp4_e2m1* ptr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) { - data[i] = ((int4*)ptr)[i]; - } - } - FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) { - ((int4*)ptr)[i] = data[i]; - } - } - FLASHINFER_INLINE void store_global_release(__nv_fp4_e2m1* addr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) { - st_global_release(*(int4*)&data[i], (int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void load_global_acquire(__nv_fp4_e2m1* addr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) { - *(int4*)&data[i] = ld_global_acquire((int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void store_global_volatile(__nv_fp4_e2m1* addr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) { - st_global_volatile(*(int4*)&data[i], (int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void load_global_volatile(__nv_fp4_e2m1* addr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) { - *(int4*)&data[i] = ld_global_volatile((int4*)(addr + i * 16)); - } - } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, const __nv_fp4_e2m1* src) { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) { - ((int4*)dst)[i] = ((int4*)src)[i]; - } - } -}; +struct vec_t<__nv_fp8_e4m3, 8> +{ + uint2 data; -#endif // FLASHINFER_ENABLE_FP4_E2M1 && CUDA_VERSION >= 12080 + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) + { + return ((__nv_fp8_e4m3*) (&data))[i]; + } -/******************* vec_t *******************/ + FLASHINFER_INLINE __nv_fp8_e4m3 const& operator[](size_t i) const + { + return ((__nv_fp8_e4m3 const*) (&data))[i]; + } -// half x 1 -template <> -struct vec_t { - half data; - - FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; } - FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; } - FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(half val); - FLASHINFER_INLINE void load(const half* ptr); - FLASHINFER_INLINE void store(half* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(half* dst, const half* src); -}; + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() + { + return reinterpret_cast<__nv_fp8_e4m3*>(&data); + } -FLASHINFER_INLINE void vec_t::fill(half val) { data = val; } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(__nv_fp8_e4m3 const* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; -FLASHINFER_INLINE void vec_t::load(const half* ptr) { data = *ptr; } + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } -FLASHINFER_INLINE void vec_t::store(half* ptr) const { *ptr = data; } + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } -FLASHINFER_INLINE void vec_t::memcpy(half* dst, const half* src) { *dst = *src; } + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } -// half x 2 -template <> -struct vec_t { - half2 data; - - FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; } - FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; } - FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(half val); - FLASHINFER_INLINE void load(const half* ptr); - FLASHINFER_INLINE void store(half* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(half* dst, const half* src); + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src); }; -FLASHINFER_INLINE void vec_t::fill(half val) { data = make_half2(val, val); } - -FLASHINFER_INLINE void vec_t::load(const half* ptr) { data = *((half2*)ptr); } - -FLASHINFER_INLINE void vec_t::store(half* ptr) const { *((half2*)ptr) = data; } - -FLASHINFER_INLINE void vec_t::memcpy(half* dst, const half* src) { - *((half2*)dst) = *((half2*)src); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::fill(__nv_fp8_e4m3 val) +{ + ((__nv_fp8x4_e4m3*) (&data.x))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) + | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*) (&data.y))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) + | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); } -// half x 4 +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::load(__nv_fp8_e4m3 const* ptr) +{ + data = *((uint2*) ptr); +} -template <> -struct vec_t { - uint2 data; - - FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; } - FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; } - FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(half val); - FLASHINFER_INLINE void load(const half* ptr); - FLASHINFER_INLINE void store(half* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(half* dst, const half* src); -}; +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::store(__nv_fp8_e4m3* ptr) const +{ + *((uint2*) ptr) = data; +} -FLASHINFER_INLINE void vec_t::fill(half val) { - *(half2*)(&data.x) = make_half2(val, val); - *(half2*)(&data.y) = make_half2(val, val); +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src) +{ + *((uint2*) dst) = *((uint2*) src); } -FLASHINFER_INLINE void vec_t::load(const half* ptr) { data = *((uint2*)ptr); } +// __nv_fp8_e4m3 x 16 or more +template +struct vec_t<__nv_fp8_e4m3, vec_size> +{ + static_assert(vec_size % 16 == 0, "Invalid vector size"); + int4 data[vec_size / 16]; -FLASHINFER_INLINE void vec_t::store(half* ptr) const { *((uint2*)ptr) = data; } + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) + { + return ((__nv_fp8_e4m3*) data)[i]; + } -FLASHINFER_INLINE void vec_t::memcpy(half* dst, const half* src) { - *((uint2*)dst) = *((uint2*)src); -} + FLASHINFER_INLINE __nv_fp8_e4m3 const& operator[](size_t i) const + { + return ((__nv_fp8_e4m3 const*) data)[i]; + } -// half x 8 or more + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() + { + return reinterpret_cast<__nv_fp8_e4m3*>(&data); + } -template -struct vec_t { - static_assert(vec_size % 8 == 0, "Invalid vector size"); - int4 data[vec_size / 8]; - FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)data)[i]; } - FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)data)[i]; } - FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(half val) { + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val) + { #pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - *(half2*)(&(data[i].x)) = make_half2(val, val); - *(half2*)(&(data[i].y)) = make_half2(val, val); - *(half2*)(&(data[i].z)) = make_half2(val, val); - *(half2*)(&(data[i].w)) = make_half2(val, val); + for (size_t i = 0; i < vec_size / 16; ++i) + { + ((__nv_fp8x4_e4m3*) (&(data[i].x)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) + | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) + | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*) (&(data[i].y)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) + | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) + | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*) (&(data[i].z)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) + | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) + | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*) (&(data[i].w)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) + | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) + | __nv_fp8x4_storage_t(val.__x); + } } - } - FLASHINFER_INLINE void load(const half* ptr) { + + FLASHINFER_INLINE void load(__nv_fp8_e4m3 const* ptr) + { #pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - data[i] = ((int4*)ptr)[i]; + for (size_t i = 0; i < vec_size / 16; ++i) + { + data[i] = ((int4*) ptr)[i]; + } } - } - FLASHINFER_INLINE void store(half* ptr) const { + + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const + { #pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - ((int4*)ptr)[i] = data[i]; + for (size_t i = 0; i < vec_size / 16; ++i) + { + ((int4*) ptr)[i] = data[i]; + } } - } - FLASHINFER_INLINE void load_global_acquire(half* addr) { + + FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e4m3* addr) + { #pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - data[i] = ld_global_acquire((int4*)(addr + i * 8)); + for (size_t i = 0; i < vec_size / 16; ++i) + { + *((int4*) (data + i)) = ld_global_acquire((int4*) (addr + i * 16)); + } } - } - FLASHINFER_INLINE void store_global_release(half* addr) const { + + FLASHINFER_INLINE void store_global_release(__nv_fp8_e4m3* addr) const + { #pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - st_global_release(data[i], (int4*)(addr + i * 8)); + for (size_t i = 0; i < vec_size / 16; ++i) + { + st_global_release(data[i], (int4*) (addr + i * 16)); + } } - } - FLASHINFER_INLINE void store_global_volatile(half* addr) const { + + FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e4m3* addr) + { #pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - st_global_volatile(data[i], (int4*)(addr + i * 8)); + for (size_t i = 0; i < vec_size / 16; ++i) + { + data[i] = ld_global_volatile((int4*) (addr + i * 16)); + } } - } - FLASHINFER_INLINE void load_global_volatile(half* addr) { + + FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e4m3* addr) const + { #pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - data[i] = ld_global_volatile((int4*)(addr + i * 8)); + for (size_t i = 0; i < vec_size / 16; ++i) + { + st_global_volatile(data[i], (int4*) (addr + i * 16)); + } } - } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(half* dst, const half* src) { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - ((int4*)dst)[i] = ((int4*)src)[i]; + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); } - } -}; -/******************* vec_t *******************/ + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } -// nv_bfloat16 x 1 -template <> -struct vec_t { - nv_bfloat16 data; - FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; } - FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { - return ((const nv_bfloat16*)(&data))[i]; - } - FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(nv_bfloat16 val); - FLASHINFER_INLINE void load(const nv_bfloat16* ptr); - FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src); -}; + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } -FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) { data = val; } + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + ((int4*) dst)[i] = ((int4*) src)[i]; + } + } +}; -FLASHINFER_INLINE void vec_t::load(const nv_bfloat16* ptr) { data = *ptr; } +/******************* vec_t<__nv_fp8_e5m2> *******************/ -FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const { *ptr = data; } +// __nv_fp8_e5m2 x 1 +template <> +struct vec_t<__nv_fp8_e5m2, 1> +{ + __nv_fp8_e5m2 data; -FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { - *dst = *src; -} + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) + { + return ((__nv_fp8_e5m2*) (&data))[i]; + } -// nv_bfloat16 x 2 -template <> -struct vec_t { - nv_bfloat162 data; - - FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; } - FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { - return ((const nv_bfloat16*)(&data))[i]; - } - FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(nv_bfloat16 val); - FLASHINFER_INLINE void load(const nv_bfloat16* ptr); - FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src); -}; + FLASHINFER_INLINE __nv_fp8_e5m2 const& operator[](size_t i) const + { + return ((__nv_fp8_e5m2 const*) (&data))[i]; + } -FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) { - data = make_bfloat162(val, val); -} + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() + { + return reinterpret_cast<__nv_fp8_e5m2*>(&data); + } -FLASHINFER_INLINE void vec_t::load(const nv_bfloat16* ptr) { - data = *((nv_bfloat162*)ptr); -} + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(__nv_fp8_e5m2 const* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; -FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const { - *((nv_bfloat162*)ptr) = data; -} + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } -FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { - *((nv_bfloat162*)dst) = *((nv_bfloat162*)src); -} + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } -// nv_bfloat16 x 4 + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } -template <> -struct vec_t { - uint2 data; - - FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; } - FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { - return ((const nv_bfloat16*)(&data))[i]; - } - FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(nv_bfloat16 val); - FLASHINFER_INLINE void load(const nv_bfloat16* ptr); - FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src); + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src); }; -FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) { - *(nv_bfloat162*)(&data.x) = make_bfloat162(val, val); - *(nv_bfloat162*)(&data.y) = make_bfloat162(val, val); +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::fill(__nv_fp8_e5m2 val) +{ + data = val; } -FLASHINFER_INLINE void vec_t::load(const nv_bfloat16* ptr) { - data = *((uint2*)ptr); +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::load(__nv_fp8_e5m2 const* ptr) +{ + data = *ptr; } -FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const { - *((uint2*)ptr) = data; +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::store(__nv_fp8_e5m2* ptr) const +{ + *ptr = data; } -FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { - *((uint2*)dst) = *((uint2*)src); +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src) +{ + *dst = *src; } -// nv_bfloat16 x 8 or more +// __nv_fp8_e5m2 x 2 +template <> +struct vec_t<__nv_fp8_e5m2, 2> +{ + __nv_fp8x2_e5m2 data; -template -struct vec_t { - static_assert(vec_size % 8 == 0, "Invalid vector size"); - int4 data[vec_size / 8]; - - FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)data)[i]; } - FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { - return ((const nv_bfloat16*)data)[i]; - } - FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(nv_bfloat16 val) { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - *(nv_bfloat162*)(&(data[i].x)) = make_bfloat162(val, val); - *(nv_bfloat162*)(&(data[i].y)) = make_bfloat162(val, val); - *(nv_bfloat162*)(&(data[i].z)) = make_bfloat162(val, val); - *(nv_bfloat162*)(&(data[i].w)) = make_bfloat162(val, val); - } - } - FLASHINFER_INLINE void load(const nv_bfloat16* ptr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - data[i] = ((int4*)ptr)[i]; - } - } - FLASHINFER_INLINE void store(nv_bfloat16* ptr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - ((int4*)ptr)[i] = data[i]; - } - } - FLASHINFER_INLINE void store_global_release(nv_bfloat16* addr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - st_global_release(data[i], (int4*)(addr + i * 8)); - } - } - FLASHINFER_INLINE void load_global_acquire(nv_bfloat16* addr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - data[i] = ld_global_acquire((int4*)(addr + i * 8)); - } - } - FLASHINFER_INLINE void store_global_volatile(nv_bfloat16* addr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - st_global_volatile(data[i], (int4*)(addr + i * 8)); - } - } - FLASHINFER_INLINE void load_global_volatile(nv_bfloat16* addr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - data[i] = ld_global_volatile((int4*)(addr + i * 8)); - } - } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) { - ((int4*)dst)[i] = ((int4*)src)[i]; - } - } -}; + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) + { + return ((__nv_fp8_e5m2*) (&data))[i]; + } -/******************* vec_t *******************/ + FLASHINFER_INLINE __nv_fp8_e5m2 const& operator[](size_t i) const + { + return ((__nv_fp8_e5m2 const*) (&data))[i]; + } -// uint8_t x 1 -template <> -struct vec_t { - uint8_t data; - - FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } - FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { - return ((const uint8_t*)(&data))[i]; - } - FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(uint8_t val); - FLASHINFER_INLINE void load(const uint8_t* ptr); - FLASHINFER_INLINE void store(uint8_t* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); -}; + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() + { + return reinterpret_cast<__nv_fp8_e5m2*>(&data); + } -FLASHINFER_INLINE void vec_t::fill(uint8_t val) { data = val; } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(__nv_fp8_e5m2 const* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; -FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *ptr; } + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } -FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *ptr = data; } + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } -FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { *dst = *src; } + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } -// uint8_t x 2 -template <> -struct vec_t { - uint16_t data; - - FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } - FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { - return ((const uint8_t*)(&data))[i]; - } - FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(uint8_t val); - FLASHINFER_INLINE void load(const uint8_t* ptr); - FLASHINFER_INLINE void store(uint8_t* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src); }; -FLASHINFER_INLINE void vec_t::fill(uint8_t val) { - data = (uint16_t(val) << 8) | uint16_t(val); +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::fill(__nv_fp8_e5m2 val) +{ + data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x); } -FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *((uint16_t*)ptr); } +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::load(__nv_fp8_e5m2 const* ptr) +{ + data = *((__nv_fp8x2_e5m2*) ptr); +} -FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *((uint16_t*)ptr) = data; } +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::store(__nv_fp8_e5m2* ptr) const +{ + *((__nv_fp8x2_e5m2*) ptr) = data; +} -FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { - *((uint16_t*)dst) = *((uint16_t*)src); +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src) +{ + *((__nv_fp8x2_e5m2*) dst) = *((__nv_fp8x2_e5m2*) src); } -// uint8_t x 4 +// __nv_fp8_e5m2 x 4 template <> -struct vec_t { - uint32_t data; - - FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } - FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { - return ((const uint8_t*)(&data))[i]; - } - FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(uint8_t val); - FLASHINFER_INLINE void load(const uint8_t* ptr); - FLASHINFER_INLINE void store(uint8_t* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); -}; +struct vec_t<__nv_fp8_e5m2, 4> +{ + __nv_fp8x4_e5m2 data; -FLASHINFER_INLINE void vec_t::fill(uint8_t val) { - data = (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); -} + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) + { + return ((__nv_fp8_e5m2*) (&data))[i]; + } -FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *((uint32_t*)ptr); } + FLASHINFER_INLINE __nv_fp8_e5m2 const& operator[](size_t i) const + { + return ((__nv_fp8_e5m2 const*) (&data))[i]; + } -FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *((uint32_t*)ptr) = data; } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() + { + return reinterpret_cast<__nv_fp8_e5m2*>(&data); + } -FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { - *((uint32_t*)dst) = *((uint32_t*)src); -} + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(__nv_fp8_e5m2 const* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; -// uint8_t x 8 + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } -template <> -struct vec_t { - uint2 data; - - FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } - FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { - return ((const uint8_t*)(&data))[i]; - } - FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(uint8_t val); - FLASHINFER_INLINE void load(const uint8_t* ptr); - FLASHINFER_INLINE void store(uint8_t* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src); }; -FLASHINFER_INLINE void vec_t::fill(uint8_t val) { - uint32_t val32 = - (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); - data.x = val32; - data.y = val32; +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::fill(__nv_fp8_e5m2 val) +{ + data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) + | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); } -FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *((uint2*)ptr); } - -FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *((uint2*)ptr) = data; } +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::load(__nv_fp8_e5m2 const* ptr) +{ + data = *((__nv_fp8x4_e5m2*) ptr); +} -FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { - *((uint2*)dst) = *((uint2*)src); +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::store(__nv_fp8_e5m2* ptr) const +{ + *((__nv_fp8x4_e5m2*) ptr) = data; } -// uint8_t x 16 or more +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src) +{ + *((__nv_fp8x4_e5m2*) dst) = *((__nv_fp8x4_e5m2*) src); +} -template -struct vec_t { - static_assert(vec_size % 16 == 0, "Invalid vector size"); - int4 data[vec_size / 16]; - - FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)data)[i]; } - FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { return ((const uint8_t*)data)[i]; } - FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(uint8_t val) { - uint32_t val32 = - (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - data[i].x = val32; - data[i].y = val32; - data[i].z = val32; - data[i].w = val32; - } - } - FLASHINFER_INLINE void load(const uint8_t* ptr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - data[i] = ((int4*)ptr)[i]; - } - } - FLASHINFER_INLINE void store(uint8_t* ptr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - ((int4*)ptr)[i] = data[i]; - } - } - FLASHINFER_INLINE void load_global_acquire(uint8_t* addr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - data[i] = ld_global_acquire((int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void store_global_release(uint8_t* addr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - st_global_release(data[i], (int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void load_global_volatile(uint8_t* addr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - data[i] = ld_global_volatile((int4*)(addr + i * 16)); - } - } - FLASHINFER_INLINE void store_global_volatile(uint8_t* addr) const { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - st_global_volatile(data[i], (int4*)(addr + i * 16)); - } - } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src) { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) { - ((int4*)dst)[i] = ((int4*)src)[i]; - } - } -}; +// __nv_fp8_e5m2 x 8 -/******************* vec_t *******************/ +template <> +struct vec_t<__nv_fp8_e5m2, 8> +{ + uint2 data; -// float x 1 + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) + { + return ((__nv_fp8_e5m2*) (&data))[i]; + } -template <> -struct vec_t { - float data; - - FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(&data))[i]; } - FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(&data))[i]; } - FLASHINFER_INLINE float* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(float val); - FLASHINFER_INLINE void load(const float* ptr); - FLASHINFER_INLINE void store(float* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(float* dst, const float* src); -}; + FLASHINFER_INLINE __nv_fp8_e5m2 const& operator[](size_t i) const + { + return ((__nv_fp8_e5m2 const*) (&data))[i]; + } -FLASHINFER_INLINE void vec_t::fill(float val) { data = val; } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() + { + return reinterpret_cast<__nv_fp8_e5m2*>(&data); + } -FLASHINFER_INLINE void vec_t::load(const float* ptr) { data = *ptr; } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(__nv_fp8_e5m2 const* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; -FLASHINFER_INLINE void vec_t::store(float* ptr) const { *ptr = data; } + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } -FLASHINFER_INLINE void vec_t::memcpy(float* dst, const float* src) { *dst = *src; } + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } -// float x 2 + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } -template <> -struct vec_t { - float2 data; - - FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(&data))[i]; } - FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(&data))[i]; } - FLASHINFER_INLINE float* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(float val); - FLASHINFER_INLINE void load(const float* ptr); - FLASHINFER_INLINE void store(float* ptr) const; - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(float* dst, const float* src); + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src); }; -FLASHINFER_INLINE void vec_t::fill(float val) { data = make_float2(val, val); } +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::fill(__nv_fp8_e5m2 val) +{ + ((__nv_fp8x4_e5m2*) (&data.x))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) + | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*) (&data.y))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) + | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +} -FLASHINFER_INLINE void vec_t::load(const float* ptr) { data = *((float2*)ptr); } +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::load(__nv_fp8_e5m2 const* ptr) +{ + data = *((uint2*) ptr); +} -FLASHINFER_INLINE void vec_t::store(float* ptr) const { *((float2*)ptr) = data; } +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::store(__nv_fp8_e5m2* ptr) const +{ + *((uint2*) ptr) = data; +} -FLASHINFER_INLINE void vec_t::memcpy(float* dst, const float* src) { - *((float2*)dst) = *((float2*)src); +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src) +{ + *((uint2*) dst) = *((uint2*) src); } -// float x 4 or more +// __nv_fp8_e5m2 x 16 or more + template -struct vec_t { - static_assert(vec_size % 4 == 0, "Invalid vector size"); - float4 data[vec_size / 4]; +struct vec_t<__nv_fp8_e5m2, vec_size> +{ + static_assert(vec_size % 16 == 0, "Invalid vector size"); + int4 data[vec_size / 16]; - FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(data))[i]; } - FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(data))[i]; } - FLASHINFER_INLINE float* ptr() { return reinterpret_cast(&data); } - FLASHINFER_INLINE void fill(float val) { -#pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) { - data[i] = make_float4(val, val, val, val); + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) + { + return ((__nv_fp8_e5m2*) data)[i]; } - } - FLASHINFER_INLINE void load(const float* ptr) { -#pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) { - data[i] = ((float4*)ptr)[i]; + + FLASHINFER_INLINE __nv_fp8_e5m2 const& operator[](size_t i) const + { + return ((__nv_fp8_e5m2 const*) data)[i]; } - } - FLASHINFER_INLINE void store(float* ptr) const { + + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() + { + return reinterpret_cast<__nv_fp8_e5m2*>(&data); + } + + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + ((__nv_fp8x4_e5m2*) (&(data[i].x)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) + | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) + | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*) (&(data[i].y)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) + | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) + | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*) (&(data[i].z)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) + | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) + | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*) (&(data[i].w)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) + | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) + | __nv_fp8x4_storage_t(val.__x); + } + } + + FLASHINFER_INLINE void load(__nv_fp8_e5m2 const* ptr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + data[i] = ((int4*) ptr)[i]; + } + } + + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + ((int4*) ptr)[i] = data[i]; + } + } + + FLASHINFER_INLINE void store_global_release(__nv_fp8_e5m2* addr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + st_global_release(data[i], (int4*) (addr + i * 16)); + } + } + + FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e5m2* addr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + data[i] = ld_global_acquire((int4*) (addr + i * 16)); + } + } + + FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e5m2* addr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + st_global_volatile(data[i], (int4*) (addr + i * 16)); + } + } + + FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e5m2* addr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + data[i] = ld_global_volatile((int4*) (addr + i * 16)); + } + } + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + ((int4*) dst)[i] = ((int4*) src)[i]; + } + } +}; + +#if defined(FLASHINFER_ENABLE_FP4_E2M1) && CUDA_VERSION >= 12080 +/******************* vec_t<__nv_fp4_e2m1> *******************/ + +// __nv_fp4_e2m1 x 2 +template <> +struct vec_t<__nv_fp4_e2m1, 2> +{ + uint8_t data; + + // index access is not supported for sub-byte data type + FLASHINFER_INLINE __nv_fp4_e2m1* ptr() + { + return reinterpret_cast<__nv_fp4_e2m1*>(&data); + } + + FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) + { + data = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); + } + + FLASHINFER_INLINE void load(__nv_fp4_e2m1 const* ptr) + { + data = *((uint8_t*) ptr); + } + + FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const + { + *((uint8_t*) ptr) = data; + } + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, __nv_fp4_e2m1 const* src) + { + *((uint8_t*) dst) = *((uint8_t*) src); + } +}; + +// __nv_fp4_e2m1 x 4 +template <> +struct vec_t<__nv_fp4_e2m1, 4> +{ + uint16_t data; + + FLASHINFER_INLINE __nv_fp4_e2m1* ptr() + { + return reinterpret_cast<__nv_fp4_e2m1*>(&data); + } + + FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) + { + __nv_fp4x2_storage_t val8 = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); + data = (uint16_t(val8) << 8) | uint16_t(val8); + } + + FLASHINFER_INLINE void load(__nv_fp4_e2m1 const* ptr) + { + data = *((uint16_t*) ptr); + } + + FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const + { + *((uint16_t*) ptr) = data; + } + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, __nv_fp4_e2m1 const* src) + { + *((uint16_t*) dst) = *((uint16_t*) src); + } +}; + +// __nv_fp4_e2m1 x 8 +template <> +struct vec_t<__nv_fp4_e2m1, 8> +{ + uint32_t data; + + FLASHINFER_INLINE __nv_fp4_e2m1* ptr() + { + return reinterpret_cast<__nv_fp4_e2m1*>(&data); + } + + FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) + { + __nv_fp4x2_storage_t val8 = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); + uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); + data = (uint32_t(val16) << 16) | uint32_t(val16); + } + + FLASHINFER_INLINE void load(__nv_fp4_e2m1 const* ptr) + { + data = *((uint32_t*) ptr); + } + + FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const + { + *((uint32_t*) ptr) = data; + } + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, __nv_fp4_e2m1 const* src) + { + *((uint32_t*) dst) = *((uint32_t*) src); + } +}; + +template <> +struct vec_t<__nv_fp4_e2m1, 16> +{ + uint2 data; + + FLASHINFER_INLINE __nv_fp4_e2m1* ptr() + { + return reinterpret_cast<__nv_fp4_e2m1*>(&data); + } + + FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) + { + __nv_fp4x2_storage_t val8 = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); + uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); + uint32_t val32 = (uint32_t(val16) << 16) | uint32_t(val16); + data.x = val32; + data.y = val32; + } + + FLASHINFER_INLINE void load(__nv_fp4_e2m1 const* ptr) + { + data = *((uint2*) ptr); + } + + FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const + { + *((uint2*) ptr) = data; + } + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, __nv_fp4_e2m1 const* src) + { + *((uint2*) dst) = *((uint2*) src); + } +}; + +// __nv_fp4_e2m1 x 32 or more +template +struct vec_t<__nv_fp4_e2m1, vec_size> +{ + static_assert(vec_size % 32 == 0, "Invalid vector size"); + int4 data[vec_size / 32]; + + FLASHINFER_INLINE __nv_fp4_e2m1* ptr() + { + return reinterpret_cast<__nv_fp4_e2m1*>(&data); + } + + FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) + { + __nv_fp4x2_storage_t val8 = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); + uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); + uint32_t val32 = (uint32_t(val16) << 16) | uint32_t(val16); +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) + { + data[i].x = val32; + data[i].y = val32; + data[i].z = val32; + data[i].w = val32; + } + } + + FLASHINFER_INLINE void load(__nv_fp4_e2m1 const* ptr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) + { + data[i] = ((int4*) ptr)[i]; + } + } + + FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) + { + ((int4*) ptr)[i] = data[i]; + } + } + + FLASHINFER_INLINE void store_global_release(__nv_fp4_e2m1* addr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) + { + st_global_release(*(int4*) &data[i], (int4*) (addr + i * 16)); + } + } + + FLASHINFER_INLINE void load_global_acquire(__nv_fp4_e2m1* addr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) + { + *(int4*) &data[i] = ld_global_acquire((int4*) (addr + i * 16)); + } + } + + FLASHINFER_INLINE void store_global_volatile(__nv_fp4_e2m1* addr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) + { + st_global_volatile(*(int4*) &data[i], (int4*) (addr + i * 16)); + } + } + + FLASHINFER_INLINE void load_global_volatile(__nv_fp4_e2m1* addr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) + { + *(int4*) &data[i] = ld_global_volatile((int4*) (addr + i * 16)); + } + } + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, __nv_fp4_e2m1 const* src) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 32; ++i) + { + ((int4*) dst)[i] = ((int4*) src)[i]; + } + } +}; + +#endif // FLASHINFER_ENABLE_FP4_E2M1 && CUDA_VERSION >= 12080 + +/******************* vec_t *******************/ + +// half x 1 +template <> +struct vec_t +{ + half data; + + FLASHINFER_INLINE half& operator[](size_t i) + { + return ((half*) (&data))[i]; + } + + FLASHINFER_INLINE half const& operator[](size_t i) const + { + return ((half const*) (&data))[i]; + } + + FLASHINFER_INLINE half* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(half val); + FLASHINFER_INLINE void load(half const* ptr); + FLASHINFER_INLINE void store(half* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(half* dst, half const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(half val) +{ + data = val; +} + +FLASHINFER_INLINE void vec_t::load(half const* ptr) +{ + data = *ptr; +} + +FLASHINFER_INLINE void vec_t::store(half* ptr) const +{ + *ptr = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(half* dst, half const* src) +{ + *dst = *src; +} + +// half x 2 +template <> +struct vec_t +{ + half2 data; + + FLASHINFER_INLINE half& operator[](size_t i) + { + return ((half*) (&data))[i]; + } + + FLASHINFER_INLINE half const& operator[](size_t i) const + { + return ((half const*) (&data))[i]; + } + + FLASHINFER_INLINE half* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(half val); + FLASHINFER_INLINE void load(half const* ptr); + FLASHINFER_INLINE void store(half* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(half* dst, half const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(half val) +{ + data = make_half2(val, val); +} + +FLASHINFER_INLINE void vec_t::load(half const* ptr) +{ + data = *((half2*) ptr); +} + +FLASHINFER_INLINE void vec_t::store(half* ptr) const +{ + *((half2*) ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(half* dst, half const* src) +{ + *((half2*) dst) = *((half2*) src); +} + +// half x 4 + +template <> +struct vec_t +{ + uint2 data; + + FLASHINFER_INLINE half& operator[](size_t i) + { + return ((half*) (&data))[i]; + } + + FLASHINFER_INLINE half const& operator[](size_t i) const + { + return ((half const*) (&data))[i]; + } + + FLASHINFER_INLINE half* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(half val); + FLASHINFER_INLINE void load(half const* ptr); + FLASHINFER_INLINE void store(half* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(half* dst, half const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(half val) +{ + *(half2*) (&data.x) = make_half2(val, val); + *(half2*) (&data.y) = make_half2(val, val); +} + +FLASHINFER_INLINE void vec_t::load(half const* ptr) +{ + data = *((uint2*) ptr); +} + +FLASHINFER_INLINE void vec_t::store(half* ptr) const +{ + *((uint2*) ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(half* dst, half const* src) +{ + *((uint2*) dst) = *((uint2*) src); +} + +// half x 8 or more + +template +struct vec_t +{ + static_assert(vec_size % 8 == 0, "Invalid vector size"); + int4 data[vec_size / 8]; + + FLASHINFER_INLINE half& operator[](size_t i) + { + return ((half*) data)[i]; + } + + FLASHINFER_INLINE half const& operator[](size_t i) const + { + return ((half const*) data)[i]; + } + + FLASHINFER_INLINE half* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(half val) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + *(half2*) (&(data[i].x)) = make_half2(val, val); + *(half2*) (&(data[i].y)) = make_half2(val, val); + *(half2*) (&(data[i].z)) = make_half2(val, val); + *(half2*) (&(data[i].w)) = make_half2(val, val); + } + } + + FLASHINFER_INLINE void load(half const* ptr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + data[i] = ((int4*) ptr)[i]; + } + } + + FLASHINFER_INLINE void store(half* ptr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + ((int4*) ptr)[i] = data[i]; + } + } + + FLASHINFER_INLINE void load_global_acquire(half* addr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + data[i] = ld_global_acquire((int4*) (addr + i * 8)); + } + } + + FLASHINFER_INLINE void store_global_release(half* addr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + st_global_release(data[i], (int4*) (addr + i * 8)); + } + } + + FLASHINFER_INLINE void store_global_volatile(half* addr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + st_global_volatile(data[i], (int4*) (addr + i * 8)); + } + } + + FLASHINFER_INLINE void load_global_volatile(half* addr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + data[i] = ld_global_volatile((int4*) (addr + i * 8)); + } + } + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(half* dst, half const* src) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + ((int4*) dst)[i] = ((int4*) src)[i]; + } + } +}; + +/******************* vec_t *******************/ + +// nv_bfloat16 x 1 +template <> +struct vec_t +{ + nv_bfloat16 data; + + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) + { + return ((nv_bfloat16*) (&data))[i]; + } + + FLASHINFER_INLINE nv_bfloat16 const& operator[](size_t i) const + { + return ((nv_bfloat16 const*) (&data))[i]; + } + + FLASHINFER_INLINE nv_bfloat16* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(nv_bfloat16 val); + FLASHINFER_INLINE void load(nv_bfloat16 const* ptr); + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, nv_bfloat16 const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) +{ + data = val; +} + +FLASHINFER_INLINE void vec_t::load(nv_bfloat16 const* ptr) +{ + data = *ptr; +} + +FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const +{ + *ptr = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, nv_bfloat16 const* src) +{ + *dst = *src; +} + +// nv_bfloat16 x 2 +template <> +struct vec_t +{ + nv_bfloat162 data; + + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) + { + return ((nv_bfloat16*) (&data))[i]; + } + + FLASHINFER_INLINE nv_bfloat16 const& operator[](size_t i) const + { + return ((nv_bfloat16 const*) (&data))[i]; + } + + FLASHINFER_INLINE nv_bfloat16* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(nv_bfloat16 val); + FLASHINFER_INLINE void load(nv_bfloat16 const* ptr); + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, nv_bfloat16 const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) +{ + data = make_bfloat162(val, val); +} + +FLASHINFER_INLINE void vec_t::load(nv_bfloat16 const* ptr) +{ + data = *((nv_bfloat162*) ptr); +} + +FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const +{ + *((nv_bfloat162*) ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, nv_bfloat16 const* src) +{ + *((nv_bfloat162*) dst) = *((nv_bfloat162*) src); +} + +// nv_bfloat16 x 4 + +template <> +struct vec_t +{ + uint2 data; + + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) + { + return ((nv_bfloat16*) (&data))[i]; + } + + FLASHINFER_INLINE nv_bfloat16 const& operator[](size_t i) const + { + return ((nv_bfloat16 const*) (&data))[i]; + } + + FLASHINFER_INLINE nv_bfloat16* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(nv_bfloat16 val); + FLASHINFER_INLINE void load(nv_bfloat16 const* ptr); + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, nv_bfloat16 const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) +{ + *(nv_bfloat162*) (&data.x) = make_bfloat162(val, val); + *(nv_bfloat162*) (&data.y) = make_bfloat162(val, val); +} + +FLASHINFER_INLINE void vec_t::load(nv_bfloat16 const* ptr) +{ + data = *((uint2*) ptr); +} + +FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const +{ + *((uint2*) ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, nv_bfloat16 const* src) +{ + *((uint2*) dst) = *((uint2*) src); +} + +// nv_bfloat16 x 8 or more + +template +struct vec_t +{ + static_assert(vec_size % 8 == 0, "Invalid vector size"); + int4 data[vec_size / 8]; + + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) + { + return ((nv_bfloat16*) data)[i]; + } + + FLASHINFER_INLINE nv_bfloat16 const& operator[](size_t i) const + { + return ((nv_bfloat16 const*) data)[i]; + } + + FLASHINFER_INLINE nv_bfloat16* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(nv_bfloat16 val) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + *(nv_bfloat162*) (&(data[i].x)) = make_bfloat162(val, val); + *(nv_bfloat162*) (&(data[i].y)) = make_bfloat162(val, val); + *(nv_bfloat162*) (&(data[i].z)) = make_bfloat162(val, val); + *(nv_bfloat162*) (&(data[i].w)) = make_bfloat162(val, val); + } + } + + FLASHINFER_INLINE void load(nv_bfloat16 const* ptr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + data[i] = ((int4*) ptr)[i]; + } + } + + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + ((int4*) ptr)[i] = data[i]; + } + } + + FLASHINFER_INLINE void store_global_release(nv_bfloat16* addr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + st_global_release(data[i], (int4*) (addr + i * 8)); + } + } + + FLASHINFER_INLINE void load_global_acquire(nv_bfloat16* addr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + data[i] = ld_global_acquire((int4*) (addr + i * 8)); + } + } + + FLASHINFER_INLINE void store_global_volatile(nv_bfloat16* addr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + st_global_volatile(data[i], (int4*) (addr + i * 8)); + } + } + + FLASHINFER_INLINE void load_global_volatile(nv_bfloat16* addr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + data[i] = ld_global_volatile((int4*) (addr + i * 8)); + } + } + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, nv_bfloat16 const* src) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) + { + ((int4*) dst)[i] = ((int4*) src)[i]; + } + } +}; + +/******************* vec_t *******************/ + +// uint8_t x 1 +template <> +struct vec_t +{ + uint8_t data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) + { + return ((uint8_t*) (&data))[i]; + } + + FLASHINFER_INLINE uint8_t const& operator[](size_t i) const + { + return ((uint8_t const*) (&data))[i]; + } + + FLASHINFER_INLINE uint8_t* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(uint8_t const* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, uint8_t const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) +{ + data = val; +} + +FLASHINFER_INLINE void vec_t::load(uint8_t const* ptr) +{ + data = *ptr; +} + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const +{ + *ptr = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, uint8_t const* src) +{ + *dst = *src; +} + +// uint8_t x 2 +template <> +struct vec_t +{ + uint16_t data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) + { + return ((uint8_t*) (&data))[i]; + } + + FLASHINFER_INLINE uint8_t const& operator[](size_t i) const + { + return ((uint8_t const*) (&data))[i]; + } + + FLASHINFER_INLINE uint8_t* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(uint8_t const* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, uint8_t const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) +{ + data = (uint16_t(val) << 8) | uint16_t(val); +} + +FLASHINFER_INLINE void vec_t::load(uint8_t const* ptr) +{ + data = *((uint16_t*) ptr); +} + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const +{ + *((uint16_t*) ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, uint8_t const* src) +{ + *((uint16_t*) dst) = *((uint16_t*) src); +} + +// uint8_t x 4 + +template <> +struct vec_t +{ + uint32_t data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) + { + return ((uint8_t*) (&data))[i]; + } + + FLASHINFER_INLINE uint8_t const& operator[](size_t i) const + { + return ((uint8_t const*) (&data))[i]; + } + + FLASHINFER_INLINE uint8_t* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(uint8_t const* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, uint8_t const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) +{ + data = (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); +} + +FLASHINFER_INLINE void vec_t::load(uint8_t const* ptr) +{ + data = *((uint32_t*) ptr); +} + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const +{ + *((uint32_t*) ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, uint8_t const* src) +{ + *((uint32_t*) dst) = *((uint32_t*) src); +} + +// uint8_t x 8 + +template <> +struct vec_t +{ + uint2 data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) + { + return ((uint8_t*) (&data))[i]; + } + + FLASHINFER_INLINE uint8_t const& operator[](size_t i) const + { + return ((uint8_t const*) (&data))[i]; + } + + FLASHINFER_INLINE uint8_t* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(uint8_t const* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, uint8_t const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) +{ + uint32_t val32 = (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); + data.x = val32; + data.y = val32; +} + +FLASHINFER_INLINE void vec_t::load(uint8_t const* ptr) +{ + data = *((uint2*) ptr); +} + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const +{ + *((uint2*) ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, uint8_t const* src) +{ + *((uint2*) dst) = *((uint2*) src); +} + +// uint8_t x 16 or more + +template +struct vec_t +{ + static_assert(vec_size % 16 == 0, "Invalid vector size"); + int4 data[vec_size / 16]; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) + { + return ((uint8_t*) data)[i]; + } + + FLASHINFER_INLINE uint8_t const& operator[](size_t i) const + { + return ((uint8_t const*) data)[i]; + } + + FLASHINFER_INLINE uint8_t* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(uint8_t val) + { + uint32_t val32 = (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + data[i].x = val32; + data[i].y = val32; + data[i].z = val32; + data[i].w = val32; + } + } + + FLASHINFER_INLINE void load(uint8_t const* ptr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + data[i] = ((int4*) ptr)[i]; + } + } + + FLASHINFER_INLINE void store(uint8_t* ptr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + ((int4*) ptr)[i] = data[i]; + } + } + + FLASHINFER_INLINE void load_global_acquire(uint8_t* addr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + data[i] = ld_global_acquire((int4*) (addr + i * 16)); + } + } + + FLASHINFER_INLINE void store_global_release(uint8_t* addr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + st_global_release(data[i], (int4*) (addr + i * 16)); + } + } + + FLASHINFER_INLINE void load_global_volatile(uint8_t* addr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + data[i] = ld_global_volatile((int4*) (addr + i * 16)); + } + } + + FLASHINFER_INLINE void store_global_volatile(uint8_t* addr) const + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + st_global_volatile(data[i], (int4*) (addr + i * 16)); + } + } + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, uint8_t const* src) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) + { + ((int4*) dst)[i] = ((int4*) src)[i]; + } + } +}; + +/******************* vec_t *******************/ + +// float x 1 + +template <> +struct vec_t +{ + float data; + + FLASHINFER_INLINE float& operator[](size_t i) + { + return ((float*) (&data))[i]; + } + + FLASHINFER_INLINE float const& operator[](size_t i) const + { + return ((float const*) (&data))[i]; + } + + FLASHINFER_INLINE float* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(float val); + FLASHINFER_INLINE void load(float const* ptr); + FLASHINFER_INLINE void store(float* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(float* dst, float const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(float val) +{ + data = val; +} + +FLASHINFER_INLINE void vec_t::load(float const* ptr) +{ + data = *ptr; +} + +FLASHINFER_INLINE void vec_t::store(float* ptr) const +{ + *ptr = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(float* dst, float const* src) +{ + *dst = *src; +} + +// float x 2 + +template <> +struct vec_t +{ + float2 data; + + FLASHINFER_INLINE float& operator[](size_t i) + { + return ((float*) (&data))[i]; + } + + FLASHINFER_INLINE float const& operator[](size_t i) const + { + return ((float const*) (&data))[i]; + } + + FLASHINFER_INLINE float* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(float val); + FLASHINFER_INLINE void load(float const* ptr); + FLASHINFER_INLINE void store(float* ptr) const; + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(float* dst, float const* src); +}; + +FLASHINFER_INLINE void vec_t::fill(float val) +{ + data = make_float2(val, val); +} + +FLASHINFER_INLINE void vec_t::load(float const* ptr) +{ + data = *((float2*) ptr); +} + +FLASHINFER_INLINE void vec_t::store(float* ptr) const +{ + *((float2*) ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(float* dst, float const* src) +{ + *((float2*) dst) = *((float2*) src); +} + +// float x 4 or more +template +struct vec_t +{ + static_assert(vec_size % 4 == 0, "Invalid vector size"); + float4 data[vec_size / 4]; + + FLASHINFER_INLINE float& operator[](size_t i) + { + return ((float*) (data))[i]; + } + + FLASHINFER_INLINE float const& operator[](size_t i) const + { + return ((float const*) (data))[i]; + } + + FLASHINFER_INLINE float* ptr() + { + return reinterpret_cast(&data); + } + + FLASHINFER_INLINE void fill(float val) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) + { + data[i] = make_float4(val, val, val, val); + } + } + + FLASHINFER_INLINE void load(float const* ptr) + { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) + { + data[i] = ((float4*) ptr)[i]; + } + } + + FLASHINFER_INLINE void store(float* ptr) const + { #pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) { - ((float4*)ptr)[i] = data[i]; + for (size_t i = 0; i < vec_size / 4; ++i) + { + ((float4*) ptr)[i] = data[i]; + } } - } - FLASHINFER_INLINE void store_global_release(float* addr) const { + + FLASHINFER_INLINE void store_global_release(float* addr) const + { #pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) { - st_global_release(*(int4*)(data + i), (int4*)(addr + i * 4)); + for (size_t i = 0; i < vec_size / 4; ++i) + { + st_global_release(*(int4*) (data + i), (int4*) (addr + i * 4)); + } } - } - FLASHINFER_INLINE void load_global_acquire(float* addr) { + + FLASHINFER_INLINE void load_global_acquire(float* addr) + { #pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) { - *((int4*)(data + i)) = ld_global_acquire((int4*)(addr + i * 4)); + for (size_t i = 0; i < vec_size / 4; ++i) + { + *((int4*) (data + i)) = ld_global_acquire((int4*) (addr + i * 4)); + } } - } - FLASHINFER_INLINE void store_global_volatile(float* addr) const { + + FLASHINFER_INLINE void store_global_volatile(float* addr) const + { #pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) { - st_global_volatile(*(int4*)(data + i), (int4*)(addr + i * 4)); + for (size_t i = 0; i < vec_size / 4; ++i) + { + st_global_volatile(*(int4*) (data + i), (int4*) (addr + i * 4)); + } } - } - FLASHINFER_INLINE void load_global_volatile(float* addr) { + + FLASHINFER_INLINE void load_global_volatile(float* addr) + { #pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) { - *((int4*)(data + i)) = ld_global_volatile((int4*)(addr + i * 4)); + for (size_t i = 0; i < vec_size / 4; ++i) + { + *((int4*) (data + i)) = ld_global_volatile((int4*) (addr + i * 4)); + } + } + + template + FLASHINFER_INLINE void cast_from(vec_t const& src) + { + cast_from_impl(*this, src); + } + + template + FLASHINFER_INLINE void cast_load(T const* ptr) + { + cast_load_impl(*this, ptr); + } + + template + FLASHINFER_INLINE void cast_store(T* ptr) const + { + cast_store_impl(ptr, *this); } - } - template - FLASHINFER_INLINE void cast_from(const vec_t& src) { - cast_from_impl(*this, src); - } - template - FLASHINFER_INLINE void cast_load(const T* ptr) { - cast_load_impl(*this, ptr); - } - template - FLASHINFER_INLINE void cast_store(T* ptr) const { - cast_store_impl(ptr, *this); - } - FLASHINFER_INLINE static void memcpy(float* dst, const float* src) { + + FLASHINFER_INLINE static void memcpy(float* dst, float const* src) + { #pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) { - ((float4*)dst)[i] = ((float4*)src)[i]; + for (size_t i = 0; i < vec_size / 4; ++i) + { + ((float4*) dst)[i] = ((float4*) src)[i]; + } } - } }; template -struct vec2_dtype { - using type = T; +struct vec2_dtype +{ + using type = T; }; template <> -struct vec2_dtype { - using type = half2; +struct vec2_dtype +{ + using type = half2; }; template <> -struct vec2_dtype<__nv_bfloat16> { - using type = __nv_bfloat162; +struct vec2_dtype<__nv_bfloat16> +{ + using type = __nv_bfloat162; }; template <> -struct vec2_dtype<__nv_fp8_e4m3> { - using type = __nv_fp8x2_e4m3; +struct vec2_dtype<__nv_fp8_e4m3> +{ + using type = __nv_fp8x2_e4m3; }; template <> -struct vec2_dtype<__nv_fp8_e5m2> { - using type = __nv_fp8x2_e5m2; +struct vec2_dtype<__nv_fp8_e5m2> +{ + using type = __nv_fp8x2_e5m2; }; template using vec2_dtype_t = typename vec2_dtype::type; template -FLASHINFER_INLINE vec2_dtype_t get_vec2_element(vec_t& vec, int i) { - static_assert(VEC_SIZE % 2 == 0, "VEC_SIZE must be a multiple of 2"); - return ((vec2_dtype_t*)&(vec[0]))[i]; +FLASHINFER_INLINE vec2_dtype_t get_vec2_element(vec_t& vec, int i) +{ + static_assert(VEC_SIZE % 2 == 0, "VEC_SIZE must be a multiple of 2"); + return ((vec2_dtype_t*) &(vec[0]))[i]; } -} // namespace flashinfer +} // namespace flashinfer -#endif // VEC_DTYPES_CUH_ +#endif // VEC_DTYPES_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 25859b8dfa4b..96b2f2892c5b 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -101,13 +101,12 @@ class ReplayCacheManager: def __init__(self): self.state_indices = [0, 3, 1, 4, 2] self.prev_num_accepted_tokens = torch.tensor( - [0, 4, 10, 11, 20], dtype=torch.int32, device="cuda") - self.cache_buf_idx = torch.tensor([0, 1, 0, 1, 0], - dtype=torch.int32, - device="cuda") + [0, 4, 10, 11, 20], dtype=torch.int32, device="cuda" + ) + self.cache_buf_idx = torch.tensor([0, 1, 0, 1, 0], dtype=torch.int32, device="cuda") def get_state_indices(self, request_ids, is_padding): - return self.state_indices[:len(request_ids)] + return self.state_indices[: len(request_ids)] def get_replay_state_update_metadata(self): return ReplayStateUpdateMetadata( @@ -127,8 +126,7 @@ def get_replay_state_update_metadata(self): kv_cache_manager=ReplayCacheManager(), request_ids=[10, 11, 12, 13, 14], kv_cache_params=SimpleNamespace( - num_cached_tokens_per_seq=torch.tensor([0], - dtype=torch.int), + num_cached_tokens_per_seq=torch.tensor([0], dtype=torch.int), ), ) @@ -146,8 +144,9 @@ def get_replay_state_update_metadata(self): ) actual = metadata.replay_work_items[:4] torch.testing.assert_close(actual, expected) - torch.testing.assert_close(metadata.replay_n_writes.cpu(), - torch.tensor([2], dtype=torch.int32)) + torch.testing.assert_close( + metadata.replay_n_writes.cpu(), torch.tensor([2], dtype=torch.int32) + ) assert actual[0, REPLAY_WORK_POSITION_IN_DECODE_BATCH] == 0 assert actual[0, REPLAY_WORK_CACHE_SLOT] == 3 assert actual[0, REPLAY_WORK_PNAT] == 11 diff --git a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py index 030fe31e86dd..41380a2e74bc 100644 --- a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py @@ -38,12 +38,18 @@ from tensorrt_llm._utils import get_sm_version -def _make_replay_work_items(prev_tokens, cache_buf_idx, T, max_window, batch, - state_batch_indices, device, explicit_order=None): +def _make_replay_work_items( + prev_tokens, + cache_buf_idx, + T, + max_window, + batch, + state_batch_indices, + device, + explicit_order=None, +): """Build the replay metadata consumed by persistent_main.""" - position_in_decode_batch = torch.arange(batch, - device=device, - dtype=torch.int32) + position_in_decode_batch = torch.arange(batch, device=device, dtype=torch.int32) if state_batch_indices is not None: cache_slot = state_batch_indices[:batch].to(torch.int32) else: @@ -55,17 +61,12 @@ def _make_replay_work_items(prev_tokens, cache_buf_idx, T, max_window, batch, n_writes = write_mask.sum().to(torch.int32).reshape(1) if explicit_order is None: - order = torch.argsort((~write_mask).to(torch.int32), - stable=True).to(torch.long) + order = torch.argsort((~write_mask).to(torch.int32), stable=True).to(torch.long) else: order = torch.tensor(explicit_order, device=device, dtype=torch.long) - replay_work_items = torch.empty(batch, - REPLAY_WORK_ITEM_WIDTH, - device=device, - dtype=torch.int32) - replay_work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = ( - position_in_decode_batch[order]) + replay_work_items = torch.empty(batch, REPLAY_WORK_ITEM_WIDTH, device=device, dtype=torch.int32) + replay_work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = position_in_decode_batch[order] replay_work_items[:, REPLAY_WORK_CACHE_SLOT] = cache_slot[order] replay_work_items[:, REPLAY_WORK_PNAT] = pnat[order] replay_work_items[:, REPLAY_WORK_CACHE_BUF_IDX] = active_cache_buf_idx[order] @@ -151,9 +152,9 @@ def _maybe_skip_dtype(state_dtype, use_sr): @pytest.mark.parametrize( "write_checkpoint,rectangle_for_nowrite", [ - (True, False), # write path (rectangle_for_nowrite is ignored) - (False, False), # nowrite path via replay-style kernels - (False, True), # nowrite path via dedicated rectangle kernels + (True, False), # write path (rectangle_for_nowrite is ignored) + (False, False), # nowrite path via replay-style kernels + (False, True), # nowrite path via dedicated rectangle kernels ], ids=["write", "no_write_replay", "no_write_rectangle"], ) @@ -163,8 +164,16 @@ def _maybe_skip_dtype(state_dtype, use_sr): ids=["persistent_dynamic", "persistent_main"], ) def test_replay_selective_state_update( - nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T, - write_checkpoint, rectangle_for_nowrite, mode, + nheads, + head_dim, + d_state, + ngroups, + state_dtype, + paged_cache, + T, + write_checkpoint, + rectangle_for_nowrite, + mode, ): """ Verify that: @@ -280,7 +289,9 @@ def test_replay_selective_state_update( old_x = torch.randn(cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) # Fill each slot's active buffer (= cache_buf_idx) with step 1's data at @@ -365,8 +376,13 @@ def test_replay_selective_state_update( # pure-nowrite cases here, all slots have the same status, so the # work-item order is identity. n_writes_t, replay_work_items_t = _make_replay_work_items( - prev_tokens, cache_buf_idx, T, max_window, batch, - state_batch_indices, device, + prev_tokens, + cache_buf_idx, + T, + max_window, + batch, + state_batch_indices, + device, ) replay_selective_state_update( test_state, @@ -396,10 +412,10 @@ def test_replay_selective_state_update( # Tolerance rationale: the replay kernel uses bf16 tl.dot for four # matmuls (dB_scaled @ old_x, C @ state, CB_scaled @ x, and C @ B in - # precompute). The reference (selective_state_update) and flashinfer - # baseline use fp32 element-wise MACs. The bf16 input casts lose the - # dt_bias/A-derived bits that the baselines keep — per-element rounding, - # not accumulating. Prefill (ssd_chunk_scan) does identical bf16 tl.dot + # precompute). The reference selective_state_update uses fp32 + # element-wise MACs. The bf16 input casts lose dt_bias/A-derived bits + # that the reference keeps — per-element rounding, not accumulating. + # Prefill (ssd_chunk_scan) does identical bf16 tl.dot # casts, so we match prefill precision exactly. Empirical: max ~1.0 at # T<=16, ~2.0 at T=32-55; mean ~0.014; <0.02% of elements exceed 0.5. # State dtype (fp16/bf16/fp32) doesn't shift the error — bf16 dot @@ -421,18 +437,23 @@ def test_replay_selective_state_update( # smaller-magnitude elements) out_atol = ( {torch.int8: 1.6, torch.int16: 1.05, torch.float8_e4m3fn: 4.0}[state_dtype] - if is_quantized else 1.0 + if is_quantized + else 1.0 ) out_rtol = ( {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] - if is_quantized else 2e-2 + if is_quantized + else 2e-2 ) out_diff = (test_out.float() - ref_out.float()).abs() out_max = out_diff.max().item() out_mean = out_diff.mean().item() try: torch.testing.assert_close( - test_out, ref_out, rtol=out_rtol, atol=out_atol, + test_out, + ref_out, + rtol=out_rtol, + atol=out_atol, msg=f"Output mismatch at k={k}", ) except AssertionError: @@ -465,15 +486,21 @@ def test_replay_selective_state_update( # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 worst-case # Atol = bf16_baseline (1.0) + quant_eps_max. state_atol = { - torch.int8: 1.1, torch.int16: 1.0, torch.float8_e4m3fn: 2.5, + torch.int8: 1.1, + torch.int16: 1.0, + torch.float8_e4m3fn: 2.5, }[state_dtype] state_rtol = { - torch.int8: 5e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 1e-1, + torch.int8: 5e-2, + torch.int16: 2e-2, + torch.float8_e4m3fn: 1e-1, }[state_dtype] try: torch.testing.assert_close( - actual_fp32, expected_fp32, - rtol=state_rtol, atol=state_atol, + actual_fp32, + expected_fp32, + rtol=state_rtol, + atol=state_atol, msg=f"State mismatch at k={k} dtype={state_dtype}", ) except AssertionError: @@ -545,39 +572,52 @@ def test_replay_selective_state_update( # --- old_x (double-buffered): write at wb, [write_offset : +T) --- written_x = old_x_w[slot, wb, write_offset : write_offset + T] torch.testing.assert_close( - written_x, x2[batch_idx], rtol=0, atol=0, + written_x, + x2[batch_idx], + rtol=0, + atol=0, msg=f"old_x written region wrong at k={k} write={write_checkpoint}", ) # Untouched ranges of old_x[slot, wb] if write_offset > 0: torch.testing.assert_close( - old_x_w[slot, wb, :write_offset], old_x[slot, wb, :write_offset], - rtol=0, atol=0, + old_x_w[slot, wb, :write_offset], + old_x[slot, wb, :write_offset], + rtol=0, + atol=0, msg=f"old_x [0:{write_offset}) modified at k={k} write={write_checkpoint}", ) if write_offset + T < max_window: torch.testing.assert_close( - old_x_w[slot, wb, write_offset + T:], old_x[slot, wb, write_offset + T:], - rtol=0, atol=0, - msg=f"old_x [{write_offset+T}:) modified at k={k} write={write_checkpoint}", + old_x_w[slot, wb, write_offset + T :], + old_x[slot, wb, write_offset + T :], + rtol=0, + atol=0, + msg=f"old_x [{write_offset + T}:) modified at k={k} write={write_checkpoint}", ) # Other-buffer (= 1-wb) untouched torch.testing.assert_close( - old_x_w[slot, 1 - wb], old_x[slot, 1 - wb], - rtol=0, atol=0, + old_x_w[slot, 1 - wb], + old_x[slot, 1 - wb], + rtol=0, + atol=0, msg=f"old_x inactive buffer modified at k={k} write={write_checkpoint}", ) # --- old_B (double-buffered): write at write_buf, [write_offset:+T) --- torch.testing.assert_close( old_B_w[slot, wb, write_offset : write_offset + T], - B2[batch_idx], rtol=0, atol=0, + B2[batch_idx], + rtol=0, + atol=0, msg=f"old_B written region wrong at k={k} write={write_checkpoint}", ) # Other-buffer (= 1-wb) untouched torch.testing.assert_close( - old_B_w[slot, 1 - wb], old_B[slot, 1 - wb], - rtol=0, atol=0, + old_B_w[slot, 1 - wb], + old_B[slot, 1 - wb], + rtol=0, + atol=0, msg=f"old_B inactive buffer modified at k={k} write={write_checkpoint}", ) @@ -585,12 +625,15 @@ def test_replay_selective_state_update( torch.testing.assert_close( old_dt_w[slot, wb, :, write_offset : write_offset + T], dt2_proc[batch_idx].T, - rtol=1e-4, atol=1e-4, + rtol=1e-4, + atol=1e-4, msg=f"old_dt written region wrong at k={k} write={write_checkpoint}", ) torch.testing.assert_close( - old_dt_w[slot, 1 - wb], old_dt[slot, 1 - wb], - rtol=0, atol=0, + old_dt_w[slot, 1 - wb], + old_dt[slot, 1 - wb], + rtol=0, + atol=0, msg=f"old_dt inactive buffer modified at k={k} write={write_checkpoint}", ) @@ -607,12 +650,15 @@ def test_replay_selective_state_update( torch.testing.assert_close( old_dA_cumsum_w[slot, wb, :, write_offset : write_offset + T], expected_dAcs, - rtol=1e-4, atol=1e-4, + rtol=1e-4, + atol=1e-4, msg=f"old_dA_cumsum written region wrong at k={k} write={write_checkpoint}", ) torch.testing.assert_close( - old_dA_cumsum_w[slot, 1 - wb], old_dA_cumsum[slot, 1 - wb], - rtol=0, atol=0, + old_dA_cumsum_w[slot, 1 - wb], + old_dA_cumsum[slot, 1 - wb], + rtol=0, + atol=0, msg=f"old_dA_cumsum inactive buf modified at k={k} write={write_checkpoint}", ) @@ -636,14 +682,21 @@ def test_replay_selective_state_update( ("mixed_auto_rect", [3, 12, 10, 15], None, True), ], ids=[ - "all_write", "all_nowrite", - "mixed_explicit", "mixed_explicit_rect", - "mixed_auto", "mixed_auto_rect", + "all_write", + "all_nowrite", + "mixed_explicit", + "mixed_explicit_rect", + "mixed_auto", + "mixed_auto_rect", ], ) @pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_replay_selective_state_update_scenarios( - scenario, pnat_per_slot_list, explicit_order, rectangle_for_nowrite, mode, + scenario, + pnat_per_slot_list, + explicit_order, + rectangle_for_nowrite, + mode, ): """ Combined scenarios test covering both kernel modes (persistent_main, @@ -680,9 +733,7 @@ def test_replay_selective_state_update_scenarios( D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) - state0 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=dtype - ) + state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=dtype) ref_input_state = state0.float() step1_T = max_window @@ -699,8 +750,14 @@ def test_replay_selective_state_update_scenarios( out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) selective_state_update( ref_input_state.clone(), - x1, dt1_input, A, B1, C1, - D=D, dt_bias=dt_bias, dt_softplus=True, + x1, + dt1_input, + A, + B1, + C1, + D=D, + dt_bias=dt_bias, + dt_softplus=True, state_batch_indices=cache_idx_for_capture, intermediate_states_buffer=states_buffer_f32, cache_steps=step1_T, @@ -711,9 +768,7 @@ def test_replay_selective_state_update_scenarios( old_x = torch.randn(batch, 2, max_window, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn( - batch, 2, nheads, max_window, device=device, dtype=torch.float32 - ) + old_dA_cumsum = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) @@ -741,48 +796,78 @@ def test_replay_selective_state_update_scenarios( ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) selective_state_update( - ref_state_f32, x2, dt2, A, B2, C2, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, out=ref_out, + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + out=ref_out, ) test_state = state0.clone() test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) n_writes_t, replay_work_items = _make_replay_work_items( - pnat_per_slot, cache_buf_idx, T, max_window, batch, None, device, + pnat_per_slot, + cache_buf_idx, + T, + max_window, + batch, + None, + device, explicit_order=explicit_order, ) replay_selective_state_update( test_state, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), cache_buf_idx.clone(), pnat_per_slot, - x=x2, dt=dt2, A=A, B=B2, C=C2, + x=x2, + dt=dt2, + A=A, + B=B2, + C=C2, out=test_out, n_writes=n_writes_t, replay_work_items=replay_work_items, - D=D, dt_bias=dt_bias, dt_softplus=True, + D=D, + dt_bias=dt_bias, + dt_softplus=True, state_batch_indices=None, mode=mode, rectangle_for_nowrite=rectangle_for_nowrite, ) torch.testing.assert_close( - test_out.float(), ref_out.float(), - atol=1.0, rtol=0.05, + test_out.float(), + ref_out.float(), + atol=1.0, + rtol=0.05, msg=f"Output mismatch (scenario={scenario})", ) for i in range(batch): if pnat_means_write[i]: torch.testing.assert_close( - test_state[i].float(), ref_state_after_replay[i].float(), - atol=1.0, rtol=0.05, + test_state[i].float(), + ref_state_after_replay[i].float(), + atol=1.0, + rtol=0.05, msg=f"Write slot {i}: state mismatch (scenario={scenario})", ) else: torch.testing.assert_close( - test_state[i], state0[i], rtol=0, atol=0, + test_state[i], + state0[i], + rtol=0, + atol=0, msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", ) @@ -803,8 +888,12 @@ def test_replay_selective_state_update_scenarios( @pytest.mark.parametrize("rectangle_for_nowrite", [True, False], ids=["rect", "norect"]) @pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_replay_selective_state_update_persistent_main_device_n_writes( - scenario, pnat_per_slot_list, n_writes_expected, work_item_order, - rectangle_for_nowrite, mode, + scenario, + pnat_per_slot_list, + n_writes_expected, + work_item_order, + rectangle_for_nowrite, + mode, ): """ Persistent_main with the device-tensor n_writes plumbing. @@ -835,9 +924,7 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( # Device-tensor n_writes. Caller mutates between iters in CUDA-graph # benchmarking; we only run one iter here so a single fill is enough. - n_writes = torch.tensor([n_writes_expected], - device=device, - dtype=torch.int32) + n_writes = torch.tensor([n_writes_expected], device=device, dtype=torch.int32) torch.manual_seed(42) A_base = -torch.rand(nheads, device=device) - 0.5 @@ -847,9 +934,7 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) - state0 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=dtype - ) + state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=dtype) ref_input_state = state0.float() step1_T = max_window @@ -866,8 +951,14 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) selective_state_update( ref_input_state.clone(), - x1, dt1_input, A, B1, C1, - D=D, dt_bias=dt_bias, dt_softplus=True, + x1, + dt1_input, + A, + B1, + C1, + D=D, + dt_bias=dt_bias, + dt_softplus=True, state_batch_indices=cache_idx_for_capture, intermediate_states_buffer=states_buffer_f32, cache_steps=step1_T, @@ -878,12 +969,16 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( old_x = torch.randn(batch, 2, max_window, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn( - batch, 2, nheads, max_window, device=device, dtype=torch.float32 - ) + old_dA_cumsum = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) _n_writes_check, replay_work_items = _make_replay_work_items( - pnat_per_slot, cache_buf_idx, T, max_window, batch, None, device, + pnat_per_slot, + cache_buf_idx, + T, + max_window, + batch, + None, + device, explicit_order=work_item_order, ) torch.testing.assert_close(_n_writes_check, n_writes) @@ -913,44 +1008,68 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) selective_state_update( - ref_state_f32, x2, dt2, A, B2, C2, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, out=ref_out, + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + out=ref_out, ) test_state = state0.clone() test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) replay_selective_state_update( test_state, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), cache_buf_idx.clone(), pnat_per_slot, - x=x2, dt=dt2, A=A, B=B2, C=C2, + x=x2, + dt=dt2, + A=A, + B=B2, + C=C2, out=test_out, n_writes=n_writes, replay_work_items=replay_work_items, - D=D, dt_bias=dt_bias, dt_softplus=True, + D=D, + dt_bias=dt_bias, + dt_softplus=True, state_batch_indices=None, mode=mode, rectangle_for_nowrite=rectangle_for_nowrite, ) torch.testing.assert_close( - test_out.float(), ref_out.float(), - atol=1.0, rtol=0.05, + test_out.float(), + ref_out.float(), + atol=1.0, + rtol=0.05, msg=f"Output mismatch (scenario={scenario})", ) for i in range(batch): if pnat_means_write[i]: torch.testing.assert_close( - test_state[i].float(), ref_state_after_replay[i].float(), - atol=1.0, rtol=0.05, + test_state[i].float(), + ref_state_after_replay[i].float(), + atol=1.0, + rtol=0.05, msg=f"Write slot {i}: state mismatch (scenario={scenario})", ) else: torch.testing.assert_close( - test_state[i], state0[i], rtol=0, atol=0, + test_state[i], + state0[i], + rtol=0, + atol=0, msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", ) @@ -965,7 +1084,14 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( @pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) @pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_replay_selective_state_update_philox( - state_dtype, nheads, head_dim, d_state, ngroups, paged_cache, T, mode, + state_dtype, + nheads, + head_dim, + d_state, + ngroups, + paged_cache, + T, + mode, ): """ Verify that Philox stochastic rounding produces correct results across @@ -1034,8 +1160,13 @@ def test_replay_selective_state_update_philox( # the philox test sets old_x's window axis = T, so max_window = T here. _max_window_philox = T _n_writes_philox, _replay_work_items_philox = _make_replay_work_items( - prev_tokens, cache_buf_idx, T, _max_window_philox, batch, - state_batch_indices, device, + prev_tokens, + cache_buf_idx, + T, + _max_window_philox, + batch, + state_batch_indices, + device, ) common_kwargs = dict( @@ -1099,14 +1230,19 @@ def test_replay_selective_state_update_philox( # fp8: amax/14 ≈ 23/14 → 6.5*1.6 ≈ 10.7 + bf16_baseline out_atol = ( {torch.int8: 1.5, torch.int16: 1.0, torch.float8_e4m3fn: 6.0}[state_dtype] - if is_quantized else 1.0 + if is_quantized + else 1.0 ) out_rtol = ( {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] - if is_quantized else 2e-2 + if is_quantized + else 2e-2 ) torch.testing.assert_close( - out_rounded, out_no_round, rtol=out_rtol, atol=out_atol, + out_rounded, + out_no_round, + rtol=out_rtol, + atol=out_atol, msg=f"Output diverged with Philox rounding ({state_dtype})", ) @@ -1124,18 +1260,14 @@ def test_replay_selective_state_update_philox( diff = (rounded_fp32 - no_round_fp32).abs() # Per-element bound = max(decode_scale_no_round, decode_scale_rounded). # decode_scale is shape (cache, nheads, dim); broadcast over dstate. - scale_bound = torch.maximum( - scales_no_round[slots], scales_rounded[slots] - ).unsqueeze(-1) + scale_bound = torch.maximum(scales_no_round[slots], scales_rounded[slots]).unsqueeze(-1) # int8 / int16: 1 cell after dequant = decode_scale exactly. # fp8_e4m3: variable grid; the largest cell within a channel scaled # to fit ±448 is at the channel's max-magnitude element, where the # cell is ~32x larger than the average. Bound = decode_scale * 32. # Apply a 1.5x slack pad for floating-point compare quirks at the # exact-cell boundary. - cell_pad = ( - 32.0 if state_dtype == torch.float8_e4m3fn else 1.0 - ) + cell_pad = 32.0 if state_dtype == torch.float8_e4m3fn else 1.0 bound = scale_bound * (cell_pad * 1.5) if not (diff <= bound).all(): offenders = (diff > bound).sum().item() @@ -1195,9 +1327,7 @@ def test_philox_rounding_unbiased(state_dtype): # fp32 reference state — replay produces values that don't fit cleanly # in the target dtype's grid, exposing the rounding bias. - state0_fp32 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) + state0_fp32 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=torch.float32) old_x = torch.randn(batch, 2, T, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(batch, 2, T, ngroups, d_state, device=device, dtype=dtype) @@ -1215,11 +1345,25 @@ def test_philox_rounding_unbiased(state_dtype): # max_window = old_x.shape[2] = T (after dbuf at axis 1) _n_writes_unb, _replay_work_items_unb = _make_replay_work_items( - prev_tokens, cache_buf_idx, T, T, batch, None, device, + prev_tokens, + cache_buf_idx, + T, + T, + batch, + None, + device, ) common_kwargs = dict( - x=x, dt=dt_val, A=A, B=B, C=C, D=D, dt_bias=dt_bias, dt_softplus=True, - n_writes=_n_writes_unb, replay_work_items=_replay_work_items_unb, + x=x, + dt=dt_val, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + n_writes=_n_writes_unb, + replay_work_items=_replay_work_items_unb, ) # 1. fp32 state — captures true post-replay fp32 state. @@ -1227,8 +1371,14 @@ def test_philox_rounding_unbiased(state_dtype): out_fp32 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) replay_selective_state_update( state_fp32, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), - cache_buf_idx.clone(), prev_tokens, out=out_fp32, **common_kwargs, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + out=out_fp32, + **common_kwargs, ) # 2. Target dtype + Philox SR. For quant we also need scales (derived @@ -1242,9 +1392,15 @@ def test_philox_rounding_unbiased(state_dtype): out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) replay_selective_state_update( state_rounded, - old_x.clone(), old_B.clone(), old_dt.clone(), old_dA_cumsum.clone(), - cache_buf_idx.clone(), prev_tokens, out=out_rounded, - rand_seed=rand_seed, philox_rounds=10, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + prev_tokens, + out=out_rounded, + rand_seed=rand_seed, + philox_rounds=10, state_scales=scales_rounded, **common_kwargs, ) @@ -1254,16 +1410,12 @@ def test_philox_rounding_unbiased(state_dtype): # both, comparing in fp32. if is_quantized: fp32_vals = state_fp32.flatten() - stochastic_residual = ( - _dequantize_state(state_rounded, scales_rounded).flatten() - fp32_vals - ) + stochastic_residual = _dequantize_state(state_rounded, scales_rounded).flatten() - fp32_vals # Deterministic reference: do the same per-channel quant on the # captured fp32 state, then dequant. This is what the kernel would # have produced with rand_seed=None. det_quant, det_scales = _quantize_state(state_fp32, state_dtype, quant_max) - deterministic_residual = ( - _dequantize_state(det_quant, det_scales).flatten() - fp32_vals - ) + deterministic_residual = _dequantize_state(det_quant, det_scales).flatten() - fp32_vals else: fp32_vals = state_fp32.flatten() stochastic_residual = state_rounded.float().flatten() - fp32_vals @@ -1286,7 +1438,7 @@ def test_philox_rounding_unbiased(state_dtype): # * fp8: residual std ~1e-1 → SE ~9e-5 (loosest, magnitude-driven) # A fixed absolute threshold is below SE for int8/fp8. Gaussian inputs # also make RN nearly unbiased, so |sr| < |det| is not reliable here. - se_sr = stochastic_std / (num_nonzero ** 0.5) + se_sr = stochastic_std / (num_nonzero**0.5) K = 4 assert abs(stochastic_mean) < K * se_sr, ( f"SR mean exceeds {K}*SE (likely biased) ({state_dtype}): " @@ -1391,9 +1543,7 @@ def test_replay_heads_per_block( old_B_init = torch.randn( cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype ) - old_dt_init = torch.randn( - cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 - ) + old_dt_init = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) old_dA_cumsum_init = torch.randn( cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 ) @@ -1410,21 +1560,21 @@ def test_replay_heads_per_block( buf = int(cache_buf_idx[slot].item()) old_x_init[slot, buf] = x1[slot] old_B_init[slot, buf] = B1[slot] - old_dt_init[slot, buf] = dt1_proc[slot].T # (nheads, max_window) - old_dA_cumsum_init[slot, buf] = dA_cumsum1[slot].T # (nheads, max_window) + old_dt_init[slot, buf] = dt1_proc[slot].T # (nheads, max_window) + old_dA_cumsum_init[slot, buf] = dA_cumsum1[slot].T # (nheads, max_window) # --- PNAT sweep -------------------------------------------------------- # Cover nowrite (PNAT+T <= max_window) and write (PNAT+T > max_window) # paths plus the boundary, with both PNAT=0 (no prefix) and PNAT=T (the # smallest prefix-load case the kernel cares about). candidate_pnats = [ - 0, # nowrite, no prefix - 1, # nowrite, smallest nontrivial prefix - T, # nowrite, prefix length one stored step - max_window - T - 1, # nowrite, largest PNAT just below threshold - max_window - T, # nowrite, exactly at threshold - max_window - T + 1, # write, smallest PNAT above threshold - max_window - 1, # write, maximum + 0, # nowrite, no prefix + 1, # nowrite, smallest nontrivial prefix + T, # nowrite, prefix length one stored step + max_window - T - 1, # nowrite, largest PNAT just below threshold + max_window - T, # nowrite, exactly at threshold + max_window - T + 1, # write, smallest PNAT above threshold + max_window - 1, # write, maximum ] seen = set() pnat_list = [] @@ -1439,8 +1589,7 @@ def test_replay_heads_per_block( has_write = any((p + T) > max_window for p in pnat_list) has_nowrite = any((p + T) <= max_window for p in pnat_list) assert has_write and has_nowrite, ( - f"PNAT sweep must cover both write and nowrite: {pnat_list}, " - f"T={T}, max_window={max_window}" + f"PNAT sweep must cover both write and nowrite: {pnat_list}, T={T}, max_window={max_window}" ) prev_tokens = torch.tensor(pnat_list, device=device, dtype=torch.int32) @@ -1466,9 +1615,17 @@ def test_replay_heads_per_block( ref_state_after_replay = ref_state_f32.clone() ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) selective_state_update( - ref_state_f32, x2, dt2, A, B2, C2, - D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, out=ref_out, + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=None, + out=ref_out, ) # Pre-call snapshots double as the "expected" baseline for untouched @@ -1489,7 +1646,13 @@ def test_replay_heads_per_block( cache_buf_idx_test = cache_buf_idx_pre.clone() n_writes_t, replay_work_items_t = _make_replay_work_items( - prev_tokens, cache_buf_idx_test, T, max_window, batch, None, device, + prev_tokens, + cache_buf_idx_test, + T, + max_window, + batch, + None, + device, ) replay_selective_state_update( @@ -1531,8 +1694,10 @@ def test_replay_heads_per_block( for slot in range(batch): if pnat_means_write[slot]: torch.testing.assert_close( - test_state[slot].float(), ref_state_after_replay[slot].float(), - rtol=2e-2, atol=1.0, + test_state[slot].float(), + ref_state_after_replay[slot].float(), + rtol=2e-2, + atol=1.0, msg=( f"Write slot {slot} (PNAT={pnat_list[slot]}): state mismatch " f"(HPB={heads_per_block}, T={T}, nheads={nheads}, " @@ -1542,8 +1707,10 @@ def test_replay_heads_per_block( ) else: torch.testing.assert_close( - test_state[slot], state_pre[slot], - rtol=0, atol=0, + test_state[slot], + state_pre[slot], + rtol=0, + atol=0, msg=( f"Nowrite slot {slot} (PNAT={pnat_list[slot]}): state HBM " f"modified (HPB={heads_per_block}, T={T}, nheads={nheads}, " @@ -1578,8 +1745,10 @@ def test_replay_heads_per_block( expected_old_x_slot = old_x_pre[slot].clone() expected_old_x_slot[target_buf, write_offset:write_end] = x2[slot] torch.testing.assert_close( - old_x_test[slot], expected_old_x_slot, - rtol=0, atol=0, + old_x_test[slot], + expected_old_x_slot, + rtol=0, + atol=0, msg=( f"old_x slot {slot} (PNAT={pnat}, is_write={is_write}, " f"target_buf={target_buf}, write_offset={write_offset}): " @@ -1592,8 +1761,10 @@ def test_replay_heads_per_block( expected_old_B_slot = old_B_pre[slot].clone() expected_old_B_slot[target_buf, write_offset:write_end] = B2[slot] torch.testing.assert_close( - old_B_test[slot], expected_old_B_slot, - rtol=0, atol=0, + old_B_test[slot], + expected_old_B_slot, + rtol=0, + atol=0, msg=( f"old_B slot {slot} (PNAT={pnat}, is_write={is_write}, " f"target_buf={target_buf}, write_offset={write_offset}): " @@ -1609,8 +1780,10 @@ def test_replay_heads_per_block( expected_old_dt_slot = old_dt_pre[slot].clone() expected_old_dt_slot[target_buf, :, write_offset:write_end] = dt2_proc[slot].T torch.testing.assert_close( - old_dt_test[slot], expected_old_dt_slot, - rtol=1e-5, atol=1e-5, + old_dt_test[slot], + expected_old_dt_slot, + rtol=1e-5, + atol=1e-5, msg=( f"old_dt slot {slot} (PNAT={pnat}, is_write={is_write}, " f"target_buf={target_buf}, write_offset={write_offset}): " @@ -1637,8 +1810,10 @@ def test_replay_heads_per_block( step_cumsum + prefix[:, None] ) torch.testing.assert_close( - old_dA_cumsum_test[slot], expected_old_dAcs_slot, - rtol=1e-5, atol=1e-5, + old_dA_cumsum_test[slot], + expected_old_dAcs_slot, + rtol=1e-5, + atol=1e-5, msg=( f"old_dA_cumsum slot {slot} (PNAT={pnat}, is_write={is_write}, " f"target_buf={target_buf}, write_offset={write_offset}): " @@ -1668,8 +1843,16 @@ def test_replay_heads_per_block( @pytest.mark.parametrize("rectangle_nowrite", [True, False], ids=["rect", "norect"]) @pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_replay_heads_per_block_multistep( - nheads, head_dim, d_state, ngroups, state_dtype, T, heads_per_block, - paged_cache, rectangle_nowrite, mode, + nheads, + head_dim, + d_state, + ngroups, + state_dtype, + T, + heads_per_block, + paged_cache, + rectangle_nowrite, + mode, ): """ Chain N decode steps with HPB > 1 and verify each step's output matches @@ -1743,18 +1926,24 @@ def test_replay_heads_per_block_multistep( if acc == 0: continue c_idx = ( - state_batch_indices[s_local].item() - if state_batch_indices is not None else s_local + state_batch_indices[s_local].item() if state_batch_indices is not None else s_local ) - s_state = ref_state[c_idx:c_idx + 1].clone() - s_x = all_x[step][s_local:s_local + 1, :acc].contiguous() - s_dt = all_dt[step][s_local:s_local + 1, :acc].contiguous() - s_B = all_B[step][s_local:s_local + 1, :acc].contiguous() - s_C = all_C[step][s_local:s_local + 1, :acc].contiguous() + s_state = ref_state[c_idx : c_idx + 1].clone() + s_x = all_x[step][s_local : s_local + 1, :acc].contiguous() + s_dt = all_dt[step][s_local : s_local + 1, :acc].contiguous() + s_B = all_B[step][s_local : s_local + 1, :acc].contiguous() + s_C = all_C[step][s_local : s_local + 1, :acc].contiguous() s_out = torch.zeros(1, acc, nheads, head_dim, device=device, dtype=dtype) selective_state_update( - s_state, s_x, s_dt, A, s_B, s_C, - D=D, dt_bias=dt_bias, dt_softplus=True, + s_state, + s_x, + s_dt, + A, + s_B, + s_C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, state_batch_indices=torch.tensor([0], device=device, dtype=torch.int32), out=s_out, ) @@ -1766,7 +1955,9 @@ def test_replay_heads_per_block_multistep( old_x = torch.zeros(cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype) old_B = torch.zeros(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.zeros(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) - old_dA_cumsum = torch.zeros(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.zeros( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) # Per-active-slot PNAT tracker, advanced by `accepted` (not T) per step. @@ -1780,8 +1971,13 @@ def test_replay_heads_per_block_multistep( prev_tokens[:] = pnat_active test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) n_writes_t, replay_work_items_t = _make_replay_work_items( - prev_tokens, cache_buf_idx, T, max_window, batch, - state_batch_indices, device, + prev_tokens, + cache_buf_idx, + T, + max_window, + batch, + state_batch_indices, + device, ) replay_selective_state_update( @@ -1814,11 +2010,14 @@ def test_replay_heads_per_block_multistep( # appends `accepted` to current PNAT. write_mask = (pnat_active + T) > max_window new_pnat_active = torch.where( - write_mask, accepted_tensor, pnat_active + accepted_tensor, + write_mask, + accepted_tensor, + pnat_active + accepted_tensor, ) pnat_active = new_pnat_active cache_active_idx = ( - state_batch_indices.long() if state_batch_indices is not None + state_batch_indices.long() + if state_batch_indices is not None else torch.arange(batch, device=device) ) write_slots = cache_active_idx[write_mask] From f3f5be4186a50c0918cd7ea6572e88fe7e3d53a9 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 11:19:23 -0700 Subject: [PATCH 70/89] Update mamba replay default tunings Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 1065 ++++++++++------- ...benchmark_replay_selective_state_update.py | 32 +- 2 files changed, 648 insertions(+), 449 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index b42bedcfa5bc..c01bbb3b7d7d 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -2286,178 +2286,159 @@ def _persistent_main_kernel( # sweep was T=6, max_window=16. Callers outside that regime silently get # the same numbers — they may be suboptimal but they're correct. # -# Source: audit_v2.py --emit-tuning. Auto-generated from per-cell search -# winners (best of pd / pm by bucket_expected_renorm). Sweep was TP=8 with -# NHEADS=128 → nheads_per_rank=16; thresholds are in effective_batch units. -# Missing dtype/SR combos (fp16/RN, int8/RN, fp8/*) fall back via the -# _resolve_tuning chain — RN→SR for same dtype, then fp8→int8/SR. +# Source: emit_tuning_from_noise.py. Auto-generated from noise-cleaned per-cell search +# winners (best of pd / pm by bucket_expected_renorm). Effective batch = raw_batch × 16. +# Missing dtype/SR combos (fp16/RN, int8/RN, fp8/RN) fall back via the +# _resolve_tuning chain — RN→SR for same dtype, then unknown fp8 cells to int8/SR. _DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { - ("fp32", "RN"): [ + ("fp16", "SR"): [ ( 16, "persistent_main", { - "_block_size_m_nowrite": 16, - "_block_size_m_write": 8, - "_cta_per_sm_nowrite": 4, - "_cta_per_sm_write": 1, + "_block_size_m_nowrite": 32, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 7, + "_cta_per_sm_write": 10, "_flatten": False, - "_heads_per_block": 2, + "_heads_per_block": 1, "_num_loop_stages_nowrite": 2, "_num_loop_stages_write": 2, "_num_stages_nowrite": 1, - "_num_stages_write": 2, - "_num_warps_nowrite": 2, - "_num_warps_write": 2, + "_num_stages_write": 3, + "_num_warps_nowrite": 4, + "_num_warps_write": 4, "_precompute_num_warps": 8, - "_use_tma_rect_load": False, + "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, "_warp_specialize": False, + "nowrite_first": False, "rectangle_for_nowrite": True, - }, - ), # raw_batch=1, score=6.22us + } + ), # raw_batch=1, score=6.66us (median of 4 noise runs; min=6.66us) ( 32, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 16, - "_block_size_m_write": 16, - "_cta_per_sm_nowrite": 9, - "_cta_per_sm_write": 7, + "_block_size_m": 4, + "_cta_per_sm": 7, "_flatten": False, - "_heads_per_block": 1, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 4, - "_num_stages_write": 4, - "_num_warps_nowrite": 2, - "_num_warps_write": 2, - "_precompute_num_warps": 8, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 5, + "_num_warps": 1, + "_precompute_num_warps": 4, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, "_warp_specialize": False, - "rectangle_for_nowrite": True, - }, - ), # raw_batch=2, score=7.17us + "rectangle_for_nowrite": False, + } + ), # raw_batch=2, score=6.84us (median of 4 noise runs; min=6.83us) ( 64, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 64, - "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 9, - "_cta_per_sm_write": 4, + "_block_size_m": 8, + "_cta_per_sm": 4, "_flatten": False, - "_heads_per_block": 2, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 3, - "_num_stages_nowrite": 1, - "_num_stages_write": 2, - "_num_warps_nowrite": 4, - "_num_warps_write": 4, - "_precompute_num_warps": 8, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 4, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": False, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_store": False, "_warp_specialize": False, - "rectangle_for_nowrite": True, - }, - ), # raw_batch=4, score=8.08us + "rectangle_for_nowrite": False, + } + ), # raw_batch=4, score=7.13us (median of 4 noise runs; min=7.12us) ( 128, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 32, - "_block_size_m_write": 64, - "_cta_per_sm_nowrite": 6, - "_cta_per_sm_write": 9, + "_block_size_m": 8, + "_cta_per_sm": 9, "_flatten": False, - "_heads_per_block": 8, - "_num_loop_stages_nowrite": 1, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 2, - "_num_stages_write": 1, - "_num_warps_nowrite": 1, - "_num_warps_write": 2, - "_precompute_num_warps": 8, + "_heads_per_block": 16, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 4, "_use_tma_rect_load": False, - "_use_tma_replay_nowrite_load": True, - "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - }, - ), # raw_batch=8, score=9.00us + } + ), # raw_batch=8, score=7.68us (median of 4 noise runs; min=7.64us) ( 256, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 64, - "_block_size_m_write": 64, - "_cta_per_sm_nowrite": 9, - "_cta_per_sm_write": 1, + "_block_size_m": 16, + "_cta_per_sm": 7, "_flatten": False, - "_heads_per_block": 4, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 3, - "_num_stages_write": 4, - "_num_warps_nowrite": 4, - "_num_warps_write": 2, - "_precompute_num_warps": 16, + "_heads_per_block": 8, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 8, "_use_tma_rect_load": False, - "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_nowrite_load": True, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, - "rectangle_for_nowrite": True, - }, - ), # raw_batch=16, score=10.92us + "rectangle_for_nowrite": False, + } + ), # raw_batch=16, score=9.04us (median of 4 noise runs; min=9.03us) ( 512, "persistent_main", { - "_block_size_m_nowrite": 32, + "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 10, + "_cta_per_sm_nowrite": 8, "_cta_per_sm_write": 4, "_flatten": False, - "_heads_per_block": 16, + "_heads_per_block": 8, "_num_loop_stages_nowrite": 1, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 2, - "_num_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 2, "_num_warps_nowrite": 1, "_num_warps_write": 2, - "_precompute_num_warps": 16, - "_use_tma_rect_load": False, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, + "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, - "rectangle_for_nowrite": False, - }, - ), # raw_batch=32, score=13.53us + "nowrite_first": True, + "rectangle_for_nowrite": True, + } + ), # raw_batch=32, score=12.51us (median of 4 noise runs; min=12.47us) ( 1024, "persistent_main", { - "_block_size_m_nowrite": 32, + "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 4, - "_cta_per_sm_write": 8, + "_cta_per_sm_nowrite": 9, + "_cta_per_sm_write": 7, "_flatten": False, "_heads_per_block": 16, - "_num_loop_stages_nowrite": 3, + "_num_loop_stages_nowrite": 1, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 3, + "_num_stages_nowrite": 2, "_num_stages_write": 4, - "_num_warps_nowrite": 1, + "_num_warps_nowrite": 2, "_num_warps_write": 1, "_precompute_num_warps": 8, "_use_tma_rect_load": True, @@ -2465,91 +2446,95 @@ def _persistent_main_kernel( "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": False, "rectangle_for_nowrite": True, - }, - ), # raw_batch=64, score=19.50us + } + ), # raw_batch=64, score=17.41us (median of 4 noise runs; min=17.41us) ( 2048, "persistent_main", { - "_block_size_m_nowrite": 32, + "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 8, + "_cta_per_sm_nowrite": 5, "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 2, + "_heads_per_block": 8, "_num_loop_stages_nowrite": 2, "_num_loop_stages_write": 1, "_num_stages_nowrite": 3, - "_num_stages_write": 3, - "_num_warps_nowrite": 1, + "_num_stages_write": 2, + "_num_warps_nowrite": 2, "_num_warps_write": 1, - "_precompute_num_warps": 2, + "_precompute_num_warps": 8, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=128, score=30.28us + } + ), # raw_batch=128, score=25.21us (median of 4 noise runs; min=25.16us) ( 4096, "persistent_main", { - "_block_size_m_nowrite": 32, + "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 4, + "_cta_per_sm_nowrite": 5, "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 4, - "_num_loop_stages_nowrite": 3, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 2, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 3, - "_num_stages_write": 4, - "_num_warps_nowrite": 1, + "_num_stages_nowrite": 1, + "_num_stages_write": 2, + "_num_warps_nowrite": 2, "_num_warps_write": 1, - "_precompute_num_warps": 2, - "_use_tma_rect_load": False, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, + "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=256, score=50.32us + } + ), # raw_batch=256, score=42.06us (median of 4 noise runs; min=42.01us) ( 8192, "persistent_main", { "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 10, + "_cta_per_sm_nowrite": 8, "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 4, - "_num_loop_stages_nowrite": 2, + "_num_loop_stages_nowrite": 3, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 4, - "_num_stages_write": 4, - "_num_warps_nowrite": 2, + "_num_stages_nowrite": 1, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, "_num_warps_write": 1, - "_precompute_num_warps": 2, + "_precompute_num_warps": 1, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, + "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=512, score=90.99us + } + ), # raw_batch=512, score=72.15us (median of 4 noise runs; min=72.11us) ( 16384, "persistent_main", { "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 9, + "_cta_per_sm_nowrite": 6, "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 8, @@ -2559,158 +2544,327 @@ def _persistent_main_kernel( "_num_stages_write": 3, "_num_warps_nowrite": 2, "_num_warps_write": 1, - "_precompute_num_warps": 2, + "_precompute_num_warps": 1, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, + "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=1024, score=171.69us + } + ), # raw_batch=1024, score=131.77us (median of 4 noise runs; min=131.63us) ], - ("fp16", "SR"): [ + ("int8", "SR"): [ ( 16, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 8, - "_block_size_m_write": 16, - "_cta_per_sm_nowrite": 5, - "_cta_per_sm_write": 7, + "_block_size_m": 8, + "_cta_per_sm": 2, "_flatten": False, - "_heads_per_block": 1, - "_num_loop_stages_nowrite": 4, - "_num_loop_stages_write": 3, - "_num_stages_nowrite": 1, - "_num_stages_write": 4, - "_num_warps_nowrite": 2, - "_num_warps_write": 4, - "_precompute_num_warps": 8, + "_heads_per_block": 8, + "_num_loop_stages": 2, + "_num_stages": 4, + "_num_warps": 2, + "_precompute_num_warps": 4, "_use_tma_rect_load": False, - "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, - "_use_tma_replay_write_store": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, "_warp_specialize": False, - "rectangle_for_nowrite": True, - }, - ), # raw_batch=1, score=6.16us + "rectangle_for_nowrite": False, + } + ), # raw_batch=1, score=6.42us (median of 4 noise runs; min=6.39us) ( 32, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 64, - "_block_size_m_write": 16, - "_cta_per_sm_nowrite": 9, - "_cta_per_sm_write": 4, + "_block_size_m": 4, + "_cta_per_sm": 10, "_flatten": False, - "_heads_per_block": 1, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 2, - "_num_stages_write": 4, - "_num_warps_nowrite": 4, - "_num_warps_write": 2, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, "_precompute_num_warps": 8, - "_use_tma_rect_load": True, + "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": False, "_use_tma_replay_write_store": True, "_warp_specialize": False, - "rectangle_for_nowrite": True, - }, - ), # raw_batch=2, score=7.01us + "rectangle_for_nowrite": False, + } + ), # raw_batch=2, score=7.13us (median of 4 noise runs; min=7.09us) ( 64, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 64, - "_block_size_m_write": 16, - "_cta_per_sm_nowrite": 5, - "_cta_per_sm_write": 8, + "_block_size_m": 8, + "_cta_per_sm": 5, "_flatten": False, - "_heads_per_block": 2, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 1, - "_num_stages_write": 4, - "_num_warps_nowrite": 2, - "_num_warps_write": 2, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 2, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + } + ), # raw_batch=4, score=7.52us (median of 4 noise runs; min=7.50us) + ( + 128, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 8, + "_flatten": False, + "_heads_per_block": 1, + "_num_loop_stages": 1, + "_num_stages": 1, + "_num_warps": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + } + ), # raw_batch=8, score=8.22us (median of 4 noise runs; min=8.21us) + ( + 256, + "persistent_dynamic", + { + "_block_size_m": 16, + "_cta_per_sm": 10, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages": 1, + "_num_stages": 1, + "_num_warps": 1, "_precompute_num_warps": 8, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": False, "_use_tma_replay_write_store": False, "_warp_specialize": False, + "rectangle_for_nowrite": False, + } + ), # raw_batch=16, score=10.06us (median of 8 noise runs; min=9.97us) + ( + 512, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 3, + "_cta_per_sm_write": 9, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=4, score=7.95us + } + ), # raw_batch=32, score=12.91us (median of 4 noise runs; min=12.90us) ( - 128, + 1024, "persistent_main", { "_block_size_m_nowrite": 64, - "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 8, - "_cta_per_sm_write": 4, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 6, + "_cta_per_sm_write": 3, "_flatten": False, - "_heads_per_block": 1, - "_num_loop_stages_nowrite": 3, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 2, "_num_loop_stages_write": 2, - "_num_stages_nowrite": 3, + "_num_stages_nowrite": 2, "_num_stages_write": 1, - "_num_warps_nowrite": 4, + "_num_warps_nowrite": 2, "_num_warps_write": 4, - "_precompute_num_warps": 4, + "_precompute_num_warps": 16, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, - "_use_tma_replay_write_store": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": False, "rectangle_for_nowrite": True, - }, - ), # raw_batch=8, score=8.87us + } + ), # raw_batch=64, score=18.05us (median of 4 noise runs; min=18.03us) ( - 256, + 2048, "persistent_main", { - "_block_size_m_nowrite": 32, - "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 10, - "_cta_per_sm_write": 2, + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 6, + "_cta_per_sm_write": 3, "_flatten": False, "_heads_per_block": 4, - "_num_loop_stages_nowrite": 1, - "_num_loop_stages_write": 1, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 3, "_num_stages_nowrite": 1, "_num_stages_write": 1, - "_num_warps_nowrite": 1, + "_num_warps_nowrite": 2, "_num_warps_write": 4, - "_precompute_num_warps": 16, - "_use_tma_rect_load": False, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, + "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": False, "_warp_specialize": False, - "rectangle_for_nowrite": False, - }, - ), # raw_batch=16, score=10.28us + "nowrite_first": False, + "rectangle_for_nowrite": True, + } + ), # raw_batch=128, score=27.66us (median of 4 noise runs; min=27.56us) ( - 512, + 4096, "persistent_main", { - "_block_size_m_nowrite": 32, - "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 7, + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 3, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 3, + "_num_stages_nowrite": 4, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, + "_num_warps_write": 4, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + } + ), # raw_batch=256, score=45.13us (median of 4 noise runs; min=45.08us) + ( + 8192, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 8, "_cta_per_sm_write": 6, "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 4, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 3, + "_num_stages_write": 4, + "_num_warps_nowrite": 1, + "_num_warps_write": 4, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + } + ), # raw_batch=512, score=77.87us (median of 4 noise runs; min=77.78us) + ( + 16384, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 3, + "_flatten": False, "_heads_per_block": 8, - "_num_loop_stages_nowrite": 1, - "_num_loop_stages_write": 1, + "_num_loop_stages_nowrite": 4, + "_num_loop_stages_write": 3, "_num_stages_nowrite": 3, - "_num_stages_write": 2, + "_num_stages_write": 1, "_num_warps_nowrite": 1, - "_num_warps_write": 2, + "_num_warps_write": 4, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + } + ), # raw_batch=1024, score=142.83us (median of 4 noise runs; min=142.79us) + ], + ("fp8", "SR"): [ + ( + 16, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 5, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages": 2, + "_num_stages": 2, + "_num_warps": 2, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + } + ), # raw_batch=1, score=6.30us (median of 4 noise runs; min=6.28us) + ( + 32, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 10, + "_flatten": False, + "_heads_per_block": 1, + "_num_loop_stages": 1, + "_num_stages": 2, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + } + ), # raw_batch=2, score=6.80us (median of 4 noise runs; min=6.79us) + ( + 64, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 9, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, "_precompute_num_warps": 8, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, @@ -2718,95 +2872,158 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - }, - ), # raw_batch=32, score=12.90us + } + ), # raw_batch=4, score=7.00us (median of 4 noise runs; min=6.98us) + ( + 128, + "persistent_dynamic", + { + "_block_size_m": 16, + "_cta_per_sm": 5, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + } + ), # raw_batch=8, score=7.67us (median of 4 noise runs; min=7.59us) + ( + 256, + "persistent_dynamic", + { + "_block_size_m": 16, + "_cta_per_sm": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + } + ), # raw_batch=16, score=8.58us (median of 4 noise runs; min=8.54us) + ( + 512, + "persistent_dynamic", + { + "_block_size_m": 32, + "_cta_per_sm": 7, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + } + ), # raw_batch=32, score=10.60us (median of 4 noise runs; min=10.54us) ( 1024, "persistent_main", { "_block_size_m_nowrite": 32, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 7, + "_cta_per_sm_nowrite": 10, "_cta_per_sm_write": 7, "_flatten": False, "_heads_per_block": 16, - "_num_loop_stages_nowrite": 3, + "_num_loop_stages_nowrite": 1, "_num_loop_stages_write": 1, "_num_stages_nowrite": 4, "_num_stages_write": 2, - "_num_warps_nowrite": 1, + "_num_warps_nowrite": 2, "_num_warps_write": 1, "_precompute_num_warps": 8, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, "_warp_specialize": False, + "nowrite_first": False, "rectangle_for_nowrite": True, - }, - ), # raw_batch=64, score=16.71us + } + ), # raw_batch=64, score=15.79us (median of 4 noise runs; min=15.77us) ( 2048, "persistent_main", { - "_block_size_m_nowrite": 32, + "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 7, + "_cta_per_sm_nowrite": 6, "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 2, - "_num_loop_stages_nowrite": 3, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 2, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 4, - "_num_stages_write": 1, - "_num_warps_nowrite": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, "_num_warps_write": 1, - "_precompute_num_warps": 2, + "_precompute_num_warps": 8, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": False, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_store": False, "_warp_specialize": False, + "nowrite_first": False, "rectangle_for_nowrite": True, - }, - ), # raw_batch=128, score=25.71us + } + ), # raw_batch=128, score=23.65us (median of 4 noise runs; min=23.65us) ( 4096, "persistent_main", { - "_block_size_m_nowrite": 32, + "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 7, + "_cta_per_sm_nowrite": 6, "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 4, - "_num_loop_stages_nowrite": 3, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 2, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 3, - "_num_stages_write": 4, - "_num_warps_nowrite": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 3, + "_num_warps_nowrite": 2, "_num_warps_write": 1, - "_precompute_num_warps": 1, + "_precompute_num_warps": 8, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, "_warp_specialize": False, + "nowrite_first": False, "rectangle_for_nowrite": True, - }, - ), # raw_batch=256, score=39.80us + } + ), # raw_batch=256, score=38.47us (median of 4 noise runs; min=38.42us) ( 8192, "persistent_main", { - "_block_size_m_nowrite": 32, + "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 7, + "_cta_per_sm_nowrite": 8, "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 4, "_num_loop_stages_nowrite": 3, - "_num_loop_stages_write": 1, + "_num_loop_stages_write": 2, "_num_stages_nowrite": 2, "_num_stages_write": 2, "_num_warps_nowrite": 1, @@ -2815,315 +3032,292 @@ def _persistent_main_kernel( "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_store": False, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=512, score=71.34us + } + ), # raw_batch=512, score=68.32us (median of 4 noise runs; min=68.28us) ( 16384, "persistent_main", { - "_block_size_m_nowrite": 32, + "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 7, + "_cta_per_sm_nowrite": 8, "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 16, + "_heads_per_block": 8, "_num_loop_stages_nowrite": 3, - "_num_loop_stages_write": 1, - "_num_stages_nowrite": 2, - "_num_stages_write": 1, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 4, + "_num_stages_write": 2, "_num_warps_nowrite": 1, "_num_warps_write": 1, "_precompute_num_warps": 1, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_store": False, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=1024, score=133.51us + } + ), # raw_batch=1024, score=123.51us (median of 4 noise runs; min=123.40us) ], - ("int8", "SR"): [ + ("fp32", "RN"): [ ( 16, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 8, - "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 5, - "_cta_per_sm_write": 8, + "_block_size_m": 8, + "_cta_per_sm": 9, "_flatten": False, - "_heads_per_block": 1, - "_num_loop_stages_nowrite": 4, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 2, - "_num_stages_write": 3, - "_num_warps_nowrite": 2, - "_num_warps_write": 4, + "_heads_per_block": 2, + "_num_loop_stages": 2, + "_num_stages": 4, + "_num_warps": 2, "_precompute_num_warps": 8, "_use_tma_rect_load": False, - "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": False, "_warp_specialize": False, - "rectangle_for_nowrite": True, - }, - ), # raw_batch=1, score=6.34us + "rectangle_for_nowrite": False, + } + ), # raw_batch=1, score=6.01us (median of 4 noise runs; min=6.00us) ( 32, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 32, - "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 10, - "_cta_per_sm_write": 8, + "_block_size_m": 8, + "_cta_per_sm": 2, "_flatten": False, "_heads_per_block": 2, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 4, - "_num_stages_write": 2, - "_num_warps_nowrite": 4, - "_num_warps_write": 4, - "_precompute_num_warps": 8, + "_num_loop_stages": 1, + "_num_stages": 5, + "_num_warps": 1, + "_precompute_num_warps": 4, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": False, "_use_tma_replay_write_store": False, "_warp_specialize": False, - "rectangle_for_nowrite": True, - }, - ), # raw_batch=2, score=7.36us + "rectangle_for_nowrite": False, + } + ), # raw_batch=2, score=6.64us (median of 4 noise runs; min=6.62us) ( 64, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 64, - "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 2, - "_cta_per_sm_write": 4, + "_block_size_m": 8, + "_cta_per_sm": 5, "_flatten": False, "_heads_per_block": 2, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 2, - "_num_stages_write": 3, - "_num_warps_nowrite": 4, - "_num_warps_write": 4, - "_precompute_num_warps": 8, - "_use_tma_rect_load": True, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, "_warp_specialize": False, - "rectangle_for_nowrite": True, - }, - ), # raw_batch=4, score=8.40us + "rectangle_for_nowrite": False, + } + ), # raw_batch=4, score=7.09us (median of 4 noise runs; min=7.08us) ( 128, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 64, - "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 10, - "_cta_per_sm_write": 10, + "_block_size_m": 8, + "_cta_per_sm": 8, "_flatten": False, "_heads_per_block": 2, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 3, - "_num_stages_write": 4, - "_num_warps_nowrite": 4, - "_num_warps_write": 4, - "_precompute_num_warps": 16, - "_use_tma_rect_load": True, - "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": True, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 2, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": False, "_use_tma_replay_write_store": True, "_warp_specialize": False, - "rectangle_for_nowrite": True, - }, - ), # raw_batch=8, score=9.37us + "rectangle_for_nowrite": False, + } + ), # raw_batch=8, score=7.96us (median of 4 noise runs; min=7.90us) ( 256, "persistent_dynamic", { - "_block_size_m": 16, + "_block_size_m": 32, "_cta_per_sm": 8, "_flatten": False, - "_heads_per_block": 16, + "_heads_per_block": 8, "_num_loop_stages": 1, - "_num_stages": 4, - "_num_warps": 1, - "_precompute_num_warps": 16, + "_num_stages": 2, + "_num_warps": 2, + "_precompute_num_warps": 8, "_use_tma_rect_load": False, - "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": False, "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - }, - ), # raw_batch=16, score=10.02us + } + ), # raw_batch=16, score=9.55us (median of 4 noise runs; min=9.50us) ( 512, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 32, - "_block_size_m_write": 16, - "_cta_per_sm_nowrite": 7, - "_cta_per_sm_write": 9, + "_block_size_m": 32, + "_cta_per_sm": 7, "_flatten": False, "_heads_per_block": 8, - "_num_loop_stages_nowrite": 1, - "_num_loop_stages_write": 1, - "_num_stages_nowrite": 2, - "_num_stages_write": 3, - "_num_warps_nowrite": 2, - "_num_warps_write": 1, + "_num_loop_stages": 1, + "_num_stages": 2, + "_num_warps": 1, "_precompute_num_warps": 8, "_use_tma_rect_load": False, - "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_nowrite_load": True, "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": False, + "_use_tma_replay_write_store": True, "_warp_specialize": False, - "rectangle_for_nowrite": True, - }, - ), # raw_batch=32, score=13.15us + "rectangle_for_nowrite": False, + } + ), # raw_batch=32, score=13.28us (median of 4 noise runs; min=13.23us) ( 1024, "persistent_main", { "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 10, - "_cta_per_sm_write": 7, + "_cta_per_sm_nowrite": 4, + "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 16, - "_num_loop_stages_nowrite": 1, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 3, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 4, - "_num_stages_write": 3, - "_num_warps_nowrite": 1, - "_num_warps_write": 1, + "_num_stages_nowrite": 1, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, + "_num_warps_write": 2, "_precompute_num_warps": 8, - "_use_tma_rect_load": False, + "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": False, + "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=64, score=17.82us + } + ), # raw_batch=64, score=19.12us (median of 4 noise runs; min=19.10us) ( 2048, "persistent_main", { "_block_size_m_nowrite": 64, - "_block_size_m_write": 64, - "_cta_per_sm_nowrite": 4, - "_cta_per_sm_write": 3, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 5, + "_cta_per_sm_write": 9, "_flatten": False, - "_heads_per_block": 4, - "_num_loop_stages_nowrite": 3, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 3, - "_num_stages_write": 4, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 1, + "_num_stages_write": 1, "_num_warps_nowrite": 2, - "_num_warps_write": 4, - "_precompute_num_warps": 1, - "_use_tma_rect_load": False, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=128, score=27.01us + } + ), # raw_batch=128, score=29.39us (median of 4 noise runs; min=29.34us) ( 4096, "persistent_main", { "_block_size_m_nowrite": 64, - "_block_size_m_write": 64, - "_cta_per_sm_nowrite": 4, - "_cta_per_sm_write": 3, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 9, + "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 4, - "_num_loop_stages_nowrite": 3, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 3, - "_num_stages_write": 2, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 1, "_num_warps_nowrite": 2, - "_num_warps_write": 4, - "_precompute_num_warps": 1, - "_use_tma_rect_load": False, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=256, score=43.23us + } + ), # raw_batch=256, score=49.36us (median of 4 noise runs; min=49.31us) ( 8192, "persistent_main", { "_block_size_m_nowrite": 32, - "_block_size_m_write": 64, + "_block_size_m_write": 32, "_cta_per_sm_nowrite": 8, - "_cta_per_sm_write": 6, + "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 4, - "_num_loop_stages_nowrite": 3, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 3, - "_num_stages_write": 1, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 4, "_num_warps_nowrite": 1, - "_num_warps_write": 4, - "_precompute_num_warps": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=512, score=77.01us + } + ), # raw_batch=512, score=87.36us (median of 4 noise runs; min=87.31us) ( 16384, "persistent_main", { "_block_size_m_nowrite": 64, - "_block_size_m_write": 64, - "_cta_per_sm_nowrite": 6, - "_cta_per_sm_write": 3, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 9, + "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 8, - "_num_loop_stages_nowrite": 3, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 1, - "_num_stages_write": 4, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 2, "_num_warps_nowrite": 2, - "_num_warps_write": 4, + "_num_warps_write": 1, "_precompute_num_warps": 1, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, + "nowrite_first": True, "rectangle_for_nowrite": True, - }, - ), # raw_batch=1024, score=140.43us + } + ), # raw_batch=1024, score=168.98us (median of 4 noise runs; min=168.83us) ], } - - -# Knob names that map between the modes' single-value (pd) and split-value -# (pm) namespaces. Used by `_bridge_tuning_knobs` when caller forces a mode -# different from the table's recommendation. _PD_TO_PM_SPLIT_MAP = { # pd unsplit knob → (pm_write_knob, pm_nowrite_knob) "_block_size_m": ("_block_size_m_write", "_block_size_m_nowrite"), "_num_warps": ("_num_warps_write", "_num_warps_nowrite"), @@ -3243,7 +3437,7 @@ def replay_selective_state_update( use_internal_pdl=True, write_checkpoint: bool = True, rectangle_for_nowrite: bool | None = None, - nowrite_first: bool = False, + nowrite_first: bool | None = None, mode: str | None = None, _block_size_m: int | None = None, _num_warps: int | None = None, @@ -3361,7 +3555,8 @@ def replay_selective_state_update( Defaults True; override for testing only. Ignored on hardware that doesn't support PDL (sm < 90). nowrite_first: benchmark/tuning knob for mode="persistent_main". - When true, launch the nowrite half before the write half. + When None, use the tuning-table value if present. When true, + launch the nowrite half before the write half. _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, _precompute_num_warps, _precompute_num_stages, _heads_per_block, @@ -3495,6 +3690,8 @@ def replay_selective_state_update( # locals() for re-read, so re-bind each kwarg explicitly. if rectangle_for_nowrite is None and "rectangle_for_nowrite" in _table_knobs: rectangle_for_nowrite = bool(_table_knobs["rectangle_for_nowrite"]) + if nowrite_first is None and "nowrite_first" in _table_knobs: + nowrite_first = bool(_table_knobs["nowrite_first"]) _block_size_m = ( _block_size_m if _block_size_m is not None else _table_knobs.get("_block_size_m") ) @@ -3595,6 +3792,8 @@ def replay_selective_state_update( mode = "persistent_dynamic" if rectangle_for_nowrite is None: rectangle_for_nowrite = False + if nowrite_first is None: + nowrite_first = False _use_tma_rect_load = bool(_use_tma_rect_load) _use_tma_replay_write_load = bool(_use_tma_replay_write_load) _use_tma_replay_write_store = bool(_use_tma_replay_write_store) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 8bc3f507ad37..5cc7964ab48e 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -2611,7 +2611,7 @@ def _bench_config( mix_samples_cpu=None, mix_label: str = "", hardcode_sort: bool = False, - nowrite_first: bool = False, + nowrite_first: bool | None = None, mix_samples_sorted_cpu=None, mix_write_frac: float | None = None, warmup_only: bool = False, @@ -3468,7 +3468,7 @@ def _emit_split(name_w, name_nw, val_w, val_nw): parts.append( f"RECT={'auto' if rectangle_for_nowrite is None else (1 if rectangle_for_nowrite else 0)}" ) - nowrite_first_list_for_tags = getattr(args, "nowrite_first_list", [False]) + nowrite_first_list_for_tags = getattr(args, "nowrite_first_list", [None]) nowrite_first_in_cell_list = "NWF" in getattr(args, "_cell_list_keys", ()) if nowrite_first or len(nowrite_first_list_for_tags) > 1 or nowrite_first_in_cell_list: parts.append(f"NWF={1 if nowrite_first else 0}") @@ -4280,7 +4280,7 @@ def _phase(label: str) -> None: rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) modes_list = getattr(args, "modes_list", ["persistent_dynamic"]) hsort_list = getattr(args, "hardcode_sort_list", [False]) - nowrite_first_list = getattr(args, "nowrite_first_list", [False]) + nowrite_first_list = getattr(args, "nowrite_first_list", [None]) # Pre-load AL distribution for mix mode. mix_al = None @@ -4862,9 +4862,10 @@ def _parse_args() -> argparse.Namespace: parser.add_argument( "--nowrite-first", type=str, - default="0", + default=None, help="Comma-separated 0/1 values for mode=persistent_main launch order. " "0 launches write before nowrite; 1 launches nowrite before write. " + "When unset (default), the wrapper resolves from _DEFAULT_TUNING. " "Ignored for persistent_dynamic.", ) parser.add_argument( @@ -5089,18 +5090,17 @@ def _round_iters_to_group(name, val): rect_list = [None] args.rectangle_for_nowrite_list = rect_list - nowrite_first_modes = [ - v.strip() - for v in (args.nowrite_first if args.nowrite_first is not None else "0").split(",") - if v.strip() - ] - nowrite_first_list = [] - for v in nowrite_first_modes: - if v not in ("0", "1"): - parser.error(f"--nowrite-first value must be 0 or 1, got {v!r}") - nowrite_first_list.append(v == "1") - if not nowrite_first_list: - nowrite_first_list = [False] + if args.nowrite_first is None: + nowrite_first_list = [None] + else: + nowrite_first_modes = [v.strip() for v in args.nowrite_first.split(",") if v.strip()] + nowrite_first_list = [] + for v in nowrite_first_modes: + if v not in ("0", "1"): + parser.error(f"--nowrite-first value must be 0 or 1, got {v!r}") + nowrite_first_list.append(v == "1") + if not nowrite_first_list: + nowrite_first_list = [None] args.nowrite_first_list = nowrite_first_list hsort_modes = [ From 6127ae3d377c413830e742ffa81184a62f575654 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 26 May 2026 16:50:35 -0700 Subject: [PATCH 71/89] [None][fix] stabilize Triton Mamba softplus Replace triton helper softplus with numerically stable version. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- tensorrt_llm/_torch/modules/mamba/softplus.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/softplus.py b/tensorrt_llm/_torch/modules/mamba/softplus.py index 7b64e96432a0..d021af75bc9d 100644 --- a/tensorrt_llm/_torch/modules/mamba/softplus.py +++ b/tensorrt_llm/_torch/modules/mamba/softplus.py @@ -1,7 +1,7 @@ # Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/softplus.py # Copyright (c) 2024, Tri Dao, Albert Gu. # -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -26,10 +26,13 @@ @triton.jit def softplus(dt): - return tl.math.log(tl.math.exp(dt) + 1) + dt_clamped = tl.minimum(dt, 20.0) + return tl.where(dt <= 20.0, tl.math.log(tl.math.exp(dt_clamped) + 1), + dt) else: @triton.jit def softplus(dt): - return tl.math.log1p(tl.exp(dt)) + dt_clamped = tl.minimum(dt, 20.0) + return tl.where(dt <= 20.0, tl.math.log1p(tl.exp(dt_clamped)), dt) From 172cc08a400ded8bdf5d1582bf07385c91cbe067 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 22 May 2026 12:54:28 -0700 Subject: [PATCH 72/89] Clean up mamba replay default tuning Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 281 +++++++----------- 1 file changed, 107 insertions(+), 174 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index c01bbb3b7d7d..2aab5971c646 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -2274,8 +2274,8 @@ def _persistent_main_kernel( # tensor shape so callers at other TP / nheads pick up the right cell. # # Schema: dict[(dtype_str, sr_str)] → list[(eff_batch_threshold, mode, knobs)] -# sorted by threshold ascending. Lookup finds the first threshold ≥ eff_b -# (so missing intermediate batches fall up to the next tuned cell). If +# sorted by threshold ascending. Lookup finds the first threshold ≥ eff_b +# (so missing intermediate batches fall up to the next tuned cell). If # eff_b exceeds the largest threshold, use the largest entry. # # Each `knobs` dict only contains keys for the chosen mode; the wrapper @@ -2288,8 +2288,8 @@ def _persistent_main_kernel( # # Source: emit_tuning_from_noise.py. Auto-generated from noise-cleaned per-cell search # winners (best of pd / pm by bucket_expected_renorm). Effective batch = raw_batch × 16. -# Missing dtype/SR combos (fp16/RN, int8/RN, fp8/RN) fall back via the -# _resolve_tuning chain — RN→SR for same dtype, then unknown fp8 cells to int8/SR. +# Missing dtype/SR combos fall back via the _resolve_tuning chain: +# RN→SR for same dtype, then bf16/int16→fp16/SR and fp8→int8/SR. _DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { ("fp16", "SR"): [ ( @@ -2316,7 +2316,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": False, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=1, score=6.66us (median of 4 noise runs; min=6.66us) ( 32, @@ -2336,7 +2336,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=2, score=6.84us (median of 4 noise runs; min=6.83us) ( 64, @@ -2356,28 +2356,28 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=4, score=7.13us (median of 4 noise runs; min=7.12us) ( 128, "persistent_dynamic", { "_block_size_m": 8, - "_cta_per_sm": 9, + "_cta_per_sm": 7, "_flatten": False, - "_heads_per_block": 16, + "_heads_per_block": 4, "_num_loop_stages": 1, - "_num_stages": 4, + "_num_stages": 2, "_num_warps": 1, - "_precompute_num_warps": 4, + "_precompute_num_warps": 2, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": False, "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } - ), # raw_batch=8, score=7.68us (median of 4 noise runs; min=7.64us) + }, + ), # raw_batch=8, score=7.83us (manual retry of pre-noise winner) ( 256, "persistent_dynamic", @@ -2396,7 +2396,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=16, score=9.04us (median of 4 noise runs; min=9.03us) ( 512, @@ -2422,7 +2422,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=32, score=12.51us (median of 4 noise runs; min=12.47us) ( 1024, @@ -2448,7 +2448,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": False, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=64, score=17.41us (median of 4 noise runs; min=17.41us) ( 2048, @@ -2474,7 +2474,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=128, score=25.21us (median of 4 noise runs; min=25.16us) ( 4096, @@ -2500,7 +2500,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=256, score=42.06us (median of 4 noise runs; min=42.01us) ( 8192, @@ -2526,7 +2526,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=512, score=72.15us (median of 4 noise runs; min=72.11us) ( 16384, @@ -2552,7 +2552,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=1024, score=131.77us (median of 4 noise runs; min=131.63us) ], ("int8", "SR"): [ @@ -2574,7 +2574,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=1, score=6.42us (median of 4 noise runs; min=6.39us) ( 32, @@ -2594,7 +2594,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=2, score=7.13us (median of 4 noise runs; min=7.09us) ( 64, @@ -2614,7 +2614,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=4, score=7.52us (median of 4 noise runs; min=7.50us) ( 128, @@ -2634,7 +2634,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=8, score=8.22us (median of 4 noise runs; min=8.21us) ( 256, @@ -2654,7 +2654,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=16, score=10.06us (median of 8 noise runs; min=9.97us) ( 512, @@ -2680,7 +2680,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=32, score=12.91us (median of 4 noise runs; min=12.90us) ( 1024, @@ -2706,7 +2706,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": False, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=64, score=18.05us (median of 4 noise runs; min=18.03us) ( 2048, @@ -2732,7 +2732,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": False, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=128, score=27.66us (median of 4 noise runs; min=27.56us) ( 4096, @@ -2758,7 +2758,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=256, score=45.13us (median of 4 noise runs; min=45.08us) ( 8192, @@ -2784,7 +2784,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=512, score=77.87us (median of 4 noise runs; min=77.78us) ( 16384, @@ -2810,7 +2810,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=1024, score=142.83us (median of 4 noise runs; min=142.79us) ], ("fp8", "SR"): [ @@ -2832,7 +2832,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=1, score=6.30us (median of 4 noise runs; min=6.28us) ( 32, @@ -2852,7 +2852,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=2, score=6.80us (median of 4 noise runs; min=6.79us) ( 64, @@ -2872,7 +2872,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=4, score=7.00us (median of 4 noise runs; min=6.98us) ( 128, @@ -2892,7 +2892,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=8, score=7.67us (median of 4 noise runs; min=7.59us) ( 256, @@ -2912,7 +2912,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=16, score=8.58us (median of 4 noise runs; min=8.54us) ( 512, @@ -2932,7 +2932,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=32, score=10.60us (median of 4 noise runs; min=10.54us) ( 1024, @@ -2958,7 +2958,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": False, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=64, score=15.79us (median of 4 noise runs; min=15.77us) ( 2048, @@ -2984,7 +2984,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": False, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=128, score=23.65us (median of 4 noise runs; min=23.65us) ( 4096, @@ -3010,7 +3010,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": False, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=256, score=38.47us (median of 4 noise runs; min=38.42us) ( 8192, @@ -3036,7 +3036,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=512, score=68.32us (median of 4 noise runs; min=68.28us) ( 16384, @@ -3062,7 +3062,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=1024, score=123.51us (median of 4 noise runs; min=123.40us) ], ("fp32", "RN"): [ @@ -3084,7 +3084,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=1, score=6.01us (median of 4 noise runs; min=6.00us) ( 32, @@ -3104,18 +3104,18 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=2, score=6.64us (median of 4 noise runs; min=6.62us) ( 64, "persistent_dynamic", { "_block_size_m": 8, - "_cta_per_sm": 5, + "_cta_per_sm": 9, "_flatten": False, "_heads_per_block": 2, "_num_loop_stages": 1, - "_num_stages": 3, + "_num_stages": 4, "_num_warps": 1, "_precompute_num_warps": 4, "_use_tma_rect_load": False, @@ -3124,8 +3124,8 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } - ), # raw_batch=4, score=7.09us (median of 4 noise runs; min=7.08us) + }, + ), # raw_batch=4, score=7.17us (manual retry of pre-noise winner) ( 128, "persistent_dynamic", @@ -3144,7 +3144,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=8, score=7.96us (median of 4 noise runs; min=7.90us) ( 256, @@ -3164,7 +3164,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=16, score=9.55us (median of 4 noise runs; min=9.50us) ( 512, @@ -3184,7 +3184,7 @@ def _persistent_main_kernel( "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, - } + }, ), # raw_batch=32, score=13.28us (median of 4 noise runs; min=13.23us) ( 1024, @@ -3210,7 +3210,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=64, score=19.12us (median of 4 noise runs; min=19.10us) ( 2048, @@ -3236,7 +3236,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=128, score=29.39us (median of 4 noise runs; min=29.34us) ( 4096, @@ -3262,7 +3262,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=256, score=49.36us (median of 4 noise runs; min=49.31us) ( 8192, @@ -3288,7 +3288,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=512, score=87.36us (median of 4 noise runs; min=87.31us) ( 16384, @@ -3314,7 +3314,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, - } + }, ), # raw_batch=1024, score=168.98us (median of 4 noise runs; min=168.83us) ], } @@ -3363,10 +3363,8 @@ def _resolve_tuning( ) -> tuple[str, dict] | None: """Look up the default mode + knobs for this (eff_batch, dt, sr) cell. - Returns (mode, knobs_dict) or None if the table has no entry covering - this dtype/sr (including the fp8→int8/SR and dtype/RN→dtype/SR fallbacks). - Returning None lets the wrapper fall back to caller-provided kwargs or - kernel-side defaults. + Returns (mode, knobs_dict) or None if the table has no entry covering this + dtype/sr after fallbacks. """ eff_b = batch * max(1, nheads_per_rank) # Lookup chain. Order: @@ -3396,7 +3394,7 @@ def _resolve_tuning( break if entries is None: return None - # Find first threshold ≥ eff_b; if none, use largest entry. + # Find first threshold >= eff_b; if none, use largest entry. for thresh, mode, knobs in entries: if eff_b <= thresh: return mode, dict(knobs) @@ -3561,7 +3559,7 @@ def replay_selective_state_update( _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, _precompute_num_warps, _precompute_num_stages, _heads_per_block, _maxnreg, _num_ctas) are benchmark-only overrides; production callers - should leave them None to use the heuristic-tuned defaults. + should leave them None to use the tuning-table defaults. """ sm_version = get_sm_version() @@ -3679,6 +3677,11 @@ def replay_selective_state_update( }.get(state.dtype, str(state.dtype)) _sr_str = "SR" if rand_seed is not None else "RN" _table_entry = _resolve_tuning(batch, nheads, _dt_str, _sr_str) + if _table_entry is None: + raise ValueError( + "replay_selective_state_update has no default tuning for " + f"state dtype {_dt_str!r} with rounding mode {_sr_str!r}." + ) if _table_entry is not None: _table_mode, _table_knobs = _table_entry if mode is None: @@ -3787,7 +3790,7 @@ def replay_selective_state_update( _use_tma_replay_nowrite_load = bool( _table_knobs.get("_use_tma_replay_nowrite_load", False) ) - # Final defaults if neither caller nor table set them (empty table case). + # Final defaults for optional mode flags if neither caller nor table set them. if mode is None: mode = "persistent_dynamic" if rectangle_for_nowrite is None: @@ -3868,124 +3871,54 @@ def replay_selective_state_update( decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) z_strides = ( - (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) if z is not None else (0, 0, 0, 0) + (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) + if z is not None else (0, 0, 0, 0) ) - # Kernel tuning: BLOCK_SIZE_M, num_warps, HEADS_PER_BLOCK, precompute_num_warps. - # Dtype-aware heuristic from B200 sweeps (batch 1-512, T=6/32, TP=8, conv1d + - # chained PDL). Keyed on total_heads, BLOCK_SIZE_T, and state dtype; 16-bit - # states prefer different tiles from fp32 due to lower bandwidth. Philox - # gets its own branch — stochastic rounding shifts compute toward CUDA cores, - # so small-batch configs want more warps to hide the extra work. - total_heads = batch * nheads heads_per_group = nheads // ngroups - state_is_16bit = state.dtype in (torch.float16, torch.bfloat16) - use_philox = rand_seed is not None - if BLOCK_SIZE_T <= 16: - if use_philox and state_is_16bit: - # Philox: more warps at small batch to hide CUDA core work. - # At large batch, converges to non-Philox fp16 config. - if total_heads <= 16: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 4, 4, 4, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - elif state_is_16bit: - if total_heads <= 16: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 4, 1 - elif total_heads <= 64: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 2, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 1, - min(2, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - else: # fp32 state (no Philox — fp32 doesn't need stochastic rounding) - if total_heads <= 32: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 4, 1 - elif total_heads <= 64: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - elif total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 2, 2, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 1, 2, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 2, - min(2, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - else: # T > 16 - if state_is_16bit: - if total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 16, - 1, - 4, - min(2, heads_per_group), - ) - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 1, - min(4, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 4, - min(2, heads_per_group), - ) - else: # fp32 state - if total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 2, - 4, - min(2, heads_per_group), - ) - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 2, - min(4, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 4, - min(2, heads_per_group), - ) - if _block_size_m is not None: - BLOCK_SIZE_M = _block_size_m - if _num_warps is not None: - num_warps = _num_warps - if _heads_per_block is not None: - heads_per_block = int(_heads_per_block) + if mode == "persistent_dynamic": + assert _block_size_m is not None, "persistent_dynamic requires _block_size_m tuning" + assert _num_warps is not None, "persistent_dynamic requires _num_warps tuning" + assert _num_stages is not None, "persistent_dynamic requires _num_stages tuning" + assert _cta_per_sm is not None, "persistent_dynamic requires _cta_per_sm tuning" + assert _num_loop_stages is not None, "persistent_dynamic requires _num_loop_stages tuning" + else: + assert _block_size_m_write is not None, ( + "persistent_main requires _block_size_m_write tuning" + ) + assert _block_size_m_nowrite is not None, ( + "persistent_main requires _block_size_m_nowrite tuning" + ) + assert _num_warps_write is not None, "persistent_main requires _num_warps_write tuning" + assert _num_warps_nowrite is not None, "persistent_main requires _num_warps_nowrite tuning" + assert _num_stages_write is not None, "persistent_main requires _num_stages_write tuning" + assert _num_stages_nowrite is not None, ( + "persistent_main requires _num_stages_nowrite tuning" + ) + assert _cta_per_sm_write is not None, "persistent_main requires _cta_per_sm_write tuning" + assert _cta_per_sm_nowrite is not None, ( + "persistent_main requires _cta_per_sm_nowrite tuning" + ) + assert _num_loop_stages_write is not None, ( + "persistent_main requires _num_loop_stages_write tuning" + ) + assert _num_loop_stages_nowrite is not None, ( + "persistent_main requires _num_loop_stages_nowrite tuning" + ) + assert _heads_per_block is not None, "replay default tuning requires _heads_per_block" + assert _precompute_num_warps is not None, ( + "replay default tuning requires _precompute_num_warps" + ) + BLOCK_SIZE_M = ( + _block_size_m if _block_size_m is not None else _block_size_m_nowrite + ) + num_warps = _num_warps if _num_warps is not None else _num_warps_nowrite + precompute_num_warps = _precompute_num_warps + heads_per_block = int(_heads_per_block) assert heads_per_block > 0, "heads_per_block must be positive" heads_per_block = min(heads_per_block, heads_per_group) while heads_per_group % heads_per_block != 0 or heads_per_block & (heads_per_block - 1) != 0: heads_per_block -= 1 - if _precompute_num_warps is not None: - precompute_num_warps = _precompute_num_warps # Per-main knob resolution: each _*_{write,nowrite} arg, if not None, # overrides the corresponding shared value for ONE main launch only. From a2b12e80bbbc56cfeded4aef315412dbd5a5356f Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 5 Jun 2026 13:15:22 -0700 Subject: [PATCH 73/89] Fix empty persistent-main PDL bridge Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../_torch/modules/mamba/replay_selective_state_update.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index c716e0045b7d..6516190a2fa4 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -1971,6 +1971,12 @@ def _persistent_main_kernel( pid = tl.program_id(axis=0) total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads + if LAUNCH_WITH_PDL and total_work == 0: + # A kernel launched with PDL must wait for its upstream dependency + # before it can finish, even if this launch has no local work and does + # not launch its own dependents. + tl.extra.cuda.gdc_wait() + # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) # with pid_m varying fastest (M-tile cache locality on state load), then # slot, then head — mirrors the existing 3D grid's axis ordering From bbed7f23ba3c899d5fe67d6ab1685c4f48709b93 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 5 Jun 2026 13:28:53 -0700 Subject: [PATCH 74/89] Zero initialize Mamba base state caches Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 45fe13134350..2fe43d5cb436 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -487,13 +487,13 @@ def __init__( ssm_state_shape = (nheads, head_dim, d_state) # create mamba conv and ssm states - conv_states = torch.empty( + conv_states = torch.zeros( size=(num_local_layers, max_batch_size) + conv_state_shape, dtype=dtype, device=device, ) - ssm_states = torch.empty( + ssm_states = torch.zeros( size=(num_local_layers, max_batch_size) + ssm_state_shape, dtype=self.mamba_ssm_cache_dtype, device=device, @@ -2127,6 +2127,8 @@ def _setup_states(self) -> None: self.local_num_mamba_layers, num_blocks_in_pool ] + self.conv_state_shape) + self.all_ssm_states.zero_() + self.all_conv_states.zero_() def _setup_mtp_intermediate_states(self, spec_config, max_batch_size) -> None: From e623f4d8e22de4d4a861eb47a0b13870f9083c32 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 5 Jun 2026 13:32:23 -0700 Subject: [PATCH 75/89] Keep Mamba dummy replay slots nowrite Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../_torch/pyexecutor/mamba_cache_manager.py | 77 +++++++++++++++++-- 1 file changed, 70 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 2fe43d5cb436..3ba2cf0d43c8 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -621,6 +621,10 @@ def __init__( # mamba cache index, maps request_id -> state indices self.mamba_cache_index: Dict[int, int] = {} + self._dummy_request_ids: set[int] = set() + self._dummy_slot_mask = torch.zeros(max_batch_size, + dtype=torch.bool, + device=device) # Permanent slot shared by every CUDA-graph padding sentinel id # (CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, one per @@ -675,6 +679,7 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]): raise RuntimeError("run out of mamba cache blocks") block = self.mamba_cache_free_blocks.pop() self.mamba_cache_index[r] = block + self._dummy_slot_mask[block] = False if (isinstance(self.mamba_cache, self.SpeculativeState) and self._use_replay_state_update): self.mamba_cache.prev_num_accepted_tokens[block] = 0 @@ -715,28 +720,44 @@ def add_dummy_requests(self, request_ids: List[int], **kwargs): # slot and are freed individually. if not request_ids: return + self._dummy_request_ids.update(request_ids) for r in request_ids: if r in self.mamba_cache_index: + block = self.mamba_cache_index[r] + self._dummy_slot_mask[block] = True + if (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 continue if self._is_padding_sentinel(r): - self.mamba_cache_index[r] = self._padding_slot + block = self._padding_slot elif (r == ATTENTION_DP_DUMMY_REQUEST_ID and self._attention_dp_dummy_slot is not None): - self.mamba_cache_index[r] = self._attention_dp_dummy_slot + block = self._attention_dp_dummy_slot else: if len(self.mamba_cache_free_blocks) == 0: raise RuntimeError("run out of mamba cache blocks") block = self.mamba_cache_free_blocks.pop() - self.mamba_cache_index[r] = block + self.mamba_cache_index[r] = block + self._dummy_slot_mask[block] = True + if (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 def free_resources(self, request: LlmRequest): request_id = request.py_request_id if request_id not in self.mamba_cache_index: return + is_dummy = request_id in self._dummy_request_ids + self._dummy_request_ids.discard(request_id) block = self.mamba_cache_index.pop(request_id) # Reserved slots must not re-enter the real-request free pool. if block != self._padding_slot and \ block != self._attention_dp_dummy_slot: + if is_dummy: + self._dummy_slot_mask[block] = False self.mamba_cache_free_blocks.append(block) def get_state_indices(self, request_ids: List[int], @@ -900,10 +921,16 @@ def update_mamba_states(self, attn_metadata: "AttentionMetadata", wrote_checkpoint, accepted_tokens, prev_num_accepted_tokens + accepted_tokens) cache_buf_idx = self.mamba_cache.cache_buf_idx[state_indices_d] + is_dummy_slot = self._dummy_slot_mask[state_indices_d] + next_num_accepted_tokens = torch.where(is_dummy_slot, + prev_num_accepted_tokens, + next_num_accepted_tokens) self.mamba_cache.prev_num_accepted_tokens[state_indices_d] = \ next_num_accepted_tokens self.mamba_cache.cache_buf_idx[state_indices_d] = \ - torch.where(wrote_checkpoint, 1 - cache_buf_idx, cache_buf_idx) + torch.where(is_dummy_slot, cache_buf_idx, + torch.where(wrote_checkpoint, 1 - cache_buf_idx, + cache_buf_idx)) else: # Legacy: copy accepted SSM state from intermediate cache. ssm_states = self.mamba_cache.temporal @@ -1634,6 +1661,8 @@ def __init__( dtype=torch.long, device="cpu") self._request_id_to_state_index = {} + self._dummy_slot_mask = None + self._dummy_slot_mask_host = None self.kv_cache_config = kv_cache_config self.is_estimating_kv_cache = is_estimating_kv_cache @@ -1723,6 +1752,8 @@ def shutdown(self): self.prev_num_accepted_tokens = None self.cache_buf_idx = None self.mamba_ssm_rand_seed = None + self._dummy_slot_mask = None + self._dummy_slot_mask_host = None self.old_x = None self.old_B = None self.old_dt = None @@ -1898,10 +1929,15 @@ def update_mamba_states(self, next_num_accepted_tokens = torch.where( wrote_checkpoint, accepted, prev_num_accepted_tokens + accepted) cache_buf_idx = self.cache_buf_idx[slots] + is_dummy_slot = self._dummy_slot_mask[slots] + next_num_accepted_tokens = torch.where(is_dummy_slot, + prev_num_accepted_tokens, + next_num_accepted_tokens) self.prev_num_accepted_tokens[slots] = next_num_accepted_tokens - self.cache_buf_idx[slots] = torch.where(wrote_checkpoint, - 1 - cache_buf_idx, - cache_buf_idx) + self.cache_buf_idx[slots] = torch.where( + is_dummy_slot, cache_buf_idx, + torch.where(wrote_checkpoint, 1 - cache_buf_idx, + cache_buf_idx)) else: # Legacy: copy the accepted SSM state from the intermediate buffer. _promote_mamba_state_triton(self.all_ssm_states, @@ -2060,6 +2096,19 @@ def _setup_state_indices(self) -> None: self.cuda_state_indices.copy_(self._host_state_indices, non_blocking=True) + if self._dummy_slot_mask is not None: + self._dummy_slot_mask_host.zero_() + for i, req in enumerate(self.requests): + if req.is_dummy: + self._dummy_slot_mask_host[ + self._host_state_indices[i].item()] = True + self._dummy_slot_mask.copy_(self._dummy_slot_mask_host, + non_blocking=True) + if (self.prev_num_accepted_tokens is not None + and self.cache_buf_idx is not None): + self.prev_num_accepted_tokens.masked_fill_( + self._dummy_slot_mask, 0) + self.cache_buf_idx.masked_fill_(self._dummy_slot_mask, 0) # Build request_id → pool block offset mapping so that # get_state_indices can return indices in arbitrary request order. @@ -2204,6 +2253,14 @@ def _setup_replay_buffers(self, spec_config) -> None: # Without spec_config or replay we still keep the seed buffer # (above) so the non-MTP flashinfer SR path has a persistent # rand_seed source. + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + self._dummy_slot_mask = None + self._dummy_slot_mask_host = None return history_size = self.replay_history_size @@ -2218,6 +2275,12 @@ def _setup_replay_buffers(self, spec_config) -> None: self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) + self._dummy_slot_mask = torch.zeros(cache_size, + dtype=torch.bool, + device=device) + self._dummy_slot_mask_host = torch.zeros(cache_size, + dtype=torch.bool, + pin_memory=prefer_pinned()) self.old_x = torch.zeros(num_local_mamba_layers, cache_size, 2, From 68ab316cfaf792f8fa6bb94aa39f8ae9e90c39bc Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Sat, 6 Jun 2026 23:54:02 -0700 Subject: [PATCH 76/89] Add memory clobber to Mamba replay PDL wait Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 6516190a2fa4..3b23fcb8e6b1 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -58,6 +58,18 @@ # - *_window: combined history + step values over the full window dimension. +@triton.jit +def _gdc_wait_with_memory_clobber(): + tl.inline_asm_elementwise( + "griddepcontrol.wait; // dummy $0", + "=r,~{memory}", + [], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + # Lazy global allocator for Triton TMA tensor descriptors. Required by any # host- or device-built tensor_descriptor; without it Triton raises at first # launch. @@ -380,7 +392,7 @@ def _replay_precompute_impl( # --- Wait for upstream kernel (external PDL) before loading B and C --- # All dt processing above is independent of conv1d outputs. if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() + _gdc_wait_with_memory_clobber() # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- group_idx = first_head // nheads_ngroups_ratio @@ -718,7 +730,7 @@ def _rectangle_precompute_impl( # ---- gdc_wait: from here on we depend on conv1d's outputs ---- if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() + _gdc_wait_with_memory_clobber() # Conv1d outputs: B and C C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group @@ -1459,7 +1471,7 @@ def _persistent_main_impl( ).to(tl.float32) if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() + _gdc_wait_with_memory_clobber() C_tile = tl.load( C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, @@ -1694,7 +1706,7 @@ def _persistent_rectangle_impl( ).to(tl.float32) if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() + _gdc_wait_with_memory_clobber() C_tile = tl.load( C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, @@ -1975,7 +1987,7 @@ def _persistent_main_kernel( # A kernel launched with PDL must wait for its upstream dependency # before it can finish, even if this launch has no local work and does # not launch its own dependents. - tl.extra.cuda.gdc_wait() + _gdc_wait_with_memory_clobber() # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) # with pid_m varying fastest (M-tile cache locality on state load), then From f4f6768f75f13502594504d3834e84481577ee3b Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 9 Jun 2026 12:10:51 -0700 Subject: [PATCH 77/89] Hoist Mamba replay PDL wait and update tunings Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../_torch/modules/mamba/mamba2_metadata.py | 5 +- .../mamba/replay_selective_state_update.py | 143 ++++++++++-------- 2 files changed, 81 insertions(+), 67 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 09c37825990e..7d37fdb8ed9d 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -316,8 +316,11 @@ def prepare(self, attn_metadata: AttentionMetadata): and hasattr(kv_cache_manager, 'get_state_indices') and request_ids is not None): batch_request_ids = request_ids[:batch_size] + max_draft_len = getattr(kv_cache_manager, + "speculative_num_draft_tokens", 0) or 0 is_padding = [ - req_id == CUDA_GRAPH_DUMMY_REQUEST_ID + CUDA_GRAPH_DUMMY_REQUEST_ID - max_draft_len <= req_id <= + CUDA_GRAPH_DUMMY_REQUEST_ID for req_id in batch_request_ids ] indices = kv_cache_manager.get_state_indices( diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 3b23fcb8e6b1..5b0de36ff791 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -1147,7 +1147,7 @@ def _persistent_main_impl( BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, BLOCK_SIZE_WINDOW: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, + WAIT_FOR_PDL_PREDECESSOR: tl.constexpr, USE_RS_ROUNDING: tl.constexpr, PHILOX_ROUNDS: tl.constexpr, QUANT_MAX: tl.constexpr, @@ -1470,7 +1470,7 @@ def _persistent_main_impl( D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 ).to(tl.float32) - if LAUNCH_WITH_PDL: + if WAIT_FOR_PDL_PREDECESSOR: _gdc_wait_with_memory_clobber() C_tile = tl.load( @@ -1616,7 +1616,7 @@ def _persistent_rectangle_impl( BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - LAUNCH_WITH_PDL: tl.constexpr, + WAIT_FOR_PDL_PREDECESSOR: tl.constexpr, QUANT_MAX: tl.constexpr, USE_TMA_LOAD: tl.constexpr = False, ): @@ -1705,7 +1705,7 @@ def _persistent_rectangle_impl( other=0.0, ).to(tl.float32) - if LAUNCH_WITH_PDL: + if WAIT_FOR_PDL_PREDECESSOR: _gdc_wait_with_memory_clobber() C_tile = tl.load( @@ -1983,10 +1983,19 @@ def _persistent_main_kernel( pid = tl.program_id(axis=0) total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads - if LAUNCH_WITH_PDL and total_work == 0: - # A kernel launched with PDL must wait for its upstream dependency - # before it can finish, even if this launch has no local work and does - # not launch its own dependents. + if LAUNCH_WITH_PDL and ( + (total_work == 0 and pid == 0) or (NUM_LOOP_STAGES > 1 and pid < total_work) + ): + # This pre-loop path reads only replay partition metadata prepared + # before the PDL chain, not conv1d/precompute outputs. + # + # Empty PDL launches still need one waiting CTA, otherwise this kernel + # can retire before its upstream dependency and a later dependent launch + # can observe producer data too early. + # + # For loop-pipelined kernels, CTAs with work wait before the loop to + # work around a Triton PDL scheduling bug that can otherwise move + # producer-dependent loads ahead of the wait. _gdc_wait_with_memory_clobber() # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) @@ -2132,7 +2141,7 @@ def _persistent_main_kernel( BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, + LAUNCH_WITH_PDL and (NUM_LOOP_STAGES == 1), # WAIT_FOR_PDL_PREDECESSOR USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, @@ -2211,7 +2220,7 @@ def _persistent_main_kernel( BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, - LAUNCH_WITH_PDL, + LAUNCH_WITH_PDL and (NUM_LOOP_STAGES == 1), # WAIT_FOR_PDL_PREDECESSOR QUANT_MAX, USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle ) @@ -2300,7 +2309,7 @@ def _persistent_main_kernel( BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, - LAUNCH_WITH_PDL, + LAUNCH_WITH_PDL and (NUM_LOOP_STAGES == 1), # WAIT_FOR_PDL_PREDECESSOR USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, @@ -2348,6 +2357,8 @@ def _persistent_main_kernel( # # Source: emit_tuning_from_noise.py. Auto-generated from noise-cleaned per-cell search # winners (best of pd / pm by bucket_expected_renorm). Effective batch = raw_batch × 16. +# The search predates a Triton PDL scheduling bug. +# Most knobs are unchanged; large PDL-hoist regressions got spot retunes. # Missing dtype/SR combos fall back via the _resolve_tuning chain: # RN→SR for same dtype, then bf16/int16→fp16/SR and fp8→int8/SR. _DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { @@ -2377,7 +2388,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=1, score=6.66us (median of 4 noise runs; min=6.66us) + ), # raw_batch=1, score=9.66us (B200 PDL-hoist default 5x200) ( 32, "persistent_dynamic", @@ -2397,7 +2408,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=2, score=6.84us (median of 4 noise runs; min=6.83us) + ), # raw_batch=2, score=7.12us (B200 PDL-hoist default 5x200) ( 64, "persistent_dynamic", @@ -2417,7 +2428,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=4, score=7.13us (median of 4 noise runs; min=7.12us) + ), # raw_batch=4, score=7.05us (B200 PDL-hoist default 5x200) ( 128, "persistent_dynamic", @@ -2437,7 +2448,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=8, score=7.83us (manual retry of pre-noise winner) + ), # raw_batch=8, score=7.76us (B200 PDL-hoist default 5x200) ( 256, "persistent_dynamic", @@ -2457,7 +2468,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=16, score=9.04us (median of 4 noise runs; min=9.03us) + ), # raw_batch=16, score=9.27us (B200 PDL-hoist default 5x200) ( 512, "persistent_main", @@ -2483,7 +2494,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=32, score=12.51us (median of 4 noise runs; min=12.47us) + ), # raw_batch=32, score=12.71us (B200 PDL-hoist default 5x200) ( 1024, "persistent_main", @@ -2509,7 +2520,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=64, score=17.41us (median of 4 noise runs; min=17.41us) + ), # raw_batch=64, score=17.51us (B200 PDL-hoist default 5x200) ( 2048, "persistent_main", @@ -2535,7 +2546,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=128, score=25.21us (median of 4 noise runs; min=25.16us) + ), # raw_batch=128, score=26.39us (B200 PDL-hoist default 5x200) ( 4096, "persistent_main", @@ -2561,7 +2572,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=256, score=42.06us (median of 4 noise runs; min=42.01us) + ), # raw_batch=256, score=43.57us (B200 PDL-hoist default 5x200) ( 8192, "persistent_main", @@ -2587,7 +2598,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=72.15us (median of 4 noise runs; min=72.11us) + ), # raw_batch=512, score=73.80us (B200 PDL-hoist default 5x200) ( 16384, "persistent_main", @@ -2613,7 +2624,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=1024, score=131.77us (median of 4 noise runs; min=131.63us) + ), # raw_batch=1024, score=134.03us (B200 PDL-hoist default 5x200) ], ("int8", "SR"): [ ( @@ -2635,7 +2646,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=6.42us (median of 4 noise runs; min=6.39us) + ), # raw_batch=1, score=8.60us (B200 PDL-hoist default 5x200) ( 32, "persistent_dynamic", @@ -2655,7 +2666,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=2, score=7.13us (median of 4 noise runs; min=7.09us) + ), # raw_batch=2, score=7.44us (B200 PDL-hoist default 5x200) ( 64, "persistent_dynamic", @@ -2675,7 +2686,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=4, score=7.52us (median of 4 noise runs; min=7.50us) + ), # raw_batch=4, score=7.97us (B200 PDL-hoist default 5x200) ( 128, "persistent_dynamic", @@ -2695,7 +2706,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=8, score=8.22us (median of 4 noise runs; min=8.21us) + ), # raw_batch=8, score=8.15us (B200 PDL-hoist default 5x200) ( 256, "persistent_dynamic", @@ -2715,7 +2726,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=16, score=10.06us (median of 8 noise runs; min=9.97us) + ), # raw_batch=16, score=9.88us (B200 PDL-hoist default 5x200) ( 512, "persistent_main", @@ -2741,7 +2752,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=32, score=12.91us (median of 4 noise runs; min=12.90us) + ), # raw_batch=32, score=12.99us (B200 PDL-hoist default 5x200) ( 1024, "persistent_main", @@ -2767,7 +2778,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=64, score=18.05us (median of 4 noise runs; min=18.03us) + ), # raw_batch=64, score=22.26us (B200 PDL-hoist default 5x200) ( 2048, "persistent_main", @@ -2793,7 +2804,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=128, score=27.66us (median of 4 noise runs; min=27.56us) + ), # raw_batch=128, score=34.03us (B200 PDL-hoist default 5x200) ( 4096, "persistent_main", @@ -2819,7 +2830,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=256, score=45.13us (median of 4 noise runs; min=45.08us) + ), # raw_batch=256, score=49.64us (B200 PDL-hoist default 5x200) ( 8192, "persistent_main", @@ -2845,7 +2856,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=77.87us (median of 4 noise runs; min=77.78us) + ), # raw_batch=512, score=82.99us (B200 PDL-hoist default 5x200) ( 16384, "persistent_main", @@ -2856,10 +2867,10 @@ def _persistent_main_kernel( "_cta_per_sm_write": 3, "_flatten": False, "_heads_per_block": 8, - "_num_loop_stages_nowrite": 4, - "_num_loop_stages_write": 3, - "_num_stages_nowrite": 3, - "_num_stages_write": 1, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 4, + "_num_stages_write": 4, "_num_warps_nowrite": 1, "_num_warps_write": 4, "_precompute_num_warps": 1, @@ -2871,7 +2882,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=1024, score=142.83us (median of 4 noise runs; min=142.79us) + ), # raw_batch=1024, score=149.25us (B200 PDL-hoist default 5x200) ], ("fp8", "SR"): [ ( @@ -2893,7 +2904,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=6.30us (median of 4 noise runs; min=6.28us) + ), # raw_batch=1, score=8.67us (B200 PDL-hoist default 5x200) ( 32, "persistent_dynamic", @@ -2913,7 +2924,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=2, score=6.80us (median of 4 noise runs; min=6.79us) + ), # raw_batch=2, score=7.05us (B200 PDL-hoist default 5x200) ( 64, "persistent_dynamic", @@ -2933,7 +2944,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=4, score=7.00us (median of 4 noise runs; min=6.98us) + ), # raw_batch=4, score=7.24us (B200 PDL-hoist default 5x200) ( 128, "persistent_dynamic", @@ -2953,7 +2964,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=8, score=7.67us (median of 4 noise runs; min=7.59us) + ), # raw_batch=8, score=7.79us (B200 PDL-hoist default 5x200) ( 256, "persistent_dynamic", @@ -2973,7 +2984,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=16, score=8.58us (median of 4 noise runs; min=8.54us) + ), # raw_batch=16, score=8.77us (B200 PDL-hoist default 5x200) ( 512, "persistent_dynamic", @@ -2993,7 +3004,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=32, score=10.60us (median of 4 noise runs; min=10.54us) + ), # raw_batch=32, score=10.71us (B200 PDL-hoist default 5x200) ( 1024, "persistent_main", @@ -3019,7 +3030,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=64, score=15.79us (median of 4 noise runs; min=15.77us) + ), # raw_batch=64, score=16.00us (B200 PDL-hoist default 5x200) ( 2048, "persistent_main", @@ -3045,7 +3056,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=128, score=23.65us (median of 4 noise runs; min=23.65us) + ), # raw_batch=128, score=26.19us (B200 PDL-hoist default 5x200) ( 4096, "persistent_main", @@ -3071,7 +3082,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=256, score=38.47us (median of 4 noise runs; min=38.42us) + ), # raw_batch=256, score=41.16us (B200 PDL-hoist default 5x200) ( 8192, "persistent_main", @@ -3097,7 +3108,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=68.32us (median of 4 noise runs; min=68.28us) + ), # raw_batch=512, score=77.62us (B200 PDL-hoist default 5x200) ( 16384, "persistent_main", @@ -3108,10 +3119,10 @@ def _persistent_main_kernel( "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 8, - "_num_loop_stages_nowrite": 3, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 4, - "_num_stages_write": 2, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 4, "_num_warps_nowrite": 1, "_num_warps_write": 1, "_precompute_num_warps": 1, @@ -3123,7 +3134,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=1024, score=123.51us (median of 4 noise runs; min=123.40us) + ), # raw_batch=1024, score=131.40us (B200 PDL-hoist default 5x200) ], ("fp32", "RN"): [ ( @@ -3145,7 +3156,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=6.01us (median of 4 noise runs; min=6.00us) + ), # raw_batch=1, score=8.58us (B200 PDL-hoist default 5x200) ( 32, "persistent_dynamic", @@ -3165,7 +3176,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=2, score=6.64us (median of 4 noise runs; min=6.62us) + ), # raw_batch=2, score=6.83us (B200 PDL-hoist default 5x200) ( 64, "persistent_dynamic", @@ -3185,7 +3196,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=4, score=7.17us (manual retry of pre-noise winner) + ), # raw_batch=4, score=7.19us (B200 PDL-hoist default 5x200) ( 128, "persistent_dynamic", @@ -3205,7 +3216,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=8, score=7.96us (median of 4 noise runs; min=7.90us) + ), # raw_batch=8, score=8.14us (B200 PDL-hoist default 5x200) ( 256, "persistent_dynamic", @@ -3225,7 +3236,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=16, score=9.55us (median of 4 noise runs; min=9.50us) + ), # raw_batch=16, score=9.66us (B200 PDL-hoist default 5x200) ( 512, "persistent_dynamic", @@ -3245,7 +3256,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=32, score=13.28us (median of 4 noise runs; min=13.23us) + ), # raw_batch=32, score=13.59us (B200 PDL-hoist default 5x200) ( 1024, "persistent_main", @@ -3271,7 +3282,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=64, score=19.12us (median of 4 noise runs; min=19.10us) + ), # raw_batch=64, score=21.92us (B200 PDL-hoist default 5x200) ( 2048, "persistent_main", @@ -3282,7 +3293,7 @@ def _persistent_main_kernel( "_cta_per_sm_write": 9, "_flatten": False, "_heads_per_block": 16, - "_num_loop_stages_nowrite": 2, + "_num_loop_stages_nowrite": 3, "_num_loop_stages_write": 1, "_num_stages_nowrite": 1, "_num_stages_write": 1, @@ -3297,7 +3308,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=128, score=29.39us (median of 4 noise runs; min=29.34us) + ), # raw_batch=128, score=35.29us (B200 PDL-hoist default 5x200) ( 4096, "persistent_main", @@ -3323,7 +3334,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=256, score=49.36us (median of 4 noise runs; min=49.31us) + ), # raw_batch=256, score=54.98us (B200 PDL-hoist default 5x200) ( 8192, "persistent_main", @@ -3349,7 +3360,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=87.36us (median of 4 noise runs; min=87.31us) + ), # raw_batch=512, score=93.58us (B200 PDL-hoist default 5x200) ( 16384, "persistent_main", @@ -3360,9 +3371,9 @@ def _persistent_main_kernel( "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 8, - "_num_loop_stages_nowrite": 2, + "_num_loop_stages_nowrite": 5, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 4, + "_num_stages_nowrite": 1, "_num_stages_write": 2, "_num_warps_nowrite": 2, "_num_warps_write": 1, @@ -3375,7 +3386,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=1024, score=168.98us (median of 4 noise runs; min=168.83us) + ), # raw_batch=1024, score=180.18us (B200 PDL-hoist default 5x200) ], } _PD_TO_PM_SPLIT_MAP = { # pd unsplit knob → (pm_write_knob, pm_nowrite_knob) From dcf8e5a2dfb6239288ac06e96d440d3b10e4577d Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 10 Jun 2026 16:54:56 -0700 Subject: [PATCH 78/89] Require explicit Mamba replay cache metadata Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 80 +++++++------------ ...benchmark_replay_selective_state_update.py | 18 +---- .../test_replay_selective_state_update.py | 51 +++++------- 3 files changed, 51 insertions(+), 98 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 5b0de36ff791..055ee51ff4f1 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -255,7 +255,6 @@ def _replay_precompute_impl( # Meta-parameters DT_SOFTPLUS: tl.constexpr, HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, @@ -273,11 +272,7 @@ def _replay_precompute_impl( pid_hg = tl.program_id(axis=1) # head-group index first_head = pid_hg * HEADS_PER_BLOCK - # Resolve cache index for writes - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - else: - cache_batch_idx = pid_b.to(tl.int64) + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) # --- Cache write semantics --- # cache_buf_idx names this step's "active" buffer — the one with the @@ -520,7 +515,6 @@ def _rectangle_precompute_impl( # Meta-parameters DT_SOFTPLUS: tl.constexpr, HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, @@ -532,10 +526,7 @@ def _rectangle_precompute_impl( pid_hg = tl.program_id(axis=1) first_head = pid_hg * HEADS_PER_BLOCK - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - else: - cache_batch_idx = pid_b.to(tl.int64) + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) @@ -887,7 +878,6 @@ def _dynamic_precompute_kernel( # Meta-parameters DT_SOFTPLUS: tl.constexpr, HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, @@ -903,10 +893,7 @@ def _dynamic_precompute_kernel( tl.extra.cuda.gdc_launch_dependents() pid_b = tl.program_id(axis=0) - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - else: - cache_batch_idx = pid_b.to(tl.int64) + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH @@ -965,7 +952,6 @@ def _dynamic_precompute_kernel( stride_old_dA_cumsum_T, DT_SOFTPLUS, HAS_DT_BIAS, - HAS_CACHE_BATCH_INDICES, BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, LAUNCH_WITH_PDL, @@ -1026,7 +1012,6 @@ def _dynamic_precompute_kernel( stride_old_dA_cumsum_T, DT_SOFTPLUS, HAS_DT_BIAS, - HAS_CACHE_BATCH_INDICES, BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, @@ -1143,7 +1128,6 @@ def _persistent_main_impl( BLOCK_SIZE_M: tl.constexpr, HAS_D: tl.constexpr, HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, BLOCK_SIZE_WINDOW: tl.constexpr, @@ -1612,7 +1596,6 @@ def _persistent_rectangle_impl( BLOCK_SIZE_M: tl.constexpr, HAS_D: tl.constexpr, HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, @@ -1780,9 +1763,6 @@ def _persistent_rectangle_impl( # Heuristics mirror those of the replay main kernel. @triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics( - {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} -) @triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) @triton.heuristics( @@ -1920,7 +1900,6 @@ def _persistent_main_kernel( BLOCK_SIZE_M: tl.constexpr, HAS_D: tl.constexpr, HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, BLOCK_SIZE_WINDOW: tl.constexpr, @@ -1952,7 +1931,6 @@ def _persistent_main_kernel( USE_TMA_LOAD_WRITE: tl.constexpr = False, USE_TMA_LOAD_NOWRITE: tl.constexpr = False, USE_TMA_STORE: tl.constexpr = False, - USE_REPLAY_CACHE_SLOT: tl.constexpr = True, ): # PDL signal: fire once at kernel entry (not per work unit). if LAUNCH_DEPENDENT_KERNELS: @@ -2015,26 +1993,16 @@ def _persistent_main_kernel( pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) work_item_idx = pid_b_local + slot_lo if IS_DYNAMIC: - if HAS_CACHE_BATCH_INDICES: - pid_b = work_item_idx - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - else: - pid_b = work_item_idx - cache_batch_idx = pid_b.to(tl.int64) + pid_b = work_item_idx + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) pnat = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) else: work_item_base = replay_work_items_ptr + work_item_idx * _REPLAY_WORK_ITEM_WIDTH - if USE_REPLAY_CACHE_SLOT: - pid_b = tl.load(work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH) - cache_batch_idx = tl.load(work_item_base + _REPLAY_WORK_CACHE_SLOT).to(tl.int64) - pnat = tl.load(work_item_base + _REPLAY_WORK_PNAT) - active_buf = tl.load(work_item_base + _REPLAY_WORK_CACHE_BUF_IDX).to(tl.int32) - else: - pid_b = tl.load(work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH) - cache_batch_idx = work_item_idx.to(tl.int64) - pnat = tl.load(work_item_base + _REPLAY_WORK_PNAT) - active_buf = tl.load(work_item_base + _REPLAY_WORK_CACHE_BUF_IDX).to(tl.int32) + pid_b = tl.load(work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH) + cache_batch_idx = tl.load(work_item_base + _REPLAY_WORK_CACHE_SLOT).to(tl.int64) + pnat = tl.load(work_item_base + _REPLAY_WORK_PNAT) + active_buf = tl.load(work_item_base + _REPLAY_WORK_CACHE_BUF_IDX).to(tl.int32) # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle # impl. `replay_work_items` carries the cache slot, PNAT and active # buffer for persistent_main; persistent_dynamic resolves those once @@ -2137,7 +2105,6 @@ def _persistent_main_kernel( BLOCK_SIZE_M, HAS_D, HAS_Z, - HAS_CACHE_BATCH_INDICES, BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, @@ -2216,7 +2183,6 @@ def _persistent_main_kernel( BLOCK_SIZE_M, HAS_D, HAS_Z, - HAS_CACHE_BATCH_INDICES, BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_K, @@ -2305,7 +2271,6 @@ def _persistent_main_kernel( BLOCK_SIZE_M, HAS_D, HAS_Z, - HAS_CACHE_BATCH_INDICES, BLOCK_SIZE_DSTATE, BLOCK_SIZE_T, BLOCK_SIZE_WINDOW, @@ -3494,11 +3459,11 @@ def replay_selective_state_update( # pd ignores it (per-slot runtime PNAT check). n_writes: torch.Tensor, replay_work_items: torch.Tensor, + state_batch_indices: torch.Tensor, D: torch.Tensor | None = None, z: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, dt_softplus: bool = False, - state_batch_indices: torch.Tensor | None = None, rand_seed: torch.Tensor | None = None, philox_rounds: int = 10, state_scales: torch.Tensor | None = None, @@ -3540,7 +3505,6 @@ def replay_selective_state_update( _use_tma_replay_write_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=True _use_tma_replay_write_store: bool | None = None, # SSM state store when WRITE_CHECKPOINT=True _use_tma_replay_nowrite_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=False - _use_replay_cache_slot: bool = True, # Persistent-mode tuning kwargs (consulted for both pd and pm; pd uses # _cta_per_sm / _num_loop_stages, pm uses the _write/_nowrite splits): # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally @@ -3605,7 +3569,7 @@ def replay_selective_state_update( D: (nheads, dim) optional feed-through parameter. z: (batch, T, nheads, dim) optional silu gate. dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). - state_batch_indices: (batch,) optional cache slot mapping. + state_batch_indices: (batch,) int32 cache slot mapping. rand_seed: optional (cache_size,) int64 CUDA tensor of per-cache-slot Philox PRNG seeds. The caller bumps this tensor in-place for each replay invocation so CUDA graph replay still gets fresh draws. The @@ -3726,6 +3690,7 @@ def replay_selective_state_update( cache_size, nheads, dim, dstate = state.shape batch, T, _, _ = x.shape + device = x.device ngroups = B.shape[2] assert nheads % ngroups == 0 @@ -3921,6 +3886,21 @@ def replay_selective_state_update( assert prev_num_accepted_tokens.dtype == torch.int32, ( f"prev_num_accepted_tokens must be int32, got {prev_num_accepted_tokens.dtype}" ) + assert isinstance(state_batch_indices, torch.Tensor), ( + f"state_batch_indices must be a torch.Tensor, " + f"got {type(state_batch_indices).__name__}" + ) + assert state_batch_indices.device == device, ( + f"state_batch_indices must be on device {device}, got {state_batch_indices.device}" + ) + assert state_batch_indices.dtype == torch.int32, ( + f"state_batch_indices must be int32, got {state_batch_indices.dtype}" + ) + assert state_batch_indices.shape == (batch,), ( + f"state_batch_indices must have shape (batch={batch},), " + f"got {tuple(state_batch_indices.shape)}" + ) + assert state_batch_indices.is_contiguous(), "state_batch_indices must be contiguous" if rand_seed is not None: assert rand_seed.dtype == torch.int64, ( f"rand_seed dtype must be int64, got {rand_seed.dtype}" @@ -3943,7 +3923,6 @@ def replay_selective_state_update( ) assert tie_hdim - device = x.device BLOCK_SIZE_T = max(triton.next_power_of_2(T), MIN_REPLAY_TILE_SIZE) # Rectangle window bound = max_window. Computed unconditionally # so the launch sites can refer to it; only used on the rectangle path. @@ -4029,8 +4008,6 @@ def replay_selective_state_update( _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages ) - HAS_CACHE_BATCH_INDICES = state_batch_indices is not None - assert nheads % heads_per_block == 0, ( f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" ) @@ -4182,7 +4159,6 @@ def launch_dynamic_precompute(rectangle: bool): old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), dt_softplus, - HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, LAUNCH_WITH_PDL=launch_with_pdl, LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, HEADS_PER_BLOCK=heads_per_block, @@ -4346,7 +4322,6 @@ def launch_persistent_main( and not write_checkpoint ), USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), - USE_REPLAY_CACHE_SLOT=bool(_use_replay_cache_slot), num_warps=launch_num_warps, **({"num_stages": launch_num_stages} if launch_num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), @@ -4475,7 +4450,6 @@ def launch_persistent_dynamic_main( _use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load ), USE_TMA_STORE=bool(_use_tma_replay_write_store), - USE_REPLAY_CACHE_SLOT=bool(_use_replay_cache_slot), num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), **({"num_ctas": _num_ctas} if _num_ctas else {}), diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 5cc7964ab48e..0a61eefb4847 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -3071,7 +3071,7 @@ def _run_pr3324_baseline(): D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=(state_batch_indices if args.use_cache_slot else None), + state_batch_indices=state_batch_indices, state_scale=state_scales_work, rand_seed=rand_seed, philox_rounds=args.philox_rounds, @@ -3365,7 +3365,6 @@ def _run_incr( extra_kwargs["_flatten"] = bool(flatten) if warp_specialize is not None: extra_kwargs["_warp_specialize"] = bool(warp_specialize) - extra_kwargs["_use_replay_cache_slot"] = bool(args.use_cache_slot) replay_selective_state_update( state_work, @@ -3384,7 +3383,7 @@ def _run_incr( D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, + state_batch_indices=state_batch_indices, rand_seed=rand_seed, philox_rounds=args.philox_rounds, use_internal_pdl=args.internal_pdl, @@ -3477,8 +3476,6 @@ def _emit_split(name_w, name_nw, val_w, val_nw): hsort_in_cell_list = "HSORT" in getattr(args, "_cell_list_keys", ()) if hardcode_sort or len(hsort_list_for_tags) > 1 or hsort_in_cell_list: parts.append(f"HSORT={1 if hardcode_sort else 0}") - if not args.use_cache_slot: - parts.append("CSLOT=0") sweep_suffix = (" " + ",".join(parts)) if parts else "" sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") @@ -3550,7 +3547,6 @@ def _emit_split(name_w, name_nw, val_w, val_nw): bool(rectangle_for_nowrite), bool(nowrite_first), bool(hardcode_sort), - bool(args.use_cache_slot), scenario_pre_iter is not None, expected_K, ) @@ -4775,14 +4771,6 @@ def _parse_args() -> argparse.Namespace: help="External PDL: conv1d launches dependents, precompute waits. " "Only relevant with --with-conv1d. --no-external-pdl disables.", ) - parser.add_argument( - "--use-cache-slot", - action=argparse.BooleanOptionalAction, - default=True, - help="Use the cache-slot field from replay_work_items in persistent_main. " - "--no-use-cache-slot keeps the old identity-cache-slot shortcut for " - "diagnostic comparisons only and requires --hardcode-sort 1.", - ) parser.add_argument( "--heads-per-block", type=str, @@ -5115,8 +5103,6 @@ def _round_iters_to_group(name, val): hsort_list.append(v == "1") if not hsort_list: hsort_list = [False] - if not args.use_cache_slot and not all(hsort_list): - parser.error("--no-use-cache-slot is a diagnostic shortcut and requires --hardcode-sort 1") args.hardcode_sort_list = hsort_list # mode=None means "let the wrapper resolve from _DEFAULT_TUNING". Same diff --git a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py index 41380a2e74bc..6f68a9e8c121 100644 --- a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py @@ -50,10 +50,7 @@ def _make_replay_work_items( ): """Build the replay metadata consumed by persistent_main.""" position_in_decode_batch = torch.arange(batch, device=device, dtype=torch.int32) - if state_batch_indices is not None: - cache_slot = state_batch_indices[:batch].to(torch.int32) - else: - cache_slot = position_in_decode_batch + cache_slot = state_batch_indices[:batch].to(torch.int32) cache_slot_long = cache_slot.to(torch.long) pnat = prev_tokens[cache_slot_long].to(torch.int32) active_cache_buf_idx = cache_buf_idx[cache_slot_long].to(torch.int32) @@ -207,7 +204,7 @@ def test_replay_selective_state_update( state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) else: cache_size = batch - state_batch_indices = None + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) torch.manual_seed(42) @@ -356,7 +353,7 @@ def test_replay_selective_state_update( D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=(state_batch_indices if paged_cache else None), + state_batch_indices=state_batch_indices, out=ref_out, ) @@ -721,6 +718,7 @@ def test_replay_selective_state_update_scenarios( batch = 4 device = "cuda" dtype = torch.bfloat16 + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) pnat_means_write = (pnat_per_slot + T > max_window).tolist() @@ -805,7 +803,7 @@ def test_replay_selective_state_update_scenarios( D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, + state_batch_indices=state_batch_indices, out=ref_out, ) @@ -817,7 +815,7 @@ def test_replay_selective_state_update_scenarios( T, max_window, batch, - None, + state_batch_indices, device, explicit_order=explicit_order, ) @@ -840,7 +838,7 @@ def test_replay_selective_state_update_scenarios( D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, + state_batch_indices=state_batch_indices, mode=mode, rectangle_for_nowrite=rectangle_for_nowrite, ) @@ -914,6 +912,7 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( batch = 4 device = "cuda" dtype = torch.bfloat16 + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) pnat_means_write = (pnat_per_slot + T > max_window).tolist() @@ -1017,7 +1016,7 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, + state_batch_indices=state_batch_indices, out=ref_out, ) @@ -1042,7 +1041,7 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, + state_batch_indices=state_batch_indices, mode=mode, rectangle_for_nowrite=rectangle_for_nowrite, ) @@ -1118,7 +1117,7 @@ def test_replay_selective_state_update_philox( state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) else: cache_size = batch - state_batch_indices = None + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) torch.manual_seed(42) @@ -1342,6 +1341,7 @@ def test_philox_rounding_unbiased(state_dtype): C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) prev_tokens = torch.full((batch,), T, device=device, dtype=torch.int32) + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) # max_window = old_x.shape[2] = T (after dbuf at axis 1) _n_writes_unb, _replay_work_items_unb = _make_replay_work_items( @@ -1350,7 +1350,7 @@ def test_philox_rounding_unbiased(state_dtype): T, T, batch, - None, + state_batch_indices, device, ) common_kwargs = dict( @@ -1362,6 +1362,7 @@ def test_philox_rounding_unbiased(state_dtype): D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=state_batch_indices, n_writes=_n_writes_unb, replay_work_items=_replay_work_items_unb, ) @@ -1614,6 +1615,7 @@ def test_replay_heads_per_block( ref_state_f32[slot] = states_buffer_f32[slot, pnat_list[slot] - 1] ref_state_after_replay = ref_state_f32.clone() ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) selective_state_update( ref_state_f32, x2, @@ -1624,7 +1626,7 @@ def test_replay_heads_per_block( D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, + state_batch_indices=state_batch_indices, out=ref_out, ) @@ -1651,7 +1653,7 @@ def test_replay_heads_per_block( T, max_window, batch, - None, + state_batch_indices, device, ) @@ -1674,7 +1676,7 @@ def test_replay_heads_per_block( D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, + state_batch_indices=state_batch_indices, mode=mode, rectangle_for_nowrite=rectangle_nowrite, _heads_per_block=heads_per_block, @@ -1880,7 +1882,7 @@ def test_replay_heads_per_block_multistep( state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) else: cache_size = batch - state_batch_indices = None + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) all_x = [] all_dt = [] @@ -1925,9 +1927,7 @@ def test_replay_heads_per_block_multistep( acc = accepted_per_slot[s_local] if acc == 0: continue - c_idx = ( - state_batch_indices[s_local].item() if state_batch_indices is not None else s_local - ) + c_idx = state_batch_indices[s_local].item() s_state = ref_state[c_idx : c_idx + 1].clone() s_x = all_x[step][s_local : s_local + 1, :acc].contiguous() s_dt = all_dt[step][s_local : s_local + 1, :acc].contiguous() @@ -1965,10 +1965,7 @@ def test_replay_heads_per_block_multistep( for step in range(n_steps): prev_tokens = torch.zeros(cache_size, device=device, dtype=torch.int32) - if state_batch_indices is not None: - prev_tokens[state_batch_indices.long()] = pnat_active - else: - prev_tokens[:] = pnat_active + prev_tokens[state_batch_indices.long()] = pnat_active test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) n_writes_t, replay_work_items_t = _make_replay_work_items( prev_tokens, @@ -2015,11 +2012,7 @@ def test_replay_heads_per_block_multistep( pnat_active + accepted_tensor, ) pnat_active = new_pnat_active - cache_active_idx = ( - state_batch_indices.long() - if state_batch_indices is not None - else torch.arange(batch, device=device) - ) + cache_active_idx = state_batch_indices.long() write_slots = cache_active_idx[write_mask] cache_buf_idx[write_slots] = 1 - cache_buf_idx[write_slots] From 63f3033f58b754a8074bb499ce09e725db146771 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 10 Jun 2026 20:12:29 -0700 Subject: [PATCH 79/89] Clean up Mamba checkpoint replay plumbing Remove dead replay tuning knobs and stale comments, tighten replay metadata invariants for AutoDeploy, and import the noise-cleaned PDL retune table. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../custom_ops/attention_interface.py | 6 +- .../mamba/flashinfer_backend_mamba.py | 63 ++- .../_torch/auto_deploy/shim/interface.py | 3 + .../_torch/modules/mamba/mamba2_metadata.py | 59 +- .../_torch/modules/mamba/mamba2_mixer.py | 31 +- .../mamba/replay_selective_state_update.py | 510 ++++++++---------- tensorrt_llm/_torch/pyexecutor/_util.py | 80 ++- .../_torch/pyexecutor/mamba_cache_manager.py | 38 +- .../_torch/pyexecutor/resource_manager.py | 4 +- .../executor/test_mamba_cache_manager.py | 41 ++ .../modules/mamba/test_mamba_ssm_rand_seed.py | 21 +- .../test_replay_selective_state_update.py | 103 +--- .../mamba/test_flashinfer_mamba_cached_op.py | 34 +- .../auto_deploy/singlegpu/shim/test_engine.py | 74 +++ 14 files changed, 605 insertions(+), 462 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index df5bab6e6a60..d3f409b6745a 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -2144,7 +2144,7 @@ def __eq__(self, other) -> bool: class ReplayOldBHandler(StateResourceHandler): """Per-layer old_B cache for the replay SSM kernel (double-buffered, bf16). - Shape: (max_batch, 2, T, n_groups, d_state) — T from manager. + Shape: (max_batch, 2, replay_history_size, n_groups, d_state). Routes to MambaHybridCacheManager via get_replay_old_B(layer_idx). """ @@ -2172,7 +2172,7 @@ def __eq__(self, other) -> bool: class ReplayOldDtHandler(StateResourceHandler): """Per-layer old_dt cache for the replay SSM kernel (double-buffered, fp32). - Shape: (max_batch, 2, num_heads, T) — T from manager. + Shape: (max_batch, 2, num_heads, replay_history_size). Routes to MambaHybridCacheManager via get_replay_old_dt(layer_idx). """ @@ -2194,7 +2194,7 @@ def __eq__(self, other) -> bool: class ReplayOldDAcumsumHandler(StateResourceHandler): """Per-layer old_dA_cumsum cache for the replay SSM kernel (double-buffered, fp32). - Shape: (max_batch, 2, num_heads, T) — T from manager. + Shape: (max_batch, 2, num_heads, replay_history_size). Routes to MambaHybridCacheManager via get_replay_old_dA_cumsum(layer_idx). """ diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index 86e2c01183d8..827cc5b3202c 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional +from typing import List, Optional, cast import torch from flashinfer.mamba import selective_state_update as _flashinfer_ssm_update @@ -97,12 +97,18 @@ def _flashinfer_cached_ssm( intermediate_ssm_state_cache: Optional[ torch.Tensor ], # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state]; None in replay mode - replay_old_x: Optional[torch.Tensor], # [max_batch, T, nheads, head_dim]; None in non-replay - replay_old_b: Optional[torch.Tensor], # [max_batch, 2, T, ngroups, dstate]; None in non-replay - replay_old_dt: Optional[torch.Tensor], # [max_batch, 2, nheads, T] fp32; None in non-replay + replay_old_x: Optional[ + torch.Tensor + ], # [max_batch, 2, history, nheads, head_dim]; None in non-replay + replay_old_b: Optional[ + torch.Tensor + ], # [max_batch, 2, history, ngroups, dstate]; None in non-replay + replay_old_dt: Optional[ + torch.Tensor + ], # [max_batch, 2, nheads, history] fp32; None in non-replay replay_old_da_cumsum: Optional[ torch.Tensor - ], # [max_batch, 2, nheads, T] fp32; None in non-replay + ], # [max_batch, 2, nheads, history] fp32; None in non-replay replay_cache_buf_idx: Optional[torch.Tensor], # [max_batch] int32; None in non-replay replay_prev_num_accepted: Optional[torch.Tensor], # [max_batch] int32; None in non-replay replay_work_items: Optional[torch.Tensor], # [max_batch, 4] int32; None in non-replay @@ -202,6 +208,34 @@ def _flashinfer_cached_ssm( use_replay = batch_info.is_use_replay() if use_replay: + missing_replay_tensors = [ + name + for name, tensor in ( + ("replay_old_x", replay_old_x), + ("replay_old_b", replay_old_b), + ("replay_old_dt", replay_old_dt), + ("replay_old_da_cumsum", replay_old_da_cumsum), + ("replay_cache_buf_idx", replay_cache_buf_idx), + ("replay_prev_num_accepted", replay_prev_num_accepted), + ("replay_work_items", replay_work_items), + ("replay_n_writes", replay_n_writes), + ) + if tensor is None + ] + if missing_replay_tensors: + raise RuntimeError( + "flashinfer_cached_ssm replay path missing required tensors: " + f"{', '.join(missing_replay_tensors)}" + ) + replay_old_x = cast(torch.Tensor, replay_old_x) + replay_old_b = cast(torch.Tensor, replay_old_b) + replay_old_dt = cast(torch.Tensor, replay_old_dt) + replay_old_da_cumsum = cast(torch.Tensor, replay_old_da_cumsum) + replay_cache_buf_idx = cast(torch.Tensor, replay_cache_buf_idx) + replay_prev_num_accepted = cast(torch.Tensor, replay_prev_num_accepted) + replay_work_items = cast(torch.Tensor, replay_work_items) + replay_n_writes = cast(torch.Tensor, replay_n_writes) + # Replay path: fast-forward SSM state via tl.dot on cached values. # State is updated in-place; no disable_state_update needed. # x_extend/B_extend/C_extend are non-contiguous views from the CUDA graph's @@ -231,6 +265,11 @@ def _flashinfer_cached_ssm( launch_with_pdl=True, # PDL chain: triton_causal_conv extend → precompute → main ) else: + if intermediate_ssm_state_cache is None: + raise RuntimeError( + "flashinfer_cached_ssm non-replay extend branch requires " + "intermediate_ssm_state_cache" + ) if intermediate_ssm_state_cache.size(1) < tokens_per_extend: raise RuntimeError( "flashinfer_cached_ssm: intermediate_ssm_state_cache is too small " @@ -348,12 +387,18 @@ def _flashinfer_cached_ssm_fake( intermediate_ssm_state_cache: Optional[ torch.Tensor ], # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state]; None in replay mode - replay_old_x: Optional[torch.Tensor], # [max_batch, T, nheads, head_dim]; None in non-replay - replay_old_b: Optional[torch.Tensor], # [max_batch, 2, T, ngroups, dstate]; None in non-replay - replay_old_dt: Optional[torch.Tensor], # [max_batch, 2, nheads, T] fp32; None in non-replay + replay_old_x: Optional[ + torch.Tensor + ], # [max_batch, 2, history, nheads, head_dim]; None in non-replay + replay_old_b: Optional[ + torch.Tensor + ], # [max_batch, 2, history, ngroups, dstate]; None in non-replay + replay_old_dt: Optional[ + torch.Tensor + ], # [max_batch, 2, nheads, history] fp32; None in non-replay replay_old_da_cumsum: Optional[ torch.Tensor - ], # [max_batch, 2, nheads, T] fp32; None in non-replay + ], # [max_batch, 2, nheads, history] fp32; None in non-replay replay_cache_buf_idx: Optional[torch.Tensor], # [max_batch] int32; None in non-replay replay_prev_num_accepted: Optional[torch.Tensor], # [max_batch] int32; None in non-replay replay_work_items: Optional[torch.Tensor], # [max_batch, 4] int32; None in non-replay diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 36f0979271d2..073d9ccc3c86 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -927,6 +927,9 @@ def prepare_replay_metadata(self) -> None: ) pnat = prev_num_accepted_tokens[cache_slot_idx].to(torch.int32) active_cache_buf_idx = cache_buf_idx[cache_slot_idx].to(torch.int32) + + # Keep field order and write-first partitioning in sync with the + # PyTorch replay metadata path in mamba2_metadata.py. writes = pnat + replay_metadata.replay_step_width > replay_metadata.replay_history_size writes_i32 = writes.to(torch.int32) write_offsets = torch.cumsum(writes_i32, dim=0) - writes_i32 diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 7d37fdb8ed9d..0c9dbe80ef6f 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -21,8 +21,7 @@ import triton.language as tl from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata -from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ - CUDA_GRAPH_DUMMY_REQUEST_ID +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm._utils import prefer_pinned REPLAY_WORK_POSITION_IN_DECODE_BATCH = 0 @@ -143,19 +142,26 @@ def cu_seqlens_to_chunk_indices_offsets( chunk_size: int) -> Tuple[torch.Tensor, torch.Tensor]: """ Args: - cu_seqlens (torch.Tensor): 1D tensor of cumulative sequence lengths, shape (num_seqs + 1,). The first element should be 0. Each entry represents the starting index of a sequence in the flattened token array. + cu_seqlens (torch.Tensor): 1D tensor of cumulative sequence lengths, + shape (num_seqs + 1,). The first element should be 0. Each entry + represents the starting index of a sequence in the flattened token + array. chunk_size (int): The size of each physical mamba chunk (number of tokens per chunk). Returns: Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - chunk_indices (torch.Tensor): 1D tensor of indices indicating the physical chunk for each logical chunk. - - chunk_offsets (torch.Tensor): 1D tensor of offsets indicating the starting index of each logical chunk within its physical chunk. + - chunk_offsets (torch.Tensor): 1D tensor of offsets indicating + the starting index of each logical chunk within its physical + chunk. This function computes the chunk indices and offsets for the given cu_seqlens and chunk_size. Both are tensors of integers with length N, where N is the number of logical (pseudo) chunks. - A logical chunk is a sequence of tokens that are all part of the same sequence and are all in the same physical mamba chunk. + A logical chunk is a sequence of tokens that are all part of the same sequence + and are all in the same physical mamba chunk. In other words, a logical chunk changes every time we cross a sequence boundary or a physical mamba chunk boundary. - Logical chunks are needed to handle batched requests with initial states (see _state_passing_fwd and _chunk_scan_fwd). + Logical chunks are needed to handle batched requests with initial states + (see _state_passing_fwd and _chunk_scan_fwd). The chunk_indices tensor contains the index of the physical chunk for each logical chunk. The chunk_offsets tensor contains the offset (AKA starting index) of the logical chunk in the physical chunk. @@ -167,9 +173,12 @@ def cu_seqlens_to_chunk_indices_offsets( In this example, we have 2 sequences, each with 5 tokens. The physical chunk size is 8 tokens. We have three logical chunks: - - the first logical chunk starts at token 0 in the first physical chunk and contains all 5 tokens from the first sequence - - the second logical chunk starts at token 5 in the first physical chunk and contains first 3 tokens from the second sequence - - the third logical chunk starts at token 0 in the second physical chunk and contains the remaining 2 tokens from the second sequence + - the first logical chunk starts at token 0 in the first physical chunk and + contains all 5 tokens from the first sequence + - the second logical chunk starts at token 5 in the first physical chunk + and contains first 3 tokens from the second sequence + - the third logical chunk starts at token 0 in the second physical chunk + and contains the remaining 2 tokens from the second sequence """ total_seqlens = cu_seqlens[-1] @@ -259,22 +268,26 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, self.replay_num_decodes = 0 if not getattr(kv_cache_manager, 'use_replay_state_update', False): return - if not hasattr(kv_cache_manager, 'get_replay_state_update_metadata'): + num_decodes = batch_size - num_contexts + self.replay_num_decodes = num_decodes + self.replay_n_writes.zero_() + if num_decodes == 0: return + if not hasattr(kv_cache_manager, 'get_replay_state_update_metadata'): + raise RuntimeError( + "Replay state update is enabled, but the KV cache manager " + "does not expose replay state update metadata.") replay_metadata = kv_cache_manager.get_replay_state_update_metadata() if replay_metadata is None: - return - self.replay_n_writes.zero_() + raise RuntimeError( + "Replay state update is enabled for a decode batch, but the " + "KV cache manager returned no replay state update metadata.") prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens cache_buf_idx = replay_metadata.cache_buf_idx replay_step_width = replay_metadata.replay_step_width replay_history_size = replay_metadata.replay_history_size - num_decodes = batch_size - num_contexts - self.replay_num_decodes = num_decodes - if num_decodes == 0: - return position_in_decode_batch = torch.arange( num_decodes, dtype=torch.int32, device=self.state_indices.device) @@ -283,6 +296,8 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, pnat = prev_num_accepted_tokens[cache_slot_idx].to(torch.int32) active_cache_buf_idx = cache_buf_idx[cache_slot_idx].to(torch.int32) + # Keep field order and write-first partitioning in sync with the + # AutoDeploy replay metadata path in shim/interface.py. writes = (pnat + replay_step_width > replay_history_size) writes_i32 = writes.to(torch.int32) write_offsets = torch.cumsum(writes_i32, dim=0) - writes_i32 @@ -333,12 +348,10 @@ def prepare(self, attn_metadata: AttentionMetadata): # cudaStreamSynchronize per element. # # Safe under CUDA graphs only when the source buffer has a - # stable data pointer across all calls (currently true for - # CppMambaHybridCacheManager.cuda_state_indices, allocated - # once in __init__). If a future cache manager reallocates - # this buffer between iterations, captured kernels would - # still read from the address seen at capture time, so we - # assert stability here. + # stable data pointer across all calls. If a cache manager + # reallocates this buffer between iterations, captured kernels + # would still read from the address seen at capture time, so + # we assert stability here. if self._state_indices_aliased_ptr is None: self._state_indices_aliased_ptr = indices.data_ptr() else: @@ -428,7 +441,7 @@ def prepare(self, attn_metadata: AttentionMetadata): # Complete any deferred recurrent-state block onboards scheduled by # CppMambaHybridCacheManager.prepare_resources(). prepare_resources # only enqueues the async cudaMemcpyAsync calls and sets a pending - # flag; we sync the onboard stream here, so the prior CPU-side prep + # flag; we sync the onboard stream here, so CPU-side prep # work in _prepare_tp_inputs overlaps with the in-flight transfers. # Cheap no-op on cache managers without this method or when no # transfers were scheduled this iteration. diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 15096f10eb5b..861807ec490b 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -22,10 +22,8 @@ from torch import nn from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata -from tensorrt_llm._torch.modules.multi_stream_utils import \ - maybe_execute_in_parallel -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import \ - use_cpp_mamba_cache_manager +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import use_cpp_mamba_cache_manager from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -35,15 +33,12 @@ from ...speculative import SpecMetadata from ..linear import Linear, TensorParallelMode from .causal_conv1d import causal_conv1d_fn, causal_conv1d_update -from .causal_conv1d_triton import \ - causal_conv1d_update as causal_conv1d_update_triton -from .fuse_elementwise_ops import (extract_transpose_xbc_prefill, - fused_split_rearrange_after_conv1d) +from .causal_conv1d_triton import causal_conv1d_update as causal_conv1d_update_triton +from .fuse_elementwise_ops import extract_transpose_xbc_prefill, fused_split_rearrange_after_conv1d from .layernorm_gated import RMSNorm as RMSNormGated from .layernorm_gated import fused_gated_rmsnorm_quant_shape_ok from .replay_selective_state_update import replay_selective_state_update -from .selective_state_update import \ - selective_state_update as selective_state_update_native +from .selective_state_update import selective_state_update as selective_state_update_native from .selective_state_update import selective_state_update_mtp_ssm_cache_trtllm from .ssd_combined import mamba_chunk_scan_combined @@ -412,11 +407,23 @@ def forward( # Speculative decoding only supported with Python path assert layer_cache is not None, \ "Speculative decoding requires Python MambaCacheManager" - # TODO: support dynamic speculation, will add current_draft_len later [TRTLLM-10319] - draft_token_num = spec_metadata.max_draft_len + 1 intermediate_conv_states = layer_cache.intermediate_conv_window use_replay = getattr(attn_metadata.kv_cache_manager, 'use_replay_state_update', False) + draft_token_num = spec_metadata.runtime_draft_len + 1 + if use_replay: + replay_metadata = ( + attn_metadata.kv_cache_manager. + get_replay_state_update_metadata()) + assert replay_metadata is not None, ( + "Mamba replay state update is enabled but replay " + "metadata was not allocated.") + replay_step_width = replay_metadata.replay_step_width + assert draft_token_num == replay_step_width, ( + "Mamba replay state update does not support dynamic " + "draft length yet. Runtime token width " + f"{draft_token_num} must match fixed replay step " + f"width {replay_step_width}.") intermediate_state_indices = _cached_arange( attn_metadata.kv_cache_manager.get_max_resource_count(), diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 055ee51ff4f1..57df41a06cb4 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -176,9 +176,8 @@ def _stochastic_round_int16_packed(x: tl.tensor, rand: tl.tensor, offs_n: tl.ten return tl.extra.cuda.libdevice.floor(x + rand01) -# Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, -# old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. -# Grid: (batch, nheads // HEADS_PER_BLOCK). +# Replay-style precompute body. Computes CB_scaled/decay_vec in T-space and +# writes this step's B/dt/dA_cumsum to the selected cache buffer. @triton.jit() @@ -261,12 +260,12 @@ def _replay_precompute_impl( HEADS_PER_BLOCK: tl.constexpr, # Checkpoint write flag — selects target buffer + offset for new-token # cache writes. See "Cache write semantics" block below. - # Runtime (not constexpr): the only WRITE_CHECKPOINT-dependent code in + # Runtime (not constexpr): the only checkpoint-dependent code in # this body is the write_buf/write_offset selection, which is plain # arithmetic — no constexpr-shaped tile or whole-block gate. Letting # it be runtime lets the dynamic dispatch kernel call us once with the # per-slot needs_write flag instead of inlining two specializations. - write_checkpoint, + needs_checkpoint_write, ): pid_b = tl.program_id(axis=0) pid_hg = tl.program_id(axis=1) # head-group index @@ -289,7 +288,7 @@ def _replay_precompute_impl( # before the caller flips to the staging buffer. buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) - if write_checkpoint: + if needs_checkpoint_write: write_buf = 1 - buf_active write_offset = 0 else: @@ -305,11 +304,8 @@ def _replay_precompute_impl( causal_mask = offs_t[:, None] >= offs_t[None, :] valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] - # --- Vectorized pre-wait phase across HEADS_PER_BLOCK heads --- - # Compute dt, dA_cumsum, decay_vec as (H, T) tiles. Pre-compute - # scale_combo = decay_matrix * dt[:, None, :] as an (H, T, T) tile that - # stays in registers across gdc_wait — eliminates the post-wait reload - # of dt + dA_cumsum and the per-head loop. + # --- Pre-wait phase across HEADS_PER_BLOCK heads --- + # Compute dt, dA_cumsum, decay_vec, and scale_combo as head-block tiles. offs_h = tl.arange(0, HEADS_PER_BLOCK) heads_block = first_head + offs_h # (H,) @@ -337,9 +333,9 @@ def _replay_precompute_impl( # this step's per-step-restarted cumsum before storing so the buffer # holds one continuous cumsum across N back-to-back nowrites. Write path # (write_buf = 1 - buf_active, write_offset = 0) starts fresh, no prefix. - # Both branches are on scalar runtime values (write_checkpoint and PNAT), - # uniform across the block — use scalar if to short-circuit the load. - if write_checkpoint or prev_num_accepted_tokens == 0: + # Both branches are on scalar runtime values (checkpoint predicate and + # PNAT), uniform across the block — use scalar if to short-circuit the load. + if needs_checkpoint_write or prev_num_accepted_tokens == 0: prev_total = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) else: last_cumsum_ptrs = ( @@ -380,12 +376,10 @@ def _replay_precompute_impl( tl.store(decay_vec_addrs, decay_vec, mask=t_mask[None, :]) # scale_combo (H, T, T) = exp(dA_cumsum[h, t1] - dA_cumsum[h, t2]) * dt[h, t2] - # Stays live across gdc_wait — used post-wait to compute CB_scaled. decay_matrix = tl.exp(dA_cumsum[:, :, None] - dA_cumsum[:, None, :]) # (H, T, T) scale_combo = decay_matrix * dt[:, None, :] # (H, T, T) - # --- Wait for upstream kernel (external PDL) before loading B and C --- - # All dt processing above is independent of conv1d outputs. + # Wait for conv1d before loading this step's B/C. if LAUNCH_WITH_PDL: _gdc_wait_with_memory_clobber() @@ -424,9 +418,9 @@ def _replay_precompute_impl( mask=t_mask[:, None] & n_mask[None, :], ) - # --- Vectorized post-wait phase: scale_combo (H, T, T) is still live in - # registers from pre-wait; multiply by raw_CB (T, T), apply causal mask, - # store as one (H, T, T) tile. --- + # --- Post-wait CB precompute --- + # Combine raw_CB with scale_combo, apply the causal mask, and store one + # (H, T, T) tile. CB_scaled_block = tl.where( valid_mask[None, :, :], raw_CB[None, :, :] * scale_combo, @@ -443,9 +437,8 @@ def _replay_precompute_impl( tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) -# Replay-style precompute kernel. Thin wrapper around _replay_precompute_impl -# that carries the @triton.heuristics for constexpr derivation; called from -# the Python wrapper on the replay-style path (write or replay-nowrite). +# Rectangle nowrite precompute body. Builds a window-space CB tile with +# history at [0, PNAT) and this step's values at [PNAT, PNAT + T). @triton.jit() def _rectangle_precompute_impl( # Input pointers @@ -455,8 +448,8 @@ def _rectangle_precompute_impl( B_ptr, C_ptr, # Output pointers - cb_scaled_ptr, # (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) — rectangle window - decay_vec_ptr, # (batch, nheads, BLOCK_SIZE_T) — total_decay * exp(cumAdt_new[t]) + cb_scaled_ptr, # rectangle window: (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) + decay_vec_ptr, # total_decay * exp(cumAdt_new[t]): (batch, nheads, BLOCK_SIZE_T) # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite # path: read from buf_active at [0, PNAT), write new tokens at # [PNAT, PNAT+T) of buf_active (same buffer). @@ -554,8 +547,6 @@ def _rectangle_precompute_impl( heads_block = first_head + offs_h # Precompute this step's (H, T) dt and continuous dA_cumsum tiles. - # Keep them in registers for the rectangle path below; reloading from - # global memory after storing would create a same-kernel write/read race. dt_addrs = ( dt_ptr + pid_b * stride_dt_batch @@ -607,9 +598,7 @@ def _rectangle_precompute_impl( mask=t_mask[None, :], ) - # ---- Work independent of conv1d ---- - # Load historical cache and build combo_block before gdc_wait so this - # work can overlap the upstream conv1d latency. + # Build the history side of combo_block from cached data. group_idx = first_head // nheads_ngroups_ratio # Group-level: history B from active buffer at [0, PNAT) of the window. @@ -627,9 +616,6 @@ def _rectangle_precompute_impl( other=0.0, ) - # combo_block stays in registers across gdc_wait and is used directly - # after the wait to compute rect_CB_scaled. - # Per-head read bases (H,) - broadcast with offs_window for 2D loads. old_dt_read_h = ( old_dt_ptr @@ -656,10 +642,7 @@ def _rectangle_precompute_impl( mask=history_mask_h, other=0.0, ).to(tl.float32) - # Use loop-1 registers for this step's newly appended tokens. These are - # exactly the values stored above at [PNAT, PNAT+T); reloading them here - # would read bytes this same kernel just wrote, which Triton does not - # guarantee to fence. + # Combine cached history with this step's values at [PNAT, PNAT+T). ht_mask = t_mask[None, :] # (1, T) dA_cumsum_new = dA_cumsum_step + dA_cumsum_prefix[:, None] # (H, T) @@ -719,11 +702,10 @@ def _rectangle_precompute_impl( exp_diff = tl.exp(neg_dA_cumsum_window[:, None, :] + dA_cumsum_new[:, :, None]) combo_block = dt_factor_window[:, None, :] * exp_diff # (H, T, window) - # ---- gdc_wait: from here on we depend on conv1d's outputs ---- + # Wait for conv1d before loading this step's B/C. if LAUNCH_WITH_PDL: _gdc_wait_with_memory_clobber() - # Conv1d outputs: B and C C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group step_B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group @@ -775,9 +757,8 @@ def _rectangle_precompute_impl( ) causal_combined = (is_history_position_2d | is_step_causal_2d) & t_mask[:, None] - # Post-wait vectorized: combo_block (H, T, window) is still live in registers. - # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store as - # one (H, T, window) tile. + # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store one + # (H, T, window) tile. rect_CB_scaled_block = tl.where( causal_combined[None, :, :], raw_rect_CB[None, :, :] * combo_block, @@ -796,9 +777,9 @@ def _rectangle_precompute_impl( tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) -# Rectangle precompute kernel. Thin wrapper around _rectangle_precompute_impl -# that carries the @triton.heuristics for constexpr derivation; called from -# the Python wrapper on the rectangle nowrite path. +# Dynamic precompute kernel. Carries constexpr heuristics and dispatches each +# slot to the replay-style or rectangle precompute body. +# Grid: (batch, nheads // HEADS_PER_BLOCK). @triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) @triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) @triton.heuristics( @@ -1021,8 +1002,8 @@ def _dynamic_precompute_kernel( ) -# Main kernel: tl.dot replay + precomputed CB output. -# Grid: (cdiv(dim, M), batch, nheads). +# Replay-style main body for one persistent work item. Used by write and +# replay-nowrite paths. @triton.jit() @@ -1475,6 +1456,7 @@ def _persistent_main_impl( step_x, mask=t_mask[:, None] & m_mask[None, :], ) + step_x_for_dot = step_x.to(tl.bfloat16) step_x = step_x.to(tl.float32) cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head @@ -1490,7 +1472,7 @@ def _persistent_main_impl( ) init_out = tl.dot(C_tile.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] - cb_out = tl.dot(CB_scaled.to(tl.bfloat16), step_x.to(tl.bfloat16)) + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), step_x_for_dot) output_tile = init_out + cb_out if HAS_D: @@ -1510,12 +1492,8 @@ def _persistent_main_impl( tl.store(output_ptrs, output_tile, mask=t_mask[:, None] & m_mask[None, :]) -# `_persistent_rectangle_impl`: rectangle nowrite path for the persistent -# kernel. Body is a copy of `_rectangle_main_impl` with `pid_m`/`pid_b`/`pid_h` -# lifted to args (same pattern as `_persistent_main_impl` vs `_replay_main_impl`). -# Called only for nowrite slots when the kernel runs with RECTANGLE=True. -# Dropped from the rect impl: LAUNCH_DEPENDENT_KERNELS / REVERSE_PERM -# (kernel-level, signalled once at top); the wrapper resolves replay metadata. +# Rectangle nowrite main body for one persistent work item. Used only for +# nowrite slots when the persistent kernel runs with RECTANGLE=True. @triton.jit() def _persistent_rectangle_impl( # Per-work-unit indices (computed by the persistent wrapper). @@ -1652,8 +1630,6 @@ def _persistent_rectangle_impl( mask=m_mask, other=1.0, ).to(tl.float32) - else: - state = state.to(tl.float32) # Group / pointer offset setup group_idx = pid_h // nheads_ngroups_ratio @@ -1679,14 +1655,14 @@ def _persistent_rectangle_impl( D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 ).to(tl.float32) - # Hoist: history x doesn't depend on conv1d/precompute; load before gdc_wait. + # History x does not depend on conv1d/precompute. history_x = tl.load( old_x_read_base + safe_history_idx[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, mask=is_history_position[:, None] & m_mask[None, :], other=0.0, - ).to(tl.float32) + ) if WAIT_FOR_PDL_PREDECESSOR: _gdc_wait_with_memory_clobber() @@ -1709,16 +1685,14 @@ def _persistent_rectangle_impl( mask=is_step_position[:, None] & m_mask[None, :], ) - step_x_in_window_f32 = step_x_in_window.to(tl.float32) - x_window = history_x + step_x_in_window_f32 + step_x_in_window_for_dot = step_x_in_window.to(tl.bfloat16) + x_window_for_dot = (history_x.to(tl.bfloat16) + step_x_in_window_for_dot).to(tl.bfloat16) - if HAS_D or HAS_Z: + if HAS_D: step_in_window_selector = offs_t[:, None] == ( offs_window[None, :] - prev_num_accepted_tokens ) - step_x = tl.dot(step_in_window_selector.to(tl.bfloat16), step_x_in_window.to(tl.bfloat16)) - else: - step_x = step_x_in_window_f32 # placeholder; unused + step_x = tl.dot(step_in_window_selector.to(tl.bfloat16), step_x_in_window_for_dot) cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head CB_scaled = tl.load( @@ -1738,7 +1712,7 @@ def _persistent_rectangle_impl( if QUANT_MAX > 0.0: state_out = state_out * decode_scale[None, :] - token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_window.to(tl.bfloat16)) + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_window_for_dot) output_tile = state_out + token_out @@ -1759,8 +1733,8 @@ def _persistent_rectangle_impl( tl.store(output_ptrs, output_tile, mask=t_mask[:, None] & m_mask[None, :]) -# Persistent main kernel: 1D grid, persistent CTA loop. -# Heuristics mirror those of the replay main kernel. +# Persistent replay main kernel: 1D grid, persistent CTA loop. Heuristics cover +# both the replay-style and rectangle main bodies. @triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) @triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) @@ -1824,9 +1798,9 @@ def _persistent_main_kernel( # device memory keeps the pointer stable across CUDA graph replay while # allowing the value to change between iterations. # When IS_DYNAMIC=True the value is unused (Triton DCEs the load). - n_writes_ptr, # int32 *: device-side count of write-mode slots - batch_total, # int32: total slot count - nheads, # int32: total head count (== _replay_main_impl's program_id axis 2 count) + n_writes_ptr, # device-side count of write-mode slots + batch_total, # total slot count + nheads, # total head count # Dimensions T: tl.constexpr, MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, @@ -1919,7 +1893,7 @@ def _persistent_main_kernel( WARP_SPECIALIZE: tl.constexpr, IS_DYNAMIC: tl.constexpr, BLOCK_SIZE_K: tl.constexpr = MIN_REPLAY_TILE_SIZE, # rectangle window dimension - RECTANGLE: tl.constexpr = False, # when True, dispatch nowrite slots to _persistent_rectangle_impl + RECTANGLE: tl.constexpr = False, # dispatch nowrite slots to _persistent_rectangle_impl when true # 3 TMA toggles per the 3 live paths per-compilation: # USE_TMA_LOAD_WRITE — SSM state load when is_write # USE_TMA_LOAD_NOWRITE — nowrite-path state load (rect when RECTANGLE, @@ -2112,8 +2086,8 @@ def _persistent_main_kernel( USE_RS_ROUNDING, PHILOX_ROUNDS, QUANT_MAX, - True, - IS_DYNAMIC, # WRITE_CHECKPOINT=True (write arm) + True, # WRITE_CHECKPOINT=True (write arm) + IS_DYNAMIC, True, # WRITE_CHECKPOINT_IS_CONSTEXPR USE_TMA_LOAD_WRITE, USE_TMA_LOAD_NOWRITE, @@ -2188,7 +2162,7 @@ def _persistent_main_kernel( BLOCK_SIZE_K, LAUNCH_WITH_PDL and (NUM_LOOP_STAGES == 1), # WAIT_FOR_PDL_PREDECESSOR QUANT_MAX, - USE_TMA_LOAD_NOWRITE, # rect-load TMA toggle + USE_TMA_LOAD_NOWRITE, # USE_TMA_LOAD ) else: _persistent_main_impl( @@ -2330,38 +2304,32 @@ def _persistent_main_kernel( ("fp16", "SR"): [ ( 16, - "persistent_main", + "persistent_dynamic", { - "_block_size_m_nowrite": 32, - "_block_size_m_write": 16, - "_cta_per_sm_nowrite": 7, - "_cta_per_sm_write": 10, + "_block_size_m": 8, + "_cta_per_sm": 6, "_flatten": False, - "_heads_per_block": 1, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 1, - "_num_stages_write": 3, - "_num_warps_nowrite": 4, - "_num_warps_write": 4, - "_precompute_num_warps": 8, - "_use_tma_rect_load": True, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, "_warp_specialize": False, - "nowrite_first": False, - "rectangle_for_nowrite": True, + "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=9.66us (B200 PDL-hoist default 5x200) + ), # raw_batch=1, score=6.57us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 32, "persistent_dynamic", { - "_block_size_m": 4, - "_cta_per_sm": 7, + "_block_size_m": 8, + "_cta_per_sm": 9, "_flatten": False, - "_heads_per_block": 4, + "_heads_per_block": 2, "_num_loop_stages": 1, "_num_stages": 5, "_num_warps": 1, @@ -2373,7 +2341,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=2, score=7.12us (B200 PDL-hoist default 5x200) + ), # raw_batch=2, score=7.0us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 64, "persistent_dynamic", @@ -2440,14 +2408,14 @@ def _persistent_main_kernel( { "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 8, - "_cta_per_sm_write": 4, + "_cta_per_sm_nowrite": 5, + "_cta_per_sm_write": 9, "_flatten": False, "_heads_per_block": 8, "_num_loop_stages_nowrite": 1, "_num_loop_stages_write": 1, "_num_stages_nowrite": 3, - "_num_stages_write": 2, + "_num_stages_write": 3, "_num_warps_nowrite": 1, "_num_warps_write": 2, "_precompute_num_warps": 8, @@ -2459,7 +2427,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=32, score=12.71us (B200 PDL-hoist default 5x200) + ), # raw_batch=32, score=12.66us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 1024, "persistent_main", @@ -2563,7 +2531,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=73.80us (B200 PDL-hoist default 5x200) + ), # raw_batch=512, score=73.8us (B200 PDL-hoist default 5x200) ( 16384, "persistent_main", @@ -2597,21 +2565,21 @@ def _persistent_main_kernel( "persistent_dynamic", { "_block_size_m": 8, - "_cta_per_sm": 2, + "_cta_per_sm": 9, "_flatten": False, - "_heads_per_block": 8, - "_num_loop_stages": 2, - "_num_stages": 4, - "_num_warps": 2, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, "_precompute_num_warps": 4, "_use_tma_rect_load": False, - "_use_tma_replay_nowrite_load": True, - "_use_tma_replay_write_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=8.60us (B200 PDL-hoist default 5x200) + ), # raw_batch=1, score=6.59us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 32, "persistent_dynamic", @@ -2723,53 +2691,53 @@ def _persistent_main_kernel( "persistent_main", { "_block_size_m_nowrite": 64, - "_block_size_m_write": 64, - "_cta_per_sm_nowrite": 6, - "_cta_per_sm_write": 3, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 3, + "_cta_per_sm_write": 10, "_flatten": False, - "_heads_per_block": 16, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 2, - "_num_stages_write": 1, - "_num_warps_nowrite": 2, - "_num_warps_write": 4, - "_precompute_num_warps": 16, - "_use_tma_rect_load": True, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 5, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_store": False, "_warp_specialize": False, - "nowrite_first": False, + "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=64, score=22.26us (B200 PDL-hoist default 5x200) + ), # raw_batch=64, score=18.51us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 2048, "persistent_main", { "_block_size_m_nowrite": 64, - "_block_size_m_write": 64, - "_cta_per_sm_nowrite": 6, - "_cta_per_sm_write": 3, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 4, - "_num_loop_stages_nowrite": 2, - "_num_loop_stages_write": 3, - "_num_stages_nowrite": 1, - "_num_stages_write": 1, - "_num_warps_nowrite": 2, - "_num_warps_write": 4, - "_precompute_num_warps": 1, - "_use_tma_rect_load": True, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 2, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": False, + "_use_tma_replay_write_store": True, "_warp_specialize": False, - "nowrite_first": False, + "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=128, score=34.03us (B200 PDL-hoist default 5x200) + ), # raw_batch=128, score=28.72us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 4096, "persistent_main", @@ -2855,21 +2823,21 @@ def _persistent_main_kernel( "persistent_dynamic", { "_block_size_m": 8, - "_cta_per_sm": 5, + "_cta_per_sm": 3, "_flatten": False, - "_heads_per_block": 2, - "_num_loop_stages": 2, - "_num_stages": 2, - "_num_warps": 2, - "_precompute_num_warps": 8, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 5, + "_num_warps": 1, + "_precompute_num_warps": 2, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": True, + "_use_tma_replay_write_load": False, "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=8.67us (B200 PDL-hoist default 5x200) + ), # raw_batch=1, score=6.46us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 32, "persistent_dynamic", @@ -2995,21 +2963,21 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=64, score=16.00us (B200 PDL-hoist default 5x200) + ), # raw_batch=64, score=16.0us (B200 PDL-hoist default 5x200) ( 2048, "persistent_main", { - "_block_size_m_nowrite": 64, + "_block_size_m_nowrite": 32, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 6, + "_cta_per_sm_nowrite": 10, "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 16, - "_num_loop_stages_nowrite": 2, + "_num_loop_stages_nowrite": 1, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 2, - "_num_stages_write": 4, + "_num_stages_nowrite": 4, + "_num_stages_write": 5, "_num_warps_nowrite": 2, "_num_warps_write": 1, "_precompute_num_warps": 8, @@ -3021,21 +2989,21 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=128, score=26.19us (B200 PDL-hoist default 5x200) + ), # raw_batch=128, score=24.5us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 4096, "persistent_main", { - "_block_size_m_nowrite": 64, + "_block_size_m_nowrite": 32, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 6, + "_cta_per_sm_nowrite": 10, "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 16, - "_num_loop_stages_nowrite": 2, + "_num_loop_stages_nowrite": 1, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 2, - "_num_stages_write": 3, + "_num_stages_nowrite": 3, + "_num_stages_write": 5, "_num_warps_nowrite": 2, "_num_warps_write": 1, "_precompute_num_warps": 8, @@ -3047,51 +3015,51 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=256, score=41.16us (B200 PDL-hoist default 5x200) + ), # raw_batch=256, score=40.04us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 8192, "persistent_main", { - "_block_size_m_nowrite": 64, + "_block_size_m_nowrite": 32, "_block_size_m_write": 32, "_cta_per_sm_nowrite": 8, "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 4, - "_num_loop_stages_nowrite": 3, - "_num_loop_stages_write": 2, - "_num_stages_nowrite": 2, - "_num_stages_write": 2, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 5, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 1, "_num_warps_nowrite": 1, "_num_warps_write": 1, - "_precompute_num_warps": 1, + "_precompute_num_warps": 2, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": False, "_warp_specialize": False, - "nowrite_first": True, + "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=77.62us (B200 PDL-hoist default 5x200) + ), # raw_batch=512, score=69.44us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 16384, "persistent_main", { "_block_size_m_nowrite": 64, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 8, + "_cta_per_sm_nowrite": 6, "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 8, - "_num_loop_stages_nowrite": 2, + "_num_loop_stages_nowrite": 5, "_num_loop_stages_write": 1, "_num_stages_nowrite": 2, - "_num_stages_write": 4, - "_num_warps_nowrite": 1, + "_num_stages_write": 3, + "_num_warps_nowrite": 2, "_num_warps_write": 1, "_precompute_num_warps": 1, - "_use_tma_rect_load": True, + "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": False, @@ -3099,7 +3067,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=1024, score=131.40us (B200 PDL-hoist default 5x200) + ), # raw_batch=1024, score=126.74us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ], ("fp32", "RN"): [ ( @@ -3107,13 +3075,13 @@ def _persistent_main_kernel( "persistent_dynamic", { "_block_size_m": 8, - "_cta_per_sm": 9, + "_cta_per_sm": 6, "_flatten": False, - "_heads_per_block": 2, - "_num_loop_stages": 2, - "_num_stages": 4, - "_num_warps": 2, - "_precompute_num_warps": 8, + "_heads_per_block": 8, + "_num_loop_stages": 1, + "_num_stages": 2, + "_num_warps": 1, + "_precompute_num_warps": 4, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": True, "_use_tma_replay_write_load": True, @@ -3121,7 +3089,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=8.58us (B200 PDL-hoist default 5x200) + ), # raw_batch=1, score=6.64us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 32, "persistent_dynamic", @@ -3229,40 +3197,40 @@ def _persistent_main_kernel( "_block_size_m_nowrite": 64, "_block_size_m_write": 32, "_cta_per_sm_nowrite": 4, - "_cta_per_sm_write": 8, + "_cta_per_sm_write": 6, "_flatten": False, - "_heads_per_block": 8, - "_num_loop_stages_nowrite": 3, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 1, "_num_loop_stages_write": 1, "_num_stages_nowrite": 1, - "_num_stages_write": 4, + "_num_stages_write": 1, "_num_warps_nowrite": 2, "_num_warps_write": 2, "_precompute_num_warps": 8, - "_use_tma_rect_load": True, - "_use_tma_replay_nowrite_load": False, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, "nowrite_first": True, - "rectangle_for_nowrite": True, + "rectangle_for_nowrite": False, }, - ), # raw_batch=64, score=21.92us (B200 PDL-hoist default 5x200) + ), # raw_batch=64, score=20.05us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 2048, "persistent_main", { - "_block_size_m_nowrite": 64, + "_block_size_m_nowrite": 32, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 5, - "_cta_per_sm_write": 9, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 16, + "_heads_per_block": 8, "_num_loop_stages_nowrite": 3, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 1, - "_num_stages_write": 1, - "_num_warps_nowrite": 2, + "_num_stages_nowrite": 3, + "_num_stages_write": 4, + "_num_warps_nowrite": 1, "_num_warps_write": 1, "_precompute_num_warps": 8, "_use_tma_rect_load": True, @@ -3273,22 +3241,22 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=128, score=35.29us (B200 PDL-hoist default 5x200) + ), # raw_batch=128, score=31.39us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 4096, "persistent_main", { - "_block_size_m_nowrite": 64, + "_block_size_m_nowrite": 32, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 9, + "_cta_per_sm_nowrite": 8, "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 8, - "_num_loop_stages_nowrite": 2, + "_num_loop_stages_nowrite": 4, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 4, - "_num_stages_write": 1, - "_num_warps_nowrite": 2, + "_num_stages_nowrite": 2, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, "_num_warps_write": 1, "_precompute_num_warps": 8, "_use_tma_rect_load": True, @@ -3299,24 +3267,24 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=256, score=54.98us (B200 PDL-hoist default 5x200) + ), # raw_batch=256, score=51.27us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 8192, "persistent_main", { "_block_size_m_nowrite": 32, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 8, + "_cta_per_sm_nowrite": 4, "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 16, - "_num_loop_stages_nowrite": 2, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 5, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 2, + "_num_stages_nowrite": 4, "_num_stages_write": 4, "_num_warps_nowrite": 1, "_num_warps_write": 1, - "_precompute_num_warps": 8, + "_precompute_num_warps": 4, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, @@ -3325,25 +3293,25 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=93.58us (B200 PDL-hoist default 5x200) + ), # raw_batch=512, score=88.87us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ( 16384, "persistent_main", { - "_block_size_m_nowrite": 64, + "_block_size_m_nowrite": 32, "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 9, + "_cta_per_sm_nowrite": 8, "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 8, - "_num_loop_stages_nowrite": 5, + "_num_loop_stages_nowrite": 2, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 1, - "_num_stages_write": 2, - "_num_warps_nowrite": 2, + "_num_stages_nowrite": 3, + "_num_stages_write": 5, + "_num_warps_nowrite": 1, "_num_warps_write": 1, "_precompute_num_warps": 1, - "_use_tma_rect_load": True, + "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, @@ -3351,10 +3319,10 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=1024, score=180.18us (B200 PDL-hoist default 5x200) + ), # raw_batch=1024, score=170.51us (B200 PDL-retune noise-cleaned 5x500) # << RETUNED ], } -_PD_TO_PM_SPLIT_MAP = { # pd unsplit knob → (pm_write_knob, pm_nowrite_knob) +_PD_TO_PM_SPLIT_MAP = { # pd unsplit knob -> (pm_write_knob, pm_nowrite_knob) "_block_size_m": ("_block_size_m_write", "_block_size_m_nowrite"), "_num_warps": ("_num_warps_write", "_num_warps_nowrite"), "_num_stages": ("_num_stages_write", "_num_stages_nowrite"), @@ -3469,7 +3437,6 @@ def replay_selective_state_update( state_scales: torch.Tensor | None = None, launch_with_pdl=False, use_internal_pdl=True, - write_checkpoint: bool = True, rectangle_for_nowrite: bool | None = None, nowrite_first: bool | None = None, mode: str | None = None, @@ -3477,13 +3444,10 @@ def replay_selective_state_update( _num_warps: int | None = None, _num_stages: int | None = None, _precompute_num_warps: int | None = None, - _precompute_num_stages: int | None = None, _heads_per_block: int | None = None, - _maxnreg: int | None = None, - _num_ctas: int | None = None, - # Per-main knobs (override shared values for one half of the dl-family / - # persistent_main launches). Default None = tied to the shared value - # (backward compat). The two main kernels (write vs nowrite) have + # Per-main knobs override shared values for one half of the persistent_main + # launches. Default None ties the half-specific value to the shared knob. + # The two main kernels (write vs nowrite) have # different per-slot work — write does a state shift + store, nowrite # just appends — so the optimum (M, W, S, H) can differ. Precompute # knobs are intentionally NOT split: shared precompute wins (cheaper @@ -3501,7 +3465,7 @@ def replay_selective_state_update( # TMA state-tensor toggles — 4 independent paths (see replay design notes # item #17 for measured perf profiles). Each is False=raw load/store, True= # use a host-built TMA tensor_descriptor for that path. - _use_tma_rect_load: bool | None = None, # rect kernel's state load (nowrite-only) + _use_tma_rect_load: bool | None = None, # rectangle path state load (nowrite-only) _use_tma_replay_write_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=True _use_tma_replay_write_store: bool | None = None, # SSM state store when WRITE_CHECKPOINT=True _use_tma_replay_nowrite_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=False @@ -3554,10 +3518,10 @@ def replay_selective_state_update( Arguments: state: (cache, nheads, dim, dstate) in-place. After the call, contains the state after replaying prev_num_accepted_tokens old tokens. - old_x: (cache, 2, T, nheads, dim) bf16 — double-buffered old x cache. - old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. - old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. - old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. + old_x: (cache, 2, max_window, nheads, dim) bf16 replay history cache. + old_B: (cache, 2, max_window, ngroups, dstate) bf16 replay history cache. + old_dt: (cache, 2, nheads, max_window) fp32 processed dt history. + old_dA_cumsum: (cache, 2, nheads, max_window) fp32 cumulative A*dt history. cache_buf_idx: (cache,) int32 — which buffer to read (0 or 1). prev_num_accepted_tokens: (cache,) int32. x: (batch, T, nheads, dim) new token inputs. @@ -3566,10 +3530,14 @@ def replay_selective_state_update( B: (batch, T, ngroups, dstate). C: (batch, T, ngroups, dstate). out: (batch, T, nheads, dim) preallocated output. + n_writes: (1,) int32 device tensor with the write-mode count. + replay_work_items: (batch, 4) int32 device tensor, sorted write-first. + Persistent-main consumes all fields; persistent-dynamic only + requires the argument for a stable wrapper signature. + state_batch_indices: (batch,) int32 cache slot mapping. D: (nheads, dim) optional feed-through parameter. z: (batch, T, nheads, dim) optional silu gate. dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). - state_batch_indices: (batch,) int32 cache slot mapping. rand_seed: optional (cache_size,) int64 CUDA tensor of per-cache-slot Philox PRNG seeds. The caller bumps this tensor in-place for each replay invocation so CUDA graph replay still gets fresh draws. The @@ -3594,10 +3562,8 @@ def replay_selective_state_update( When None, use the tuning-table value if present. When true, launch the nowrite half before the write half. - _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, - _precompute_num_warps, _precompute_num_stages, _heads_per_block, - _maxnreg, _num_ctas) are benchmark-only overrides; production callers - should leave them None to use the tuning-table defaults. + _-prefixed kwargs are tuning overrides; production callers should + leave them None to use the tuning-table defaults. """ sm_version = get_sm_version() @@ -3612,14 +3578,15 @@ def replay_selective_state_update( # mode="persistent_dynamic": single persistent-CTA kernel covering the # full batch. Each work-item dispatches via runtime PNAT check # (is_write = (pnat + T) > MAX). No write/nowrite split. - # replay_work_items is ignored. write_checkpoint is ignored. + # The wrapper requires replay_work_items for a uniform signature, but + # the dynamic kernel ignores its contents. # mode="persistent_main": persistent-CTA kernel with two launches # (write half + nowrite half). Caller MUST pre-sort replay_work_items # write-first; the n_writes tensor partitions the persistent loop # into the two halves with the right WRITE_CHECKPOINT constexpr # each time. RECTANGLE constexpr (= rectangle_for_nowrite) picks # rect vs replay for the nowrite half. nowrite_first controls - # launch order only. write_checkpoint is ignored. + # launch order only. # Note: mode-and-knob resolution from the default-tuning table happens # below, after we have `batch` and `nheads`. @@ -3749,11 +3716,6 @@ def replay_selective_state_update( if _precompute_num_warps is not None else _table_knobs.get("_precompute_num_warps") ) - _precompute_num_stages = ( - _precompute_num_stages - if _precompute_num_stages is not None - else _table_knobs.get("_precompute_num_stages") - ) _block_size_m_write = ( _block_size_m_write if _block_size_m_write is not None @@ -3887,8 +3849,7 @@ def replay_selective_state_update( f"prev_num_accepted_tokens must be int32, got {prev_num_accepted_tokens.dtype}" ) assert isinstance(state_batch_indices, torch.Tensor), ( - f"state_batch_indices must be a torch.Tensor, " - f"got {type(state_batch_indices).__name__}" + f"state_batch_indices must be a torch.Tensor, got {type(state_batch_indices).__name__}" ) assert state_batch_indices.device == device, ( f"state_batch_indices must be on device {device}, got {state_batch_indices.device}" @@ -3928,7 +3889,7 @@ def replay_selective_state_update( # so the launch sites can refer to it; only used on the rectangle path. # If this differs from BLOCK_SIZE_T, rectangle precompute uses a slower # one-hot fallback because tl.gather requires matching padded dimension sizes. - # Production uses matching padded T and window sizes; mismatches are for tests/debugging. + # Production uses matching padded T and window sizes. BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), MIN_REPLAY_TILE_SIZE) rectangle_use_gather = BLOCK_SIZE_T == BLOCK_SIZE_K @@ -4028,7 +3989,7 @@ def replay_selective_state_update( state_scales.stride(2), ) else: - state_scales_arg = state # any valid ptr — gated by QUANT_MAX==0 + state_scales_arg = state # any valid pointer; gated by QUANT_MAX == 0 state_scales_strides = (0, 0, 0) # Per-path TMA descriptors for state — write-side and nowrite-side. Each @@ -4070,8 +4031,8 @@ def replay_selective_state_update( block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], ) else: - state_tma_descriptor_write = state # dummy; all consuming constexprs False - state_tma_descriptor_nowrite = state # dummy; all consuming constexprs False + state_tma_descriptor_write = state # dummy; all consuming constexprs are false + state_tma_descriptor_nowrite = state # dummy; all consuming constexprs are false # Work items are sorted write-first for persistent_main. Each row carries # decode-batch position, cache slot, PNAT, and active cache buffer index. @@ -4104,7 +4065,7 @@ def replay_selective_state_update( # ---- Launch helpers (close over locals) ------------------------------- # Each helper is a thin closure that calls one Triton kernel with the # full positional + kwarg argument list. Mode-dependent constexprs - # (write_checkpoint, early_out, rectangle) are passed in. + # (write_checkpoint_mode, rectangle) are passed in. def launch_dynamic_precompute(rectangle: bool): _dynamic_precompute_kernel[precomp_grid]( @@ -4165,7 +4126,6 @@ def launch_dynamic_precompute(rectangle: bool): RECTANGLE=rectangle, RECTANGLE_USE_GATHER=rectangle_use_gather, num_warps=precompute_num_warps, - **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), launch_pdl=launch_with_pdl, ) @@ -4192,22 +4152,26 @@ def launch_dynamic_precompute(rectangle: bool): # multiple tiles). NUM_PERSISTENT is now a runtime int (see kernel def # docstring at _persistent_main_kernel) so changing cta_per_sm does NOT # trigger a new Triton compile — same kernel binary, different loop step. - # (Named UPPERCASE for historical Triton-style consistency only; not - # constexpr.) + # Runtime value, despite the Triton-style uppercase name. _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M def launch_persistent_main( - write_checkpoint: bool, *, launch_dependent_kernels: bool = False, rectangle: bool = False + write_checkpoint_mode: bool, + *, + launch_dependent_kernels: bool = False, + rectangle: bool = False, ): # `n_writes` is the (1,) int32 device tensor with the write count. # Both halves always launch; an empty half has a zero-length slot # range and the persistent loop does no work. - block_size_m = BLOCK_SIZE_M_WRITE if write_checkpoint else BLOCK_SIZE_M_NOWRITE - launch_num_warps = NUM_WARPS_WRITE if write_checkpoint else NUM_WARPS_NOWRITE - launch_num_stages = NUM_STAGES_WRITE if write_checkpoint else NUM_STAGES_NOWRITE - ctas_per_sm = CTA_PER_SM_WRITE if write_checkpoint else CTA_PER_SM_NOWRITE + block_size_m = BLOCK_SIZE_M_WRITE if write_checkpoint_mode else BLOCK_SIZE_M_NOWRITE + launch_num_warps = NUM_WARPS_WRITE if write_checkpoint_mode else NUM_WARPS_NOWRITE + launch_num_stages = NUM_STAGES_WRITE if write_checkpoint_mode else NUM_STAGES_NOWRITE + ctas_per_sm = CTA_PER_SM_WRITE if write_checkpoint_mode else CTA_PER_SM_NOWRITE ctas_per_sm = ctas_per_sm if ctas_per_sm else 1 - num_loop_stages = NUM_LOOP_STAGES_WRITE if write_checkpoint else NUM_LOOP_STAGES_NOWRITE + num_loop_stages = ( + NUM_LOOP_STAGES_WRITE if write_checkpoint_mode else NUM_LOOP_STAGES_NOWRITE + ) num_loop_stages = num_loop_stages if num_loop_stages else 2 num_persistent = ctas_per_sm * _num_sms num_pid_m_local = (dim + block_size_m - 1) // block_size_m @@ -4219,7 +4183,7 @@ def launch_persistent_main( grid = (min(num_persistent, total_work_launch),) # Per-path TMA descriptor — block_shape[0] must match block_size_m. selected_state_tma_descriptor = ( - state_tma_descriptor_write if write_checkpoint else state_tma_descriptor_nowrite + state_tma_descriptor_write if write_checkpoint_mode else state_tma_descriptor_nowrite ) _persistent_main_kernel[grid]( state, @@ -4303,7 +4267,7 @@ def launch_persistent_main( LAUNCH_WITH_PDL=use_internal_pdl, PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, QUANT_MAX=quant_max, - WRITE_CHECKPOINT=write_checkpoint, + WRITE_CHECKPOINT=write_checkpoint_mode, LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, NUM_PERSISTENT=num_persistent, NUM_LOOP_STAGES=num_loop_stages, @@ -4316,16 +4280,14 @@ def launch_persistent_main( # NOWRITE_LOAD is dummy False; when WRITE_CHECKPOINT=False, WRITE_LOAD/STORE # dummy False. NOWRITE_LOAD picks rect-load (RECTANGLE) or # replay-nowrite-load. - USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint), + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint_mode), USE_TMA_LOAD_NOWRITE=bool( (_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load) - and not write_checkpoint + and not write_checkpoint_mode ), - USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint), + USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint_mode), num_warps=launch_num_warps, **({"num_stages": launch_num_stages} if launch_num_stages else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), launch_pdl=use_internal_pdl, ) @@ -4452,8 +4414,6 @@ def launch_persistent_dynamic_main( USE_TMA_STORE=bool(_use_tma_replay_write_store), num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), - **({"num_ctas": _num_ctas} if _num_ctas else {}), - **({"maxnreg": _maxnreg} if _maxnreg else {}), launch_pdl=use_internal_pdl, ) @@ -4482,7 +4442,7 @@ def launch_persistent_dynamic_main( # pre-sorted write-first. def launch_nowrite(launch_dependent_kernels: bool): launch_persistent_main( - write_checkpoint=False, + write_checkpoint_mode=False, launch_dependent_kernels=launch_dependent_kernels, rectangle=rectangle_for_nowrite, ) @@ -4491,15 +4451,15 @@ def launch_nowrite(launch_dependent_kernels: bool): if nowrite_first: launch_nowrite(launch_dependent_kernels=True) launch_persistent_main( - write_checkpoint=True, + write_checkpoint_mode=True, launch_dependent_kernels=False, - rectangle=False, # write always replay-style + rectangle=False, # write always uses the replay-style path ) else: launch_persistent_main( - write_checkpoint=True, + write_checkpoint_mode=True, launch_dependent_kernels=True, - rectangle=False, # write always replay-style + rectangle=False, # write always uses the replay-style path ) launch_nowrite(launch_dependent_kernels=False) else: diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index addf98611597..e4f246309a1f 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -7,10 +7,13 @@ import tensorrt_llm import tensorrt_llm.bindings.executor as trtllm -from tensorrt_llm._torch.models.modeling_utils import \ - MODEL_CLASS_VISION_ENCODER_MAPPING -from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version, - str_dtype_to_binding, torch_dtype_to_str) +from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_VISION_ENCODER_MAPPING +from tensorrt_llm._utils import ( + confidential_compute_enabled, + get_sm_version, + str_dtype_to_binding, + torch_dtype_to_str, +) from tensorrt_llm.bindings.executor import DecodingMode # isort: off @@ -21,38 +24,55 @@ WaitingQueuePolicy) # isort: on from tensorrt_llm.logger import logger -from tensorrt_llm.lora_helper import (LoraConfig, - get_default_trtllm_modules_to_hf_modules) +from tensorrt_llm.lora_helper import LoraConfig, get_default_trtllm_modules_to_hf_modules from tensorrt_llm.lora_manager import load_torch_lora from tensorrt_llm.mapping import CpType, Mapping from ..attention_backend import get_sparse_attn_kv_cache_manager from ..model_config import ModelConfig -from ..speculative import (get_num_extra_kv_tokens, get_num_spec_layers, - get_spec_decoder, should_use_separate_draft_kv_cache) -from .config_utils import (extract_mamba_kv_cache_params, is_gemma4_hybrid, - is_hybrid_linear, is_mla, is_nemotron_hybrid, - is_qwen3_hybrid) +from ..speculative import ( + get_num_extra_kv_tokens, + get_num_spec_layers, + get_spec_decoder, + should_use_separate_draft_kv_cache, +) +from .config_utils import ( + extract_mamba_kv_cache_params, + is_gemma4_hybrid, + is_hybrid_linear, + is_mla, + is_nemotron_hybrid, + is_qwen3_hybrid, +) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder from .kv_cache_transceiver import AttentionTypeCpp, create_kv_cache_transceiver from .llm_request import ExecutorResponse -from .mamba_cache_manager import (BaseMambaCacheManager, - CppMambaHybridCacheManager, - MixedMambaHybridCacheManager, - use_cpp_mamba_cache_manager, - use_py_mamba_cache_manager) +from .mamba_cache_manager import ( + BaseMambaCacheManager, + CppMambaHybridCacheManager, + MixedMambaHybridCacheManager, + use_cpp_mamba_cache_manager, + use_py_mamba_cache_manager, +) from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor -from .resource_manager import (KVCacheManager, KVCacheManagerV2, - PeftCacheManager, ResourceManager, - ResourceManagerType) -from .sampler import (EarlyStopSampler, EarlyStopWithMMResult, TorchSampler, - TRTLLMSampler) -from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, - KVCacheV2Scheduler, SimpleScheduler, - SimpleUnifiedScheduler) +from .resource_manager import ( + KVCacheManager, + KVCacheManagerV2, + PeftCacheManager, + ResourceManager, + ResourceManagerType, +) +from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler, TRTLLMSampler +from .scheduler import ( + BindCapacityScheduler, + BindMicroBatchScheduler, + KVCacheV2Scheduler, + SimpleScheduler, + SimpleUnifiedScheduler, +) from .seq_slot_manager import SeqSlotManager GB = 1 << 30 @@ -97,9 +117,19 @@ def get_kv_cache_manager_cls( logger.info("Hybrid linear model has 0 mamba layers; using " "KVCacheManager without mamba caching") return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) + if use_py_mamba_cache_manager(): + if kv_cache_config.enable_block_reuse: + raise ValueError( + "TRTLLM_USE_PY_MAMBA=1 forces " + "MixedMambaHybridCacheManager, which does not support " + "block reuse. Disable block reuse or unset " + "TRTLLM_USE_PY_MAMBA to use CppMambaHybridCacheManager.") + logger.info( + "Using MixedMambaHybridCacheManager for hybrid mamba model") + return MixedMambaHybridCacheManager if kv_cache_config.enable_block_reuse: return CppMambaHybridCacheManager - if use_cpp_mamba_cache_manager() or use_py_mamba_cache_manager(): + if use_cpp_mamba_cache_manager(): logger.info( "Using MixedMambaHybridCacheManager for hybrid mamba model") return MixedMambaHybridCacheManager diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 3ba2cf0d43c8..eb44997cfb8c 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -29,16 +29,18 @@ from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig -from tensorrt_llm._torch.pyexecutor.llm_request import ( - ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest) +from tensorrt_llm._torch.pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import ( - BaseResourceManager, CacheTypeCpp, DataType, KVCacheManager, - PoolConfiguration, get_pp_layers) + BaseResourceManager, + CacheTypeCpp, + DataType, + KVCacheManager, + PoolConfiguration, + get_pp_layers, +) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._utils import (nvtx_range, prefer_pinned, - torch_dtype_to_binding) -from tensorrt_llm.bindings.internal.batch_manager import ( - LinearAttentionMetadata, LinearCacheType) +from tensorrt_llm._utils import nvtx_range, prefer_pinned, torch_dtype_to_binding +from tensorrt_llm.bindings.internal.batch_manager import LinearAttentionMetadata, LinearCacheType from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -643,12 +645,15 @@ def __init__( dtype=torch.int32, device=device) - # Store max_batch_size for resource management + # Physical tensor rows include reserved dummy slots. Resource capacity + # is the number of real request slots remaining after those + # reservations. self._max_batch_size = max_batch_size + self._max_resource_count = len(self.mamba_cache_free_blocks) def get_max_resource_count(self) -> int: - """Return the maximum number of sequences that can be cached.""" - return self._max_batch_size + """Return the maximum number of real requests that can be cached.""" + return self._max_resource_count def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: """Return the prefix of *context_requests* that fits in the @@ -683,6 +688,7 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]): if (isinstance(self.mamba_cache, self.SpeculativeState) and self._use_replay_state_update): self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 if self._mamba_ssm_rand_seed is not None: # Deterministic per-slot rotation on fresh assignment. # `block` is pulled from mamba_cache_free_blocks, which @@ -708,8 +714,7 @@ def _is_padding_sentinel(self, request_id: int) -> bool: # cuda_graph_runner caches one dummy per runtime_draft_len value # (see _get_padded_batch), so any id in the range of dummy request IDs # may be live concurrently. - from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ - CUDA_GRAPH_DUMMY_REQUEST_ID + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID max_dl = self.speculative_num_draft_tokens or 0 return (CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl <= request_id <= CUDA_GRAPH_DUMMY_REQUEST_ID) @@ -1479,8 +1484,8 @@ def __init__( # on ranks with no local mamba layers. self._use_replay_state_update = use_replay_state_update self.replay_step_width: Optional[int] = ( - spec_config.max_draft_len + - 1 if spec_config is not None and use_replay_state_update else None) + spec_config.tokens_per_gen_step + if spec_config is not None and use_replay_state_update else None) self.replay_history_size: Optional[int] = ( max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) if self.replay_step_width is not None else None) @@ -1693,8 +1698,7 @@ def get_cache_size_per_token( ``T = budget // bytes_per_token``. """ # Lazy import to avoid pulling config_utils into module import order. - from tensorrt_llm._torch.pyexecutor.config_utils import \ - extract_mamba_kv_cache_params + from tensorrt_llm._torch.pyexecutor.config_utils import extract_mamba_kv_cache_params # Attention slope from the parent's existing formula. attention_slope = KVCacheManager.get_cache_size_per_token( diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 8f46d88794cd..b29cb9911fb6 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -329,6 +329,7 @@ class BaseResourceManager(ABC): @abstractmethod def get_max_resource_count(self) -> int: + """Return the maximum number of real requests this manager can admit.""" raise NotImplementedError @abstractmethod @@ -945,8 +946,7 @@ def probe_prefix_match_length(self, input_tokens, lora_task_id=None): return 0 from tensorrt_llm.bindings import SamplingConfig from tensorrt_llm.bindings.internal.batch_manager import BlockKey - from tensorrt_llm.bindings.internal.batch_manager import \ - LlmRequest as CppLlmRequest + from tensorrt_llm.bindings.internal.batch_manager import LlmRequest as CppLlmRequest block_key = BlockKey(tokens=input_tokens, lora_task_id=lora_task_id) unique_tokens = block_key.unique_tokens dummy_req = CppLlmRequest(request_id=0, diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index e4d1bd1fc62d..78da5e14e2d9 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -50,6 +50,20 @@ def _make_mgr( ) +@skip_no_cuda +@pytest.mark.parametrize("enable_attention_dp", [False, True]) +def test_python_mamba_resource_count_excludes_reserved_dummy_slots(enable_attention_dp): + max_batch_size = 4 + mgr = _make_mgr( + max_batch_size=max_batch_size, + max_draft_len=2, + enable_attention_dp=enable_attention_dp, + ) + + assert mgr.get_max_resource_count() == max_batch_size + assert len(mgr.mamba_cache_free_blocks) == max_batch_size + + @skip_no_cuda def test_padding_slot_not_held_by_parked_real(): """Padding must not resolve to a slot owned by a parked real @@ -165,6 +179,33 @@ def test_replay_update_mamba_states_uses_history_window(): assert torch.all(mgr.mamba_cache.conv[:, slot_checkpointed] == 13.0) +@skip_no_cuda +def test_replay_update_mamba_states_skips_dummy_slots(): + mgr = _make_mgr(max_batch_size=2, max_draft_len=5, use_replay_state_update=True) + mgr._prepare_mamba_cache_blocks([100]) + mgr.add_dummy_requests([CUDA_GRAPH_DUMMY_REQUEST_ID]) + + real_slot = mgr.mamba_cache_index[100] + dummy_slot = mgr.mamba_cache_index[CUDA_GRAPH_DUMMY_REQUEST_ID] + mgr.mamba_cache.prev_num_accepted_tokens[real_slot] = 13 + mgr.mamba_cache.prev_num_accepted_tokens[dummy_slot] = 13 + mgr.mamba_cache.cache_buf_idx[real_slot] = 1 + mgr.mamba_cache.cache_buf_idx[dummy_slot] = 1 + + state_indices = torch.tensor([real_slot, dummy_slot], dtype=torch.int32, device="cuda") + attn = SimpleNamespace(num_seqs=2, num_contexts=0) + mgr.update_mamba_states( + attn, + torch.tensor([3, 3], dtype=torch.int32, device="cuda"), + state_indices=state_indices, + ) + + assert mgr.mamba_cache.prev_num_accepted_tokens[real_slot].item() == 3 + assert mgr.mamba_cache.prev_num_accepted_tokens[dummy_slot].item() == 13 + assert mgr.mamba_cache.cache_buf_idx[real_slot].item() == 0 + assert mgr.mamba_cache.cache_buf_idx[dummy_slot].item() == 1 + + @skip_no_cuda def test_attention_dp_dummy_has_reserved_slot_with_batch_size_one(): mgr = _make_mgr(max_batch_size=1, max_draft_len=0, enable_attention_dp=True) diff --git a/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py b/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py index f94a5f7a28cf..04f1ff00bcf3 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py @@ -1,13 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for the Mamba SSM stochastic-rounding Philox seed plumbing. - -The Mamba SSM SR path previously generated `rand_seed` tensors via -`torch.randint(..., (1,))` on every decode forward. The cache manager now -owns a persistent per-cache-slot int64 buffer that is deterministically -initialized and rewritten on fresh request assignment. These tests pin the -contract: pure-function seed generation, deterministic allocation, and -per-slot reset without `torch.randint`. +"""Unit tests for Mamba SSM stochastic-rounding Philox seed plumbing. + +The cache manager owns a persistent per-cache-slot int64 seed buffer that is +deterministically initialized and rewritten on fresh request assignment. These +tests pin pure-function seed generation, deterministic allocation, and per-slot +reset without per-forward `torch.randint`. """ import pytest @@ -149,8 +147,7 @@ def test_padding_sentinel_does_not_churn_seeds(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_replay_path_still_allocates_seed_buffer(): - # Backward-compatibility: the replay path used to allocate the seed - # buffer; the new wiring must not regress that. + # Replay stochastic rounding uses the persistent per-slot seed buffer. mgr = _make_python_manager(sr=False, replay=True) seed_buf = mgr.get_mamba_ssm_rand_seed() assert seed_buf is not None @@ -205,9 +202,7 @@ def test_cpp_hybrid_non_replay_mtp_layer_cache_carries_rand_seed(): mamba_ssm_rand_seed on the returned SpeculativeState. The mixer's non-replay MTP SR branch (mamba2_mixer.py) reads - `layer_cache.mamba_ssm_rand_seed` and asserts non-None. Iter5 review - caught the regression where the seed was only forwarded inside the - replay branch of mamba_layer_cache; this test pins both paths.""" + `layer_cache.mamba_ssm_rand_seed` and asserts non-None.""" from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig spec_config = MTPDecodingConfig(max_draft_len=2) diff --git a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py index 6f68a9e8c121..be13c409ace5 100644 --- a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py @@ -142,14 +142,13 @@ def _maybe_skip_dtype(state_dtype, use_sr): ], ids=["fp16", "bf16", "fp32", "int8", "int16", "fp8"], ) -@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) @pytest.mark.parametrize( "T", [6, 10, 16, 27, 32, 55], ids=["T6", "T10", "T16", "T27", "T32", "T55"] ) @pytest.mark.parametrize( "write_checkpoint,rectangle_for_nowrite", [ - (True, False), # write path (rectangle_for_nowrite is ignored) + (True, False), # write path; rectangle_for_nowrite is ignored (False, False), # nowrite path via replay-style kernels (False, True), # nowrite path via dedicated rectangle kernels ], @@ -166,7 +165,6 @@ def test_replay_selective_state_update( d_state, ngroups, state_dtype, - paged_cache, T, write_checkpoint, rectangle_for_nowrite, @@ -199,12 +197,8 @@ def test_replay_selective_state_update( # max_window); for larger T it scales with np2(T). max_window = max(triton.next_power_of_2(T), 16) - if paged_cache: - cache_size = 4 - state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) - else: - cache_size = batch - state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) torch.manual_seed(42) @@ -239,9 +233,7 @@ def test_replay_selective_state_update( state0_scales = None ref_input_state = state0.float() - # Old inputs: up to `max_window` tokens per batch request, so the test - # loop can probe PNAT > T-1 (which the prior T-token setup couldn't - # reach). step1_T = max_window covers the full PNAT range we sweep. + # Seed enough history to cover every PNAT value swept below. step1_T = max_window x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) @@ -254,11 +246,7 @@ def test_replay_selective_state_update( states_buffer_f32 = torch.zeros( cache_size, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 ) - cache_idx_for_capture = ( - state_batch_indices - if paged_cache - else torch.arange(batch, device=device, dtype=torch.int32) - ) + cache_idx_for_capture = state_batch_indices out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) selective_state_update( ref_input_state.clone(), @@ -280,8 +268,8 @@ def test_replay_selective_state_update( # Build cache tensors for the replay kernel. # old_x: (cache, 2, max_window, nheads, dim) bf16 — double-buffered # old_B: (cache, 2, max_window, ngroups, dstate) bf16 — double-buffered - # old_dt: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous - # old_dA_cumsum: (cache, 2, nheads, max_window) fp32 — double-buffered, T contiguous + # old_dt: (cache, 2, nheads, max_window) fp32 — double-buffered, window contiguous + # old_dA_cumsum: (cache, 2, nheads, max_window) fp32 — double-buffered, window contiguous # cache_buf_idx: random 0s and 1s to verify indexing correctness old_x = torch.randn(cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) @@ -295,16 +283,13 @@ def test_replay_selective_state_update( # positions [0:step1_T) = [0:max_window). Whole buffer covered so PNAT # values up to max_window are exercised. Inactive buffer has random # garbage to catch indexing bugs. - slots = state_batch_indices if paged_cache else slice(None) - # old_x active-buffer fill is done in the per-slot loop below (same pattern - # as old_B/old_dt/old_dA_cumsum since old_x is now double-buffered too). # Compute processed dt and dA_cumsum for step 1 dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) # Write to each slot's active buffer based on its cache_buf_idx - slot_indices = state_batch_indices.tolist() if paged_cache else list(range(cache_size)) + slot_indices = state_batch_indices.tolist() for i, slot in enumerate(slot_indices): buf = cache_buf_idx[slot].item() batch_idx = i # maps slot back to the batch index @@ -338,6 +323,7 @@ def test_replay_selective_state_update( # Reference (fp32, starting from the same lossy-or-not state the # kernel sees). + slots = state_batch_indices ref_state_f32 = ref_input_state.clone() if k > 0: ref_state_f32[slots] = states_buffer_f32[slots, k - 1] @@ -402,36 +388,14 @@ def test_replay_selective_state_update( dt_softplus=True, state_batch_indices=state_batch_indices, state_scales=test_scales, - write_checkpoint=write_checkpoint, rectangle_for_nowrite=rectangle_for_nowrite, mode=mode, ) - # Tolerance rationale: the replay kernel uses bf16 tl.dot for four - # matmuls (dB_scaled @ old_x, C @ state, CB_scaled @ x, and C @ B in - # precompute). The reference selective_state_update uses fp32 - # element-wise MACs. The bf16 input casts lose dt_bias/A-derived bits - # that the reference keeps — per-element rounding, not accumulating. - # Prefill (ssd_chunk_scan) does identical bf16 tl.dot - # casts, so we match prefill precision exactly. Empirical: max ~1.0 at - # T<=16, ~2.0 at T=32-55; mean ~0.014; <0.02% of elements exceed 0.5. - # State dtype (fp16/bf16/fp32) doesn't shift the error — bf16 dot - # inputs dominate, not state storage. - # - # Quantized states add a per-element state quant error eps that - # propagates through C @ state in the output dot. With dstate=128 - # and C ~ N(0,1), the output channel std from this noise is roughly - # eps * sqrt(128/3) ≈ 6.5 * eps. Stack with the bf16 baseline: - # out_atol = bf16_atol + 6.5 * eps_max - # where eps_max is the worst-case per-element error at the - # post-replay SSM state magnitude (T=55 → amax ≈ 23). - # - # Per-element error (eps_max for T=55): - # int8 (uniform grid): amax/(2*127) ≈ 0.091 - # int16 (uniform grid): amax/(2*32767) ≈ 3.5e-4 - # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 (worst-case - # cell at top of channel; smaller for - # smaller-magnitude elements) + # Tolerance rationale: replay uses bf16 tl.dot while the reference + # selective_state_update uses fp32 element-wise MACs. This matches + # prefill precision but needs a small absolute tolerance. Quantized + # states add decode-grid error that propagates through C @ state. out_atol = ( {torch.int8: 1.6, torch.int16: 1.05, torch.float8_e4m3fn: 4.0}[state_dtype] if is_quantized @@ -473,15 +437,8 @@ def test_replay_selective_state_update( ref_input_state[slots] if k == 0 else states_buffer_f32[slots, k - 1] ) actual_fp32 = _dequantize_state(test_state[slots], test_scales[slots]) - # State diff = bf16_replay_error + quant_error (per element). - # The bf16 component is the SAME error source the non-quant - # test absorbs in its atol=1.0 baseline (replay's tl.dot is - # bf16-input fp32-accum; per-element error ~ 2^-7 * amax, - # empirically ≤ ~0.2 at T=55 amax≈23). Quant adds: - # int8: amax/(2*127) ≈ 0.091 worst-case - # int16: amax/(2*32767) ≈ 3.5e-4 (negligible vs bf16) - # fp8_e4m3 (variable grid): amax/16 ≈ 1.44 worst-case - # Atol = bf16_baseline (1.0) + quant_eps_max. + # State tolerance covers bf16 replay dot error plus one + # per-element quantization step for the state dtype. state_atol = { torch.int8: 1.1, torch.int16: 1.0, @@ -976,7 +933,7 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( T, max_window, batch, - None, + state_batch_indices, device, explicit_order=work_item_order, ) @@ -1079,7 +1036,6 @@ def test_replay_selective_state_update_persistent_main_device_n_writes( [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], ids=["fp16", "int8", "int16", "fp8"], ) -@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) @pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) @pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_replay_selective_state_update_philox( @@ -1088,7 +1044,6 @@ def test_replay_selective_state_update_philox( head_dim, d_state, ngroups, - paged_cache, T, mode, ): @@ -1096,7 +1051,7 @@ def test_replay_selective_state_update_philox( Verify that Philox stochastic rounding produces correct results across all SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). - Runs our kernel twice with identical inputs — once without rand_seed + Runs the replay kernel twice with identical inputs — once without rand_seed (deterministic RN), once with rand_seed (Philox SR) — and confirms: - Outputs are within bf16-dot tolerance (state perturbation ≤ 1 ULP). - State dtype is preserved. @@ -1112,12 +1067,8 @@ def test_replay_selective_state_update_philox( dtype = torch.bfloat16 assert nheads % ngroups == 0 - if paged_cache: - cache_size = 4 - state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) - else: - cache_size = batch - state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) torch.manual_seed(42) @@ -1139,7 +1090,7 @@ def test_replay_selective_state_update_philox( ) state0_scales = None - # Cache tensors (old_x now double-buffered like the others) + # Replay history cache tensors. old_x = torch.randn(cache_size, 2, T, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) @@ -1252,7 +1203,7 @@ def test_replay_selective_state_update_philox( # Per-channel decode_scale varies by 10x+ across channels (amax depends # on randn extremes), so a single flat atol can't bound it accurately — # use per-channel ULP-aware comparison. - slots = state_batch_indices if paged_cache else slice(None) + slots = state_batch_indices if is_quantized: rounded_fp32 = _dequantize_state(state_rounded[slots], scales_rounded[slots]) no_round_fp32 = _dequantize_state(state_no_round[slots], scales_no_round[slots]) @@ -1841,7 +1792,6 @@ def test_replay_heads_per_block( @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("T", [6, 16], ids=["T6", "T16"]) @pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) -@pytest.mark.parametrize("paged_cache", [False, True], ids=["contig", "paged"]) @pytest.mark.parametrize("rectangle_nowrite", [True, False], ids=["rect", "norect"]) @pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_replay_heads_per_block_multistep( @@ -1852,7 +1802,6 @@ def test_replay_heads_per_block_multistep( state_dtype, T, heads_per_block, - paged_cache, rectangle_nowrite, mode, ): @@ -1877,12 +1826,8 @@ def test_replay_heads_per_block_multistep( D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) - if paged_cache: - cache_size = 4 - state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) - else: - cache_size = batch - state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) all_x = [] all_dt = [] @@ -2031,7 +1976,7 @@ def test_replay_heads_per_block_multistep( msg=f"Output mismatch at step {step}, slot {s_local} " f"(acc={acc}) with HPB={heads_per_block}, T={T}, " f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}, " - f"paged_cache={paged_cache}, rectangle={rectangle_nowrite}", + f"rectangle={rectangle_nowrite}", ) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py index b8db6132dc0b..89beba153c45 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py @@ -20,6 +20,13 @@ import tensorrt_llm._torch.auto_deploy # noqa: F401 from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import BatchInfo +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_ITEM_WIDTH, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, +) from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( replay_selective_state_update as _real_replay_selective_state_update, ) @@ -116,6 +123,8 @@ def test_flashinfer_decode_matches_triton(mamba_env): None, # replay_old_da_cumsum None, # replay_cache_buf_idx None, # replay_prev_num_accepted + None, # replay_work_items + None, # replay_n_writes # CONSTANTS time_step_limit, chunk_size, @@ -166,26 +175,41 @@ def test_flashinfer_extend_replay_calls_replay_kernel(mamba_env, head_dim): slot_idx = torch.tensor([0], device=device, dtype=torch.int32) # Replay buffers: all zeros (first step, nothing cached yet; kernel still runs). + replay_history_size = 16 replay_old_x = torch.zeros( - max_batch_size, tokens_per_extend, num_heads, head_dim, device=device, dtype=torch.bfloat16 + max_batch_size, + 2, + replay_history_size, + num_heads, + head_dim, + device=device, + dtype=torch.bfloat16, ) replay_old_b = torch.zeros( max_batch_size, 2, - tokens_per_extend, + replay_history_size, n_groups, ssm_state_size, device=device, dtype=torch.bfloat16, ) replay_old_dt = torch.zeros( - max_batch_size, 2, num_heads, tokens_per_extend, device=device, dtype=torch.float32 + max_batch_size, 2, num_heads, replay_history_size, device=device, dtype=torch.float32 ) replay_old_da_cumsum = torch.zeros( - max_batch_size, 2, num_heads, tokens_per_extend, device=device, dtype=torch.float32 + max_batch_size, 2, num_heads, replay_history_size, device=device, dtype=torch.float32 ) replay_cache_buf_idx = torch.zeros(max_batch_size, device=device, dtype=torch.int32) replay_prev_num_accepted = torch.zeros(max_batch_size, device=device, dtype=torch.int32) + replay_work_items = torch.zeros( + max_batch_size, REPLAY_WORK_ITEM_WIDTH, device=device, dtype=torch.int32 + ) + replay_work_items[0, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = 0 + replay_work_items[0, REPLAY_WORK_CACHE_SLOT] = slot_idx[0] + replay_work_items[0, REPLAY_WORK_PNAT] = 0 + replay_work_items[0, REPLAY_WORK_CACHE_BUF_IDX] = 0 + replay_n_writes = torch.zeros(1, device=device, dtype=torch.int32) # Extend-only batch with replay mode enabled. _bi = BatchInfo() @@ -229,6 +253,8 @@ def test_flashinfer_extend_replay_calls_replay_kernel(mamba_env, head_dim): replay_old_da_cumsum, replay_cache_buf_idx, replay_prev_num_accepted, + replay_work_items, + replay_n_writes, # CONSTANTS time_step_limit, chunk_size, diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py b/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py index e68233861935..cd5f327b9632 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py @@ -12,6 +12,7 @@ # 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. +from types import SimpleNamespace from typing import List, Optional, Type import pytest @@ -24,6 +25,13 @@ from tensorrt_llm._torch.auto_deploy.shim.ad_executor import ADEngine from tensorrt_llm._torch.auto_deploy.shim.demollm import DemoEngine from tensorrt_llm._torch.auto_deploy.shim.interface import CachedSequenceInterface +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_ITEM_WIDTH, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, +) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -599,6 +607,72 @@ def get_tokens(self, _beam: int) -> List[int]: return self._tokens +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cached_sequence_interface_prepare_replay_metadata_write_first(): + device = torch.device("cuda") + cache_seq_interface = CachedSequenceInterface( + max_seq_len=64, + max_batch_size=4, + max_num_tokens=64, + device=device, + kv_cache_config=KvCacheConfig(tokens_per_block=16), + ) + cache_seq_interface.to(device) + + prev_num_accepted_tokens = torch.zeros(6, dtype=torch.int32, device=device) + cache_buf_idx = torch.zeros(6, dtype=torch.int32, device=device) + slot_idx = torch.tensor([3, 1, 5, 0], dtype=torch.long, device=device) + prev_num_accepted_tokens[slot_idx] = torch.tensor( + [13, 7, 14, 2], dtype=torch.int32, device=device + ) + cache_buf_idx[slot_idx] = torch.tensor([1, 0, 1, 0], dtype=torch.int32, device=device) + + cache_seq_interface._replay_work_items = torch.empty( + cache_seq_interface.info.max_num_state_slots, + REPLAY_WORK_ITEM_WIDTH, + dtype=torch.int32, + device=device, + ) + cache_seq_interface._replay_n_writes = torch.zeros(1, dtype=torch.int32, device=device) + cache_seq_interface._kv_cache_manager = SimpleNamespace( + shutdown=lambda: None, + get_replay_state_update_metadata=lambda: SimpleNamespace( + prev_num_accepted_tokens=prev_num_accepted_tokens, + cache_buf_idx=cache_buf_idx, + replay_step_width=6, + replay_history_size=16, + ), + ) + + cache_seq_interface.info.batch_info.update([0, 0, 4, 4, 0, 0]) + cache_seq_interface.info.batch_info.update_use_replay(True) + cache_seq_interface.info._input_buffer.copy_("slot_idx", slot_idx) + + cache_seq_interface.prepare_replay_metadata() + + expected = torch.tensor( + [ + [0, 3, 13, 1], + [2, 5, 14, 1], + [1, 1, 7, 0], + [3, 0, 2, 0], + ], + dtype=torch.int32, + device=device, + ) + actual = cache_seq_interface._replay_work_items[:4] + assert cache_seq_interface._replay_n_writes.item() == 2 + assert torch.equal( + actual[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH], + expected[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH], + ) + assert torch.equal(actual[:, REPLAY_WORK_CACHE_SLOT], expected[:, REPLAY_WORK_CACHE_SLOT]) + assert torch.equal(actual[:, REPLAY_WORK_PNAT], expected[:, REPLAY_WORK_PNAT]) + assert torch.equal(actual[:, REPLAY_WORK_CACHE_BUF_IDX], expected[:, REPLAY_WORK_CACHE_BUF_IDX]) + + cache_seq_interface.shutdown() + + def test_ad_engine_prepare_inputs_with_hybrid_cache_manager(): """Test ADEngine _prepare_inputs uses mamba_cache_index when available.""" seed = 42 From 63ad831013c2f8375074d7580ed1a4f6750f45a0 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Thu, 11 Jun 2026 22:16:56 -0700 Subject: [PATCH 80/89] Clean up Mamba replay seed handling and precompute tuning Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../_torch/modules/mamba/mamba2_mixer.py | 13 ++--- .../mamba/replay_selective_state_update.py | 51 +++++++++---------- 2 files changed, 29 insertions(+), 35 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index fcc341042cb9..03f9d8068679 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -523,11 +523,9 @@ def convert_dt(): philox_kwargs = {} if use_stochastic_rounding: - # Both replay and flashinfer read from the cache manager's - # persistent per-slot Philox seed buffer; replay indexes by - # cache_batch_idx, flashinfer reads slot 0 from a (1,) - # view. In-place add_(1) keeps CUDA-graph replay fresh - # without allocating any new CUDA tensors per forward. + # Both replay and flashinfer use a single Philox seed. The + # cache manager owns the persistent buffer; passing a (1,) + # view avoids allocating CUDA tensors per forward. rand_seed = layer_cache.mamba_ssm_rand_seed assert rand_seed is not None, ( "Mamba SSM stochastic rounding is enabled but the " @@ -535,10 +533,7 @@ def convert_dt(): "_util.py passes mamba_ssm_stochastic_rounding=True " "to the cache manager.") rand_seed.add_(1) - if use_replay: - philox_kwargs['rand_seed'] = rand_seed - else: - philox_kwargs['rand_seed'] = rand_seed[:1] + philox_kwargs['rand_seed'] = rand_seed[:1] philox_kwargs['philox_rounds'] = self._philox_rounds if use_replay: diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 965f71d844a0..11bfd883aac9 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -1331,7 +1331,7 @@ def _persistent_main_impl( else: RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized - rand_seed = tl.load(rand_seed_ptr + cache_batch_idx) + rand_seed = tl.load(rand_seed_ptr) base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head # Number of unique randoms per row = dstate / RAND_DIVISOR. # randint4x emits 4 randoms per offset, so use that / 4 offsets. @@ -1456,7 +1456,6 @@ def _persistent_main_impl( step_x, mask=t_mask[:, None] & m_mask[None, :], ) - step_x_for_dot = step_x.to(tl.bfloat16) step_x = step_x.to(tl.float32) cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head @@ -1472,7 +1471,7 @@ def _persistent_main_impl( ) init_out = tl.dot(C_tile.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] - cb_out = tl.dot(CB_scaled.to(tl.bfloat16), step_x_for_dot) + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), step_x.to(tl.bfloat16)) output_tile = init_out + cb_out if HAS_D: @@ -1630,6 +1629,8 @@ def _persistent_rectangle_impl( mask=m_mask, other=1.0, ).to(tl.float32) + else: + state = state.to(tl.float32) # Group / pointer offset setup group_idx = pid_h // nheads_ngroups_ratio @@ -1662,7 +1663,7 @@ def _persistent_rectangle_impl( + offs_m[None, :] * stride_old_x_dim, mask=is_history_position[:, None] & m_mask[None, :], other=0.0, - ) + ).to(tl.float32) if WAIT_FOR_PDL_PREDECESSOR: _gdc_wait_with_memory_clobber() @@ -1686,7 +1687,7 @@ def _persistent_rectangle_impl( ) step_x_in_window_for_dot = step_x_in_window.to(tl.bfloat16) - x_window_for_dot = (history_x.to(tl.bfloat16) + step_x_in_window_for_dot).to(tl.bfloat16) + x_window_for_dot = history_x + step_x_in_window.to(tl.float32) if HAS_D: step_in_window_selector = offs_t[:, None] == ( @@ -1712,7 +1713,7 @@ def _persistent_rectangle_impl( if QUANT_MAX > 0.0: state_out = state_out * decode_scale[None, :] - token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_window_for_dot) + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_window_for_dot.to(tl.bfloat16)) output_tile = state_out + token_out @@ -2309,7 +2310,7 @@ def _persistent_main_kernel( "_block_size_m": 8, "_cta_per_sm": 6, "_flatten": False, - "_heads_per_block": 4, + "_heads_per_block": 1, "_num_loop_stages": 1, "_num_stages": 4, "_num_warps": 1, @@ -2321,7 +2322,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=6.57us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=1, score=6.83us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX ( 32, "persistent_dynamic", @@ -2849,7 +2850,7 @@ def _persistent_main_kernel( "_num_loop_stages": 1, "_num_stages": 2, "_num_warps": 1, - "_precompute_num_warps": 8, + "_precompute_num_warps": 2, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": False, @@ -2857,7 +2858,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=2, score=7.05us (B200 PDL-hoist default 5x200) + ), # raw_batch=2, score=6.72us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX ( 64, "persistent_dynamic", @@ -3025,14 +3026,14 @@ def _persistent_main_kernel( "_cta_per_sm_nowrite": 8, "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 8, + "_heads_per_block": 16, "_num_loop_stages_nowrite": 5, "_num_loop_stages_write": 1, "_num_stages_nowrite": 4, "_num_stages_write": 1, "_num_warps_nowrite": 1, "_num_warps_write": 1, - "_precompute_num_warps": 2, + "_precompute_num_warps": 4, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, @@ -3041,7 +3042,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=69.44us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=512, score=67.94us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX ( 16384, "persistent_main", @@ -3277,7 +3278,7 @@ def _persistent_main_kernel( "_cta_per_sm_nowrite": 4, "_cta_per_sm_write": 8, "_flatten": False, - "_heads_per_block": 8, + "_heads_per_block": 16, "_num_loop_stages_nowrite": 5, "_num_loop_stages_write": 1, "_num_stages_nowrite": 4, @@ -3293,7 +3294,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=88.87us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=512, score=88.57us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX ( 16384, "persistent_main", @@ -3538,13 +3539,13 @@ def replay_selective_state_update( D: (nheads, dim) optional feed-through parameter. z: (batch, T, nheads, dim) optional silu gate. dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). - rand_seed: optional (cache_size,) int64 CUDA tensor of per-cache-slot - Philox PRNG seeds. The caller bumps this tensor in-place for each - replay invocation so CUDA graph replay still gets fresh draws. The - kernel indexes it by cache_batch_idx. When provided, state is - stochastically rounded on store. Supported for state.dtype in - (fp16, int8, int16, fp8_e4m3fn). fp16+SR and fp8+SR both require - sm_100a (Blackwell B200+) — wrapper asserts this loudly. + rand_seed: optional single-element int64 CUDA tensor of Philox PRNG + seed. The caller bumps this tensor in-place for each replay + invocation so CUDA graph replay still gets fresh draws. When + provided, state is stochastically rounded on store. Supported for + state.dtype in (fp16, int8, int16, fp8_e4m3fn). fp16+SR and + fp8+SR both require sm_100a (Blackwell B200+) — wrapper asserts + this loudly. When None, standard deterministic rounding is used. philox_rounds: number of Philox PRNG rounds (default 10). state_scales: required when state.dtype in (int8, int16, fp8_e4m3fn). @@ -3869,10 +3870,8 @@ def replay_selective_state_update( assert rand_seed.dim() == 1, ( f"rand_seed must be a 1D tensor; got shape {tuple(rand_seed.shape)}" ) - if rand_seed.shape[0] == 1 and cache_size > 1: - rand_seed = rand_seed.expand(cache_size).contiguous() - assert rand_seed.shape[0] >= cache_size, ( - f"rand_seed must have length 1 or >= cache_size ({cache_size}); " + assert rand_seed.shape[0] == 1, ( + "rand_seed must have length 1; " f"got shape {tuple(rand_seed.shape)}" ) From c1365f505a5dfeb897025f035cbe335a81adbab6 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 12 Jun 2026 01:11:18 -0700 Subject: [PATCH 81/89] Fix Mamba replay gating without spec decode Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 3 +- tensorrt_llm/_torch/pyexecutor/_util.py | 6 +- .../_torch/pyexecutor/mamba_cache_manager.py | 143 +++++++++++------- .../executor/test_mamba_cache_manager.py | 18 ++- 4 files changed, 115 insertions(+), 55 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 11bfd883aac9..030fa06c9be0 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -3871,8 +3871,7 @@ def replay_selective_state_update( f"rand_seed must be a 1D tensor; got shape {tuple(rand_seed.shape)}" ) assert rand_seed.shape[0] == 1, ( - "rand_seed must have length 1; " - f"got shape {tuple(rand_seed.shape)}" + f"rand_seed must have length 1; got shape {tuple(rand_seed.shape)}" ) tie_hdim = ( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 42216f83c3c3..d40e5ddc8d18 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1292,7 +1292,11 @@ def _create_kv_cache_manager( quant_config, 'mamba_ssm_stochastic_rounding', False) if quant_config is not None else False - use_replay = sm >= 80 + use_replay = spec_config is not None and sm >= 80 + if spec_config is None: + logger.info( + "Replay kernel requires speculative decoding; using non-replay path" + ) # Block reuse (prefix caching): replay leaves SSM state at a # checkpoint after speculation. The next decode step replays forward diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 3512ad337c86..a43b8cd20178 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -622,9 +622,14 @@ def __init__( # mamba cache index, maps request_id -> state indices self.mamba_cache_index: Dict[int, int] = {} self._dummy_request_ids: set[int] = set() - self._dummy_slot_mask = torch.zeros(max_batch_size, - dtype=torch.bool, - device=device) + # Batch-order mask aligned with state_indices; duplicate dummy request + # IDs mark every batch row even when they share one cache slot. + self._dummy_request_mask = torch.zeros(max_batch_size, + dtype=torch.bool, + device=device) + self._dummy_request_mask_host = torch.zeros(max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned()) # Permanent slot shared by every CUDA-graph padding sentinel id # (CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, one per @@ -682,7 +687,6 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]): raise RuntimeError("run out of mamba cache blocks") block = self.mamba_cache_free_blocks.pop() self.mamba_cache_index[r] = block - self._dummy_slot_mask[block] = False if (isinstance(self.mamba_cache, self.SpeculativeState) and self._use_replay_state_update): self.mamba_cache.prev_num_accepted_tokens[block] = 0 @@ -699,6 +703,8 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]): self._seed_rank_offset)) def prepare_resources(self, scheduled_batch: ScheduledRequests): + requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) context_ids = [ i.py_request_id for i in scheduled_batch.context_requests ] @@ -707,6 +713,7 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): ] request_ids = context_ids + generation_ids self._prepare_mamba_cache_blocks(request_ids) + self._refresh_dummy_request_mask([req.is_dummy for req in requests]) def _is_padding_sentinel(self, request_id: int) -> bool: # cuda_graph_runner caches one dummy per runtime_draft_len value @@ -728,7 +735,6 @@ def add_dummy_requests(self, request_ids: List[int], **kwargs): for r in request_ids: if r in self.mamba_cache_index: block = self.mamba_cache_index[r] - self._dummy_slot_mask[block] = True if (isinstance(self.mamba_cache, self.SpeculativeState) and self._use_replay_state_update): self.mamba_cache.prev_num_accepted_tokens[block] = 0 @@ -744,7 +750,6 @@ def add_dummy_requests(self, request_ids: List[int], **kwargs): raise RuntimeError("run out of mamba cache blocks") block = self.mamba_cache_free_blocks.pop() self.mamba_cache_index[r] = block - self._dummy_slot_mask[block] = True if (isinstance(self.mamba_cache, self.SpeculativeState) and self._use_replay_state_update): self.mamba_cache.prev_num_accepted_tokens[block] = 0 @@ -754,19 +759,33 @@ def free_resources(self, request: LlmRequest): request_id = request.py_request_id if request_id not in self.mamba_cache_index: return - is_dummy = request_id in self._dummy_request_ids self._dummy_request_ids.discard(request_id) block = self.mamba_cache_index.pop(request_id) # Reserved slots must not re-enter the real-request free pool. if block != self._padding_slot and \ block != self._attention_dp_dummy_slot: - if is_dummy: - self._dummy_slot_mask[block] = False self.mamba_cache_free_blocks.append(block) def get_state_indices(self, request_ids: List[int], is_padding: List[bool]) -> List[int]: - return [self.mamba_cache_index[rid] for rid in request_ids] + assert len(request_ids) == len(is_padding) + indices = [self.mamba_cache_index[rid] for rid in request_ids] + is_dummy = [ + rid in self._dummy_request_ids or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_( + torch.as_tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask.copy_(self._dummy_request_mask_host, + non_blocking=True) def get_conv_states(self, layer_idx: int) -> torch.Tensor: layer_offset = self.mamba_layer_offsets[layer_idx] @@ -850,7 +869,7 @@ def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: @property def use_replay_state_update(self) -> bool: - return self._use_replay_state_update + return self.get_replay_state_update_metadata() is not None def get_replay_state_update_metadata( self) -> Optional[ReplayStateUpdateMetadata]: @@ -925,14 +944,15 @@ def update_mamba_states(self, attn_metadata: "AttentionMetadata", wrote_checkpoint, accepted_tokens, prev_num_accepted_tokens + accepted_tokens) cache_buf_idx = self.mamba_cache.cache_buf_idx[state_indices_d] - is_dummy_slot = self._dummy_slot_mask[state_indices_d] - next_num_accepted_tokens = torch.where(is_dummy_slot, + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + next_num_accepted_tokens = torch.where(is_dummy_request, prev_num_accepted_tokens, next_num_accepted_tokens) self.mamba_cache.prev_num_accepted_tokens[state_indices_d] = \ next_num_accepted_tokens self.mamba_cache.cache_buf_idx[state_indices_d] = \ - torch.where(is_dummy_slot, cache_buf_idx, + torch.where(is_dummy_request, cache_buf_idx, torch.where(wrote_checkpoint, 1 - cache_buf_idx, cache_buf_idx)) else: @@ -1073,7 +1093,7 @@ def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: @property def use_replay_state_update(self) -> bool: - return getattr(self._impl, 'use_replay_state_update', False) + return self.get_replay_state_update_metadata() is not None def get_replay_state_update_metadata( self) -> Optional[ReplayStateUpdateMetadata]: @@ -1665,8 +1685,11 @@ def __init__( dtype=torch.long, device="cpu") self._request_id_to_state_index = {} - self._dummy_slot_mask = None - self._dummy_slot_mask_host = None + self._request_id_to_is_dummy = {} + # Batch-order mask aligned with state_indices; duplicate dummy request + # IDs mark every batch row even when they share one cache slot. + self._dummy_request_mask = None + self._dummy_request_mask_host = None self.kv_cache_config = kv_cache_config self.is_estimating_kv_cache = is_estimating_kv_cache @@ -1756,8 +1779,8 @@ def shutdown(self): self.prev_num_accepted_tokens = None self.cache_buf_idx = None self.mamba_ssm_rand_seed = None - self._dummy_slot_mask = None - self._dummy_slot_mask_host = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None self.old_x = None self.old_B = None self.old_dt = None @@ -1823,8 +1846,6 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests): # we skip refresh_blocks entirely when nothing was scheduled. self._pending_state_transfers = self.impl.copy_linear_attention_block_batch( self.requests) - if self._pending_state_transfers: - logger.info("Need to transfer mamba state blocks") self._setup_state_indices() # Reset replay double-buffer state for fresh context blocks. A reused # block (prefix-cache hit or block recycled across requests) may carry @@ -1933,13 +1954,15 @@ def update_mamba_states(self, next_num_accepted_tokens = torch.where( wrote_checkpoint, accepted, prev_num_accepted_tokens + accepted) cache_buf_idx = self.cache_buf_idx[slots] - is_dummy_slot = self._dummy_slot_mask[slots] - next_num_accepted_tokens = torch.where(is_dummy_slot, + assert self._dummy_request_mask is not None + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + next_num_accepted_tokens = torch.where(is_dummy_request, prev_num_accepted_tokens, next_num_accepted_tokens) self.prev_num_accepted_tokens[slots] = next_num_accepted_tokens self.cache_buf_idx[slots] = torch.where( - is_dummy_slot, cache_buf_idx, + is_dummy_request, cache_buf_idx, torch.where(wrote_checkpoint, 1 - cache_buf_idx, cache_buf_idx)) else: # Legacy: copy the accepted SSM state from the intermediate buffer. @@ -1956,6 +1979,19 @@ def update_mamba_states(self, src_state_indices, num_accepted_draft_tokens, state_indices_d) + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + if self._dummy_request_mask is None: + return + + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_( + torch.tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask.copy_(self._dummy_request_mask_host, + non_blocking=True) + def get_num_available_tokens(self, token_num_upper_bound: int, max_num_draft_tokens: int = 0, @@ -2045,6 +2081,7 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): if request in self.requests: self.requests.remove(request) self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) super().free_resources(request, pin_on_release) def _setup_state_indices(self) -> None: @@ -2099,25 +2136,15 @@ def _setup_state_indices(self) -> None: self.cuda_state_indices.copy_(self._host_state_indices, non_blocking=True) - if self._dummy_slot_mask is not None: - self._dummy_slot_mask_host.zero_() - for i, req in enumerate(self.requests): - if req.is_dummy: - self._dummy_slot_mask_host[ - self._host_state_indices[i].item()] = True - self._dummy_slot_mask.copy_(self._dummy_slot_mask_host, - non_blocking=True) - if (self.prev_num_accepted_tokens is not None - and self.cache_buf_idx is not None): - self.prev_num_accepted_tokens.masked_fill_( - self._dummy_slot_mask, 0) - self.cache_buf_idx.masked_fill_(self._dummy_slot_mask, 0) + self._refresh_dummy_request_mask( + [req.is_dummy for req in self.requests]) # Build request_id → pool block offset mapping so that # get_state_indices can return indices in arbitrary request order. for i, req in enumerate(self.requests): self._request_id_to_state_index[ req.py_request_id] = self._host_state_indices[i].item() + self._request_id_to_is_dummy[req.py_request_id] = req.is_dummy def get_state_indices(self, request_ids: Optional[List[int]] = None, @@ -2127,7 +2154,18 @@ def get_state_indices(self, # not the internal self.requests order. This is critical when # the batch is reordered after prepare_resources (e.g. disagg # serving sorts generation_requests by py_batch_idx). - return [self._request_id_to_state_index[rid] for rid in request_ids] + indices = [ + self._request_id_to_state_index[rid] for rid in request_ids + ] + if is_padding is None: + is_padding = [False] * len(request_ids) + assert len(request_ids) == len(is_padding) + is_dummy = [ + self._request_id_to_is_dummy.get(rid, False) or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices return self.cuda_state_indices def calc_next_context_chunk_size(self, request: LlmRequest) -> int: @@ -2262,8 +2300,8 @@ def _setup_replay_buffers(self, spec_config) -> None: self.old_B = None self.old_dt = None self.old_dA_cumsum = None - self._dummy_slot_mask = None - self._dummy_slot_mask_host = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None return history_size = self.replay_history_size @@ -2278,12 +2316,12 @@ def _setup_replay_buffers(self, spec_config) -> None: self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) - self._dummy_slot_mask = torch.zeros(cache_size, - dtype=torch.bool, - device=device) - self._dummy_slot_mask_host = torch.zeros(cache_size, - dtype=torch.bool, - pin_memory=prefer_pinned()) + self._dummy_request_mask = torch.zeros(self.max_batch_size, + dtype=torch.bool, + device=device) + self._dummy_request_mask_host = torch.zeros(self.max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned()) self.old_x = torch.zeros(num_local_mamba_layers, cache_size, 2, @@ -2318,18 +2356,21 @@ def _setup_replay_buffers(self, spec_config) -> None: @property def use_replay_state_update(self) -> bool: - return self._use_replay_state_update + return self.get_replay_state_update_metadata() is not None def get_replay_state_update_metadata( self) -> Optional[ReplayStateUpdateMetadata]: + prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', + None) + cache_buf_idx = getattr(self, 'cache_buf_idx', None) if (not self._use_replay_state_update - or self.prev_num_accepted_tokens is None - or self.cache_buf_idx is None or self.replay_step_width is None + or prev_num_accepted_tokens is None or cache_buf_idx is None + or self.replay_step_width is None or self.replay_history_size is None): return None return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=self.prev_num_accepted_tokens, - cache_buf_idx=self.cache_buf_idx, + prev_num_accepted_tokens=prev_num_accepted_tokens, + cache_buf_idx=cache_buf_idx, replay_step_width=self.replay_step_width, replay_history_size=self.replay_history_size) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index c14c065f528f..5f73a1f86881 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -64,6 +64,18 @@ def test_python_mamba_resource_count_excludes_reserved_dummy_slots(enable_attent assert len(mgr.mamba_cache_free_blocks) == max_batch_size +@skip_no_cuda +def test_replay_inactive_without_spec_config(): + mgr = _make_mgr( + max_batch_size=2, + max_draft_len=None, + use_replay_state_update=True, + ) + + assert mgr.use_replay_state_update is False + assert mgr.get_replay_state_update_metadata() is None + + @skip_no_cuda def test_padding_slot_not_held_by_parked_real(): """Padding must not resolve to a slot owned by a parked real @@ -192,7 +204,11 @@ def test_replay_update_mamba_states_skips_dummy_slots(): mgr.mamba_cache.cache_buf_idx[real_slot] = 1 mgr.mamba_cache.cache_buf_idx[dummy_slot] = 1 - state_indices = torch.tensor([real_slot, dummy_slot], dtype=torch.int32, device="cuda") + state_indices = torch.tensor( + mgr.get_state_indices([100, CUDA_GRAPH_DUMMY_REQUEST_ID], [False, True]), + dtype=torch.int32, + device="cuda", + ) attn = SimpleNamespace(num_seqs=2, num_contexts=0) mgr.update_mamba_states( attn, From 12a46429c8173fcdd17adb93e82a36538017bd17 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 12 Jun 2026 09:41:55 -0700 Subject: [PATCH 82/89] Fix AutoDeploy replay resource validation Treat replay work-item and n_writes resources as manager-backed when speculative replay is enabled, so AD disaggregated cache validation does not reject them as unmanaged persistent state. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- tensorrt_llm/_torch/auto_deploy/shim/interface.py | 6 ++++++ .../singlegpu/shim/test_cached_sequence_interface.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 543ae1710920..12855b2028a6 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -1025,6 +1025,8 @@ def _validate_no_unmanaged_persistent_caches( replay_old_dA_cumsum: list, replay_cache_buf_idx: list, replay_prev_num_accepted: list, + replay_work_items: list, + replay_n_writes: list, ) -> None: """Validate persistent cache resources are cache-manager backed. @@ -1050,6 +1052,8 @@ def _validate_no_unmanaged_persistent_caches( replay_old_dA_cumsum, replay_cache_buf_idx, replay_prev_num_accepted, + replay_work_items, + replay_n_writes, ): managed_names.update(name for name, _ in replay_resources) @@ -1305,6 +1309,8 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: replay_old_dA_cumsum, replay_cache_buf_idx, replay_prev_num_accepted, + replay_work_items, + replay_n_writes, ) # 8. Allocate remaining unmanaged resources diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py b/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py index 2dab50197afc..57ee76d03d10 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py @@ -33,11 +33,13 @@ IntermediateSSMStateHandler, KVPagedResourceHandler, ReplayCacheBufIdxHandler, + ReplayNWritesHandler, ReplayOldBHandler, ReplayOldDAcumsumHandler, ReplayOldDtHandler, ReplayOldXHandler, ReplayPrevNumAcceptedHandler, + ReplayWorkItemsHandler, SequenceInfo, SSMResourceHandler, StateResourceHandler, @@ -1390,6 +1392,10 @@ def _add_managed_spec_replay_resources(interface, num_layers=2): replay_names.append( interface.add_resource(f"replay_prev_num_accepted_{i}", ReplayPrevNumAcceptedHandler()) ) + replay_names.append( + interface.add_resource(f"replay_work_items_{i}", ReplayWorkItemsHandler()) + ) + replay_names.append(interface.add_resource(f"replay_n_writes_{i}", ReplayNWritesHandler())) return replay_names From a90d5cc4e63c32025ad08e9926b1362d444d92b1 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 12 Jun 2026 12:50:41 -0700 Subject: [PATCH 83/89] Fix replay benchmark kwargs and int8 tunings The benchmark still passed precompute_num_stages, maxnreg, and num_ctas after those kernel kwargs were removed, causing every bench invocation to fail with TypeError. Remove the stale plumbing and fix _gen_from_cell_list yield to match the reduced unpack count. Also import the B200 precompute retunes for the int8/SR batch 256 and 512 default tuning cells. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 30 ++++----- ...benchmark_replay_selective_state_update.py | 64 +------------------ 2 files changed, 16 insertions(+), 78 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 030fa06c9be0..85425028e64f 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -2748,15 +2748,15 @@ def _persistent_main_kernel( "_cta_per_sm_nowrite": 8, "_cta_per_sm_write": 3, "_flatten": False, - "_heads_per_block": 8, - "_num_loop_stages_nowrite": 3, - "_num_loop_stages_write": 3, + "_heads_per_block": 2, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, "_num_stages_nowrite": 4, - "_num_stages_write": 3, + "_num_stages_write": 1, "_num_warps_nowrite": 1, "_num_warps_write": 4, - "_precompute_num_warps": 8, - "_use_tma_rect_load": True, + "_precompute_num_warps": 1, + "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, @@ -2764,33 +2764,33 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=256, score=49.64us (B200 PDL-hoist default 5x200) + ), # raw_batch=256, score=47.98us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX ( 8192, "persistent_main", { "_block_size_m_nowrite": 64, "_block_size_m_write": 64, - "_cta_per_sm_nowrite": 8, + "_cta_per_sm_nowrite": 6, "_cta_per_sm_write": 6, "_flatten": False, - "_heads_per_block": 4, - "_num_loop_stages_nowrite": 4, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 2, "_num_loop_stages_write": 2, - "_num_stages_nowrite": 3, + "_num_stages_nowrite": 5, "_num_stages_write": 4, - "_num_warps_nowrite": 1, + "_num_warps_nowrite": 2, "_num_warps_write": 4, - "_precompute_num_warps": 1, + "_precompute_num_warps": 4, "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, - "nowrite_first": True, + "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=82.99us (B200 PDL-hoist default 5x200) + ), # raw_batch=512, score=83.00us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX ( 16384, "persistent_main", diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 0a61eefb4847..80167ff2048a 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -2472,10 +2472,7 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): "num_loop_stages", "num_loop_stages_write", "num_loop_stages_nowrite" ) pw_vals = _ps(args.precompute_num_warps) - ps_vals = _ps(args.precompute_num_stages) h_vals = _ps(args.heads_per_block) - mr_vals = _ps(args.maxnreg) - ct_vals = _ps(args.num_ctas) fl_vals = _ps(args.flatten) wsp_vals = _ps(args.warp_specialize) trl_vals = _ps(args.use_tma_rect_load) @@ -2488,17 +2485,14 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): for (mw, mnw), (ww, wnw), (sw, snw), (cw, cnw), ( lw, lnw, - ), pw, ps_, h, mr, ct, fl, wsp, trl, twl, tnl, tws in _it.product( + ), pw, h, fl, wsp, trl, twl, tnl, tws in _it.product( m_pairs, w_pairs, ns_pairs, cps_pairs, ls_pairs, pw_vals, - ps_vals, h_vals, - mr_vals, - ct_vals, fl_vals, wsp_vals, trl_vals, @@ -2519,10 +2513,7 @@ def _split_or_pair(shared_attr, w_attr, nw_attr): ("num_loop_stages_write", lw), ("num_loop_stages_nowrite", lnw), ("precompute_num_warps", pw), - ("precompute_num_stages", ps_), ("heads_per_block", h), - ("maxnreg", mr), - ("num_ctas", ct), ("flatten", fl), ("warp_specialize", wsp), ("use_tma_rect_load", trl), @@ -2798,10 +2789,7 @@ def _parse_sweep(val): num_warps_values = _parse_sweep(args.num_warps) num_stages_values = _parse_sweep(args.num_stages) precompute_num_warps_values = _parse_sweep(args.precompute_num_warps) - precompute_num_stages_values = _parse_sweep(args.precompute_num_stages) heads_per_block_values = _parse_sweep(args.heads_per_block) - maxnreg_values = _parse_sweep(args.maxnreg) - num_ctas_values = _parse_sweep(args.num_ctas) # Persistent-only sweep dims; ignored when the cell's mode != persistent_main. cta_per_sm_values = _parse_sweep(args.cta_per_sm) num_loop_stages_values = _parse_sweep(args.num_loop_stages) @@ -3142,10 +3130,7 @@ def _run_pr3324_baseline(): num_stages_write_values, num_stages_nowrite_values, precompute_num_warps_values, - precompute_num_stages_values, heads_per_block_values, - maxnreg_values, - num_ctas_values, cta_per_sm_write_values, cta_per_sm_nowrite_values, num_loop_stages_write_values, @@ -3168,10 +3153,7 @@ def _run_pr3324_baseline(): num_stages_values, [None], precompute_num_warps_values, - precompute_num_stages_values, heads_per_block_values, - maxnreg_values, - num_ctas_values, cta_per_sm_values, [None], num_loop_stages_values, @@ -3214,10 +3196,7 @@ def _gen_from_cell_list(): d.get("Sw"), d.get("Snw"), d.get("pW"), - d.get("pS"), d.get("H"), - d.get("R"), - d.get("CT"), d.get("CPSw"), d.get("CPSnw"), d.get("LSw"), @@ -3242,10 +3221,7 @@ def _gen_from_cell_list(): num_stages_w, num_stages_nw, precompute_num_warps, - precompute_num_stages, heads_per_block, - maxnreg, - num_ctas, cta_per_sm_w, cta_per_sm_nw, num_loop_stages_w, @@ -3321,10 +3297,7 @@ def _run_incr( num_warps=num_warps, num_stages=num_stages, precompute_num_warps=precompute_num_warps, - precompute_num_stages=precompute_num_stages, heads_per_block=heads_per_block, - maxnreg=maxnreg, - num_ctas=num_ctas, cta_per_sm=cta_per_sm, num_loop_stages=num_loop_stages, flatten=flatten, @@ -3391,10 +3364,7 @@ def _run_incr( _num_warps=num_warps, _num_stages=num_stages, _precompute_num_warps=precompute_num_warps, - _precompute_num_stages=precompute_num_stages, _heads_per_block=heads_per_block, - _maxnreg=maxnreg, - _num_ctas=num_ctas, # Per-main overrides (None = tied to shared above; explicit # only when the inner loop is iterating split axes). _block_size_m_write=block_size_m_w if _any_split else None, @@ -3438,14 +3408,6 @@ def _emit_split(name_w, name_nw, val_w, val_nw): _emit_split("Sw", "Snw", num_stages_w, num_stages_nw) parts.append(f"pW={_val(precompute_num_warps)}") parts.append(f"H={_val(heads_per_block)}") - optional_tag_parts = ( - ("pS", precompute_num_stages), - ("R", maxnreg), - ("CT", num_ctas), - ) - for _name, _value in optional_tag_parts: - if _value is not None: - parts.append(f"{_name}={_value}") # Persistent-only knobs (only meaningful when MODE=persistent_main; # printed unconditionally so output rows are uniformly comparable # across modes when the user passed these sweeps). @@ -3841,10 +3803,7 @@ def _submit_result_job( "LSw": "num_loop_stages_write", "LSnw": "num_loop_stages_nowrite", "pW": "precompute_num_warps", - "pS": "precompute_num_stages", "H": "heads_per_block", - "R": "maxnreg", - "CT": "num_ctas", "FL": "flatten", "WS": "warp_specialize", "TMARL": "use_tma_rect_load", @@ -3958,10 +3917,7 @@ def _load_cell_list_into_args(args) -> None: "LSw": "num_loop_stages_w", "LSnw": "num_loop_stages_nw", "pW": "precompute_num_warps", - "pS": "precompute_num_stages", "H": "heads_per_block", - "R": "maxnreg", - "CT": "num_ctas", "FL": "flatten", "WS": "warp_specialize", "TMARL": "use_tma_rect_load", @@ -4731,12 +4687,6 @@ def _parse_args() -> argparse.Namespace: default=None, help="Override num_warps for precompute kernel (comma-separated sweep).", ) - parser.add_argument( - "--precompute-num-stages", - type=str, - default=None, - help="Override num_stages for precompute kernel (comma-separated sweep).", - ) parser.add_argument( "--max-window", type=int, @@ -4777,18 +4727,6 @@ def _parse_args() -> argparse.Namespace: default=None, help="Override HEADS_PER_BLOCK for precompute kernel (comma-separated sweep).", ) - parser.add_argument( - "--maxnreg", - type=str, - default=None, - help="Override maxnreg for the main kernel (comma-separated sweep).", - ) - parser.add_argument( - "--num-ctas", - type=str, - default=None, - help="Override num_ctas for the main kernel (comma-separated sweep).", - ) parser.add_argument( "--cta-per-sm", type=str, From 061ce29c02e5f948b187770d2d8ff0b0b18a44a1 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Fri, 12 Jun 2026 21:03:33 -0700 Subject: [PATCH 84/89] Handle strided TMA state layout for replay Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 91 +++++++++--- ...benchmark_replay_selective_state_update.py | 94 ++++++++++-- .../test_replay_selective_state_update.py | 140 ++++++++++++++++++ 3 files changed, 293 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 85425028e64f..4c0ecdd3edd2 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -3470,6 +3470,7 @@ def replay_selective_state_update( _use_tma_replay_write_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=True _use_tma_replay_write_store: bool | None = None, # SSM state store when WRITE_CHECKPOINT=True _use_tma_replay_nowrite_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=False + _require_tma_state_layout: bool = False, # Persistent-mode tuning kwargs (consulted for both pd and pm; pd uses # _cta_per_sm / _num_loop_stages, pm uses the _write/_nowrite splits): # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally @@ -3590,6 +3591,12 @@ def replay_selective_state_update( # launch order only. # Note: mode-and-knob resolution from the default-tuning table happens # below, after we have `batch` and `nheads`. + caller_forced_tma_state = ( + _use_tma_rect_load is True + or _use_tma_replay_write_load is True + or _use_tma_replay_write_store is True + or _use_tma_replay_nowrite_load is True + ) # --- Hardware support gates --- # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX @@ -3990,42 +3997,80 @@ def replay_selective_state_update( state_scales_arg = state # any valid pointer; gated by QUANT_MAX == 0 state_scales_strides = (0, 0, 0) - # Per-path TMA descriptors for state — write-side and nowrite-side. Each - # kernel launch consumes the descriptor whose block_shape[0] matches its - # BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need - # distinct descriptors; otherwise the descriptor's block_shape[0] would - # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot / arithmetic - # on the loaded tile fails shape inference at compile time - # ("Cannot make_shape_compatible: incompatible dimensions"). When Mw == - # Mnw (tied, the common case) the two descriptors are the same object. - # Same memory (state's flat 2D view, shape (cache*nheads*dim, dstate)) - # and same dstate block_shape — only block_shape[0] differs. - # When no TMA flag is on, both variables hold the raw `state` tensor as a - # dummy; kernels never reference it because their constexprs are all - # False (Triton DCEs the dead branches). - # `triton.set_allocator()` must run before any descriptor-using launch. - if ( + # Per-path TMA descriptors for state — write-side and nowrite-side. + # Kernels see state as a 2D row-space: [cache/head/dim row, dstate]. + # Dense tensors use the original flat view. Block-reuse cache tensors may + # have a gap between slots (conv state packed after SSM state), but the + # per-slot SSM rows are still dense and can use the same 2D row-space with + # explicit strides. Unsupported layouts fall back to raw loads/stores. + use_tma_state = ( _use_tma_rect_load or _use_tma_replay_write_load or _use_tma_replay_write_store or _use_tma_replay_nowrite_load - ): + ) + if use_tma_state: + state_row_stride = state.stride(2) + tma_state_supported = ( + state.stride(-1) == 1 + and state_row_stride > 0 + and state.stride(1) == dim * state_row_stride + and state.stride(0) % state_row_stride == 0 + and (state_row_stride * state.element_size()) % 16 == 0 + ) + if not tma_state_supported: + if _require_tma_state_layout or caller_forced_tma_state: + raise AssertionError( + "TMA state layout requires inner stride 1, dense dim rows " + "within each head, cache stride aligned to dim-row stride, " + "and 16-byte-aligned row stride; got " + f"shape={tuple(state.shape)} strides={state.stride()}" + ) + _use_tma_rect_load = False + _use_tma_replay_write_load = False + _use_tma_replay_write_store = False + _use_tma_replay_nowrite_load = False + use_tma_state = False + + # Each kernel launch consumes the descriptor whose block_shape[0] matches + # its BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need + # distinct descriptors; otherwise the descriptor's block_shape[0] would + # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot/arithmetic on + # the loaded tile fails shape inference at compile time. When no TMA flag + # is on, both variables hold raw `state` as a dummy; kernels never + # reference it because their constexprs are all False. + if use_tma_state: from triton.tools.tensor_descriptor import TensorDescriptor _ensure_tma_allocator() - assert state.is_contiguous(), "TMA state requires contiguous state" - assert state.stride(-1) == 1, "TMA state requires inner stride 1" - _state_flat = state.view(-1, state.shape[-1]) + if state.is_contiguous(): + state_tma_base = state.view(-1, state.shape[-1]) + make_state_tma_descriptor = TensorDescriptor.from_tensor + else: + slot_stride_rows = state.stride(0) // state_row_stride + state_rows = (cache_size - 1) * slot_stride_rows + nheads * dim + state_tma_base = state + state_tma_shape = [state_rows, dstate] + state_tma_strides = [state_row_stride, state.stride(3)] + + def make_state_tma_descriptor(_state, block_shape): + return TensorDescriptor( + _state, + state_tma_shape, + state_tma_strides, + block_shape=block_shape, + ) + _dstate_pow2 = triton.next_power_of_2(dstate) - state_tma_descriptor_write = TensorDescriptor.from_tensor( - _state_flat, + state_tma_descriptor_write = make_state_tma_descriptor( + state_tma_base, block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], ) if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: state_tma_descriptor_nowrite = state_tma_descriptor_write else: - state_tma_descriptor_nowrite = TensorDescriptor.from_tensor( - _state_flat, + state_tma_descriptor_nowrite = make_state_tma_descriptor( + state_tma_base, block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], ) else: diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py index 80167ff2048a..36d1949cdcf7 100644 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py @@ -515,6 +515,34 @@ def _build_replay_work_items_cpu( _TENSOR_CACHE: dict = {} +def _empty_state_with_slot_gap( + batch: int, + nheads: int, + head_dim: int, + d_state: int, + state_dtype: torch.dtype, + device: str, + slot_gap_rows: int, +): + if slot_gap_rows == 0: + return torch.empty(batch, nheads, head_dim, d_state, device=device, dtype=state_dtype) + + ssm_rows_per_slot = nheads * head_dim + rows_per_slot = ssm_rows_per_slot + slot_gap_rows + backing = torch.empty(batch, rows_per_slot, d_state, device=device, dtype=state_dtype) + return backing[:, :ssm_rows_per_slot].view(batch, nheads, head_dim, d_state) + + +def _clone_state_preserving_layout(state: torch.Tensor): + if state.is_contiguous(): + return state.clone() + clone = torch.empty_strided( + tuple(state.shape), state.stride(), device=state.device, dtype=state.dtype + ) + clone.copy_(state) + return clone + + def _build_tensors( batch: int, mtp_len: int, @@ -525,6 +553,7 @@ def _build_tensors( d_state: int, ngroups: int, max_window: int | None = None, + strided_state_cache: bool = False, ): """ Build all tensors for one benchmark configuration. @@ -544,7 +573,17 @@ def _build_tensors( device = "cuda" # Cache lookup — grow batch in place if needed; else return views. - cache_key = (state_dtype, act_dtype, max_window, mtp_len, nheads, head_dim, d_state, ngroups) + cache_key = ( + state_dtype, + act_dtype, + max_window, + mtp_len, + nheads, + head_dim, + d_state, + ngroups, + strided_state_cache, + ) cached = _TENSOR_CACHE.get(cache_key) if cached is not None and cached["max_batch"] >= batch: # Hit — return slices for current batch. @@ -587,6 +626,20 @@ def _build_tensors( torch.manual_seed(42) + d_inner = nheads * head_dim + conv_dim = d_inner + 2 * ngroups * d_state + d_conv = 4 # conv kernel width for Nemotron/Mamba2 + slot_gap_rows = 0 + if strided_state_cache: + conv_state_elems_per_slot = conv_dim * d_conv + if conv_state_elems_per_slot % d_state != 0: + raise ValueError( + "strided state cache requires conv state to be an integer " + f"number of d_state rows, got {conv_state_elems_per_slot=} " + f"and {d_state=}" + ) + slot_gap_rows = conv_state_elems_per_slot // d_state + # --- SSM parameters (float32, tie_hdim strides) --- A_base = -torch.rand(nheads, device=device) - 0.5 A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) # stride(-1)=0, stride(-2)=0 @@ -618,11 +671,20 @@ def _build_tensors( state_scales0 = (1.0 / encode_scale).to(torch.float32) # decode scale scaled = state_fp32 * encode_scale.unsqueeze(-1) if state_dtype == torch.float8_e4m3fn: - state0 = scaled.clamp(-quant_max, quant_max).to(state_dtype) + state0_dense = scaled.clamp(-quant_max, quant_max).to(state_dtype) else: - state0 = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) + state0_dense = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) + state0 = _empty_state_with_slot_gap( + batch, nheads, head_dim, d_state, state_dtype, device, slot_gap_rows + ) + state0.copy_(state0_dense) else: - state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=state_dtype) + state0 = _empty_state_with_slot_gap( + batch, nheads, head_dim, d_state, state_dtype, device, slot_gap_rows + ) + state0.copy_( + torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=state_dtype) + ) state_scales0 = None # --- Cache tensors for replay kernel --- @@ -662,10 +724,6 @@ def _build_tensors( out_base = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) # --- Conv1d tensors (for --with-conv1d mode) --- - d_inner = nheads * head_dim - conv_dim = d_inner + 2 * ngroups * d_state - d_conv = 4 # conv kernel width for Nemotron/Mamba2 - # xbc_input: (batch, conv_dim, mtp_len) — "hot" input from in_proj. # Match production layout: in_proj output is (batch*mtp_len, conv_dim) # contiguous, then .view(batch, mtp_len, conv_dim).transpose(1, 2) @@ -2658,6 +2716,7 @@ def _bench_config( args.d_state, args.tp_ngroups, max_window=getattr(args, "max_window", None) or None, + strided_state_cache=getattr(args, "strided_state_cache", False), ) nheads = args.tp_nheads @@ -2680,7 +2739,7 @@ def _bench_config( if mode != "persistent_main" and nowrite_first: return - state_work = state0.clone() + state_work = _clone_state_preserving_layout(state0) state_scales_work = state_scales0.clone() if state_scales0 is not None else None old_x_work = old_x0.clone() old_B_work = old_B0.clone() @@ -3330,6 +3389,8 @@ def _run_incr( extra_kwargs["_use_tma_replay_nowrite_load"] = bool(use_tma_replay_nowrite_load) if use_tma_replay_write_store is not None: extra_kwargs["_use_tma_replay_write_store"] = bool(use_tma_replay_write_store) + if getattr(args, "require_tma_state_layout", False): + extra_kwargs["_require_tma_state_layout"] = True if cta_per_sm is not None: extra_kwargs["_cta_per_sm"] = cta_per_sm if num_loop_stages is not None: @@ -4826,6 +4887,21 @@ def _parse_args() -> argparse.Namespace: help="Comma-separated 0/1 sweep. TMA state STORE in replay main " "for the checkpoint/write half. Independent from all load TMA flags.", ) + parser.add_argument( + "--strided-state-cache", + action=argparse.BooleanOptionalAction, + default=False, + help="Allocate SSM state as a view with block-reuse-like gaps between " + "cache slots. Used to validate TMA descriptor handling for recurrent " + "state pools that pack SSM and conv state together.", + ) + parser.add_argument( + "--require-tma-state-layout", + action=argparse.BooleanOptionalAction, + default=False, + help="Raise instead of falling back if a TMA state path is requested " + "but the state layout cannot be represented by the 2D descriptor.", + ) parser.add_argument( "--modes", type=str, diff --git a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py index be13c409ace5..dadebbd996db 100644 --- a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py @@ -120,6 +120,22 @@ def _dequantize_state(state_quant: torch.Tensor, decode_scale: torch.Tensor): return state_quant.to(torch.float32) * decode_scale.unsqueeze(-1) +def _make_strided_state_with_slot_gap(state: torch.Tensor, slot_gap_rows: int): + cache_size, nheads, head_dim, d_state = state.shape + ssm_rows_per_slot = nheads * head_dim + backing = torch.empty( + cache_size, + ssm_rows_per_slot + slot_gap_rows, + d_state, + device=state.device, + dtype=state.dtype, + ) + backing[:, ssm_rows_per_slot:].fill_(float("nan")) + strided_state = backing[:, :ssm_rows_per_slot].view_as(state) + strided_state.copy_(state) + return strided_state, backing + + def _maybe_skip_dtype(state_dtype, use_sr): """Skip on insufficient SM. fp8 e4m3fn (any) needs SM 89+; fp16/fp8 SR needs SM 100+; int8/int16 (RN or SR) runs anywhere.""" @@ -129,6 +145,130 @@ def _maybe_skip_dtype(state_dtype, use_sr): pytest.skip(f"{state_dtype} stochastic rounding requires SM 100+ (Blackwell B200+)") +@pytest.mark.skipif(get_sm_version() < 90, reason="TMA descriptor path requires SM 90+") +@pytest.mark.parametrize("block_size_m", [8, 64], ids=["M8", "M64"]) +@pytest.mark.parametrize("rectangle_for_nowrite", [False, True], ids=["replay_nowrite", "rect"]) +def test_replay_tma_strided_state_layout_matches_contiguous(rectangle_for_nowrite, block_size_m): + """Block-reuse packs conv state after each slot's SSM state.""" + torch.manual_seed(123) + + cache_size = 4 + batch = 2 + T = 8 + max_window = 16 + nheads = 16 + head_dim = 64 + d_state = 128 + ngroups = 1 + device = "cuda" + dtype = torch.bfloat16 + state_dtype = torch.float32 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + + state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + d_inner = nheads * head_dim + conv_dim = d_inner + 2 * ngroups * d_state + slot_gap_rows = conv_dim * 4 // d_state + assert (conv_dim * 4) % d_state == 0 + + old_x = torch.randn(cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) + prev_tokens = torch.zeros(cache_size, device=device, dtype=torch.int32) + prev_tokens[state_batch_indices[0]] = 4 + prev_tokens[state_batch_indices[1]] = 12 + n_writes, replay_work_items = _make_replay_work_items( + prev_tokens, + cache_buf_idx, + T, + max_window, + batch, + state_batch_indices, + device, + ) + + x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + B = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + D = repeat(torch.randn(nheads, device=device, dtype=dtype), "h -> h p", p=head_dim) + dt_bias = repeat(torch.randn(nheads, device=device, dtype=dtype), "h -> h p", p=head_dim) + + dense_state = state0.clone() + strided_state, strided_backing = _make_strided_state_with_slot_gap(state0, slot_gap_rows) + strided_gap = strided_backing[:, nheads * head_dim :] + dense_out = torch.empty(batch, T, nheads, head_dim, device=device, dtype=dtype) + strided_out = torch.empty_like(dense_out) + dense_old_x = old_x.clone() + dense_old_B = old_B.clone() + dense_old_dt = old_dt.clone() + dense_old_dA_cumsum = old_dA_cumsum.clone() + strided_old_x = old_x.clone() + strided_old_B = old_B.clone() + strided_old_dt = old_dt.clone() + strided_old_dA_cumsum = old_dA_cumsum.clone() + + common_kwargs = dict( + prev_num_accepted_tokens=prev_tokens, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + n_writes=n_writes, + replay_work_items=replay_work_items, + use_internal_pdl=False, + rectangle_for_nowrite=rectangle_for_nowrite, + mode="persistent_dynamic", + _block_size_m=block_size_m, + _use_tma_rect_load=True, + _use_tma_replay_write_load=True, + _use_tma_replay_write_store=True, + _use_tma_replay_nowrite_load=True, + _require_tma_state_layout=True, + ) + + replay_selective_state_update( + dense_state, + dense_old_x, + dense_old_B, + dense_old_dt, + dense_old_dA_cumsum, + cache_buf_idx.clone(), + out=dense_out, + **common_kwargs, + ) + replay_selective_state_update( + strided_state, + strided_old_x, + strided_old_B, + strided_old_dt, + strided_old_dA_cumsum, + cache_buf_idx.clone(), + out=strided_out, + **common_kwargs, + ) + + torch.testing.assert_close(strided_out, dense_out, rtol=0, atol=0) + torch.testing.assert_close(strided_state, dense_state, rtol=0, atol=0) + torch.testing.assert_close(strided_old_x, dense_old_x, rtol=0, atol=0) + torch.testing.assert_close(strided_old_B, dense_old_B, rtol=0, atol=0) + torch.testing.assert_close(strided_old_dt, dense_old_dt, rtol=0, atol=0) + torch.testing.assert_close(strided_old_dA_cumsum, dense_old_dA_cumsum, rtol=0, atol=0) + assert torch.isnan(strided_gap).all() + + @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) @pytest.mark.parametrize( "state_dtype", From 3436b76d2258caf2eff7f7b0c4643cf312c7c7b3 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 15 Jun 2026 13:06:45 -0700 Subject: [PATCH 85/89] Update Mamba replay precompute tunings Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../mamba/replay_selective_state_update.py | 151 +++++++++--------- 1 file changed, 76 insertions(+), 75 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 4c0ecdd3edd2..a59aaeaf1b0e 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -2298,7 +2298,7 @@ def _persistent_main_kernel( # Source: emit_tuning_from_noise.py. Auto-generated from noise-cleaned per-cell search # winners (best of pd / pm by bucket_expected_renorm). Effective batch = raw_batch × 16. # The search predates a Triton PDL scheduling bug. -# Most knobs are unchanged; large PDL-hoist regressions got spot retunes. +# Most knobs are unchanged; large PDL-hoist/precompute regressions got spot retunes. # Missing dtype/SR combos fall back via the _resolve_tuning chain: # RN→SR for same dtype, then bf16/int16→fp16/SR and fp8→int8/SR. _DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { @@ -2307,32 +2307,32 @@ def _persistent_main_kernel( 16, "persistent_dynamic", { - "_block_size_m": 8, - "_cta_per_sm": 6, + "_block_size_m": 4, + "_cta_per_sm": 4, "_flatten": False, - "_heads_per_block": 1, + "_heads_per_block": 2, "_num_loop_stages": 1, "_num_stages": 4, "_num_warps": 1, - "_precompute_num_warps": 4, + "_precompute_num_warps": 8, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": False, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_store": False, "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=6.83us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=1, score=6.40us (B200 precompute-retune noise-cleaned 5x500) ( 32, "persistent_dynamic", { "_block_size_m": 8, - "_cta_per_sm": 9, + "_cta_per_sm": 4, "_flatten": False, - "_heads_per_block": 2, + "_heads_per_block": 1, "_num_loop_stages": 1, - "_num_stages": 5, + "_num_stages": 3, "_num_warps": 1, "_precompute_num_warps": 4, "_use_tma_rect_load": False, @@ -2342,17 +2342,17 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=2, score=7.0us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=2, score=6.99us (B200 precompute-retune noise-cleaned 5x500) ( 64, "persistent_dynamic", { "_block_size_m": 8, - "_cta_per_sm": 4, + "_cta_per_sm": 10, "_flatten": False, - "_heads_per_block": 4, + "_heads_per_block": 2, "_num_loop_stages": 1, - "_num_stages": 4, + "_num_stages": 2, "_num_warps": 1, "_precompute_num_warps": 4, "_use_tma_rect_load": False, @@ -2362,7 +2362,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=4, score=7.05us (B200 PDL-hoist default 5x200) + ), # raw_batch=4, score=7.05us (B200 precompute-retune noise-cleaned 5x500) ( 128, "persistent_dynamic", @@ -2408,27 +2408,27 @@ def _persistent_main_kernel( "persistent_main", { "_block_size_m_nowrite": 64, - "_block_size_m_write": 32, - "_cta_per_sm_nowrite": 5, - "_cta_per_sm_write": 9, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 7, "_flatten": False, "_heads_per_block": 8, "_num_loop_stages_nowrite": 1, "_num_loop_stages_write": 1, "_num_stages_nowrite": 3, - "_num_stages_write": 3, + "_num_stages_write": 2, "_num_warps_nowrite": 1, - "_num_warps_write": 2, + "_num_warps_write": 1, "_precompute_num_warps": 8, - "_use_tma_rect_load": True, + "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, "_warp_specialize": False, - "nowrite_first": True, + "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=32, score=12.66us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=32, score=13.08us (B200 precompute-retune noise-cleaned 5x500) ( 1024, "persistent_main", @@ -2565,22 +2565,22 @@ def _persistent_main_kernel( 16, "persistent_dynamic", { - "_block_size_m": 8, - "_cta_per_sm": 9, + "_block_size_m": 4, + "_cta_per_sm": 7, "_flatten": False, - "_heads_per_block": 4, + "_heads_per_block": 2, "_num_loop_stages": 1, - "_num_stages": 3, + "_num_stages": 2, "_num_warps": 1, - "_precompute_num_warps": 4, + "_precompute_num_warps": 8, "_use_tma_rect_load": False, - "_use_tma_replay_nowrite_load": False, - "_use_tma_replay_write_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=6.59us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=1, score=6.54us (B200 precompute-retune noise-cleaned 5x500) ( 32, "persistent_dynamic", @@ -2712,51 +2712,51 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=64, score=18.51us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=64, score=18.51us (B200 PDL-retune noise-cleaned 5x500) ( 2048, "persistent_main", { "_block_size_m_nowrite": 64, - "_block_size_m_write": 16, - "_cta_per_sm_nowrite": 8, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 5, "_cta_per_sm_write": 8, "_flatten": False, "_heads_per_block": 8, "_num_loop_stages_nowrite": 1, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 4, + "_num_stages_nowrite": 2, "_num_stages_write": 2, "_num_warps_nowrite": 1, "_num_warps_write": 1, "_precompute_num_warps": 8, - "_use_tma_rect_load": False, + "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": True, + "_use_tma_replay_write_store": False, "_warp_specialize": False, "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=128, score=28.72us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=128, score=28.83us (B200 precompute-retune noise-cleaned 5x500) ( 4096, "persistent_main", { "_block_size_m_nowrite": 64, "_block_size_m_write": 64, - "_cta_per_sm_nowrite": 8, + "_cta_per_sm_nowrite": 6, "_cta_per_sm_write": 3, "_flatten": False, - "_heads_per_block": 2, - "_num_loop_stages_nowrite": 1, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 2, "_num_loop_stages_write": 1, - "_num_stages_nowrite": 4, - "_num_stages_write": 1, - "_num_warps_nowrite": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, "_num_warps_write": 4, - "_precompute_num_warps": 1, - "_use_tma_rect_load": False, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": True, "_use_tma_replay_write_store": True, @@ -2764,7 +2764,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=256, score=47.98us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=256, score=46.60us (B200 precompute-retune noise-cleaned 5x500) ( 8192, "persistent_main", @@ -2775,10 +2775,10 @@ def _persistent_main_kernel( "_cta_per_sm_write": 6, "_flatten": False, "_heads_per_block": 16, - "_num_loop_stages_nowrite": 2, + "_num_loop_stages_nowrite": 5, "_num_loop_stages_write": 2, - "_num_stages_nowrite": 5, - "_num_stages_write": 4, + "_num_stages_nowrite": 2, + "_num_stages_write": 2, "_num_warps_nowrite": 2, "_num_warps_write": 4, "_precompute_num_warps": 4, @@ -2790,7 +2790,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=83.00us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=512, score=81.07us (B200 precompute-retune noise-cleaned 5x500) ( 16384, "persistent_main", @@ -2819,6 +2819,7 @@ def _persistent_main_kernel( ), # raw_batch=1024, score=149.25us (B200 PDL-hoist default 5x200) ], ("fp8", "SR"): [ + # --- ALL UNCHANGED (not retuned this round) --- ( 16, "persistent_dynamic", @@ -2838,7 +2839,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=6.46us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=1, score=6.46us (B200 PDL-retune noise-cleaned 5x500) ( 32, "persistent_dynamic", @@ -2858,7 +2859,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=2, score=6.72us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=2, score=6.72us (B200 precompute retune 5x100) ( 64, "persistent_dynamic", @@ -2990,7 +2991,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=128, score=24.5us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=128, score=24.5us (B200 PDL-retune noise-cleaned 5x500) ( 4096, "persistent_main", @@ -3016,7 +3017,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=256, score=40.04us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=256, score=40.04us (B200 PDL-retune noise-cleaned 5x500) ( 8192, "persistent_main", @@ -3042,7 +3043,7 @@ def _persistent_main_kernel( "nowrite_first": False, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=67.94us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=512, score=67.94us (B200 precompute retune 5x100) ( 16384, "persistent_main", @@ -3068,7 +3069,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=1024, score=126.74us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=1024, score=126.74us (B200 PDL-retune noise-cleaned 5x500) ], ("fp32", "RN"): [ ( @@ -3078,19 +3079,19 @@ def _persistent_main_kernel( "_block_size_m": 8, "_cta_per_sm": 6, "_flatten": False, - "_heads_per_block": 8, + "_heads_per_block": 2, "_num_loop_stages": 1, - "_num_stages": 2, + "_num_stages": 1, "_num_warps": 1, - "_precompute_num_warps": 4, + "_precompute_num_warps": 8, "_use_tma_rect_load": False, - "_use_tma_replay_nowrite_load": True, - "_use_tma_replay_write_load": True, - "_use_tma_replay_write_store": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=1, score=6.64us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=1, score=6.42us (B200 precompute-retune noise-cleaned 5x500) ( 32, "persistent_dynamic", @@ -3116,13 +3117,13 @@ def _persistent_main_kernel( "persistent_dynamic", { "_block_size_m": 8, - "_cta_per_sm": 9, + "_cta_per_sm": 7, "_flatten": False, "_heads_per_block": 2, "_num_loop_stages": 1, - "_num_stages": 4, + "_num_stages": 1, "_num_warps": 1, - "_precompute_num_warps": 4, + "_precompute_num_warps": 8, "_use_tma_rect_load": False, "_use_tma_replay_nowrite_load": False, "_use_tma_replay_write_load": False, @@ -3130,7 +3131,7 @@ def _persistent_main_kernel( "_warp_specialize": False, "rectangle_for_nowrite": False, }, - ), # raw_batch=4, score=7.19us (B200 PDL-hoist default 5x200) + ), # raw_batch=4, score=7.10us (B200 precompute-retune noise-cleaned 5x500) ( 128, "persistent_dynamic", @@ -3216,7 +3217,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": False, }, - ), # raw_batch=64, score=20.05us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=64, score=20.05us (B200 PDL-retune noise-cleaned 5x500) ( 2048, "persistent_main", @@ -3242,7 +3243,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=128, score=31.39us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=128, score=31.39us (B200 PDL-retune noise-cleaned 5x500) ( 4096, "persistent_main", @@ -3268,7 +3269,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=256, score=51.27us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=256, score=51.27us (B200 PDL-retune noise-cleaned 5x500) ( 8192, "persistent_main", @@ -3294,7 +3295,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=512, score=88.57us (B200 precompute retune 5x100) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=512, score=88.57us (B200 precompute retune 5x100) ( 16384, "persistent_main", @@ -3320,7 +3321,7 @@ def _persistent_main_kernel( "nowrite_first": True, "rectangle_for_nowrite": True, }, - ), # raw_batch=1024, score=170.51us (B200 PDL-retune noise-cleaned 5x500) # << TUNED_AFTER_PDL_FIX + ), # raw_batch=1024, score=170.51us (B200 PDL-retune noise-cleaned 5x500) ], } _PD_TO_PM_SPLIT_MAP = { # pd unsplit knob -> (pm_write_knob, pm_nowrite_knob) From 199082822b210b2213e77d2d8dce5add8f4b3e54 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Mon, 15 Jun 2026 21:18:52 -0700 Subject: [PATCH 86/89] Refresh Mamba dummy mask under inference mode. The dummy mask refactor broke some paths. Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index a43b8cd20178..43535f838160 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -777,6 +777,7 @@ def get_state_indices(self, request_ids: List[int], self._refresh_dummy_request_mask(is_dummy) return indices + @torch.inference_mode() def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: n = len(is_dummy) assert n <= self._dummy_request_mask_host.shape[0] @@ -1979,6 +1980,7 @@ def update_mamba_states(self, src_state_indices, num_accepted_draft_tokens, state_indices_d) + @torch.inference_mode() def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: if self._dummy_request_mask is None: return From 80b1f7a5a86d859435f145b4919c4436a361f7f0 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Tue, 16 Jun 2026 09:33:17 -0700 Subject: [PATCH 87/89] Keep AutoDeploy replay metadata standalone Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- .../custom_ops/attention_interface.py | 3 +-- .../custom_ops/mamba/replay_metadata.py | 22 +++++++++++++++++++ .../_torch/auto_deploy/shim/interface.py | 15 ++++++------- 3 files changed, 30 insertions(+), 10 deletions(-) create mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/replay_metadata.py diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index 321ff3c80fdd..c5903f9bbb42 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -35,11 +35,10 @@ from torch.fx import GraphModule, Node from torch.types import Number -from tensorrt_llm._torch.modules.mamba.mamba2_metadata import REPLAY_WORK_ITEM_WIDTH - from .._compat import KvCacheConfig, nvtx_range, prefer_pinned, str_dtype_to_torch from ..utils.logger import ad_logger from ..utils.node_utils import extract_op_args, get_op_schema +from .mamba.replay_metadata import REPLAY_WORK_ITEM_WIDTH Constant = Union[int, float, str, None] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/replay_metadata.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/replay_metadata.py new file mode 100644 index 000000000000..ce924b2215af --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/replay_metadata.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Replay metadata layout shared by AutoDeploy Mamba descriptors.""" + +REPLAY_WORK_POSITION_IN_DECODE_BATCH = 0 +REPLAY_WORK_CACHE_SLOT = 1 +REPLAY_WORK_PNAT = 2 +REPLAY_WORK_CACHE_BUF_IDX = 3 +REPLAY_WORK_ITEM_WIDTH = 4 diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 12855b2028a6..30ec0b8ae616 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -48,14 +48,6 @@ Mapping = None torch_dtype_to_binding = None -from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( - REPLAY_WORK_CACHE_BUF_IDX, - REPLAY_WORK_CACHE_SLOT, - REPLAY_WORK_ITEM_WIDTH, - REPLAY_WORK_PNAT, - REPLAY_WORK_POSITION_IN_DECODE_BATCH, -) - from ..custom_ops.attention_interface import ( AttentionType, CausalConvResourceHandler, @@ -77,6 +69,13 @@ SSMResourceHandler, StateResourceHandler, ) +from ..custom_ops.mamba.replay_metadata import ( + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_ITEM_WIDTH, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, +) from ..distributed.common import all_gather_object, get_world_size from ..distributed.common import is_initialized as is_distributed_initialized from ..utils.cuda_mem_tracker import bytes_to, get_mem_info From d41775976f3746f04fef402a430b082200ae981a Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 17 Jun 2026 12:08:40 -0700 Subject: [PATCH 88/89] Lower Ultra ADP MTP KV cache fraction Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 8b4d2c499860..af3339132b81 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7435,7 +7435,7 @@ def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, enable_block_reuse=True, mamba_ssm_cache_dtype="float16", mamba_state_cache_interval=mamba_state_cache_interval, - free_gpu_memory_fraction=0.6, + free_gpu_memory_fraction=0.5, ), max_batch_size=max_batch_size, tensor_parallel_size=tp_size, From ca3249496f77589f354836280de302f82cffcf83 Mon Sep 17 00:00:00 2001 From: Harris Nover <249353502+hnover-nv@users.noreply.github.com> Date: Wed, 20 May 2026 17:26:58 -0700 Subject: [PATCH 89/89] remove benchmark and flashinfer imports for PR Signed-off-by: Harris Nover <249353502+hnover-nv@users.noreply.github.com> --- ...benchmark_replay_selective_state_update.py | 5209 ----------------- .../flashinfer_checkpointing_ssu_pr3324.py | 325 - .../csrc/checkpointing_ssu.cu | 554 -- .../checkpointing_ssu_customize_config.jinja | 38 - .../csrc/checkpointing_ssu_jit_binding.cu | 51 - .../csrc/checkpointing_ssu_kernel_inst.cu | 14 - .../include/flashinfer/exception.h | 126 - .../flashinfer/mamba/checkpointing_ssu.cuh | 153 - .../include/flashinfer/mamba/common.cuh | 234 - .../include/flashinfer/mamba/conversion.cuh | 555 -- .../mamba/kernel_checkpointing_ssu.cuh | 1114 ---- .../mamba/kernel_checkpointing_ssu_8bit.cuh | 1481 ----- .../mamba/kernel_checkpointing_ssu_common.cuh | 1733 ------ .../mamba/launch_checkpointing_ssu.cuh | 192 - .../flashinfer/mamba/ssu_mtp_common.cuh | 160 - .../include/flashinfer/utils.cuh | 648 -- .../include/flashinfer/vec_dtypes.cuh | 3201 ---------- 17 files changed, 15788 deletions(-) delete mode 100644 tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_customize_config.jinja delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh delete mode 100644 tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py deleted file mode 100644 index 36d1949cdcf7..000000000000 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ /dev/null @@ -1,5209 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. -"""Standalone benchmark for replay_selective_state_update (Triton kernel). - -Suitable for nsight-compute (ncu) and nsight-systems (nsys) capture. - -Fixed model config: NVIDIA-Nemotron-3-Super-120B-A12B at TP=8 - nheads=16, head_dim=64, d_state=128, ngroups=1 - -mtp_len is the per-request sequence length processed by replay: in MTP it -equals num_draft_tokens + 1 target token, so --mtp-lengths 6 models 5 drafts -+ 1 target. - -Baseline kernel (--baseline flashinfer_pr3324): - Calls FlashInfer PR 3324's checkpointing_ssu kernel against the same replay - cache tensors as the replay kernel. - -Timing methodology -================== - -All in-bench timing comes from CUPTI's Activity API (1 ns kernel -timestamps from the GPU profiling fabric). cudaEvent.elapsed_time() was -removed — its ~0.5 us resolution overshoots CUPTI by ~50% on short kernels -in graphs, and we have no other use for it here. See the CUPTI block -lower in this file for the timer source. - -Three modes: - - --cupti --cuda-graph (default) - Capture a small CUDA graph for the cell, replay it for warmup + timed - iterations, and read kernel start/end from CUPTI. Raw CUPTI buffers are - parsed out-of-process on the timed path, with a cached ordinal plan used - to keep only the kernels we care about. - - --cupti --no-cuda-graph - Eager loop with CUPTI. Per-kernel timestamps are still accurate, but - the per-iter SPAN (max(end) - min(start)) now includes the Python - launch latency BETWEEN consecutive kernels in run_fn (~100 µs on - Hopper/Blackwell). Graph capture and PDL hide that latency; eager - mode honestly reports it. For per-kernel timing in eager mode, look - at per_kernel.start_us/end_us in --json-detailed output rather than - the span percentiles. Useful when graph capture is undesirable. - - --no-cupti (with or without --cuda-graph) - No in-bench timing — just runs the kernels for an external profiler - (nsys / ncu) to time. In-process CUPTI conflicts with nsys's own - subscriber, so disable ours when wrapping in nsys. Bench output - reports zeros for median/p95/p99; trust the external trace. - -JSON output schema (--json-output PATH) -======================================= - -Designed to be parsed by collect.py / report.py without touching sqlite or -NVTX traces. Future agents: prefer reading this JSON over re-running nsys. - - { - "metadata": {timestamp, cmd, tp_size, warmup, iters, cupti}, - "results": { - "": {median, p95, p99, n, iters_us, [n_writes_per_iter], - [kmix_bucket_score], [per_kernel]} - } - } - -Key format mirrors collect.py's kernel_data.json convention: - incremental/{batch}/{mtp}/{sd}/k{prev_k}/{sweep_parts}/tp{tp} - flashinfer_pr3324/{batch}/{mtp}/{sd}/k{prev_k}/{sweep_parts}/tp{tp} - - - is normalized: bf16 / fp16 / fp32 / int8 / int16 / fp8. - - is e.g. "M16_W1_S3_SR0_RECT0" — flags concatenated by - underscore in canonical order. pS/R/CT appear only when explicitly swept. - - All numeric values in microseconds (us). - -Per-record fields: - - median, p95, p99: span statistics (us). Span = max(kernel_end_ns) - - min(kernel_start_ns) across the iter's kernels — same convention as - nsys-derived collect.py used to use. - - n: number of timed iters that contributed. - - iters_us: list of length n, raw per-iter spans. - - n_writes_per_iter: for mix rows, list of length n with the number of - write-path slots in each timed iteration. - - kmix_bucket_score: for mix rows, binomial-PDF-weighted expected time - from per-n_writes bucket medians. - - per_kernel: {: {start_us: [...], end_us: [...]}} where - timestamps are RELATIVE to that iter's first kernel start, in us. Lets - you see PDL overlap directly without an external profiler. Only with - --json-detailed. - -Example usage: - # Basic sweep (default = --cupti, just summary stats) - python benchmark_replay_selective_state_update.py \\ - --batch-sizes 1,2,4 --mtp-lengths 1,4,8 --warmup 5 --iters 20 - - # JSON output, summary stats only (compact) - python benchmark_replay_selective_state_update.py \\ - --batch-sizes 16 --mtp-lengths 6 --json-output /tmp/out.json - - # JSON output, full per-iter / per-kernel data (for PDL analysis etc.) - python benchmark_replay_selective_state_update.py \\ - --batch-sizes 16 --mtp-lengths 6 \\ - --json-output /tmp/out.json --json-detailed - - # nsys capture (--no-cupti so our subscriber doesn't conflict) - nsys profile --capture-range=cudaProfilerApi \\ - python benchmark_replay_selective_state_update.py --profile --no-cupti - - # ncu capture (--no-cupti --no-cuda-graph: each kernel replayable solo) - ncu --target-processes all \\ - python benchmark_replay_selective_state_update.py --profile \\ - --no-cupti --no-cuda-graph \\ - --batch-sizes 1 --mtp-lengths 4 --warmup 5 --iters 5 -""" - -import argparse -import atexit -import csv -import ctypes -import importlib -import itertools -import json -import math -import multiprocessing as mp -import os -import queue -import statistics -import sys -import threading -import time -from datetime import datetime -from multiprocessing import shared_memory -from pathlib import Path - -import numpy as np -import torch -from einops import repeat - -REPLAY_WORK_POSITION_IN_DECODE_BATCH = 0 -REPLAY_WORK_CACHE_SLOT = 1 -REPLAY_WORK_PNAT = 2 -REPLAY_WORK_CACHE_BUF_IDX = 3 -REPLAY_WORK_ITEM_WIDTH = 4 - -DEFAULT_PMIX_T = 6 -DEFAULT_PMIX_LABEL = "pmix_T6" -DEFAULT_PMIX_REPLAY_COUNTS = { - 1: 539622, - 2: 329144, - 3: 336224, - 4: 495832, - 5: 444398, - 6: 1007008, -} - - -def _import_mamba_kernels_fast(): - """Load kernel modules directly (~40s faster than a full tensorrt_llm init). - Use --full-import as the fallback if module dependencies change. - - Strategy: stub the parent packages (tensorrt_llm, tensorrt_llm._torch, - tensorrt_llm._torch.modules) in sys.modules with __path__ set, but do - NOT execute their __init__.py. Then load the leaf kernel modules. - When a kernel body imports e.g. tensorrt_llm._utils.get_sm_version, - Python's machinery resolves it against our stub's __path__ and loads - only _utils.py — skipping the heavy tensorrt_llm package init. - """ - import types - - repo_root = Path(__file__).resolve().parents[5] - trtllm_dir = repo_root / "tensorrt_llm" - mamba_pkg = "tensorrt_llm._torch.modules.mamba" - mamba_dir = trtllm_dir / "_torch" / "modules" / "mamba" - - def _stub_pkg(fqn: str, pkg_dir: Path): - """Register a stub package in sys.modules without running its - __init__.py. Sets __path__ so Python can resolve submodule imports - against the real directory on disk.""" - if fqn in sys.modules: - return - stub = types.ModuleType(fqn) - stub.__path__ = [str(pkg_dir)] - sys.modules[fqn] = stub - - # Stub the parent chain so `from tensorrt_llm._utils import ...` (and - # similar) work without triggering tensorrt_llm/__init__.py. - _stub_pkg("tensorrt_llm", trtllm_dir) - _stub_pkg("tensorrt_llm._torch", trtllm_dir / "_torch") - _stub_pkg("tensorrt_llm._torch.modules", trtllm_dir / "_torch" / "modules") - - utils_stub = types.ModuleType("tensorrt_llm._utils") - - def _fast_get_sm_version(): - prop = torch.cuda.get_device_properties(0) - return prop.major * 10 + prop.minor - - utils_stub.get_sm_version = _fast_get_sm_version - sys.modules["tensorrt_llm._utils"] = utils_stub - - def _load(mod_name: str, file_name: str): - fqn = f"{mamba_pkg}.{mod_name}" if mod_name else mamba_pkg - if fqn in sys.modules: - return sys.modules[fqn] - kwargs = {} - if file_name == "__init__.py": - kwargs["submodule_search_locations"] = [str(mamba_dir)] - spec = importlib.util.spec_from_file_location(fqn, mamba_dir / file_name, **kwargs) - mod = importlib.util.module_from_spec(spec) - sys.modules[fqn] = mod - spec.loader.exec_module(mod) - return mod - - # 1. Package __init__ (defines replay work-item constants) - _load("", "__init__.py") - # 2. softplus helper (used by both kernel modules) - _load("softplus", "softplus.py") - # 3. replay_selective_state_update only needs replay-work-item constants - # from mamba2_metadata at import time. Stub those constants instead of - # importing the full metadata module and its scheduler/attention deps. - metadata_stub = types.ModuleType(f"{mamba_pkg}.mamba2_metadata") - metadata_stub.REPLAY_WORK_POSITION_IN_DECODE_BATCH = REPLAY_WORK_POSITION_IN_DECODE_BATCH - metadata_stub.REPLAY_WORK_CACHE_SLOT = REPLAY_WORK_CACHE_SLOT - metadata_stub.REPLAY_WORK_PNAT = REPLAY_WORK_PNAT - metadata_stub.REPLAY_WORK_CACHE_BUF_IDX = REPLAY_WORK_CACHE_BUF_IDX - metadata_stub.REPLAY_WORK_ITEM_WIDTH = REPLAY_WORK_ITEM_WIDTH - sys.modules[f"{mamba_pkg}.mamba2_metadata"] = metadata_stub - # 4. The actual kernels - replay_mod = _load("replay_selective_state_update", "replay_selective_state_update.py") - conv1d_mod = _load("causal_conv1d_triton", "causal_conv1d_triton.py") - - return ( - replay_mod.replay_selective_state_update, - replay_mod._resolve_tuning, - conv1d_mod.causal_conv1d_update, - ) - - -def _import_mamba_kernels_full(): - """Import via the standard tensorrt_llm package (slow but safe).""" - replay_mod = importlib.import_module( - "tensorrt_llm._torch.modules.mamba.replay_selective_state_update" - ) - from tensorrt_llm._torch.modules.mamba.causal_conv1d_triton import causal_conv1d_update - - return ( - replay_mod.replay_selective_state_update, - replay_mod._resolve_tuning, - causal_conv1d_update, - ) - - -# Use fast import by default; --full-import parsed later but we need the -# functions at module level. Check sys.argv early. -if "--full-import" in sys.argv: - ( - replay_selective_state_update, - resolve_replay_tuning, - causal_conv1d_update, - ) = _import_mamba_kernels_full() -else: - try: - ( - replay_selective_state_update, - resolve_replay_tuning, - causal_conv1d_update, - ) = _import_mamba_kernels_fast() - except Exception as e: # noqa: BLE001 - exit loudly; don't hide a fast-import regression - print( - f"ERROR: fast import failed ({type(e).__name__}: {e})\n" - "Re-run with --full-import for the slow but stable path, " - "then file a bug or fix _import_mamba_kernels_fast.", - file=sys.stderr, - ) - sys.exit(1) - -# Model config defaults (Nemotron-3-Super-120B full model). -# --tp-size divides nheads and ngroups to get the per-GPU slice. -# TP=1: nheads=128, ngroups=8 -# TP=4: nheads=32, ngroups=2 -# TP=8: nheads=16, ngroups=1 (default) -NHEADS = 128 -HEAD_DIM = 64 -D_STATE = 128 -NGROUPS = 8 -TP_SIZE = 8 # default; overridden by --tp-size - -# L2 flush buffer: ~128 MB — larger than L2 on A100/H100/B200 -_L2_FLUSH_SIZE = 32 * 1024 * 1024 # float32 elements → 128 MB -_l2_flush: torch.Tensor | None = None - - -def _init_l2_flush() -> None: - global _l2_flush - _l2_flush = torch.empty(_L2_FLUSH_SIZE, dtype=torch.float32, device="cuda") - - -def _flush_l2() -> None: - """Evict L2 by writing to a large buffer then synchronising.""" - assert _l2_flush is not None - _l2_flush.fill_(0.0) - torch.cuda.synchronize() - - -def _resolve_prev_ks(args, mtp_len: int) -> list[int]: - """Resolve prev_k values for one mtp_len cell. - - Two input modes (mutually exclusive in spirit; absolute wins if both given): - --prev-tokens-int "0,10,11,16" → use literal integers, clamped to - [0, max_window] (where max_window is the cache T-axis capacity). - --prev-tokens-fracs "0,0.5,1.0" → fractions of mtp_len, clamped to - [0, mtp_len] (current behavior). - - For replay-style checkpointing the cache holds up to max_window old - tokens, so absolute integers are the right knob. Fractions are kept - for back-compat with prior placeholder runs. - """ - upper = getattr(args, "max_window", 0) or mtp_len - if getattr(args, "prev_tokens_int", None): - return sorted(set(max(0, min(upper, int(v))) for v in args.prev_tokens_int)) - return sorted(set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs)) - - -def _al_counts_to_distribution( - al_to_count: dict[int, float], - T: int, - label: str, -) -> np.ndarray: - """Return a length-(T + 1) probability vector indexed by accepted length.""" - if not al_to_count: - raise SystemExit(f"no numeric AL rows found in {label}") - if max(al_to_count) > T: - raise SystemExit( - f"AL histogram {label} contains accepted length {max(al_to_count)} > T={T}" - ) - if min(al_to_count) < 0: - raise SystemExit(f"AL histogram {label} contains negative accepted lengths") - if al_to_count.get(0, 0.0) > 0.0: - print( - f"[WARN] {label} contains AL=0 mass; spec decoding normally " - "accepts at least the target token.", - file=sys.stderr, - ) - - dist = np.zeros(T + 1, dtype=np.float64) - for accepted_length, count in al_to_count.items(): - dist[accepted_length] = count - total = dist.sum() - if total == 0.0: - raise SystemExit(f"AL distribution sums to zero in {label}") - return dist / total - - -def _load_al_distribution(path: Path, T: int, column: int = 1) -> np.ndarray: - with path.open(newline="") as f: - rows = list(csv.reader(f)) - if not rows: - raise SystemExit(f"empty CSV: {path}") - - al_to_count: dict[int, float] = {} - for row in rows: - if not row or not row[0].strip(): - continue - try: - accepted_length = int(float(row[0])) - count = float(row[column]) - except (IndexError, ValueError): - continue - al_to_count[accepted_length] = al_to_count.get(accepted_length, 0.0) + count - - return _al_counts_to_distribution(al_to_count, T, str(path)) - - -def _load_builtin_pmix_distribution(T: int) -> np.ndarray: - if T != DEFAULT_PMIX_T: - raise SystemExit( - f"--pmix uses the built-in T{DEFAULT_PMIX_T} histogram; " - f"use --mtp-lengths {DEFAULT_PMIX_T} or pass --mix-csv for another T." - ) - return _al_counts_to_distribution( - DEFAULT_PMIX_REPLAY_COUNTS, - T, - f"built-in {DEFAULT_PMIX_LABEL}", - ) - - -def _markov_stationary(al_dist: np.ndarray, T: int, window: int) -> np.ndarray: - """Stationary PNAT distribution for the replay-window Markov chain.""" - n_states = window + 1 - transition = np.zeros((n_states, n_states), dtype=np.float64) - for pnat in range(n_states): - is_write = pnat + T > window - for accepted_length in range(1, T + 1): - prob = al_dist[accepted_length] - if prob == 0.0: - continue - next_pnat = accepted_length if is_write else pnat + accepted_length - assert 0 <= next_pnat <= window, ( - f"unreachable transition: pnat={pnat} al={accepted_length} " - f"write={is_write} window={window}" - ) - transition[pnat, next_pnat] += prob - - eigvals, eigvecs = np.linalg.eig(transition.T) - idx = int(np.argmin(np.abs(eigvals - 1.0))) - if abs(eigvals[idx] - 1.0) > 1e-6: - raise SystemExit( - f"stationary distribution eigensolve failed: closest eigenvalue to 1 is {eigvals[idx]}" - ) - - pi = np.real(eigvecs[:, idx]) - pi = np.maximum(pi, 0.0) - if pi.sum() == 0.0: - pi = np.ones(n_states, dtype=np.float64) / n_states - for _ in range(2000): - new_pi = pi @ transition - if np.allclose(new_pi, pi, atol=1e-12, rtol=0): - pi = new_pi - break - pi = new_pi - return pi / pi.sum() - - -def _sample_steady_state_pnat( - al_dist: np.ndarray, - T: int, - window: int, - batch: int, - K: int, - seed: int = 42, -) -> np.ndarray: - pi = _markov_stationary(al_dist, T, window) - rng = np.random.default_rng(seed) - return rng.choice( - np.arange(window + 1, dtype=np.int64), - size=(K, batch), - p=pi, - ).astype(np.int32) - - -def _build_replay_work_items_cpu( - pnat_samples: np.ndarray, - T: int, - window: int, - cache_buf_idx_samples: np.ndarray | None = None, -) -> tuple[np.ndarray, np.ndarray]: - """Build write-first replay work items for one or more PNAT sample rows.""" - samples = np.asarray(pnat_samples, dtype=np.int32) - squeeze = samples.ndim == 1 - if squeeze: - samples = samples[None, :] - if samples.ndim != 2: - raise ValueError(f"pnat_samples must be 1D or 2D, got shape {samples.shape}") - - n_samples, batch = samples.shape - write_mask = samples + T > window - order = np.argsort(-write_mask.astype(np.int8), kind="stable", axis=1) - positions = np.broadcast_to(np.arange(batch, dtype=np.int32), (n_samples, batch)) - - if cache_buf_idx_samples is None: - cache_buf_idx_samples = np.zeros_like(samples, dtype=np.int32) - else: - cache_buf_idx_samples = np.asarray(cache_buf_idx_samples, dtype=np.int32) - if cache_buf_idx_samples.ndim == 1: - cache_buf_idx_samples = np.broadcast_to(cache_buf_idx_samples[None, :], samples.shape) - if cache_buf_idx_samples.shape != samples.shape: - raise ValueError( - "cache_buf_idx_samples shape must match pnat_samples, got " - f"{cache_buf_idx_samples.shape} and {samples.shape}" - ) - - work_items = np.empty((n_samples, batch, REPLAY_WORK_ITEM_WIDTH), dtype=np.int32) - work_items[:, :, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = np.take_along_axis( - positions, order, axis=1 - ) - work_items[:, :, REPLAY_WORK_CACHE_SLOT] = work_items[ - :, :, REPLAY_WORK_POSITION_IN_DECODE_BATCH - ] - work_items[:, :, REPLAY_WORK_PNAT] = np.take_along_axis(samples, order, axis=1) - work_items[:, :, REPLAY_WORK_CACHE_BUF_IDX] = np.take_along_axis( - cache_buf_idx_samples, order, axis=1 - ) - n_writes = write_mask.sum(axis=1).astype(np.int32) - - if squeeze: - return n_writes[:1], work_items[0] - return n_writes, work_items - - -# Tensor construction helpers - -# Module-level cache for tensor buffers shared across cells. Keyed by all -# the "fixed" dimensions (state_dtype, act_dtype, max_window, mtp_len, -# nheads, head_dim, d_state, ngroups). Within a key, the batch dim grows -# in place: if a new cell requests a batch <= cached max_batch, we return -# views (slices) of the existing tensors; if batch > cached max_batch, we -# realloc at the new batch (which becomes the new max). Tensors never shrink. -# -# Rationale: torch.randn/zeros for these tensor shapes at b=512 takes -# ~10-30ms per call. At ~895 cells/min with 5 different batch sizes, -# we were re-allocating every cell. Caching saves the bulk of that per-cell -# overhead, raising GPU util in the timing phase. -# -# Reset state lives in caller (state_work = state0.copy_), so cached state0 -# is purely a reference whose contents stay fixed once allocated. This is -# fine: it's only read by the reset path. -_TENSOR_CACHE: dict = {} - - -def _empty_state_with_slot_gap( - batch: int, - nheads: int, - head_dim: int, - d_state: int, - state_dtype: torch.dtype, - device: str, - slot_gap_rows: int, -): - if slot_gap_rows == 0: - return torch.empty(batch, nheads, head_dim, d_state, device=device, dtype=state_dtype) - - ssm_rows_per_slot = nheads * head_dim - rows_per_slot = ssm_rows_per_slot + slot_gap_rows - backing = torch.empty(batch, rows_per_slot, d_state, device=device, dtype=state_dtype) - return backing[:, :ssm_rows_per_slot].view(batch, nheads, head_dim, d_state) - - -def _clone_state_preserving_layout(state: torch.Tensor): - if state.is_contiguous(): - return state.clone() - clone = torch.empty_strided( - tuple(state.shape), state.stride(), device=state.device, dtype=state.dtype - ) - clone.copy_(state) - return clone - - -def _build_tensors( - batch: int, - mtp_len: int, - state_dtype: torch.dtype, - act_dtype: torch.dtype, - nheads: int, - head_dim: int, - d_state: int, - ngroups: int, - max_window: int | None = None, - strided_state_cache: bool = False, -): - """ - Build all tensors for one benchmark configuration. - - nheads/ngroups are already TP-split (i.e. full_nheads // tp_size). - - Returns: - state0 : (batch, nheads, head_dim, d_state) – initial SSM state - x, dt, B, C : (batch, mtp_len, ...) – token inputs for both kernels - A, dt_bias, D : SSM parameters (float32, tie_hdim strides) - prev_tokens : (batch,) - state_batch_indices : identity cache-slot indirection (batch,) - replay_work_items : packed per-slot replay metadata (batch, 4) - out_incr : pre-allocated output for replay kernel (batch, mtp_len, nheads, head_dim) - out_base : pre-allocated output for baseline kernel (batch, mtp_len, nheads, head_dim) - """ - device = "cuda" - - # Cache lookup — grow batch in place if needed; else return views. - cache_key = ( - state_dtype, - act_dtype, - max_window, - mtp_len, - nheads, - head_dim, - d_state, - ngroups, - strided_state_cache, - ) - cached = _TENSOR_CACHE.get(cache_key) - if cached is not None and cached["max_batch"] >= batch: - # Hit — return slices for current batch. - b = batch - return ( - cached["state0"][:b], - cached["state_scales0"][:b] if cached["state_scales0"] is not None else None, - cached["old_x"][:b], - cached["old_B"][:b], - cached["old_dt"][:b], - cached["old_dA_cumsum"][:b], - cached["cache_buf_idx"][:b], - cached["x"][:b], - cached["dt"][:b], - cached["B"][:b], - cached["C"][:b], - cached["A"], - cached["dt_bias"], - cached["D"], - cached["prev_tokens"][:b], - cached["state_batch_indices"][:b], - cached["replay_work_items"][:b], - cached["out_incr"][:b], - cached["out_base"][:b], - cached["xbc_input"][:b], - cached["conv_state"][:b], - cached["conv_weight"], - cached["conv_bias"], - cached["d_inner"], - cached["conv_dim"], - ) - - # Miss or grow. Allocate at new max_batch (existing data, if any, is - # released — caller code re-fills via reset paths anyway). Rebind - # `batch` locally to alloc_batch so the existing allocation code below - # uses the larger size; keep request_batch for the final slice. - request_batch = batch - alloc_batch = batch if cached is None else max(batch, cached["max_batch"]) - batch = alloc_batch - - torch.manual_seed(42) - - d_inner = nheads * head_dim - conv_dim = d_inner + 2 * ngroups * d_state - d_conv = 4 # conv kernel width for Nemotron/Mamba2 - slot_gap_rows = 0 - if strided_state_cache: - conv_state_elems_per_slot = conv_dim * d_conv - if conv_state_elems_per_slot % d_state != 0: - raise ValueError( - "strided state cache requires conv state to be an integer " - f"number of d_state rows, got {conv_state_elems_per_slot=} " - f"and {d_state=}" - ) - slot_gap_rows = conv_state_elems_per_slot // d_state - - # --- SSM parameters (float32, tie_hdim strides) --- - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) # stride(-1)=0, stride(-2)=0 - - dt_bias_base = torch.randn(nheads, device=device, dtype=torch.float32) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) # stride(-1)=0 - - D_base = torch.randn(nheads, device=device, dtype=torch.float32) - D = repeat(D_base, "h -> h p", p=head_dim) - - # --- SSM state --- - # Quantized dtypes need their own initializer (torch.randn doesn't accept - # int) and a parallel fp32 scales tensor (per-(head, dim) channel decode - # scale, broadcast over dstate). Quant state is filled with realistic- - # range values via fp32 → quant; scales are derived consistently so the - # initial state isn't garbage on dequant. - _QUANT_BENCH = { - torch.int8: 127.0, - torch.int16: 32767.0, - torch.float8_e4m3fn: 448.0, - } - if state_dtype in _QUANT_BENCH: - quant_max = _QUANT_BENCH[state_dtype] - state_fp32 = torch.randn( - batch, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - amax = state_fp32.abs().amax(dim=-1) # (batch, nheads, head_dim) - encode_scale = quant_max / amax.clamp(min=1e-30) - state_scales0 = (1.0 / encode_scale).to(torch.float32) # decode scale - scaled = state_fp32 * encode_scale.unsqueeze(-1) - if state_dtype == torch.float8_e4m3fn: - state0_dense = scaled.clamp(-quant_max, quant_max).to(state_dtype) - else: - state0_dense = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) - state0 = _empty_state_with_slot_gap( - batch, nheads, head_dim, d_state, state_dtype, device, slot_gap_rows - ) - state0.copy_(state0_dense) - else: - state0 = _empty_state_with_slot_gap( - batch, nheads, head_dim, d_state, state_dtype, device, slot_gap_rows - ) - state0.copy_( - torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=state_dtype) - ) - state_scales0 = None - - # --- Cache tensors for replay kernel --- - # max_window is the cache T-axis capacity; defaults to mtp_len (the - # placeholder/degenerate case where every step is a checkpoint step). - # For real replay-style checkpointing, max_window > mtp_len. - cache_T = max_window if max_window is not None else mtp_len - # old_x: double-buffered (cache, 2, max_window, nheads, dim) - old_x = torch.randn(batch, 2, cache_T, nheads, head_dim, device=device, dtype=act_dtype) - # old_B: double-buffered (cache, 2, max_window, ngroups, dstate) - old_B = torch.randn(batch, 2, cache_T, ngroups, d_state, device=device, dtype=act_dtype) - # old_dt: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous - old_dt = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) - # old_dA_cumsum: double-buffered (cache, 2, nheads, max_window) fp32 — T contiguous - old_dA_cumsum = torch.randn(batch, 2, nheads, cache_T, device=device, dtype=torch.float32) - # cache_buf_idx: which buffer to read (0 or 1) - cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) - - # --- Token inputs (used by both replay and baseline kernels) --- - x = torch.randn(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - # dt must match D's dtype (fp32) for flashinfer — force it for all paths. - dt_base = torch.randn(batch, mtp_len, nheads, device=device, dtype=torch.float32) - dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) # tie_hdim - B = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) - C = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) - - # prev_tokens placeholder — overwritten per-run - prev_tokens = torch.zeros(batch, device=device, dtype=torch.int32) - state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) - replay_work_items = torch.empty(batch, REPLAY_WORK_ITEM_WIDTH, device=device, dtype=torch.int32) - replay_work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = state_batch_indices - replay_work_items[:, REPLAY_WORK_CACHE_SLOT] = state_batch_indices - replay_work_items[:, REPLAY_WORK_PNAT] = 0 - replay_work_items[:, REPLAY_WORK_CACHE_BUF_IDX] = 0 - - out_incr = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - out_base = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - - # --- Conv1d tensors (for --with-conv1d mode) --- - # xbc_input: (batch, conv_dim, mtp_len) — "hot" input from in_proj. - # Match production layout: in_proj output is (batch*mtp_len, conv_dim) - # contiguous, then .view(batch, mtp_len, conv_dim).transpose(1, 2) - # gives strides (mtp_len*conv_dim, 1, conv_dim) — NOT the standard - # (conv_dim*mtp_len, mtp_len, 1) of a freshly allocated 3D tensor. - # Conv1d preserves input strides in its output, so downstream split - # + view inherits the correct layout without needing .contiguous(). - xbc_input_flat = torch.randn(batch * mtp_len, conv_dim, device=device, dtype=act_dtype) - xbc_input = xbc_input_flat.view(batch, mtp_len, conv_dim).transpose(1, 2) - # conv_state: (batch, conv_dim, d_conv) — "cold" cache - conv_state = torch.randn(batch, conv_dim, d_conv, device=device, dtype=act_dtype) - # conv_weight: (conv_dim, d_conv) — parameter - conv_weight = torch.randn(conv_dim, d_conv, device=device, dtype=act_dtype) - # conv_bias: (conv_dim,) — parameter - conv_bias = torch.randn(conv_dim, device=device, dtype=act_dtype) - - # Store full-batch buffers in cache and return slices at request_batch. - _TENSOR_CACHE[cache_key] = { - "max_batch": alloc_batch, - "state0": state0, - "state_scales0": state_scales0, - "old_x": old_x, - "old_B": old_B, - "old_dt": old_dt, - "old_dA_cumsum": old_dA_cumsum, - "cache_buf_idx": cache_buf_idx, - "x": x, - "dt": dt, - "B": B, - "C": C, - "A": A, - "dt_bias": dt_bias, - "D": D, - "prev_tokens": prev_tokens, - "state_batch_indices": state_batch_indices, - "replay_work_items": replay_work_items, - "out_incr": out_incr, - "out_base": out_base, - "xbc_input": xbc_input, - "conv_state": conv_state, - "conv_weight": conv_weight, - "conv_bias": conv_bias, - "d_inner": d_inner, - "conv_dim": conv_dim, - } - rb = request_batch - return ( - state0[:rb], - state_scales0[:rb] if state_scales0 is not None else None, - old_x[:rb], - old_B[:rb], - old_dt[:rb], - old_dA_cumsum[:rb], - cache_buf_idx[:rb], - x[:rb], - dt[:rb], - B[:rb], - C[:rb], - A, - dt_bias, - D, - prev_tokens[:rb], - state_batch_indices[:rb], - replay_work_items[:rb], - out_incr[:rb], - out_base[:rb], - xbc_input[:rb], - conv_state[:rb], - conv_weight, - conv_bias, - d_inner, - conv_dim, - ) - - -# ============================================================================= -# CUPTI in-process kernel timing -# -# Self-contained module-in-a-file. Reads kernel start/end timestamps directly -# from the GPU profiling fabric via CUPTI's Activity API (1 ns -# resolution), avoiding two pitfalls of the cuda-events path: -# -# 1. cudaEvent.elapsed_time() resolution (~0.5 us) is too coarse for the -# short kernels we care about, especially with PDL + cuda graphs at -# small batch — events recorded inside a graph have proven noisy. -# 2. nsys is the only known accurate alternative, but the -# profile-export-sqlite-parse pipeline is heavy and out-of-process. -# -# This is functionally equivalent to wrapping each cell in nsys, except it -# runs in the same benchmark process and sends raw activity buffers to a -# parser process instead of materializing Python objects in the CUPTI callback. -# ============================================================================= - - -# Substring match: kernels run_fn launches that we want to time. Mirrors -# the parser in scripts/.../collect.py so cupti and nsys-based outputs agree. -_CUPTI_KEEP_KERNEL_SUBSTRINGS = ( - "_dynamic_precompute", - "_persistent_main", - "selective_scan_update", - "selective_state_update", - "checkpointing_ssu", - "causal_conv1d_update", -) - -_SR_SUPPORTED_DTYPES = ( - torch.float16, - torch.int8, - torch.int16, - torch.float8_e4m3fn, -) - - -def _sr_modes_for_dtype(state_dtype: torch.dtype, requested_modes: list[str]) -> list[str]: - """Return the rounding modes that should run for one state dtype.""" - if state_dtype == torch.float32: - return ["RN"] - if state_dtype not in _SR_SUPPORTED_DTYPES: - return [mode for mode in requested_modes if mode == "RN"] - return requested_modes - - -def _kernels_per_iter_incremental( - mode: str, - with_conv1d: bool, -) -> int: - """Expected number of CUPTI-tracked kernels per iter for the incremental - kernel chain, given the dispatch mode and the conv1d flag. - - Used to validate CUPTI record counts (no auto-inference — silent - mis-timing is the failure mode we're guarding against). - - `persistent_main` always launches both write and nowrite halves. The - half ranges are derived from device `n_writes` inside the kernel so CUDA - graphs can replay with changing write counts. - `persistent_dynamic` always launches 1 main; not affected by the flag. - """ - if mode == "persistent_dynamic": - k = 2 # 1 dynamic_precomp + 1 persistent_main - elif mode == "persistent_main": - k = 3 # 1 dynamic_precomp + write-main + nowrite-main - else: - raise ValueError(f"mode must be resolved before CUPTI kernel counting, got {mode!r}") - if with_conv1d: - k += 1 - return k - - -def _dtype_key_for_replay_tuning(dtype: torch.dtype) -> str: - return { - torch.float32: "fp32", - torch.bfloat16: "bf16", - torch.float16: "fp16", - torch.int8: "int8", - torch.int16: "int16", - torch.float8_e4m3fn: "fp8", - }.get(dtype, str(dtype)) - - -def _resolve_effective_replay_mode( - args, - batch: int, - state_dtype: torch.dtype, - use_philox: bool, - mode: str | None, -) -> str: - if mode is not None: - return mode - table_entry = resolve_replay_tuning( - batch, - args.tp_nheads, - _dtype_key_for_replay_tuning(state_dtype), - "SR" if use_philox else "RN", - ) - if table_entry is None: - return "persistent_dynamic" - table_mode, _ = table_entry - return table_mode - - -def _kernels_per_iter_baseline(with_conv1d: bool) -> int: - """Expected kernels per iter for the FlashInfer PR3324 baseline. - - The baseline runs a single state-update kernel; `--with-conv1d` prepends - one conv1d kernel. - """ - return 2 if with_conv1d else 1 - - -def _is_flashinfer_pr3324_baseline(args) -> bool: - return getattr(args, "baseline", None) == "flashinfer_pr3324" - - -def _prepare_flashinfer_jit_workspace() -> None: - # The container home cache can be read-only under sandboxed runs. Set a - # writable default before importing flashinfer.jit, which resolves this at - # import time. - os.environ.setdefault("FLASHINFER_WORKSPACE_BASE", "/tmp/flashinfer") - - -_LIBCUPTI_CANDIDATES = ( - os.environ.get("CUPTI_LIBRARY_PATH"), - "/usr/local/lib/python3.12/dist-packages/nvidia/cu13/lib/libcupti.so.13", - "libcupti.so.13", - "libcupti.so", -) -_CUPTI_SUCCESS = 0 -_CUPTI_ERROR_MAX_LIMIT_REACHED = 12 -_CUPTI_ERROR_INVALID_KIND = 21 -_CUPTI_ACTIVITY_KIND_KERNEL = 3 -_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL = 10 -_CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER = 5 -_CUPTI_HOST_BUFFER_BYTES = 1024 * 1024 -_CUPTI_HOST_BUFFER_COUNT = 16 - -# Multiprocessing start method for compile-warmup + CUPTI parser children. -# Set in __main__ from --mp-start-method. "spawn" (default) is robust; each -# child re-imports torch/triton/etc (~15s). "forkserver" preloads once and -# forks cheaply (~1s/child) — see __main__ block for the preload setup. -_MP_START_METHOD = "spawn" -_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE = 1 -_DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX = 4 - - -def _load_libcupti() -> ctypes.CDLL: - errors = [] - for candidate in _LIBCUPTI_CANDIDATES: - if not candidate: - continue - try: - return ctypes.CDLL(candidate) - except OSError as exc: - errors.append(f"{candidate}: {exc}") - raise ImportError("Unable to load libcupti: " + "; ".join(errors)) - - -class _CuptiActivityKernel11Prefix(ctypes.Structure): - _pack_ = 1 - _fields_ = [ - ("kind", ctypes.c_int), - ("cache_config", ctypes.c_uint8), - ("shared_memory_config", ctypes.c_uint8), - ("registers_per_thread", ctypes.c_uint16), - ("partitioned_global_cache_requested", ctypes.c_int), - ("partitioned_global_cache_executed", ctypes.c_int), - ("start", ctypes.c_uint64), - ("end", ctypes.c_uint64), - ("completed", ctypes.c_uint64), - ("device_id", ctypes.c_uint32), - ("context_id", ctypes.c_uint32), - ("stream_id", ctypes.c_uint32), - ("grid_x", ctypes.c_int32), - ("grid_y", ctypes.c_int32), - ("grid_z", ctypes.c_int32), - ("block_x", ctypes.c_int32), - ("block_y", ctypes.c_int32), - ("block_z", ctypes.c_int32), - ("static_shared_memory", ctypes.c_int32), - ("dynamic_shared_memory", ctypes.c_int32), - ("local_memory_per_thread", ctypes.c_uint32), - ("local_memory_total", ctypes.c_uint32), - ("correlation_id", ctypes.c_uint32), - ("grid_id", ctypes.c_int64), - ("name", ctypes.c_void_p), - ("reserved0", ctypes.c_void_p), - ("queued", ctypes.c_uint64), - ("submitted", ctypes.c_uint64), - ("launch_type", ctypes.c_uint8), - ("is_shared_memory_carveout_requested", ctypes.c_uint8), - ("shared_memory_carveout_requested", ctypes.c_uint8), - ("padding", ctypes.c_uint8), - ("shared_memory_executed", ctypes.c_uint32), - ("graph_node_id", ctypes.c_uint64), - ] - - -def _configure_cupti_get_next_record(libcupti) -> None: - libcupti.cuptiActivityGetNextRecord.argtypes = [ - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.POINTER(ctypes.c_void_p), - ] - libcupti.cuptiActivityGetNextRecord.restype = ctypes.c_int - - -def _parse_cupti_buffer_ptr(libcupti, buffer_ptr: int, valid_size: int, *, include_names: bool): - records = [] - zero_ts_count = 0 - zero_ts_names: dict[str, int] = {} - record_ptr = ctypes.c_void_p(None) - while True: - result = libcupti.cuptiActivityGetNextRecord( - ctypes.c_void_p(buffer_ptr), - valid_size, - ctypes.byref(record_ptr), - ) - if result == _CUPTI_SUCCESS: - kind = ctypes.cast(record_ptr, ctypes.POINTER(ctypes.c_int)).contents.value - if kind not in (_CUPTI_ACTIVITY_KIND_KERNEL, _CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL): - continue - kernel = ctypes.cast(record_ptr, ctypes.POINTER(_CuptiActivityKernel11Prefix)).contents - name = None - if include_names: - if kernel.name: - name = ctypes.string_at(kernel.name).decode("utf-8", errors="replace") - else: - name = "?" - if kernel.start == 0 or kernel.end == 0: - zero_ts_count += 1 - if name is not None: - zero_ts_names[name] = zero_ts_names.get(name, 0) + 1 - continue - if include_names: - records.append( - ( - name, - int(kernel.start), - int(kernel.end), - int(kernel.correlation_id), - 0, - int(kernel.graph_node_id), - int(kernel.stream_id), - ) - ) - else: - records.append( - ( - int(kernel.start), - int(kernel.end), - int(kernel.correlation_id), - int(kernel.graph_node_id), - int(kernel.stream_id), - ) - ) - elif result == _CUPTI_ERROR_MAX_LIMIT_REACHED: - break - elif result == _CUPTI_ERROR_INVALID_KIND: - break - else: - raise RuntimeError(f"cuptiActivityGetNextRecord failed with CUptiResult={result}") - return records, zero_ts_count, zero_ts_names - - -def _apply_cupti_filter_plan(numeric_records, filter_plan): - if not filter_plan: - return [ - (None, start, end, corr, 0, graph_node_id, stream_id) - for start, end, corr, graph_node_id, stream_id in sorted(numeric_records) - ] - - filtered = [] - replay_idx = 0 - record_idx = 0 - for start, end, corr, graph_node_id, stream_id in sorted(numeric_records): - if replay_idx >= len(filter_plan): - break - records_per_replay, ordinal_names = filter_plan[replay_idx] - if record_idx < len(ordinal_names): - name = ordinal_names[record_idx] - if name is not None: - filtered.append((name, start, end, corr, 0, graph_node_id, stream_id)) - record_idx += 1 - if record_idx >= records_per_replay: - replay_idx += 1 - record_idx = 0 - return filtered - - -def _cupti_parser_worker(input_queue, output_queue, ready_event) -> None: - libcupti = _load_libcupti() - _configure_cupti_get_next_record(libcupti) - shared_blocks: dict[str, shared_memory.SharedMemory] = {} - records_by_generation: dict[int, list[tuple[int, int, int, int, int]]] = {} - zero_ts_by_generation: dict[int, int] = {} - ready_event.set() - while True: - item = input_queue.get() - if item is None: - break - kind = item[0] - if kind == "buffer": - _, generation, buffer_id, name, valid_size = item - shm = shared_blocks.get(name) - if shm is None: - shm = shared_memory.SharedMemory(name=name) - shared_blocks[name] = shm - shared_char = ctypes.c_char.from_buffer(shm.buf) - try: - parser_ptr = ctypes.addressof(shared_char) - records, zero_ts_count, _ = _parse_cupti_buffer_ptr( - libcupti, - parser_ptr, - valid_size, - include_names=False, - ) - records_by_generation.setdefault(generation, []).extend(records) - zero_ts_by_generation[generation] = ( - zero_ts_by_generation.get(generation, 0) + zero_ts_count - ) - ctypes.memset(parser_ptr, 0, len(shm.buf)) - except Exception as exc: # pragma: no cover - diagnostic worker path - output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) - finally: - del shared_char - output_queue.put( - {"kind": "buffer_done", "generation": generation, "buffer_id": buffer_id} - ) - elif kind == "finish": - if len(item) == 4: - _, generation, filter_plan, stats_request = item - else: - _, generation, filter_plan = item - stats_request = None - try: - raw_records = records_by_generation.pop(generation, []) - zero_ts_count = zero_ts_by_generation.pop(generation, 0) - filtered_records = _apply_cupti_filter_plan(raw_records, filter_plan) - stats = None - parser_stats_ms = 0.0 - stats_ready = stats_request is not None - if stats_request is not None: - stats_start_s = time.perf_counter() - stats = _stats_from_cupti_records( - filtered_records, - int(stats_request["warmup"]), - int(stats_request["iters"]), - str(stats_request["tag"]), - int(stats_request["expected_K"]), - zero_ts_count=zero_ts_count, - zero_ts_names={}, - include_details=bool(stats_request.get("include_details", True)), - ) - parser_stats_ms = 1000.0 * (time.perf_counter() - stats_start_s) - filtered_records = [] - output_queue.put( - { - "kind": "finish_done", - "generation": generation, - "records": filtered_records, - "zero_ts_count": zero_ts_count, - "zero_ts_names": {}, - "raw_record_count": len(raw_records), - "stats": stats, - "stats_ready": stats_ready, - "parser_stats_ms": parser_stats_ms, - } - ) - except Exception as exc: # pragma: no cover - diagnostic worker path - output_queue.put({"kind": "error", "generation": generation, "error": repr(exc)}) - else: - output_queue.put( - {"kind": "error", "generation": -1, "error": f"unknown parser message {kind!r}"} - ) - for shm in shared_blocks.values(): - shm.close() - - -class CuptiKernelTimer: - """Raw CUPTI Activity timer with out-of-process parsing for timed runs. - - CUPTI's callback gives us raw activity buffers. The callback only hands - shared-memory buffer metadata to a parser process, so the main process - avoids the cupti-python per-record object creation cost during the timed - path. A single local calibration replay may parse names in-process to - build an ordinal filter plan for a just-captured CUDA graph. - """ - - _instance = None - _import_error = None - - _request_callback_type = ctypes.CFUNCTYPE( - None, - ctypes.POINTER(ctypes.c_void_p), - ctypes.POINTER(ctypes.c_size_t), - ctypes.POINTER(ctypes.c_size_t), - ) - _complete_callback_type = ctypes.CFUNCTYPE( - None, - ctypes.c_void_p, - ctypes.c_uint32, - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.c_size_t, - ) - - @classmethod - def get(cls) -> "CuptiKernelTimer": - if cls._instance is not None: - return cls._instance - if cls._import_error is not None: - raise cls._import_error - try: - cls._instance = cls() - return cls._instance - except ImportError as exc: # pragma: no cover - env-dependent - cls._import_error = exc - raise - - def __init__(self) -> None: - self._libcupti = _load_libcupti() - self._configure_functions() - self._lock = threading.Lock() - self._shared_buffers: dict[int, shared_memory.SharedMemory] = {} - self._buffer_id_by_ptr: dict[int, int] = {} - self._free_buffer_ids: list[int] = [] - self._local_completed: list[tuple[int, int]] = [] - self._mode = "drop" - self._generation = 0 - self._finish_results: dict[int, dict] = {} - self._parser_errors: list[str] = [] - self._filter_plan = () - self._last_start_timing: dict[str, float] = {} - self._last_stop_timing: dict[str, float] = {} - self._current_flush_period_ms = 0 - self._mp_ctx = mp.get_context(_MP_START_METHOD) - # Retry parser-process spawn: concurrent bench instances on the same - # node race on POSIX named semaphores in /dev/shm — child can die in - # pickle.load with FileNotFoundError in SemLock._rebuild before - # signalling ready_event. Detect early-dead child via is_alive() so - # we don't waste the full timeout, and retry up to 3x with jitter. - last_err = None - for _spawn_attempt in range(3): - self._parse_input_queue = self._mp_ctx.Queue() - self._parse_output_queue = self._mp_ctx.Queue() - ready_event = self._mp_ctx.Event() - self._parse_process = self._mp_ctx.Process( - target=_cupti_parser_worker, - args=(self._parse_input_queue, self._parse_output_queue, ready_event), - ) - self._parse_process.start() - deadline = time.time() + 30.0 - spawn_ok = False - while time.time() < deadline: - if ready_event.wait(timeout=0.5): - spawn_ok = True - break - if not self._parse_process.is_alive(): - break - if spawn_ok: - last_err = None - break - last_err = ( - f"attempt {_spawn_attempt + 1}: " - f"alive={self._parse_process.is_alive()}, " - f"exitcode={self._parse_process.exitcode}" - ) - try: - if self._parse_process.is_alive(): - self._parse_process.terminate() - self._parse_process.join(timeout=2.0) - except Exception: - pass - time.sleep(0.5 + 0.5 * _spawn_attempt) - if last_err is not None: - raise RuntimeError( - f"CUPTI parser process did not initialize after 3 attempts: {last_err}" - ) - - self._set_zeroed_host_buffer_attr() - for _ in range(_CUPTI_HOST_BUFFER_COUNT): - self._free_buffer_ids.append(self._allocate_shared_buffer()) - - self._request_callback = self._request_callback_type(self._request_buffer) - self._complete_callback = self._complete_callback_type(self._complete_buffer) - self._check( - self._libcupti.cuptiActivityRegisterCallbacks( - self._request_callback, - self._complete_callback, - ) - ) - self._check(self._libcupti.cuptiActivityEnable(_CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL)) - atexit.register(self.close) - - def _configure_functions(self) -> None: - self._libcupti.cuptiActivityRegisterCallbacks.argtypes = [ - self._request_callback_type, - self._complete_callback_type, - ] - self._libcupti.cuptiActivityRegisterCallbacks.restype = ctypes.c_int - self._libcupti.cuptiActivityEnable.argtypes = [ctypes.c_int] - self._libcupti.cuptiActivityEnable.restype = ctypes.c_int - self._libcupti.cuptiActivityFlushAll.argtypes = [ctypes.c_uint32] - self._libcupti.cuptiActivityFlushAll.restype = ctypes.c_int - self._libcupti.cuptiActivityFlushPeriod.argtypes = [ctypes.c_uint32] - self._libcupti.cuptiActivityFlushPeriod.restype = ctypes.c_int - self._libcupti.cuptiActivitySetAttribute.argtypes = [ - ctypes.c_int, - ctypes.POINTER(ctypes.c_size_t), - ctypes.c_void_p, - ] - self._libcupti.cuptiActivitySetAttribute.restype = ctypes.c_int - _configure_cupti_get_next_record(self._libcupti) - - def _set_zeroed_host_buffer_attr(self) -> None: - value_obj = ctypes.c_uint8(1) - size_obj = ctypes.c_size_t(ctypes.sizeof(value_obj)) - result = self._libcupti.cuptiActivitySetAttribute( - _CUPTI_ACTIVITY_ATTR_ZEROED_OUT_ACTIVITY_BUFFER, - ctypes.byref(size_obj), - ctypes.byref(value_obj), - ) - if result != _CUPTI_SUCCESS: - print( - "[WARN] CUPTI zeroed host-buffer attribute failed; " - f"continuing with default CUPTI buffer handling (CUptiResult={result}).", - file=sys.stderr, - ) - - def _check(self, result: int) -> None: - if result != _CUPTI_SUCCESS: - raise RuntimeError(f"CUPTI call failed with CUptiResult={result}") - - def _allocate_shared_buffer(self) -> int: - buffer_id = len(self._shared_buffers) - shm = shared_memory.SharedMemory(create=True, size=_CUPTI_HOST_BUFFER_BYTES) - shared_char = ctypes.c_char.from_buffer(shm.buf) - try: - ptr = ctypes.addressof(shared_char) - finally: - del shared_char - if ptr % 8 != 0: - shm.close() - shm.unlink() - raise RuntimeError("CUPTI shared-memory activity buffer was not 8-byte aligned") - self._shared_buffers[buffer_id] = shm - self._buffer_id_by_ptr[ptr] = buffer_id - return buffer_id - - def _buffer_ptr(self, buffer_id: int) -> int: - shm = self._shared_buffers[buffer_id] - shared_char = ctypes.c_char.from_buffer(shm.buf) - try: - return ctypes.addressof(shared_char) - finally: - del shared_char - - def _request_buffer(self, buffer, size, max_num_records) -> None: - with self._lock: - if self._free_buffer_ids: - buffer_id = self._free_buffer_ids.pop() - else: - buffer_id = self._allocate_shared_buffer() - ptr = self._buffer_ptr(buffer_id) - buffer[0] = ptr - size[0] = _CUPTI_HOST_BUFFER_BYTES - max_num_records[0] = 0 - - def _complete_buffer(self, context, stream_id, buffer, size, valid_size) -> None: - del context, stream_id, size - buffer_ptr = int(buffer) - valid_size_int = int(valid_size) - with self._lock: - mode = self._mode - generation = self._generation - buffer_id = self._buffer_id_by_ptr[buffer_ptr] - if valid_size_int == 0 or mode == "drop": - self._free_buffer_ids.append(buffer_id) - return - if mode == "local": - self._local_completed.append((buffer_id, valid_size_int)) - return - shm = self._shared_buffers[buffer_id] - self._parse_input_queue.put(("buffer", generation, buffer_id, shm.name, valid_size_int)) - - def _handle_parser_result(self, result: dict) -> None: - kind = result.get("kind") - if kind == "buffer_done": - with self._lock: - self._free_buffer_ids.append(int(result["buffer_id"])) - elif kind == "finish_done": - self._finish_results[int(result["generation"])] = result - elif kind == "error": - self._parser_errors.append(str(result.get("error"))) - - def _drain_parser_results(self) -> None: - while True: - try: - result = self._parse_output_queue.get_nowait() - except queue.Empty: - break - self._handle_parser_result(result) - - def is_generation_ready(self, generation: int) -> bool: - self._drain_parser_results() - return generation in self._finish_results or bool(self._parser_errors) - - def _flush(self, flag: int) -> None: - self._check(self._libcupti.cuptiActivityFlushAll(flag)) - - def _set_flush_period_ms(self, period_ms: int) -> None: - if period_ms == self._current_flush_period_ms: - return - self._check(self._libcupti.cuptiActivityFlushPeriod(period_ms)) - self._current_flush_period_ms = period_ms - - def _begin( - self, - mode: str, - filter_plan=(), - flush_period_ms: int = 0, - collect_timing: bool = False, - ) -> int: - start_timing: dict[str, float] = {} - with self._lock: - self._mode = "drop" - phase_start_s = time.perf_counter() if collect_timing else 0.0 - self._flush(1) - if collect_timing: - start_timing["forced_flush_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) - phase_start_s = time.perf_counter() if collect_timing else 0.0 - self._drain_parser_results() - if collect_timing: - start_timing["drain_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) - with self._lock: - self._generation += 1 - generation = self._generation - self._mode = mode - self._local_completed = [] - self._filter_plan = filter_plan - if flush_period_ms > 0: - phase_start_s = time.perf_counter() if collect_timing else 0.0 - self._set_flush_period_ms(flush_period_ms) - if collect_timing: - start_timing["period_enable_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) - self._last_start_timing = start_timing - return generation - - def capture_names(self, replay_fn) -> tuple[list[tuple], int, dict]: - """Run a small calibration replay and parse kernel names locally.""" - self._begin("local") - replay_fn() - torch.cuda.synchronize() - self._flush(0) - records: list[tuple] = [] - zero_ts_count = 0 - zero_ts_names: dict[str, int] = {} - with self._lock: - completed = list(self._local_completed) - self._local_completed = [] - self._mode = "drop" - for buffer_id, valid_size in completed: - ptr = self._buffer_ptr(buffer_id) - recs, zeros, zero_names = _parse_cupti_buffer_ptr( - self._libcupti, - ptr, - valid_size, - include_names=True, - ) - records.extend(recs) - zero_ts_count += zeros - for name, count in zero_names.items(): - zero_ts_names[name] = zero_ts_names.get(name, 0) + count - ctypes.memset(ptr, 0, _CUPTI_HOST_BUFFER_BYTES) - with self._lock: - self._free_buffer_ids.append(buffer_id) - records.sort(key=lambda r: r[1]) - return records, zero_ts_count, zero_ts_names - - def start( - self, - filter_plan=(), - flush_period_ms: int = 0, - collect_timing: bool = False, - ) -> None: - self._begin("parser", filter_plan, flush_period_ms, collect_timing) - - def stop_async( - self, - collect_timing: bool = False, - stats_request: dict | None = None, - ) -> tuple[int, dict[str, float]]: - stop_timing: dict[str, float] = {} - generation = self._generation - phase_start_s = time.perf_counter() if collect_timing else 0.0 - self._set_flush_period_ms(0) - if collect_timing: - stop_timing["period_disable_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) - phase_start_s = time.perf_counter() if collect_timing else 0.0 - self._flush(0) - if collect_timing: - stop_timing["flush_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) - with self._lock: - self._mode = "drop" - filter_plan = self._filter_plan - self._parse_input_queue.put(("finish", generation, filter_plan, stats_request)) - self._last_stop_timing = stop_timing - return generation, stop_timing - - def wait_for_generation_result( - self, - generation: int, - stop_timing: dict[str, float] | None = None, - collect_timing: bool = False, - ) -> dict: - if stop_timing is None: - stop_timing = {} - phase_start_s = time.perf_counter() if collect_timing else 0.0 - deadline = time.perf_counter() + 10.0 - while time.perf_counter() < deadline: - result = self._finish_results.pop(generation, None) - if result is not None: - if collect_timing: - stop_timing["parser_wait_ms"] = 1000.0 * (time.perf_counter() - phase_start_s) - stop_timing["total_ms"] = ( - stop_timing.get("period_disable_ms", 0.0) - + stop_timing.get("flush_ms", 0.0) - + stop_timing["parser_wait_ms"] - ) - self._last_stop_timing = stop_timing - return result - timeout_s = max(0.0, min(0.01, deadline - time.perf_counter())) - try: - parser_result = self._parse_output_queue.get(timeout=timeout_s) - except queue.Empty: - continue - self._handle_parser_result(parser_result) - if self._parser_errors: - raise RuntimeError("CUPTI parser process failed: " + "; ".join(self._parser_errors)) - raise TimeoutError("Timed out waiting for CUPTI parser process") - - def wait_for_generation( - self, - generation: int, - stop_timing: dict[str, float] | None = None, - collect_timing: bool = False, - ) -> tuple[list[tuple], int, dict, int]: - result = self.wait_for_generation_result(generation, stop_timing, collect_timing) - return ( - list(result["records"]), - int(result["zero_ts_count"]), - dict(result["zero_ts_names"]), - int(result["raw_record_count"]), - ) - - def stop(self, collect_timing: bool = False) -> tuple[list[tuple], int, dict, int]: - generation, stop_timing = self.stop_async(collect_timing) - return self.wait_for_generation(generation, stop_timing, collect_timing) - - def last_start_timing(self) -> dict[str, float]: - return dict(self._last_start_timing) - - def last_stop_timing(self) -> dict[str, float]: - return dict(self._last_stop_timing) - - def close(self) -> None: - parse_process = getattr(self, "_parse_process", None) - if parse_process is not None and parse_process.is_alive(): - self._parse_input_queue.put(None) - parse_process.join(timeout=5.0) - if parse_process.is_alive(): - parse_process.terminate() - parse_process.join(timeout=1.0) - for shm in getattr(self, "_shared_buffers", {}).values(): - try: - shm.close() - shm.unlink() - except FileNotFoundError: - pass - - -# ============================================================================= -# Timing helpers -# ============================================================================= - - -def _stats_from_spans(spans_us: list[float]) -> dict: - """Compute median / p95 / p99 / n from a per-iter span list.""" - s = sorted(spans_us) - return { - "median": statistics.median(s), - "p95": s[int(0.95 * len(s))], - "p99": s[int(0.99 * len(s))], - "n": len(s), - } - - -def _binomial_pmf(n: int, p: float) -> tuple[float, ...]: - return tuple(math.comb(n, k) * (p**k) * ((1.0 - p) ** (n - k)) for k in range(n + 1)) - - -def _kmix_bucket_score( - iters_us, - n_writes_per_iter, - batch: int, - write_frac: float | None, -) -> float | None: - """Compute the search-driver kmix score from per-iteration timings.""" - if ( - write_frac is None - or not iters_us - or not n_writes_per_iter - or len(iters_us) != len(n_writes_per_iter) - ): - return None - - buckets: dict[int, list[float]] = {} - for span, n_writes in zip(iters_us, n_writes_per_iter): - k = int(n_writes) - if 0 <= k <= batch: - buckets.setdefault(k, []).append(float(span)) - if not buckets: - return None - - pmf = _binomial_pmf(batch, write_frac) - numerator = 0.0 - denominator = 0.0 - for k, spans in buckets.items(): - numerator += pmf[k] * statistics.median(spans) - denominator += pmf[k] - return numerator / denominator if denominator > 0.0 else None - - -def _stats_from_cupti_records( - records, - warmup, - iters, - tag, - expected_K, - zero_ts_count: int = 0, - zero_ts_names: dict | None = None, - include_details: bool = True, -): - """Bin a flat CUPTI kernel record stream into per-iter spans + per-kernel - relative timestamps. Used by both graph and eager CUPTI paths. - - `records` are tuples (name, start_ns, end_ns, ...) — see CuptiKernelTimer. - `expected_K` is the kernels-per-iter count the caller declares; we - validate the CUPTI total matches `expected_K * (warmup + iters)` exactly. - On mismatch we dump per-name record counts so missing or extra kernels - are obvious (most common cause: a new dispatch mode whose kernels lack - a matching entry in `_CUPTI_KEEP_KERNEL_SUBSTRINGS`, silently filtering - them out). - """ - records = [ - r - for r in records - if r[0] is not None and any(s in r[0] for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS) - ] - records.sort(key=lambda r: r[1]) # by start_ns - - total = len(records) - expected_iters = warmup + iters - expected_total = expected_K * expected_iters - if total != expected_total: - from collections import Counter - - name_counts = dict(Counter(r[0] for r in records)) - # Non-fatal: skip this cell instead of killing the whole sweep. - # Mismatch may be a CUPTI dropped-records issue (rare configs), - # not necessarily a K-table bug. Log so the user can investigate - # the specific cell post-hoc; return None so the caller can skip - # writing a JSON row. - zero_msg = "" - if zero_ts_count: - zero_msg = ( - f" + {zero_ts_count} records with start/end=0 " - f"(dropped by callback, breakdown {zero_ts_names}). " - f"Total observed kernel records (timed + zero-ts) = " - f"{total + zero_ts_count} / {expected_total}." - ) - print( - f"[WARN] CUPTI capture mismatch for {tag!r}: expected " - f"{expected_K} kernels/iter × {expected_iters} iters " - f"(warmup+iters) = {expected_total} records, got {total}. " - f"Kernel record counts: {name_counts}.{zero_msg} SKIPPING cell.", - file=sys.stderr, - flush=True, - ) - # Per-record dump: (name, start_ns_rel, end_ns_rel, corr_id, graph_id, stream_id). - # Times relative to first record so absolute ns isn't drowning output. - # Limit dump to first 30 records to avoid flooding logs at high K. - if records: - t0_ns = records[0][1] - for i, r in enumerate(records[:30]): - # r = (name, start_ns, end_ns, corr_id, graph_id, graph_node_id, stream_id) - rel_start = (r[1] - t0_ns) / 1000.0 # us - rel_end = (r[2] - t0_ns) / 1000.0 - print( - f" rec[{i:3d}] name={r[0]!r} start={rel_start:.2f}us " - f"end={rel_end:.2f}us corr={r[3]} graph={r[4]} stream={r[6]}", - file=sys.stderr, - flush=True, - ) - if len(records) > 30: - print( - f" ... ({len(records) - 30} more records elided)", file=sys.stderr, flush=True - ) - return None - K = expected_K - timed = records[warmup * K :] - - spans_us: list[float] = [] - per_kernel: dict[str, dict[str, list[float]]] = {} - for i in range(iters): - chunk = timed[i * K : (i + 1) * K] - iter_start_ns = min(r[1] for r in chunk) - iter_end_ns = max(r[2] for r in chunk) - spans_us.append((iter_end_ns - iter_start_ns) / 1000.0) - if include_details: - for r in chunk: - name = r[0] - slot = per_kernel.setdefault(name, {"start_us": [], "end_us": []}) - slot["start_us"].append((r[1] - iter_start_ns) / 1000.0) - slot["end_us"].append((r[2] - iter_start_ns) / 1000.0) - - out = _stats_from_spans(spans_us) - out["iters_us"] = spans_us - if include_details: - out["per_kernel"] = per_kernel - return out - - -_PRE_GRAPH_WARMUP_ITERS = 1 -_CUPTI_FILTER_PLAN_CACHE: dict[tuple, tuple[int, tuple[str | None, ...]]] = {} - - -class _HostTiming: - def __init__(self, enabled: bool) -> None: - self.enabled = enabled - self.values: dict[str, float | int | bool] = {} - self._total_start_s = time.perf_counter() if enabled else 0.0 - self._phase_start_s = 0.0 - - def start(self) -> None: - if self.enabled: - self._phase_start_s = time.perf_counter() - - def stop(self, key: str) -> None: - if self.enabled: - self.values[key] = 1000.0 * (time.perf_counter() - self._phase_start_s) - - def add(self, key: str, value: float | int | bool) -> None: - if self.enabled: - self.values[key] = value - - def stop_total(self) -> None: - if self.enabled: - self.values["total_ms"] = 1000.0 * (time.perf_counter() - self._total_start_s) - - def attach(self, stats: dict | None) -> None: - if self.enabled and stats is not None: - stats["host_timing"] = self.values - - -class _PendingCuptiStats: - def __init__( - self, - timer: CuptiKernelTimer, - generation: int, - stop_timing: dict[str, float], - host_timing: _HostTiming, - *, - warmup: int, - iters: int, - tag: str, - expected_K: int, - expected_raw_record_count: int, - ) -> None: - self._timer = timer - self._generation = generation - self._stop_timing = stop_timing - self._host_timing = host_timing - self._warmup = warmup - self._iters = iters - self._tag = tag - self._expected_K = expected_K - self._expected_raw_record_count = expected_raw_record_count - - def is_ready(self) -> bool: - return self._timer.is_generation_ready(self._generation) - - def resolve(self) -> dict | None: - result = self._timer.wait_for_generation_result( - self._generation, - self._stop_timing, - collect_timing=self._host_timing.enabled, - ) - for key, value in self._timer.last_stop_timing().items(): - self._host_timing.add(f"cupti_stop_{key}", value) - raw_record_count = int(result["raw_record_count"]) - if raw_record_count != self._expected_raw_record_count: - print( - f"[WARN] CUPTI raw-record mismatch for {self._tag!r}: expected " - f"{self._expected_raw_record_count}, got {raw_record_count}. SKIPPING cell.", - file=sys.stderr, - ) - return None - - if result.get("stats_ready"): - stats = result.get("stats") - self._host_timing.add("stats_ms", 0.0) - self._host_timing.add("parser_stats_ms", float(result.get("parser_stats_ms", 0.0))) - else: - self._host_timing.start() - stats = _stats_from_cupti_records( - list(result["records"]), - self._warmup, - self._iters, - self._tag, - self._expected_K, - zero_ts_count=int(result["zero_ts_count"]), - zero_ts_names=dict(result["zero_ts_names"]), - ) - self._host_timing.stop("stats_ms") - self._host_timing.attach(stats) - return stats - - -def _target_name_or_none(name: str | None) -> str | None: - if name is None: - return None - if any(s in name for s in _CUPTI_KEEP_KERNEL_SUBSTRINGS): - return name - return None - - -def _capture_group_graph( - args, - run_fn, - reset_fn, - group_iters: int, - graph_pre_iter_fn=None, -) -> torch.cuda.CUDAGraph: - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - for j in range(group_iters): - if graph_pre_iter_fn is not None: - graph_pre_iter_fn(j) - reset_fn() - if args.l2_flush: - _l2_flush.fill_(0.0) - run_fn() - return graph - - -def _graph_group_iters(args, total_iters: int, pre_iter_fn, pre_iter_group_factory) -> int: - """Pick the graph-group size unconditionally; the caller is expected to - round total_iters up to a multiple of this so all iters fit in clean - replays. Sample arrays are pre-padded at allocation (see _sample_pnat - call site) so the per-replay window can index past the user-requested - iter count by up to group_iters-1 extra samples. - """ - if pre_iter_fn is not None and pre_iter_group_factory is None: - # Per-iter callback without a group-factory: can't batch. - return 1 - requested = getattr(args, "cuda_graph_group_iters", None) - if requested is None: - return ( - _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX - if pre_iter_group_factory is not None - else _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE - ) - return max(1, int(requested)) - - -def _get_cupti_filter_plan( - timer: CuptiKernelTimer, graph, cache_key: tuple | None, group_iters: int -) -> tuple[int, tuple[str | None, ...]]: - full_cache_key = None if cache_key is None else (cache_key, group_iters) - if full_cache_key is not None: - cached = _CUPTI_FILTER_PLAN_CACHE.get(full_cache_key) - if cached is not None: - return cached - - records, zero_ts_count, zero_ts_names = timer.capture_names(graph.replay) - if zero_ts_count: - print( - f"[WARN] CUPTI calibration saw {zero_ts_count} zero-timestamp records " - f"(breakdown {zero_ts_names}); continuing with nonzero records.", - file=sys.stderr, - ) - ordinal_names = tuple(_target_name_or_none(r[0]) for r in records) - target_count = sum(name is not None for name in ordinal_names) - if target_count == 0: - raise RuntimeError("CUPTI calibration did not find any target kernel records") - plan = (len(records), ordinal_names) - if full_cache_key is not None: - _CUPTI_FILTER_PLAN_CACHE[full_cache_key] = plan - return plan - - -def _time_kernel_cuda_graph( - args, - run_fn, - reset_fn, - tag: str, - *, - expected_K: int, - pre_iter_fn=None, - pre_iter_group_factory=None, - iters_override: int | None = None, - cupti_plan_key: tuple | None = None, -) -> dict: - """CUDA-graph CUPTI timer (graph-per-iter design). - - Captures one CUDA graph holding a small group of logical iterations - (per-iter setup + reset + l2_flush + run_fn) and replays it enough - times to cover `warmup + iters`. - - Why graph-per-iter (vs the older "one giant graph holding all iters" - design): instantiating a CUDA graph is expensive — proportional to - graph size — so a single small graph instantiated once is much - cheaper than one big graph instantiated for each cell of a sweep. - Replays are cheap regardless. - - Mix cells use a per-replay device window: an outside-graph copy loads - the next group of PNAT/n_writes samples, then graph-captured per-iter - copies update kernel inputs before each reset + L2 flush + run. - - Pre-graph eager warmup: forces PyTorch's caching allocator - + Triton's autotune cache to settle before capture so the graph - doesn't bake in init-only allocations. - - ``iters_override`` (if not None) overrides ``args.iters`` for this - call. Used to give mix scenarios a higher iter count than pure - (more iters = more independent mix draws averaged in). - """ - host_timing = _HostTiming(bool(getattr(args, "host_timing", False))) - timer = CuptiKernelTimer.get() - warmup = args.warmup - iters = iters_override if iters_override is not None else args.iters - - # Pre-graph eager warmup: full per-iter chain once. This settles - # Triton/PyTorch setup and wrapper-side intermediate allocations; - # skipping it risks lazy work leaking into graph capture. - warmup_iters = _PRE_GRAPH_WARMUP_ITERS - host_timing.add("pre_graph_warmup_iters", warmup_iters) - host_timing.start() - for _ in range(warmup_iters): - reset_fn() - if pre_iter_fn is not None: - pre_iter_fn(0) - run_fn() - if warmup_iters > 0: - torch.cuda.synchronize() - host_timing.stop("pre_graph_warmup_ms") - - total_iters = warmup + iters - group_iters = _graph_group_iters(args, total_iters, pre_iter_fn, pre_iter_group_factory) - # Args are rounded at argparse-time so warmup+iters/mix_iters are already - # multiples of the relevant group_iters. Assert here to catch any caller - # bypassing argparse. - assert total_iters % group_iters == 0, ( - f"total_iters={total_iters} not a multiple of group_iters={group_iters}; " - f"args.warmup/iters/mix_iters should be rounded post-argparse." - ) - pre_replay_fn = None - graph_pre_iter_fn = None - if pre_iter_group_factory is not None and group_iters > 1: - pre_replay_fn, graph_pre_iter_fn = pre_iter_group_factory(group_iters) - - # Reset just before capture so warmup state changes don't bleed in. - host_timing.start() - reset_fn() - torch.cuda.synchronize() - host_timing.stop("pre_capture_reset_ms") - - # Capture a small group of identical logical iterations. Mix/pre_iter - # cells can group when they provide a graph-side pre-iter updater backed - # by a per-replay device window. - host_timing.start() - g = _capture_group_graph(args, run_fn, reset_fn, group_iters, graph_pre_iter_fn) - host_timing.stop("graph_capture_ms") - - if pre_replay_fn is not None: - host_timing.start() - pre_replay_fn(0) - torch.cuda.synchronize() - host_timing.stop("graph_preload_ms") - - plan_cache_key = None if cupti_plan_key is None else (cupti_plan_key, group_iters) - host_timing.add( - "cupti_plan_cached", - (plan_cache_key is not None and plan_cache_key in _CUPTI_FILTER_PLAN_CACHE), - ) - host_timing.start() - records_per_replay, ordinal_names = _get_cupti_filter_plan( - timer, - g, - cupti_plan_key, - group_iters, - ) - host_timing.stop("cupti_plan_ms") - target_count = sum(name is not None for name in ordinal_names) - expected_targets_per_replay = expected_K * group_iters - if target_count != expected_targets_per_replay: - print( - f"[WARN] CUPTI calibration mismatch for {tag!r}: expected " - f"{expected_targets_per_replay} target records in a {group_iters}-iter graph replay, " - f"got {target_count} target records out of {records_per_replay} total records.", - file=sys.stderr, - ) - - # Time: replay the grouped graph enough times to cover warmup+iters. - # Mix cells preload one device window per replay on the same stream. - # CUPTI records every kernel launch; _stats_from_cupti_records - # validates against expected_K and slices warmup off the front. - graph_replays = total_iters // group_iters - filter_plan = ((records_per_replay, ordinal_names),) * graph_replays - cupti_flush_period_ms = max(0, int(getattr(args, "cupti_flush_period_ms", 0))) - host_timing.start() - timer.start( - filter_plan, - flush_period_ms=cupti_flush_period_ms, - collect_timing=host_timing.enabled, - ) - host_timing.stop("cupti_start_ms") - for key, value in timer.last_start_timing().items(): - host_timing.add(f"cupti_start_{key}", value) - torch.cuda.nvtx.range_push(tag) - host_timing.start() - for i in range(graph_replays): - if pre_replay_fn is not None: - pre_replay_fn(i) - elif pre_iter_fn is not None: - pre_iter_fn(i) - g.replay() - host_timing.stop("graph_enqueue_ms") - host_timing.start() - torch.cuda.synchronize() - host_timing.stop("graph_sync_ms") - torch.cuda.nvtx.range_pop() - expected_raw_record_count = records_per_replay * graph_replays - host_timing.start() - if int(getattr(args, "cupti_defer_depth", 1)) > 1: - generation, stop_timing = timer.stop_async( - collect_timing=host_timing.enabled, - stats_request={ - "warmup": warmup, - "iters": iters, - "tag": tag, - "expected_K": expected_K, - "include_details": bool(getattr(args, "json_detailed", False)), - }, - ) - host_timing.stop("cupti_stop_ms") - for key, value in timer.last_stop_timing().items(): - host_timing.add(f"cupti_stop_{key}", value) - host_timing.stop_total() - host_timing.add("graph_group_iters", group_iters) - host_timing.add("graph_replays", graph_replays) - host_timing.add("cupti_records_per_replay", records_per_replay) - host_timing.add("cupti_target_records_per_replay", target_count) - host_timing.add("cupti_raw_records_expected", expected_raw_record_count) - host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) - return _PendingCuptiStats( - timer, - generation, - stop_timing, - host_timing, - warmup=warmup, - iters=iters, - tag=tag, - expected_K=expected_K, - expected_raw_record_count=expected_raw_record_count, - ) - - records, zero_ts_count, zero_ts_names, raw_record_count = timer.stop( - collect_timing=host_timing.enabled, - ) - host_timing.stop("cupti_stop_ms") - for key, value in timer.last_stop_timing().items(): - host_timing.add(f"cupti_stop_{key}", value) - if raw_record_count != expected_raw_record_count: - print( - f"[WARN] CUPTI raw-record mismatch for {tag!r}: expected " - f"{records_per_replay} total records/replay × {graph_replays} replays " - f"= {expected_raw_record_count}, got {raw_record_count}. SKIPPING cell.", - file=sys.stderr, - ) - return None - - host_timing.start() - stats = _stats_from_cupti_records( - records, - warmup, - iters, - tag, - expected_K, - zero_ts_count=zero_ts_count, - zero_ts_names=zero_ts_names, - include_details=bool(getattr(args, "json_detailed", False)), - ) - host_timing.stop("stats_ms") - host_timing.stop_total() - host_timing.add("graph_group_iters", group_iters) - host_timing.add("graph_replays", graph_replays) - host_timing.add("cupti_records_per_replay", records_per_replay) - host_timing.add("cupti_target_records_per_replay", target_count) - host_timing.add("cupti_raw_records", raw_record_count) - host_timing.add("cupti_raw_records_expected", expected_raw_record_count) - host_timing.add("cupti_flush_period_ms", cupti_flush_period_ms) - host_timing.attach(stats) - return stats - - -def _time_kernel_eager( - args, - run_fn, - reset_fn, - tag: str, - *, - expected_K: int, - pre_iter_fn=None, - iters_override: int | None = None, - cupti_plan_key: tuple | None = None, -) -> dict: - """Non-graph CUPTI timer (for ncu wrapping, debugging, etc.). - - Each iter runs serially with sync between, but kernel start/end still - come from CUPTI — same accuracy as the graph path, just slower per-iter - (extra Python + sync overhead). - """ - host_timing = _HostTiming(bool(getattr(args, "host_timing", False))) - timer = CuptiKernelTimer.get() - warmup = args.warmup - iters = iters_override if iters_override is not None else args.iters - - del cupti_plan_key - - def _run_eager_loop(): - torch.cuda.nvtx.range_push(tag) - # Unified warmup+iters loop; CUPTI filters by warmup count internally. - for i in range(warmup + iters): - reset_fn() - if args.l2_flush: - _flush_l2() # includes synchronize - if pre_iter_fn is not None: - pre_iter_fn(i) - run_fn() - torch.cuda.synchronize() - torch.cuda.nvtx.range_pop() - - host_timing.start() - records, zero_ts_count, zero_ts_names = timer.capture_names(_run_eager_loop) - host_timing.stop("timed_loop_and_cupti_parse_ms") - - host_timing.start() - stats = _stats_from_cupti_records( - records, - warmup, - iters, - tag, - expected_K, - zero_ts_count=zero_ts_count, - zero_ts_names=zero_ts_names, - include_details=bool(getattr(args, "json_detailed", False)), - ) - host_timing.stop("stats_ms") - host_timing.stop_total() - host_timing.attach(stats) - return stats - - -def _run_kernel_untimed(args, run_fn, reset_fn, tag: str) -> dict: - """No in-bench timing: just run the kernels for an external profiler - (nsys / ncu) to time externally. Returns a stats dict full of zeros so - downstream code (table, JSON) doesn't break. - - Note: pre_iter_fn / iters_override aren't plumbed here yet — mix-mode - benchmarking relies on CUPTI. Add when a use-case lands. - """ - warmup = args.warmup - iters = args.iters - - if args.cuda_graph: - # Eager warmup before capture (Triton autotune) - reset_fn() - run_fn() - torch.cuda.synchronize() - reset_fn() - torch.cuda.synchronize() - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - for _ in range(warmup + iters): - reset_fn() - if args.l2_flush: - _l2_flush.fill_(0.0) - run_fn() - torch.cuda.synchronize() - torch.cuda.nvtx.range_push(tag) - g.replay() - torch.cuda.synchronize() - torch.cuda.nvtx.range_pop() - else: - torch.cuda.nvtx.range_push(tag) - for _ in range(warmup + iters): - reset_fn() - if args.l2_flush: - _flush_l2() - run_fn() - torch.cuda.synchronize() - torch.cuda.nvtx.range_pop() - - spans_us = [0.0] * iters - out = _stats_from_spans(spans_us) - out["iters_us"] = spans_us - out["per_kernel"] = {} - return out - - -def _time_kernel( - args, - run_fn, - reset_fn, - tag: str, - *, - expected_K: int, - pre_iter_fn=None, - pre_iter_group_factory=None, - iters_override: int | None = None, - cupti_plan_key: tuple | None = None, -) -> dict: - """Dispatch to graph-CUPTI / eager-CUPTI / no-timer path. - - --cupti: in-process CUPTI Activity API timing (default). Use --no-cupti - when running under nsys (in-process CUPTI conflicts with nsys's own - subscriber); the bench then runs the kernels for nsys to time externally. - - `expected_K` is the kernels-per-iter count the caller declares - (computed via _kernels_per_iter_*). CUPTI paths validate against it - explicitly; the no-timer fallback ignores it (no records to validate). - """ - if not getattr(args, "cupti", True): - if pre_iter_fn is not None: - raise RuntimeError( - "_time_kernel: pre_iter_fn requires CUPTI (mix-mode); " - "got --no-cupti. Re-run with CUPTI on or plumb pre_iter_fn " - "through _run_kernel_untimed." - ) - return _run_kernel_untimed(args, run_fn, reset_fn, tag) - if args.cuda_graph: - return _time_kernel_cuda_graph( - args, - run_fn, - reset_fn, - tag, - expected_K=expected_K, - pre_iter_fn=pre_iter_fn, - pre_iter_group_factory=pre_iter_group_factory, - iters_override=iters_override, - cupti_plan_key=cupti_plan_key, - ) - return _time_kernel_eager( - args, - run_fn, - reset_fn, - tag, - expected_K=expected_K, - pre_iter_fn=pre_iter_fn, - iters_override=iters_override, - cupti_plan_key=cupti_plan_key, - ) - - -# Per-config benchmark (consolidated baseline + replay) - - -def _warm_one_config(args, cfg, baseline_fn) -> None: - """Module-level worker for the compile-warmup process pool. - - Module-level so ProcessPoolExecutor can pickle it (nested functions - aren't picklable). Each worker process holds its own GIL → no - serialization between concurrent compiles. - - ``cfg`` is a tuple of (outer_cfg, inner_overrides_or_list): - * outer_cfg = (batch, mtp_len, prev_ks, state_dtype, act_dtype, - sr_mode, rect, mode, hardcode_sort) - * inner_overrides_or_list = dict of args attribute name -> value-string, - OR a list of such dicts. In the list form (CPS-grouped task) the - worker compiles each entry sequentially within the same process so - Triton's in-process kernel cache catches value-spec hits across - related entries (e.g. CPS={1,2} and {4,8} each form a `div_by_16` - spec bucket; the second compile in a bucket short-circuits). - - ``baseline_fn`` is optional — when ``None``, only the replay kernel is - warmed (the baseline-selection kernel can be warmed once in - the parent if needed). This lets us avoid pickling C-extension - function references across processes. - """ - outer_cfg, inner_overrides_or_list = cfg - overrides_list = ( - inner_overrides_or_list - if isinstance(inner_overrides_or_list, list) - else [inner_overrides_or_list] - ) - (batch, mtp_len, prev_ks, state_dtype, act_dtype, sr_mode, rect, mode, hardcode_sort) = ( - outer_cfg - ) - import argparse as _ap - - for inner_overrides in overrides_list: - # Fresh clone per entry: prevents knob-value leakage between - # consecutive cells in a CPS-grouped task (entries may set - # different non-CPS knobs in degenerate edge cases). - args_copy = _ap.Namespace(**vars(args)) - for k, v in inner_overrides.items(): - setattr(args_copy, k, v) - _bench_config( - args_copy, - batch, - mtp_len, - prev_ks, - state_dtype, - act_dtype, - baseline_fn, - sr_mode=sr_mode, - rectangle_for_nowrite=rect, - mode=mode, - hardcode_sort=hardcode_sort, - warmup_only=True, - ) - - -def _compile_warmup_phase( - args, batch_sizes, mtp_lengths, state_dtypes, act_dtypes, baseline_fn, max_workers: int -) -> None: - _cw_t0 = time.perf_counter() - - def _cw(label: str) -> None: - dt = time.perf_counter() - _cw_t0 - print(f"[compile-warmup] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) - - _cw("entered _compile_warmup_phase") - """Parallel compile-warmup using a ProcessPoolExecutor with `spawn` - start method. - - Each worker process holds its own GIL and its own CUDA context, so - Triton compiles (Python AST/codegen + LLVM/ptxas) run truly in - parallel. Previous ThreadPoolExecutor design hit GIL contention - in the Python codegen phase, capping throughput at ~1-2 cores even - with 28 threads (observed: 4 R threads vs 28 in pool). - - Compiled binaries land in Triton's on-disk cache (TRITON_CACHE_DIR - or default ~/.triton/cache). Workers share the cache via filesystem - — first to write any given (kernel_source × constexpr_set) hash - wins; concurrent writes to the SAME hash are wasteful but not - corrupting. - - spawn start method avoids inheriting parent CUDA state (which is - unsafe after fork on Linux with active CUDA contexts). Per-worker - import + CUDA init costs ~10s, amortized over each worker's many - compiles. baseline_fn is intentionally NOT passed to workers to - avoid pickling complications; the parent compiles the baseline - kernel itself before launching the pool when applicable. - """ - import multiprocessing - from concurrent.futures import ProcessPoolExecutor - - sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) - - rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) - modes_list = getattr(args, "modes_list", ["persistent_dynamic"]) - hsort_list = getattr(args, "hardcode_sort_list", [False]) - - # Compile-warmup task enumeration: outer × inner cartesian. - # CRITICAL: only enumerate axes that change the kernel's COMPILE signature. - # Drop runtime axes (batch, prev_k) that produce identical kernel hashes — - # otherwise we'd pay ~50-100ms of bench setup per redundant cache-hit task. - # - # Batches collapse to first only when the caller has forced an explicit - # mode/knob set. With mode=None, _DEFAULT_TUNING depends on effective - # batch and can choose different constexpr knobs or even a different mode. - # prev_k is already a list passed into _bench_config (not enumerated here). - configs = [] - _default_tuning_requested = any(m is None for m in modes_list) - _compile_batches = batch_sizes if _default_tuning_requested else batch_sizes[:1] - for batch in _compile_batches: - for mtp_len in mtp_lengths: - prev_ks = _resolve_prev_ks(args, mtp_len) - for state_dtype in state_dtypes: - for act_dtype in act_dtypes: - for sr_mode in _sr_modes_for_dtype(state_dtype, sr_modes_list): - for mode in modes_list: - for rect in rect_list: - can_sort = args.mix_csv is not None or getattr(args, "pmix", False) - effective_hsort_list = hsort_list if can_sort else [False] - for hardcode_sort in effective_hsort_list: - configs.append( - ( - batch, - mtp_len, - prev_ks, - state_dtype, - act_dtype, - sr_mode, - rect, - mode, - hardcode_sort, - ) - ) - - # Enumerate inner-knob signatures. Two paths: - # (1) --cell-list mode (preferred when set): pull exactly the cells - # that will be timed from args._cell_list_set. No synthetic - # cartesian — we only pre-compile what will run. - # (2) Sweep-args mode: cartesian over split-aware axes that read - # BOTH unsplit (args.X) and per-half (args.X_write/_nowrite) - # knob settings. Older code read only args.X and silently - # enumerated 1 inner combo when callers set only the per-half - # versions (all cell-list usage, plus any --block-size-m-write/ - # _nowrite CLI invocation), causing massive in-process JIT - # compile tax for persistent_main especially. - # - # In BOTH paths we GROUP tasks by non-CPS signature so each worker - # process compiles all CPS values for its group sequentially. - # NUM_PERSISTENT = CPS * num_sms is a runtime int but Triton auto- - # specializes on `div_by_16`, partitioning {CPS=1,2} (132,264) from - # {CPS=4,8} (528,1056) into two distinct compiled variants. By - # keeping all CPS variants for one (M,W,S,LS,TMA,...) signature in - # the same worker, the second compile in each spec bucket hits the - # in-process Triton cache (no disk-cache round trip). - - def _ps(val): - if val is None or (isinstance(val, str) and not val): - return [None] - if isinstance(val, str): - return [v.strip() for v in val.split(",") if v.strip()] - return [val] - - def _split_or_pair(shared_attr, w_attr, nw_attr): - """Return list of (write_val, nowrite_val) strings. - - Reads shared (args.X), write-side (args.X_write), and nowrite-side - (args.X_nowrite) values. If both per-half attrs are None, emits - tied pairs (v,v) over the shared values. If either per-half is - set, cartesian-iterates per-half values, falling back to shared - for whichever side is None. - """ - w = _ps(getattr(args, w_attr, None)) - nw = _ps(getattr(args, nw_attr, None)) - s = _ps(getattr(args, shared_attr, None)) - if w == [None] and nw == [None]: - return [(v, v) for v in s] - if w == [None]: - w = s - if nw == [None]: - nw = s - return [(a, b) for a in w for b in nw] - - _cw(f"built {len(configs)} outer configs") - cell_set = getattr(args, "_cell_list_set", set()) - cell_keys = getattr(args, "_cell_list_keys", ()) - # CPS keys are runtime ints (kernel value-specializes on `div_by_16`); - # cells differing only on CPS values can SHARE a worker so the second - # CPS value in a div_by_16 bucket hits the in-process Triton cache. - _cps_keys = ("cta_per_sm_write", "cta_per_sm_nowrite", "cta_per_sm") - - if cell_set: - # ============== CELL-LIST PATH ============== - # Build tasks DIRECTLY from cells. Each cell carries its OWN - # outer-axis values (RECT, MODE, SR, HSORT) so we - # pair each cell with its specific outer config — NOT the union- - # cartesian of all cells' outer values. Previously the OUTER × - # CELL cartesian doubled task count when a cell-list spanned both - # RECT=0 and RECT=1 (or any other outer-axis split); half the - # tasks then failed the cell-list filter inside the worker and - # wasted dispatch overhead. This path is O(|unique cell groups|). - from collections import defaultdict as _dd - - cell_groups: dict = _dd(list) - for tup in cell_set: - d = dict(zip(cell_keys, tup)) - cell_outer = ( - "SR" if d.get("SR", 0) else "RN", # sr_mode - bool(d.get("RECT", 0)), # rect - d.get("MODE", "persistent_dynamic"), # mode - bool(d.get("HSORT", 0)), # hardcode_sort - ) - inner = {} - for k, v in d.items(): - if k in _CELL_LIST_KEY_TO_ARG: - inner[_CELL_LIST_KEY_TO_ARG[k]] = str(v) - non_cps_sig = tuple(sorted((k, v) for k, v in inner.items() if k not in _cps_keys)) - cell_groups[(cell_outer, non_cps_sig)].append(inner) - - # CLI-runtime axes (batch/mtp/dtype) are NOT in cell-list — they - # come from CLI args and cartesian here (typically just 1 combo). - cli_outers = [] - for _b in _compile_batches: - for _m in mtp_lengths: - _pk = _resolve_prev_ks(args, _m) - for _sd in state_dtypes: - for _ad in act_dtypes: - cli_outers.append((_b, _m, _pk, _sd, _ad)) - - tasks = [] - for cli_outer in cli_outers: - for (cell_outer, _sig), inner_list in cell_groups.items(): - outer_cfg = (*cli_outer, *cell_outer) - tasks.append((outer_cfg, inner_list)) - n_groups = len(cell_groups) - n_total_cells = sum(len(g) for g in cell_groups.values()) - n_outer_used = len(cli_outers) - else: - # ============== SWEEP-ARGS PATH ============== - # Build inner_dicts via cartesian over knob axes, then cross with - # the `configs` outer cartesian. Existing behavior. - m_pairs = _split_or_pair("block_size_m", "block_size_m_write", "block_size_m_nowrite") - w_pairs = _split_or_pair("num_warps", "num_warps_write", "num_warps_nowrite") - ns_pairs = _split_or_pair("num_stages", "num_stages_write", "num_stages_nowrite") - cps_pairs = _split_or_pair("cta_per_sm", "cta_per_sm_write", "cta_per_sm_nowrite") - ls_pairs = _split_or_pair( - "num_loop_stages", "num_loop_stages_write", "num_loop_stages_nowrite" - ) - pw_vals = _ps(args.precompute_num_warps) - h_vals = _ps(args.heads_per_block) - fl_vals = _ps(args.flatten) - wsp_vals = _ps(args.warp_specialize) - trl_vals = _ps(args.use_tma_rect_load) - twl_vals = _ps(args.use_tma_replay_write_load) - tnl_vals = _ps(args.use_tma_replay_nowrite_load) - tws_vals = _ps(args.use_tma_replay_write_store) - import itertools as _it - - inner_dicts = [] - for (mw, mnw), (ww, wnw), (sw, snw), (cw, cnw), ( - lw, - lnw, - ), pw, h, fl, wsp, trl, twl, tnl, tws in _it.product( - m_pairs, - w_pairs, - ns_pairs, - cps_pairs, - ls_pairs, - pw_vals, - h_vals, - fl_vals, - wsp_vals, - trl_vals, - twl_vals, - tnl_vals, - tws_vals, - ): - d = {} - for k, v in ( - ("block_size_m_write", mw), - ("block_size_m_nowrite", mnw), - ("num_warps_write", ww), - ("num_warps_nowrite", wnw), - ("num_stages_write", sw), - ("num_stages_nowrite", snw), - ("cta_per_sm_write", cw), - ("cta_per_sm_nowrite", cnw), - ("num_loop_stages_write", lw), - ("num_loop_stages_nowrite", lnw), - ("precompute_num_warps", pw), - ("heads_per_block", h), - ("flatten", fl), - ("warp_specialize", wsp), - ("use_tma_rect_load", trl), - ("use_tma_replay_write_load", twl), - ("use_tma_replay_nowrite_load", tnl), - ("use_tma_replay_write_store", tws), - ): - if v is not None: - d[k] = str(v) - inner_dicts.append(d) - - groups: dict = {} - for d in inner_dicts: - sig = tuple(sorted((k, v) for k, v in d.items() if k not in _cps_keys)) - groups.setdefault(sig, []).append(d) - tasks = [] - for outer in configs: - for sig, group in groups.items(): - tasks.append((outer, group)) - n_groups = len(groups) - n_total_cells = sum(len(g) for g in groups.values()) - n_outer_used = len(configs) - - # Shuffle ACROSS tasks (preserve within-group CPS sequence for in-process - # cache adjacency — within-group order is intentional, not shuffled). - import random as _r - - _r.shuffle(tasks) - - _cw(f"built {len(tasks)} tasks covering {n_total_cells} cells in {n_groups} groups") - print( - f"[compile-warmup] {len(tasks)} compile tasks " - f"({n_outer_used} outer × {n_groups} cell-groups " - f"covering {n_total_cells} cells, CPS-grouped" - + (", per-cell outer" if cell_set else "") - + f") across {max_workers} processes (ProcessPoolExecutor, {_MP_START_METHOD} start)" - ) - t0 = time.perf_counter() - - ctx = multiprocessing.get_context(_MP_START_METHOD) - errors = [] - _cw("about to create ProcessPoolExecutor") - with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as ex: - _cw("ProcessPoolExecutor created, about to submit tasks") - # baseline_fn=None: workers compile only the replay kernel. - # Baseline kernels (if any) get compiled lazily in the parent during - # the timing phase — usually just one extra compile, negligible. - futures = {ex.submit(_warm_one_config, args, task, None): task for task in tasks} - _cw(f"submitted {len(futures)} tasks, waiting for results") - _n_done = 0 - for fut in futures: - try: - fut.result() - except Exception as e: - errors.append((futures[fut], e)) - _n_done += 1 - # Progress beacons at 10/25/50/75/100% to gauge effective parallelism. - if _n_done in ( - max(1, len(futures) // 10), - max(1, len(futures) // 4), - max(1, len(futures) // 2), - max(1, (3 * len(futures)) // 4), - len(futures), - ): - _cw(f"{_n_done}/{len(futures)} tasks complete") - - if errors: - for cfg, e in errors: - print(f"[compile-warmup] FAILED config {cfg}: {type(e).__name__}: {e}", file=sys.stderr) - raise errors[0][1] - - print(f"[compile-warmup] done in {time.perf_counter() - t0:.1f}s") - - -def _bench_config( - args, - batch: int, - mtp_len: int, - prev_ks: list[int], - state_dtype: torch.dtype, - act_dtype: torch.dtype, - baseline_fn, - sr_mode: str = "RN", - rectangle_for_nowrite: bool = False, - mode: str = "persistent_dynamic", - mix_samples_cpu=None, - mix_label: str = "", - hardcode_sort: bool = False, - nowrite_first: bool | None = None, - mix_samples_sorted_cpu=None, - mix_write_frac: float | None = None, - warmup_only: bool = False, -) -> None: - """ - Benchmark one (batch, mtp_len, dtype) configuration. - - Runs the baseline kernel (if baseline_fn is not None) followed by the - replay kernel for each prev_k value. Tensors are built once and - shared across all runs in this config. - - When ``warmup_only`` is True, calls each kernel exactly once instead of - timing it. Used by the parallel-warmup phase to populate Triton's - persistent compile cache across all configs concurrently. No timing - output is produced. - """ - state_dtype_name = str(state_dtype).split(".")[-1] - act_dtype_name = str(act_dtype).split(".")[-1] - - ( - state0, - state_scales0, - old_x0, - old_B0, - old_dt0, - old_dA_cumsum0, - cache_buf_idx0, - x, - dt, - B, - C, - A, - dt_bias, - D, - prev_tokens, - state_batch_indices, - replay_work_items_buf, - out_incr, - out_base, - xbc_input0, - conv_state0, - conv_weight, - conv_bias, - d_inner, - conv_dim, - ) = _build_tensors( - batch, - mtp_len, - state_dtype, - act_dtype, - args.tp_nheads, - args.head_dim, - args.d_state, - args.tp_ngroups, - max_window=getattr(args, "max_window", None) or None, - strided_state_cache=getattr(args, "strided_state_cache", False), - ) - - nheads = args.tp_nheads - ngroups = args.tp_ngroups - head_dim = args.head_dim - d_state = args.d_state - with_conv1d = getattr(args, "with_conv1d", False) - use_philox = sr_mode == "SR" - - # SR rounding: allow fp16 and the quantized dtypes (int8/int16/fp8). - # bf16/fp32 SR is not supported (no PTX path for bf16; fp32 doesn't need - # rounding). _sr_modes_for_dtype maps fp32 to RN so production sweeps can - # request SR once while still getting fp32/RN. - rand_seed = None - if use_philox: - if state_dtype not in _SR_SUPPORTED_DTYPES: - return - rand_seed = torch.randint(0, 2**62, (1,), device="cuda", dtype=torch.int64) - mode = _resolve_effective_replay_mode(args, batch, state_dtype, use_philox, mode) - if mode != "persistent_main" and nowrite_first: - return - - state_work = _clone_state_preserving_layout(state0) - state_scales_work = state_scales0.clone() if state_scales0 is not None else None - old_x_work = old_x0.clone() - old_B_work = old_B0.clone() - old_dt_work = old_dt0.clone() - old_dA_cumsum_work = old_dA_cumsum0.clone() - cache_buf_idx_work = cache_buf_idx0.clone() - xbc_input_work = xbc_input0.clone() - conv_state_work = conv_state0.clone() - - def _reset(): - state_work.copy_(state0) - if state_scales_work is not None: - state_scales_work.copy_(state_scales0) - old_x_work.copy_(old_x0) - old_B_work.copy_(old_B0) - old_dt_work.copy_(old_dt0) - old_dA_cumsum_work.copy_(old_dA_cumsum0) - cache_buf_idx_work.copy_(cache_buf_idx0) - if with_conv1d: - conv_state_work.copy_(conv_state0) - - def _reset_conv1d_realistic(): - """Realistic reset: cold cache, L2 flush, then hot in_proj output.""" - # 1. Reset cold state (cache tensors, SSM state) - state_work.copy_(state0) - if state_scales_work is not None: - state_scales_work.copy_(state_scales0) - old_x_work.copy_(old_x0) - old_B_work.copy_(old_B0) - old_dt_work.copy_(old_dt0) - old_dA_cumsum_work.copy_(old_dA_cumsum0) - cache_buf_idx_work.copy_(cache_buf_idx0) - conv_state_work.copy_(conv_state0) - # 2. L2 flush (evicts cold state from cache) - if _l2_flush is not None: - _l2_flush.fill_(0.0) - # 3. Write hot tensors (simulates in_proj output landing in L2) - xbc_input_work.copy_(xbc_input0) - - is_pr3324_baseline = _is_flashinfer_pr3324_baseline(args) - - # Silently skip the baseline row for any (baseline, state_dtype, SR) - # combo it can't run. Better than erroring on a partial sweep — our - # kernel rows still print. Compatibility: - # * flashinfer_pr3324: enable the dtypes that match this benchmark's - # semantics. Skip int16 because PR3324 treats it as raw int16 state, - # without our per-channel state_scale path; skip bf16 because we do not - # use bf16 state. - def _baseline_supports() -> bool: - if baseline_fn is None: - return False - assert is_pr3324_baseline, args.baseline - if state_dtype not in ( - torch.float32, - torch.float16, - torch.int8, - torch.float8_e4m3fn, - ): - return False - return not (use_philox and state_dtype == torch.float32) - - if baseline_fn is not None and not _baseline_supports(): - if not warmup_only: - sr_tag = " + SR" if use_philox else "" - print( - f"# Skipping {args.baseline} baseline for " - f"state_dtype={state_dtype_name}{sr_tag} (unsupported)." - ) - baseline_fn = None - - show_kernel_col = args.baseline is not None - - def _conv1d_split(xbc_in, conv_st, launch_dependent_kernels=False): - """Run conv1d update and split output into (x, B, C) views. - - The input tensor's strides are preserved through conv1d and the - transpose+view chain. With the production-matching layout - (contiguous (batch*T, conv_dim) viewed as (batch, conv_dim, T)), - the output after transpose+view has stride(-1)==1 and - stride(1)==dim, satisfying both our kernel and flashinfer. - """ - xbc_result = causal_conv1d_update( - xbc_in, - conv_st, - conv_weight, - conv_bias, - activation="silu", - launch_dependent_kernels=launch_dependent_kernels, - ) - xbc_flat = xbc_result.transpose(1, 2).view(batch * mtp_len, conv_dim) - x_flat, B_flat, C_flat = torch.split( - xbc_flat, [d_inner, ngroups * d_state, ngroups * d_state], dim=-1 - ) - x_conv = x_flat.view(batch, mtp_len, nheads, head_dim) - B_conv = B_flat.view(batch, mtp_len, ngroups, d_state) - C_conv = C_flat.view(batch, mtp_len, ngroups, d_state) - return x_conv, B_conv, C_conv - - # --- Sweep parameter parsing (invariant across prev_k) --- - def _parse_sweep(val): - if val is None: - return [None] - return [int(v) for v in val.split(",")] - - block_size_m_values = _parse_sweep(args.block_size_m) - num_warps_values = _parse_sweep(args.num_warps) - num_stages_values = _parse_sweep(args.num_stages) - precompute_num_warps_values = _parse_sweep(args.precompute_num_warps) - heads_per_block_values = _parse_sweep(args.heads_per_block) - # Persistent-only sweep dims; ignored when the cell's mode != persistent_main. - cta_per_sm_values = _parse_sweep(args.cta_per_sm) - num_loop_stages_values = _parse_sweep(args.num_loop_stages) - flatten_values = _parse_sweep(args.flatten) - warp_specialize_values = _parse_sweep(args.warp_specialize) - - # Per-main split-knob sweeps. Default = same as the shared sweep (so each - # combo is tied). When set independently, the inner loop sweeps the - # cross-product (write × nowrite); --skip-diagonal drops the tied subset. - def _split_or_share(split_csv, shared_values): - return _parse_sweep(split_csv) if split_csv else shared_values - - block_size_m_write_values = _split_or_share(args.block_size_m_write, block_size_m_values) - block_size_m_nowrite_values = _split_or_share(args.block_size_m_nowrite, block_size_m_values) - num_warps_write_values = _split_or_share(args.num_warps_write, num_warps_values) - num_warps_nowrite_values = _split_or_share(args.num_warps_nowrite, num_warps_values) - num_stages_write_values = _split_or_share(args.num_stages_write, num_stages_values) - num_stages_nowrite_values = _split_or_share(args.num_stages_nowrite, num_stages_values) - cta_per_sm_write_values = _split_or_share(args.cta_per_sm_write, cta_per_sm_values) - cta_per_sm_nowrite_values = _split_or_share(args.cta_per_sm_nowrite, cta_per_sm_values) - num_loop_stages_write_values = _split_or_share( - args.num_loop_stages_write, num_loop_stages_values - ) - num_loop_stages_nowrite_values = _split_or_share( - args.num_loop_stages_nowrite, num_loop_stages_values - ) - # Whether any *_write / *_nowrite knob was independently set — used by - # --skip-diagonal to know if the cross-product is non-trivial. Without - # any split, the per-main values == shared values and skip-diagonal is - # a no-op (which is correct). - _any_split = any( - getattr(args, name) - for name in ( - "block_size_m_write", - "block_size_m_nowrite", - "num_warps_write", - "num_warps_nowrite", - "num_stages_write", - "num_stages_nowrite", - "cta_per_sm_write", - "cta_per_sm_nowrite", - "num_loop_stages_write", - "num_loop_stages_nowrite", - ) - ) - # TMA toggles — independent 0/1 sweep per path. The skip-dupe at the - # top of the inner loop body collapses cells where a flag's path is - # unreachable for the current rectangle_for_nowrite setting. - use_tma_rect_load_values = _parse_sweep(args.use_tma_rect_load) - use_tma_replay_write_load_values = _parse_sweep(args.use_tma_replay_write_load) - use_tma_replay_nowrite_load_values = _parse_sweep(args.use_tma_replay_nowrite_load) - use_tma_replay_write_store_values = _parse_sweep(args.use_tma_replay_write_store) - - # --- Replay kernel --- - # Cache T-axis capacity (for prev_k validity check on the nowrite path). - max_window = getattr(args, "max_window", 0) or mtp_len - - # Build the list of scenarios to time. A scenario is one cell in the - # output: pure-mode scenarios fill prev_tokens with one constant before - # the timing loop; mix-mode scenarios feed a pre-baked per-iter samples - # tensor, with the per-iter copy captured inside the CUDA graph. Pure - # and mix can coexist in one call so a single nsys trace covers both. - scenarios = [] - if not (getattr(args, "mix_only", False) and mix_samples_cpu is not None): - for prev_k in prev_ks: - # Persistent modes dispatch per-slot from PNAT, so any prev_k - # <= max_window is valid. - scenarios.append( - { - "label": f"k{prev_k}", - "print_label": prev_k, - "fill": prev_k, - "pre_iter": None, - "iters": None, # use args.iters - } - ) - # Mix scenario: bench pre-bakes per-iter PNAT samples, n_writes, and - # replay_work_items. Grouped graph capture copies window rows into the - # persistent kernel-input tensors before each in-graph L2 flush, so timed - # kernels read the metadata cold. - if mix_samples_cpu is not None: - device = state_work.device - # Hardcode-sort is a precluster diagnostic: sort PNAT samples before - # building replay_work_items. Production keeps decode rows unsorted - # and only sorts the secondary replay metadata. - src = ( - mix_samples_sorted_cpu - if (hardcode_sort and mix_samples_sorted_cpu is not None) - else mix_samples_cpu - ) - samples_gpu = torch.from_numpy(src).to(device=device, dtype=torch.int32) - - n_writes_per_iter_all, replay_work_items_samples_cpu = _build_replay_work_items_cpu( - src, mtp_len, max_window - ) - n_writes_samples_gpu = torch.from_numpy(n_writes_per_iter_all).to( - device=device, dtype=torch.int32 - ) - replay_work_items_samples_gpu = torch.from_numpy(replay_work_items_samples_cpu).to( - device=device, dtype=torch.int32 - ) - n_writes_mix = torch.zeros(1, dtype=torch.int32, device=device) - - def _mix_pre_iter( - i, - _s=samples_gpu, - _ns=n_writes_samples_gpu, - _wi=replay_work_items_samples_gpu, - _pt=prev_tokens, - _nw=n_writes_mix, - _rwi=replay_work_items_buf, - ): - _pt.copy_(_s[i]) - _nw.copy_(_ns[i : i + 1]) - _rwi.copy_(_wi[i]) - - def _mix_pre_iter_group_factory( - group_iters, - _s=samples_gpu, - _ns=n_writes_samples_gpu, - _wi=replay_work_items_samples_gpu, - _pt=prev_tokens, - _nw=n_writes_mix, - _rwi=replay_work_items_buf, - ): - sample_window = torch.empty( - (group_iters, _s.shape[1]), - device=_s.device, - dtype=_s.dtype, - ) - n_writes_window = torch.empty((group_iters,), device=_ns.device, dtype=_ns.dtype) - work_items_window = torch.empty( - (group_iters, _wi.shape[1], _wi.shape[2]), - device=_wi.device, - dtype=_wi.dtype, - ) - - def _pre_replay(replay_idx): - start = replay_idx * group_iters - end = start + group_iters - sample_window.copy_(_s[start:end]) - n_writes_window.copy_(_ns[start:end]) - work_items_window.copy_(_wi[start:end]) - - def _graph_pre_iter(j): - _pt.copy_(sample_window[j]) - _nw.copy_(n_writes_window[j : j + 1]) - _rwi.copy_(work_items_window[j]) - - return _pre_replay, _graph_pre_iter - - # Mix iters override: if --mix-iters set, use it; else use args.iters. - mix_iters = getattr(args, "mix_iters", None) - scenarios.append( - { - "label": f"mix{mix_label}", - "print_label": "mix", - "fill": None, - "pre_iter": _mix_pre_iter, - "pre_iter_group_factory": _mix_pre_iter_group_factory, - "iters": mix_iters, # None => use args.iters - "n_writes": n_writes_mix, - # Full per-iter n_writes array (size = warmup + iters). Used by - # the JSON-detailed output to pair each iter's span with its - # mix composition for post-hoc bucketing analysis. - "n_writes_per_iter": n_writes_per_iter_all, - } - ) - - # Pure scenarios use one constant n_writes/work-items row. Mix scenarios - # update both tensors per iter. persistent_main always launches both - # halves because n_writes is device-resident for graph replay. - for scn in scenarios: - scenario_n_writes = scn.get("n_writes") - if scn["fill"] is not None: - prev_tokens.fill_(scn["fill"]) - n_writes_cpu, work_items_cpu = _build_replay_work_items_cpu( - np.full((batch,), scn["fill"], dtype=np.int32), - mtp_len, - max_window, - ) - scenario_n_writes = torch.from_numpy(n_writes_cpu).to( - device=state_work.device, dtype=torch.int32 - ) - replay_work_items_buf.copy_( - torch.from_numpy(work_items_cpu).to(device=state_work.device, dtype=torch.int32) - ) - prev_k_for_print = scn["print_label"] - scenario_pre_iter = scn["pre_iter"] - scenario_pre_iter_group_factory = scn.get("pre_iter_group_factory") - scenario_iters = scn.get("iters") # None => use args.iters - tag = f"incr_b{batch}_mtp{mtp_len}_{scn['label']}_s{state_dtype_name}_a{act_dtype_name}" - - scenario_per_iter_nw = None - if getattr(args, "json_detailed", False) or scn["fill"] is None: - eff_iters = scenario_iters if scenario_iters is not None else args.iters - if scn["fill"] is not None: - if getattr(args, "json_detailed", False): - is_write = scn["fill"] + mtp_len > max_window - scenario_per_iter_nw = [batch if is_write else 0] * eff_iters - else: - nw_full = scn.get("n_writes_per_iter") - if nw_full is not None: - scenario_per_iter_nw = nw_full[args.warmup : args.warmup + eff_iters].tolist() - - if baseline_fn is not None and is_pr3324_baseline: - baseline_suffix_parts = [f"SR={int(use_philox)}"] - hsort_list_for_tags = getattr(args, "hardcode_sort_list", [False]) - hsort_in_cell_list = "HSORT" in getattr(args, "_cell_list_keys", ()) - emit_hsort_tag = hardcode_sort or len(hsort_list_for_tags) > 1 or hsort_in_cell_list - if scn["fill"] is None and emit_hsort_tag: - baseline_suffix_parts.append(f"HSORT={1 if hardcode_sort else 0}") - baseline_sweep_suffix = ",".join(baseline_suffix_parts) - baseline_key = _build_json_key( - args.baseline, - batch, - mtp_len, - prev_k_for_print, - state_dtype_name, - baseline_sweep_suffix, - args.tp_size, - ) - run_baseline = True - if not warmup_only: - baseline_seen_keys = getattr(args, "_baseline_seen_keys", None) - if baseline_seen_keys is None: - baseline_seen_keys = set() - args._baseline_seen_keys = baseline_seen_keys - if ( - baseline_key in getattr(args, "_done_keys", set()) - or baseline_key in baseline_seen_keys - ): - run_baseline = False - else: - baseline_seen_keys.add(baseline_key) - - base_tag = ( - f"base_pr3324_b{batch}_mtp{mtp_len}_{scn['label']}_" - f"s{state_dtype_name}_a{act_dtype_name}" - ) - - def _run_pr3324_baseline(): - if with_conv1d: - x_call, B_call, C_call = _conv1d_split( - xbc_input_work, - conv_state_work, - launch_dependent_kernels=args.external_pdl, - ) - else: - x_call, B_call, C_call = x, B, C - # PR3324 predates double-buffered old_x. The benchmark keeps - # cache_buf_idx at zero, so buffer 0 is the active baseline view. - baseline_fn( - state_work, - old_x_work[:, 0], - old_B_work, - old_dt_work, - old_dA_cumsum_work, - cache_buf_idx_work, - prev_tokens, - x=x_call, - dt=dt, - A=A, - B=B_call, - C=C_call, - out=out_base, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=state_batch_indices, - state_scale=state_scales_work, - rand_seed=rand_seed, - philox_rounds=args.philox_rounds, - enable_pdl=with_conv1d and args.external_pdl, - ) - - reset_fn = _reset_conv1d_realistic if with_conv1d else _reset - if warmup_only: - reset_fn() - if scenario_pre_iter is not None: - scenario_pre_iter(0) - _run_pr3324_baseline() - torch.cuda.synchronize() - elif run_baseline: - stats = _time_kernel( - args, - _run_pr3324_baseline, - reset_fn, - base_tag, - expected_K=_kernels_per_iter_baseline(with_conv1d), - pre_iter_fn=scenario_pre_iter, - pre_iter_group_factory=scenario_pre_iter_group_factory, - iters_override=scenario_iters, - cupti_plan_key=( - "baseline", - args.baseline, - batch, - mtp_len, - prev_k_for_print, - state_dtype_name, - act_dtype_name, - with_conv1d, - bool(args.l2_flush), - bool(args.external_pdl), - bool(use_philox), - _kernels_per_iter_baseline(with_conv1d), - ), - ) - _submit_result_job( - args, - stats, - show_kernel_col=show_kernel_col, - kernel_name=args.baseline, - batch=batch, - mtp_len=mtp_len, - prev_k=prev_k_for_print, - state_dtype_name=state_dtype_name, - act_dtype_name=act_dtype_name, - sweep_suffix=baseline_sweep_suffix, - per_iter_nw=scenario_per_iter_nw, - kmix_bucket_write_frac=mix_write_frac if scn["fill"] is None else None, - skipped_tag=base_tag, - ) - - # Iteration over per-cell knob combos. - # When NO per-main split is requested (_any_split=False), each row in - # the cross-product gives the same value to both write_main and - # nowrite_main (current behavior — backward-compat). When ANY split - # IS requested, we iterate the write and nowrite axes independently - # (cross-product blowup is the user's responsibility — they typically - # pair this with --skip-diagonal to drop the tied subset). - if _any_split: - _iter_axes = ( - block_size_m_write_values, - block_size_m_nowrite_values, - num_warps_write_values, - num_warps_nowrite_values, - num_stages_write_values, - num_stages_nowrite_values, - precompute_num_warps_values, - heads_per_block_values, - cta_per_sm_write_values, - cta_per_sm_nowrite_values, - num_loop_stages_write_values, - num_loop_stages_nowrite_values, - flatten_values, - warp_specialize_values, - use_tma_rect_load_values, - use_tma_replay_write_load_values, - use_tma_replay_nowrite_load_values, - use_tma_replay_write_store_values, - ) - else: - # Tied: one value per shared knob. Wrap in single-element list for - # uniform iteration; the body sets w/nw both to the shared value. - _iter_axes = ( - block_size_m_values, - [None], - num_warps_values, - [None], - num_stages_values, - [None], - precompute_num_warps_values, - heads_per_block_values, - cta_per_sm_values, - [None], - num_loop_stages_values, - [None], - flatten_values, - warp_specialize_values, - use_tma_rect_load_values, - use_tma_replay_write_load_values, - use_tma_replay_nowrite_load_values, - use_tma_replay_write_store_values, - ) - # Iteration source: when --cell-list is active AND this is the main - # timing path (not a compile-warmup worker), iterate the cell set - # DIRECTLY (one yield per cell). The earlier design iterated the - # full inner cartesian and filtered each iteration via membership in - # args._cell_list_set — that's O(cartesian) which blows up to - # billions of iterations when the cell-list spans wide split-knob - # values (CPS, LS, M, W, S each contributing a Wx*Wnw factor on top - # of TMA flags), producing 50+ min of CPU spin per bench call before - # any actual timing. Direct iteration is O(|cell_list|). - # - # IMPORTANT exception for workers (warmup_only=True): _warm_one_config - # clamps args.*_write/_nowrite via inner_overrides to single values, - # making the cartesian 1×1×...×1 = 1 iter, which is exactly the one - # cell that worker was given. If we used cell-list-direct iteration - # here, every worker would iterate ALL 2884 cells instead of just - # its assigned one — turning compile-warmup into 28-way duplication. - # (Observed: 256 tasks in 233s under that bug vs ~18s correct.) - if getattr(args, "_cell_list_set", None) and not warmup_only: - - def _gen_from_cell_list(): - keys = args._cell_list_keys - for tup in args._cell_list_set: - d = dict(zip(keys, tup)) - yield ( - d.get("Mw"), - d.get("Mnw"), - d.get("Ww"), - d.get("Wnw"), - d.get("Sw"), - d.get("Snw"), - d.get("pW"), - d.get("H"), - d.get("CPSw"), - d.get("CPSnw"), - d.get("LSw"), - d.get("LSnw"), - d.get("FL"), - d.get("WS"), - d.get("TMARL"), - d.get("TMAWL"), - d.get("TMANL"), - d.get("TMAWS"), - ) - - _iter_source = _gen_from_cell_list() - else: - _iter_source = itertools.product(*_iter_axes) - - for ( - block_size_m_w, - block_size_m_nw, - num_warps_w, - num_warps_nw, - num_stages_w, - num_stages_nw, - precompute_num_warps, - heads_per_block, - cta_per_sm_w, - cta_per_sm_nw, - num_loop_stages_w, - num_loop_stages_nw, - flatten, - warp_specialize, - use_tma_rect_load, - use_tma_replay_write_load, - use_tma_replay_nowrite_load, - use_tma_replay_write_store, - ) in _iter_source: - # When tied, _nw values were placeholder None; fill from _w (the - # shared value). When split, _w and _nw came from independent lists. - if not _any_split: - block_size_m_nw = block_size_m_w - num_warps_nw = num_warps_w - num_stages_nw = num_stages_w - cta_per_sm_nw = cta_per_sm_w - num_loop_stages_nw = num_loop_stages_w - # Skip-diagonal: when split is on, drop the tied subset (same as a - # prior shared-knob sweep would cover). - if ( - _any_split - and args.skip_diagonal - and ( - block_size_m_w == block_size_m_nw - and num_warps_w == num_warps_nw - and num_stages_w == num_stages_nw - and cta_per_sm_w == cta_per_sm_nw - and num_loop_stages_w == num_loop_stages_nw - ) - ): - continue - # Backward-compat aliases used by the existing body below. When - # tied, these are simply the shared value. When split, the - # _write copy is used for sweep_tag and grouping (a stable choice - # so the tag is unique per (write, nowrite) combo). - block_size_m = block_size_m_w - num_warps = num_warps_w - num_stages = num_stages_w - cta_per_sm = cta_per_sm_w - num_loop_stages = num_loop_stages_w - # Skip-dupe for TMA flag sweeps: a flag whose code path isn't - # reachable in this cell produces identical timing for value=0 - # and value=1. We canonicalize by skipping value=1 cells when - # the flag's path is unreachable. Path reachability rules: - # * write path (replay write-load + write-store): always true. - # * rect path (rect-load): rectangle_for_nowrite=True. - # * replay-nowrite path (nowrite-load): rect isn't taking it. - _write_path = True # both halves exist for persistent modes - _rect_path = rectangle_for_nowrite - _replay_nowrite_path = not rectangle_for_nowrite - - def _set(v): # flag set to a non-zero sweep value - return v is not None and v != 0 - - if ( - _set(use_tma_rect_load) - and not _rect_path - or _set(use_tma_replay_write_load) - and not _write_path - or _set(use_tma_replay_nowrite_load) - and not _replay_nowrite_path - or _set(use_tma_replay_write_store) - and not _write_path - ): - continue - - assert scenario_n_writes is not None - - def _run_incr( - block_size_m=block_size_m, - num_warps=num_warps, - num_stages=num_stages, - precompute_num_warps=precompute_num_warps, - heads_per_block=heads_per_block, - cta_per_sm=cta_per_sm, - num_loop_stages=num_loop_stages, - flatten=flatten, - warp_specialize=warp_specialize, - use_tma_rect_load=use_tma_rect_load, - use_tma_replay_write_load=use_tma_replay_write_load, - use_tma_replay_nowrite_load=use_tma_replay_nowrite_load, - use_tma_replay_write_store=use_tma_replay_write_store, - ): - if with_conv1d: - x_call, B_call, C_call = _conv1d_split( - xbc_input_work, conv_state_work, launch_dependent_kernels=args.external_pdl - ) - extra_kwargs = {"launch_with_pdl": args.external_pdl} - else: - x_call, B_call, C_call = x, B, C - extra_kwargs = {} - extra_kwargs["rectangle_for_nowrite"] = rectangle_for_nowrite - extra_kwargs["nowrite_first"] = nowrite_first - extra_kwargs["mode"] = mode - extra_kwargs["n_writes"] = scenario_n_writes - extra_kwargs["replay_work_items"] = replay_work_items_buf - if state_scales_work is not None: - extra_kwargs["state_scales"] = state_scales_work - if use_tma_rect_load is not None: - extra_kwargs["_use_tma_rect_load"] = bool(use_tma_rect_load) - if use_tma_replay_write_load is not None: - extra_kwargs["_use_tma_replay_write_load"] = bool(use_tma_replay_write_load) - if use_tma_replay_nowrite_load is not None: - extra_kwargs["_use_tma_replay_nowrite_load"] = bool(use_tma_replay_nowrite_load) - if use_tma_replay_write_store is not None: - extra_kwargs["_use_tma_replay_write_store"] = bool(use_tma_replay_write_store) - if getattr(args, "require_tma_state_layout", False): - extra_kwargs["_require_tma_state_layout"] = True - if cta_per_sm is not None: - extra_kwargs["_cta_per_sm"] = cta_per_sm - if num_loop_stages is not None: - extra_kwargs["_num_loop_stages"] = num_loop_stages - if flatten is not None: - extra_kwargs["_flatten"] = bool(flatten) - if warp_specialize is not None: - extra_kwargs["_warp_specialize"] = bool(warp_specialize) - - replay_selective_state_update( - state_work, - old_x_work, - old_B_work, - old_dt_work, - old_dA_cumsum_work, - cache_buf_idx_work, - prev_tokens, - x=x_call, - dt=dt, - A=A, - B=B_call, - C=C_call, - out=out_incr, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=state_batch_indices, - rand_seed=rand_seed, - philox_rounds=args.philox_rounds, - use_internal_pdl=args.internal_pdl, - _block_size_m=block_size_m, - _num_warps=num_warps, - _num_stages=num_stages, - _precompute_num_warps=precompute_num_warps, - _heads_per_block=heads_per_block, - # Per-main overrides (None = tied to shared above; explicit - # only when the inner loop is iterating split axes). - _block_size_m_write=block_size_m_w if _any_split else None, - _block_size_m_nowrite=block_size_m_nw if _any_split else None, - _num_warps_write=num_warps_w if _any_split else None, - _num_warps_nowrite=num_warps_nw if _any_split else None, - _num_stages_write=num_stages_w if _any_split else None, - _num_stages_nowrite=num_stages_nw if _any_split else None, - _cta_per_sm_write=cta_per_sm_w if _any_split else None, - _cta_per_sm_nowrite=cta_per_sm_nw if _any_split else None, - _num_loop_stages_write=num_loop_stages_w if _any_split else None, - _num_loop_stages_nowrite=num_loop_stages_nw if _any_split else None, - **extra_kwargs, - ) - - parts = [] - - # Tuned wrapper knobs emit "auto" when unset, meaning the wrapper - # resolves them from _DEFAULT_TUNING per cell. pS/R/CT are not - # tuning-table knobs today, so leave them out unless explicitly - # swept. - def _val(v): - return "auto" if v is None else v - - # When tied (not _any_split), emit the shared single-value tag - # (M=8 etc). When split, emit explicit Mw / Mnw tags so cells - # with the same shared value but different per-main values get - # unique JSON keys. - def _emit_split(name_w, name_nw, val_w, val_nw): - if val_w is None and val_nw is None: - parts.append(f"{name_w[:-1]}=auto") # tied form, both auto - return - if not _any_split or val_w == val_nw: - parts.append(f"{name_w[:-1]}={_val(val_w)}") - else: - parts.append(f"{name_w}={_val(val_w)}") - parts.append(f"{name_nw}={_val(val_nw)}") - - _emit_split("Mw", "Mnw", block_size_m_w, block_size_m_nw) - _emit_split("Ww", "Wnw", num_warps_w, num_warps_nw) - _emit_split("Sw", "Snw", num_stages_w, num_stages_nw) - parts.append(f"pW={_val(precompute_num_warps)}") - parts.append(f"H={_val(heads_per_block)}") - # Persistent-only knobs (only meaningful when MODE=persistent_main; - # printed unconditionally so output rows are uniformly comparable - # across modes when the user passed these sweeps). - _emit_split("CPSw", "CPSnw", cta_per_sm_w, cta_per_sm_nw) - _emit_split("LSw", "LSnw", num_loop_stages_w, num_loop_stages_nw) - parts.append(f"FL={_val(flatten)}") - parts.append(f"WS={_val(warp_specialize)}") - # TMA sweep tags. Four wrapper-level flags map to three - # kernel-level constexprs (rect-load and replay-nowrite-load - # share `USE_TMA_LOAD_NOWRITE`, picked by the wrapper based on - # RECTANGLE). TMARL specifically gates the rectangle path's - # state load; TMANL specifically gates the replay-style - # nowrite path's state load. - parts.append(f"TMARL={_val(use_tma_rect_load)}") - parts.append(f"TMAWL={_val(use_tma_replay_write_load)}") - parts.append(f"TMANL={_val(use_tma_replay_nowrite_load)}") - parts.append(f"TMAWS={_val(use_tma_replay_write_store)}") - parts.append(f"SR={1 if use_philox else 0}") - parts.append( - f"RECT={'auto' if rectangle_for_nowrite is None else (1 if rectangle_for_nowrite else 0)}" - ) - nowrite_first_list_for_tags = getattr(args, "nowrite_first_list", [None]) - nowrite_first_in_cell_list = "NWF" in getattr(args, "_cell_list_keys", ()) - if nowrite_first or len(nowrite_first_list_for_tags) > 1 or nowrite_first_in_cell_list: - parts.append(f"NWF={1 if nowrite_first else 0}") - parts.append(f"MODE={_val(mode)}") - hsort_list_for_tags = getattr(args, "hardcode_sort_list", [False]) - hsort_in_cell_list = "HSORT" in getattr(args, "_cell_list_keys", ()) - if hardcode_sort or len(hsort_list_for_tags) > 1 or hsort_in_cell_list: - parts.append(f"HSORT={1 if hardcode_sort else 0}") - sweep_suffix = (" " + ",".join(parts)) if parts else "" - sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") - - reset_fn = _reset_conv1d_realistic if with_conv1d else _reset - # --cell-list filter: only time cells whose (canonical-knob-values) - # tuple is in the loaded set. Robust to bench gaining new knobs - # (old cell-list files keep working: any keys they don't list - # become wildcards that retain CLI defaults). - if args._cell_list_keys: - _tup = _current_cell_tuple(args, locals()) - if _tup is None or _tup not in args._cell_list_set: - continue - # Resume from JSONL: skip cells already recorded. Built the same - # way _print_row builds JSON keys; must stay in sync. - done_keys = getattr(args, "_done_keys", None) - if done_keys: - # One key per scenario (k=6, k=11, mix) — skip the whole cell - # only if ALL of its scenarios are already done. We don't - # know which scenarios will be emitted here without - # re-evaluating the inner scenario loop; conservatively skip - # only when the prev_k_for_print's specific key is done. - _resume_key = _build_json_key( - "replay", - batch, - mtp_len, - prev_k_for_print, - state_dtype_name, - sweep_suffix, - args.tp_size, - ) - if _resume_key in done_keys: - continue - if warmup_only: - reset_fn() - if scenario_pre_iter is not None: - scenario_pre_iter(0) - _run_incr() - torch.cuda.synchronize() - else: - # Inline retry: CUPTI sometimes loses records under PDL + - # high cell count; retrying the SAME cell often catches it - # because the failure is transient at the kernel-launch level. - # Per --cupti-retry budget. On final failure, append tag to - # the skipped list for an external rerun in a fresh process. - defer_results = ( - args.cuda_graph - and getattr(args, "cupti", True) - and int(getattr(args, "cupti_defer_depth", 1)) > 1 - ) - retry_budget = 0 if defer_results else max(0, getattr(args, "cupti_retry", 1)) - stats = None - expected_K = _kernels_per_iter_incremental( - mode, - with_conv1d=with_conv1d, - ) - plan_key = ( - "incremental", - "replay", - mode, - batch, - mtp_len, - state_dtype_name, - act_dtype_name, - with_conv1d, - bool(args.l2_flush), - bool(args.external_pdl), - bool(args.internal_pdl), - bool(use_philox), - bool(rectangle_for_nowrite), - bool(nowrite_first), - bool(hardcode_sort), - scenario_pre_iter is not None, - expected_K, - ) - for attempt in range(retry_budget + 1): - stats = _time_kernel( - args, - _run_incr, - reset_fn, - sweep_tag, - expected_K=expected_K, - pre_iter_fn=scenario_pre_iter, - pre_iter_group_factory=scenario_pre_iter_group_factory, - iters_override=scenario_iters, - cupti_plan_key=plan_key, - ) - if stats is not None: - break - if attempt < retry_budget: - print( - f"[retry] CUPTI mismatch on {sweep_tag!r}; " - f"retrying ({attempt + 1}/{retry_budget})", - file=sys.stderr, - flush=True, - ) - if stats is None: - args._skipped_cells.append(sweep_tag) - - per_iter_nw = scenario_per_iter_nw if stats is not None else None - - if stats is not None: - _submit_result_job( - args, - stats, - show_kernel_col=show_kernel_col, - kernel_name="replay", - batch=batch, - mtp_len=mtp_len, - prev_k=prev_k_for_print, - state_dtype_name=state_dtype_name, - act_dtype_name=act_dtype_name, - sweep_suffix=sweep_suffix, - per_iter_nw=per_iter_nw, - kmix_bucket_write_frac=mix_write_frac if scn["fill"] is None else None, - skipped_tag=sweep_tag, - ) - - -# Map full torch dtype name → short tag used in JSON keys (matches collect.py). -_DTYPE_SHORT = { - "float32": "fp32", - "bfloat16": "bf16", - "float16": "fp16", - "int8": "int8", - "int16": "int16", - "float8_e4m3fn": "fp8", -} - - -def _build_json_key(kernel_name, batch, mtp_len, prev_k, state_dtype_name, sweep_suffix, tp_size): - """Build a key matching collect.py's kernel_data.json convention: - - incremental/{batch}/{mtp}/{sd}/k{k}/{sweep_parts}/tp{tp} - flashinfer_pr3324/{batch}/{mtp}/{sd}/k{k}/{sweep_parts}/tp{tp} - - Replay rows collapse to "incremental"; baseline rows keep their baseline - name. - """ - if kernel_name == "replay": - kind = "incremental" - else: - kind = kernel_name - - sd = _DTYPE_SHORT.get(state_dtype_name, state_dtype_name) - parts = [kind, str(batch), str(mtp_len), sd] - if prev_k != "N/A": - parts.append(f"k{prev_k}") - if sweep_suffix: - # sweep_suffix format: " M=4,W=1,S=1,SR=0,RECT=0" - # collect.py format: "M4_W1_S1_SR0_RECT0" - # Strip leading/trailing whitespace, drop '=', commas → underscores. - parts.append(sweep_suffix.strip().replace("=", "").replace(",", "_")) - parts.append(f"tp{tp_size}") - return "/".join(parts) - - -def _print_row( - show_kernel_col, - kernel_name, - batch, - mtp_len, - prev_k, - state_dtype_name, - act_dtype_name, - stats, - sweep_suffix="", - tp_size=None, - json_detailed=False, - jsonl_path=None, - jsonl_host=None, - jsonl_gpu=None, -): - """Print one summary row and append the result to the JSONL sidecar. - - `stats` is a dict from _time_kernel: {median, p95, p99, n, iters_us, - [n_writes_per_iter], [kmix_bucket_score], [per_kernel]}. The summary - table shows kmix_bucket_score when present, otherwise median. JSONL - captures the compact per-iter spans + - n_writes_per_iter by default; with json_detailed=True it also captures - per-kernel data. - - When `jsonl_path` is provided, appends one JSON line per row to the - JSONL sidecar (crash-safe incremental persistence; lets a killed sweep - resume from the last completed cell on rerun, even across hosts). Open - per-write because `args` is pickled to ProcessPoolExecutor workers and - file handles aren't picklable. JSONL is the canonical artifact — the - bench no longer writes a final `.json` summary; use `jsonl_to_json.py` - if a one-shot `.json` snapshot is needed. - """ - kernel_col = f"{kernel_name:>11} | " if show_kernel_col else "" - headline_us = stats.get("kmix_bucket_score", stats["median"]) - print( - f"| {kernel_col}{batch:>5} | {mtp_len:>7} | {str(prev_k):>6} | " - f"{state_dtype_name:>11} | {act_dtype_name:>9} | " - f"{headline_us:>9.2f} | {stats['p95']:>7.2f} | {stats['p99']:>7.2f} |" - f"{sweep_suffix}" - ) - if jsonl_path is not None: - key = _build_json_key( - kernel_name, - batch, - mtp_len, - prev_k, - state_dtype_name, - sweep_suffix, - tp_size, - ) - if json_detailed: - row_stats = stats - else: - row_stats = { - k: stats[k] - for k in ( - "median", - "p95", - "p99", - "n", - "iters_us", - "n_writes_per_iter", - "kmix_bucket_score", - ) - if k in stats - } - if "host_timing" in stats: - row_stats["host_timing"] = stats["host_timing"] - # Append to JSONL sidecar if a path is set (incremental persistence). - # Open per-write because args is pickled to ProcessPoolExecutor - # workers, and file handles aren't picklable. A clean SIGTERM or - # Python exception will leave the file consistent up to the last - # newline; catastrophic kills can leave a partial last line, which - # the resume reader tolerates via json.JSONDecodeError pass. - if jsonl_path is not None: - # Wall-clock timestamp (float seconds since UNIX epoch) at write - # time. Lets post-hoc analysis diff consecutive rows to derive - # per-cell wall budget and identify startup-bound vs steady-state - # segments (cells/sec, downtime between bench invocations) without - # needing to instrument the bench's outer loops separately. - import time as _time - - rec = {"key": key, "stats": row_stats, "t": _time.time()} - if jsonl_host is not None: - rec["host"] = jsonl_host - if jsonl_gpu is not None: - rec["gpu"] = jsonl_gpu - with open(jsonl_path, "a") as f: - f.write(json.dumps(rec) + "\n") - - -def _finish_result_job(args, job: dict) -> None: - result = job["result"] - if isinstance(result, _PendingCuptiStats): - stats = result.resolve() - else: - stats = result - - if stats is None: - skipped_tag = job.get("skipped_tag") - if skipped_tag is not None: - args._skipped_cells.append(skipped_tag) - return - - per_iter_nw = job.get("per_iter_nw") - if per_iter_nw is not None: - stats["n_writes_per_iter"] = per_iter_nw - kmix_bucket_score = _kmix_bucket_score( - stats.get("iters_us"), - per_iter_nw, - job["batch"], - job.get("kmix_bucket_write_frac"), - ) - if kmix_bucket_score is not None: - stats["kmix_bucket_score"] = kmix_bucket_score - - _print_row( - job["show_kernel_col"], - job["kernel_name"], - job["batch"], - job["mtp_len"], - job["prev_k"], - job["state_dtype_name"], - job["act_dtype_name"], - stats, - job.get("sweep_suffix", ""), - tp_size=args.tp_size, - json_detailed=getattr(args, "json_detailed", False), - jsonl_path=getattr(args, "_jsonl_path", None), - jsonl_host=getattr(args, "_jsonl_host", None), - jsonl_gpu=getattr(args, "_jsonl_gpu", None), - ) - - -def _drain_pending_results(args, *, force: bool = False) -> None: - pending_results = getattr(args, "_pending_results", None) - if not pending_results: - return - - max_pending = max(1, int(getattr(args, "cupti_defer_depth", 1))) - while pending_results: - first_result = pending_results[0]["result"] - should_block = force or len(pending_results) >= max_pending - if ( - not should_block - and isinstance(first_result, _PendingCuptiStats) - and not first_result.is_ready() - ): - break - job = pending_results.pop(0) - _finish_result_job(args, job) - - -def _submit_result_job( - args, - result, - *, - show_kernel_col, - kernel_name, - batch, - mtp_len, - prev_k, - state_dtype_name, - act_dtype_name, - sweep_suffix="", - per_iter_nw=None, - kmix_bucket_write_frac=None, - skipped_tag=None, -) -> None: - job = { - "result": result, - "show_kernel_col": show_kernel_col, - "kernel_name": kernel_name, - "batch": batch, - "mtp_len": mtp_len, - "prev_k": prev_k, - "state_dtype_name": state_dtype_name, - "act_dtype_name": act_dtype_name, - "sweep_suffix": sweep_suffix, - "per_iter_nw": per_iter_nw, - "kmix_bucket_write_frac": kmix_bucket_write_frac, - "skipped_tag": skipped_tag, - } - if isinstance(result, _PendingCuptiStats): - args._pending_results.append(job) - _drain_pending_results(args) - else: - _finish_result_job(args, job) - - -# Cell-list mode — canonical knob-key mapping to argparse args + local -# loop variable. See _load_cell_list_into_args / inner-loop filter. -# -# Each entry: cell-key → (args attribute name, comma-separated string flag) -# For split (write/nowrite) knobs, we use Xw / Xnw keys. Tied forms (M, W, -# S, CPS, LS) accepted on load and expanded to their w/nw variants. -_CELL_LIST_KEY_TO_ARG = { - "Mw": "block_size_m_write", - "Mnw": "block_size_m_nowrite", - "Ww": "num_warps_write", - "Wnw": "num_warps_nowrite", - "Sw": "num_stages_write", - "Snw": "num_stages_nowrite", - "CPSw": "cta_per_sm_write", - "CPSnw": "cta_per_sm_nowrite", - "LSw": "num_loop_stages_write", - "LSnw": "num_loop_stages_nowrite", - "pW": "precompute_num_warps", - "H": "heads_per_block", - "FL": "flatten", - "WS": "warp_specialize", - "TMARL": "use_tma_rect_load", - "TMAWL": "use_tma_replay_write_load", - "TMANL": "use_tma_replay_nowrite_load", - "TMAWS": "use_tma_replay_write_store", - "RECT": "rectangle_for_nowrite", - "NWF": "nowrite_first", - "HSORT": "hardcode_sort", - # MODE and SR get special handling (string values): - # MODE → args.modes (single mode name) - # SR → args.sr_modes ("RN" if 0, "SR" if 1) -} - -# Split-knob tied form: "M" expands to both "Mw" and "Mnw". -_CELL_LIST_TIED_EXPANSIONS = { - "M": ("Mw", "Mnw"), - "W": ("Ww", "Wnw"), - "S": ("Sw", "Snw"), - "CPS": ("CPSw", "CPSnw"), - "LS": ("LSw", "LSnw"), -} - - -def _normalize_cell(cell: dict) -> dict: - """Expand tied-form keys (M, W, S, CPS, LS) to their w/nw variants. - Returns a new dict with only canonical split-or-plain keys. - """ - out = dict(cell) - for tied, (w_key, nw_key) in _CELL_LIST_TIED_EXPANSIONS.items(): - if tied in out: - v = out.pop(tied) - out.setdefault(w_key, v) - out.setdefault(nw_key, v) - return out - - -def _load_cell_list_into_args(args) -> None: - """Read --cell-list JSON, normalize, override args.* knob ranges, and - populate args._cell_list_keys + args._cell_list_set for the inner-loop - filter. Errors out if cells aren't uniform (different key sets). - """ - with open(args.cell_list) as f: - raw = json.load(f) - if not isinstance(raw, list): - sys.exit(f"--cell-list: expected JSON list, got {type(raw).__name__}") - cells = [_normalize_cell(c) for c in raw] - if not cells: - print("[cell-list] empty list — nothing to time", file=sys.stderr) - return - allowed_keys = set(_CELL_LIST_KEY_TO_ARG) | {"MODE", "SR"} - unknown_keys = sorted({key for cell in cells for key in cell if key not in allowed_keys}) - if unknown_keys: - sys.exit(f"--cell-list: unknown keys {unknown_keys}") - # All cells must share the same key set (uniform schema) - keys0 = frozenset(cells[0].keys()) - for i, c in enumerate(cells[1:], start=1): - if frozenset(c.keys()) != keys0: - sys.exit( - f"--cell-list: cells must have uniform key sets; cell[0] " - f"has {sorted(keys0)} but cell[{i}] has {sorted(c.keys())}" - ) - - # Auto-cover: collect per-knob value set across all cells - cover: dict = {} - for c in cells: - for k, v in c.items(): - cover.setdefault(k, set()).add(v) - # Apply overrides - for key, vals in cover.items(): - if key in _CELL_LIST_KEY_TO_ARG: - arg_name = _CELL_LIST_KEY_TO_ARG[key] - vals_str = ",".join(str(v) for v in sorted(vals)) - setattr(args, arg_name, vals_str) - elif key == "MODE": - args.modes = ",".join(sorted({str(v) for v in vals})) - elif key == "SR": - args.sr_modes = ",".join(sorted({"SR" if v else "RN" for v in vals})) - else: - print( - f"[cell-list] WARNING: unknown key {key!r} in cells; " - f"will not override any args.* attribute (the value will " - f"still be matched in the filter if a matching local var " - f"is in scope)", - file=sys.stderr, - ) - - # Canonical key order (sorted) for tuple matching in the inner loop - args._cell_list_keys = tuple(sorted(keys0)) - args._cell_list_set = {tuple(c[k] for k in args._cell_list_keys) for c in cells} - print( - f"[cell-list] loaded {len(cells)} cells with keys " - f"{list(args._cell_list_keys)}; overrode args.* to auto-cover", - file=sys.stderr, - ) - - -# Maps cell-list key → name of the local variable in _bench_config's inner -# loop. Used to extract the "current cell" tuple for the filter check. -# Keep in sync with the loop-variable names; the filter is lenient about -# missing names (it picks them up from the inner scope at runtime). -_CELL_LIST_KEY_TO_LOCAL = { - "Mw": "block_size_m_w", - "Mnw": "block_size_m_nw", - "Ww": "num_warps_w", - "Wnw": "num_warps_nw", - "Sw": "num_stages_w", - "Snw": "num_stages_nw", - "CPSw": "cta_per_sm_w", - "CPSnw": "cta_per_sm_nw", - "LSw": "num_loop_stages_w", - "LSnw": "num_loop_stages_nw", - "pW": "precompute_num_warps", - "H": "heads_per_block", - "FL": "flatten", - "WS": "warp_specialize", - "TMARL": "use_tma_rect_load", - "TMAWL": "use_tma_replay_write_load", - "TMANL": "use_tma_replay_nowrite_load", - "TMAWS": "use_tma_replay_write_store", - "RECT": "rectangle_for_nowrite", - "NWF": "nowrite_first", - "MODE": "mode", - "HSORT": "hardcode_sort", - "SR": "use_philox", -} - - -def _current_cell_tuple(args, locals_dict: dict) -> tuple | None: - """Build the (key1=val1, key2=val2, ...) tuple for the current inner-loop - iteration, matching args._cell_list_keys' order. Used by the inner-loop - filter to check membership in args._cell_list_set. Returns None if any - expected local is missing (the bench evolved a knob name — caller skips). - """ - if not args._cell_list_keys: - return None - vals = [] - for k in args._cell_list_keys: - local_name = _CELL_LIST_KEY_TO_LOCAL.get(k, k) - if local_name not in locals_dict: - return None - v = locals_dict[local_name] - # Coerce bools to ints to match cell-list JSON (1/0) - if isinstance(v, bool): - v = int(v) - vals.append(v) - return tuple(vals) - - -# Main benchmark loop - - -def _run_benchmark(args) -> None: - # Phase-timing markers — emit timestamped checkpoints so a captured-stdout - # run can later attribute wall time to setup vs compile-warmup vs prewarm - # vs timing. Single-line format makes log-grepping trivial. - _phase_t0 = time.perf_counter() - - def _phase(label: str) -> None: - dt = time.perf_counter() - _phase_t0 - print(f"[phase] t={dt:7.2f}s {label}", file=sys.stderr, flush=True) - - _phase("enter _run_benchmark") - - # Pending-results FIFO for srxl's deferred CUPTI parsing pipeline. Each - # entry holds a _PendingCuptiStats handle; _drain_pending_results pulls - # ready entries and routes them to _print_row (which appends to JSONL). - args._pending_results = [] - - # JSONL incremental sidecar. Path = `.jsonl`. Each completed - # cell appends one line `{"key": , "stats": {...}, "host": }` - # to this file as it finishes timing. On startup we read this sidecar (if - # present) and populate _done_keys so a killed bench can resume without - # redoing already-timed cells. Crash-safe by construction: append-only - # writes survive SIGTERM/SIGKILL/reboot mid-sweep. - # - # Resume is host-blind: _done_keys includes records from any host, so a - # bench restarted on a different node fills in the missing cells without - # redoing cells already covered elsewhere. Cross-host *timings* aren't - # directly comparable, but each JSONL record carries its `host` stamp so - # the analyzer can group/compare per host. This bench no longer writes a - # final `.json` summary — the JSONL is the canonical artifact; use the - # `jsonl_to_json.py` helper if a one-shot `.json` snapshot is needed. - # - # Note: we store only paths/strings on `args` because args is pickled to - # ProcessPoolExecutor workers during compile-warmup, and file handles - # (TextIOWrapper) aren't picklable. _print_row open-appends per cell. - args._jsonl_path = None - args._done_keys: set[str] = set() - args._baseline_seen_keys: set[str] = set() - args._jsonl_host = None # hostname stamp for the current run - args._jsonl_gpu = None # GPU device id stamp (current process visibility) - if getattr(args, "json_output", None): - import socket - - args._jsonl_host = socket.gethostname() - # Capture GPU id once at startup. Used by the oracle-cache layer in - # search_driver to attribute timings to a specific (host, gpu) pair - # for cross-process pruning. os.environ['CUDA_VISIBLE_DEVICES'] - # is the right source pre-torch-init (it's what the harness sets); - # post-init we could use torch.cuda.current_device() but we keep it - # to env to avoid forcing a CUDA init at this point in startup. - args._jsonl_gpu = os.environ.get("CUDA_VISIBLE_DEVICES", "") - args._jsonl_path = args.json_output + ".jsonl" - # Read existing JSONL if present: load every record's key into the - # skip set regardless of host (gap-fill on a new node). - if os.path.exists(args._jsonl_path): - n_loaded = 0 - host_counts: dict[str, int] = {} - with open(args._jsonl_path) as f: - for line in f: - line = line.strip() - if not line: - continue - try: - rec = json.loads(line) - except json.JSONDecodeError: - # Tolerate partial last line from a crash mid-write. - continue - k = rec.get("key") - if k is None: - continue - args._done_keys.add(k) - n_loaded += 1 - rec_host = rec.get("host") - if rec_host: - host_counts[rec_host] = host_counts.get(rec_host, 0) + 1 - if n_loaded: - host_summary = ( - ", ".join(f"{h}={n}" for h, n in sorted(host_counts.items())) - if host_counts - else "(no host stamps)" - ) - print( - f"[resume] {args._jsonl_path}: loaded {n_loaded} prior " - f"cell results across hosts [{host_summary}]; sweep will " - f"skip them. New cells stamp host={args._jsonl_host}.", - file=sys.stderr, - ) - - # Sidecar metadata: cmd, host, tp_size, cupti, etc. Written - # once at startup; helps later analysis identify how this JSONL was - # produced even though there's no top-level .json wrapper anymore. - meta_path = args.json_output + ".meta.json" - meta_payload = { - "timestamp": datetime.now().isoformat(), - "host": args._jsonl_host, - "cmd": " ".join(sys.argv), - "tp_size": getattr(args, "tp_size", None), - "warmup": getattr(args, "warmup", None), - "iters": getattr(args, "iters", None), - "cupti": getattr(args, "cupti", False), - } - # Append to a list so successive runs (gap-fill, retry) keep history. - existing_meta = [] - if os.path.exists(meta_path): - try: - with open(meta_path) as f: - existing_meta = json.load(f) - if not isinstance(existing_meta, list): - existing_meta = [existing_meta] - except (OSError, json.JSONDecodeError): - existing_meta = [] - existing_meta.append(meta_payload) - # Bench is sometimes invoked with --json-output pointing into a dir - # the caller hasn't created (subprocess driver, search loop, etc.). - # Ensure the dir exists before writing the meta sidecar OR the JSONL. - os.makedirs(os.path.dirname(os.path.abspath(meta_path)), exist_ok=True) - tmp = meta_path + ".tmp" - with open(tmp, "w") as f: - json.dump(existing_meta, f, indent=2) - os.replace(tmp, meta_path) - - # Skipped cells accumulator — populated by _bench_config when CUPTI capture - # mismatch causes a cell to be skipped. Written to args.skipped_output - # (or derived from json_output) at end of run. - args._skipped_cells = [] - - # Cell-list filter (replaces the old --retry-cells tag-string filter). - # When set, the sweep iterates ONLY the cells described in the list. - # - # Each entry in the JSON file is a dict of canonical knob keys → values, - # using the same names that appear in the sweep_tag (Mw/Mnw, Ww/Wnw, - # Sw/Snw, pW, pS, H, R, CT, CPSw/CPSnw, LSw/LSnw, FL, WS, TMARL, - # TMAWL, TMANL, TMAWS, SR, RECT, NWF, MODE). Optional diagnostic HSORT cells - # are accepted for compatibility. Each cell may also use the tied forms - # M / W / S / CPS / LS (single value applied to both write and nowrite - # halves). - # - # On load we: - # - Override the bench's CLI knob args (`args.block_size_m_write`, - # etc.) with the union of values present across all cells per knob, - # so the cartesian iteration auto-covers the list. - # - Build `args._cell_list_keys` (the canonical key order used by - # every cell — must be uniform across the list) and - # `args._cell_list_set` (frozen tuples for O(1) membership check - # inside the inner loop). - # - # In the inner loop, we build the current iteration's tuple and skip - # cells not in the set. Dict-matching is robust to bench gaining new - # knobs (old cell-list files keep working — newly-added knobs simply - # aren't matched on, so they retain CLI defaults). - # Cell-list state may already have been populated by main() (so that - # the args.*_list derivations downstream see the override). Default to - # empty if not. - _phase(f"done loading _done_keys ({len(args._done_keys)} entries)") - - if not hasattr(args, "_cell_list_keys"): - args._cell_list_keys: tuple = () - args._cell_list_set: set = set() - if getattr(args, "cell_list", None): - _load_cell_list_into_args(args) - _phase(f"done loading cell-list ({len(args._cell_list_set)} cells)") - - assert args.nheads % args.tp_size == 0, ( - f"nheads ({args.nheads}) must be divisible by tp_size ({args.tp_size})" - ) - assert args.ngroups % args.tp_size == 0, ( - f"ngroups ({args.ngroups}) must be divisible by tp_size ({args.tp_size})" - ) - args.tp_nheads = args.nheads // args.tp_size - args.tp_ngroups = args.ngroups // args.tp_size - - batch_sizes = [int(x) for x in args.batch_sizes.split(",")] - mtp_lengths = [int(x) for x in args.mtp_lengths.split(",")] - if args.mix_csv is not None and len(mtp_lengths) != 1: - sys.exit("--mix-csv requires exactly one --mtp-lengths value") - - dtype_map = { - "bf16": torch.bfloat16, - "fp32": torch.float32, - "fp16": torch.float16, - "int8": torch.int8, - "int16": torch.int16, - "fp8": torch.float8_e4m3fn, - } - state_dtypes = [dtype_map[s] for s in args.state_dtypes.split(",")] - act_dtypes = [dtype_map[s] for s in args.act_dtypes.split(",")] - if args.json_output and len(act_dtypes) != 1: - sys.exit("--json-output requires exactly one --act-dtypes value") - - # Resolve baseline function. - if args.baseline == "flashinfer_pr3324": - _prepare_flashinfer_jit_workspace() - from flashinfer_checkpointing_ssu_pr3324 import checkpointing_ssu as baseline_fn - else: - baseline_fn = None - - # --with-conv1d uses its own realistic L2 flush (cold cache flush then - # hot in_proj write). Override the generic l2_flush to avoid double-flushing. - if args.with_conv1d: - args.l2_flush = False - _init_l2_flush() # still needed for the realistic reset's flush step - elif args.l2_flush: - _init_l2_flush() - - _phase("about to enter compile-warmup") - if args.compile_threads > 0: - _compile_warmup_phase( - args, - batch_sizes, - mtp_lengths, - state_dtypes, - act_dtypes, - baseline_fn, - max_workers=args.compile_threads, - ) - _phase("returned from compile-warmup") - - # Pre-warm the per-(state_dtype, act_dtype, mtp_len, ...) tensor cache at - # the largest requested batch size. Without this, the timing loop would - # progressively grow the cache as it encounters larger batches (e.g., - # iterate 1 -> 16 -> 64 -> 128 -> 512 = 5 separate growth allocations, - # each freeing the previous buffers). Pre-warming at max-batch up front - # makes every subsequent timing cell a view-slice (zero alloc cost). - _max_batch = max(batch_sizes) - for state_dtype in state_dtypes: - for act_dtype in act_dtypes: - for mtp_len in mtp_lengths: - _build_tensors( - _max_batch, - mtp_len, - state_dtype, - act_dtype, - args.tp_nheads, - args.head_dim, - args.d_state, - args.tp_ngroups, - max_window=getattr(args, "max_window", None) or None, - ) - _phase("done tensor prewarm — entering timing") - - if args.profile: - torch.cuda.cudart().cudaProfilerStart() - - # Print header - mix_enabled = args.mix_csv is not None or getattr(args, "pmix", False) - headline_name = "score_us" if mix_enabled else "median_us" - print(f"conv1d: {'enabled' if args.with_conv1d else 'disabled'}") - if mix_enabled: - print( - "kmix bucket score: mix rows report bucket-weighted score_us; " - "non-mix rows report median_us." - ) - if baseline_fn is not None: - print( - f"| {'kernel':>11} | {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " - f"{'state_dtype':>11} | {'act_dtype':>9} | " - f"{headline_name:>9} | {'p95_us':>7} | {'p99_us':>7} |" - ) - print( - f"|{'-' * 13}|{'-' * 7}|{'-' * 9}|{'-' * 8}|" - f"{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" - ) - else: - print( - f"| {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " - f"{'state_dtype':>11} | {'act_dtype':>9} | " - f"{headline_name:>9} | {'p95_us':>7} | {'p99_us':>7} |" - ) - print( - f"|{'-' * 7}|{'-' * 9}|{'-' * 8}|{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" - ) - - sr_modes_list = getattr(args, "sr_modes_list", ["RN"]) - rect_list = getattr(args, "rectangle_for_nowrite_list", [False]) - modes_list = getattr(args, "modes_list", ["persistent_dynamic"]) - hsort_list = getattr(args, "hardcode_sort_list", [False]) - nowrite_first_list = getattr(args, "nowrite_first_list", [None]) - - # Pre-load AL distribution for mix mode. - mix_al = None - mix_label = "" - if getattr(args, "pmix", False): - mix_label = DEFAULT_PMIX_LABEL - mix_al = _load_builtin_pmix_distribution(T=max(mtp_lengths)) - elif args.mix_csv is not None: - mix_csv = Path(args.mix_csv) - mix_label = mix_csv.stem - # T (= mtp_len) varies per cell; load once with the LARGEST mtp so - # we have enough columns; the loader normalizes the dist anyway. - mix_al = _load_al_distribution(mix_csv, T=max(mtp_lengths), column=args.mix_csv_column) - - for batch in batch_sizes: - for mtp_len in mtp_lengths: - # Resolve prev_k fractions → clamped integers in [0, mtp_len] - prev_ks = _resolve_prev_ks(args, mtp_len) - - # Pre-generate mix samples once per (batch, mtp_len) cell so all - # tuning configs see the same per-iter prev_tokens vectors — - # tuning differences become signal, mix-noise is shared. - # Size the sample buffer for the LARGER of args.iters and - # args.mix_iters since mix scenarios use mix_iters. - mix_samples_cpu = None - mix_samples_sorted_cpu = None # per-iter prev_tokens, write-first - mix_write_frac = None - if mix_al is not None: - _max_window = getattr(args, "max_window", 0) or mtp_len - mix_pi = _markov_stationary(mix_al, mtp_len, _max_window) - mix_write_frac = float(mix_pi[_max_window - mtp_len + 1 :].sum()) - _max_iters = max(args.iters, getattr(args, "mix_iters", None) or args.iters) - mix_samples_cpu = _sample_steady_state_pnat( - mix_al, - T=mtp_len, - window=_max_window, - batch=batch, - K=args.warmup + _max_iters, - seed=args.mix_seed, - ) - if any(hsort_list): - # write-first stable argsort: kind='stable' preserves - # original-slot order within each mode group. - write_mask = (mix_samples_cpu + mtp_len > _max_window).astype( - np.int8 - ) # 1 = write, 0 = nowrite - perm_idx = np.argsort(-write_mask, kind="stable", axis=-1).astype(np.int32) - # Apply the perm to the prev_tokens samples themselves. - # Result row i = mix_samples_cpu[i] reordered such - # that write-mode entries come first. - mix_samples_sorted_cpu = np.take_along_axis( - mix_samples_cpu, perm_idx, axis=-1 - ).astype(mix_samples_cpu.dtype) - - for state_dtype in state_dtypes: - for act_dtype in act_dtypes: - for sr_mode in _sr_modes_for_dtype(state_dtype, sr_modes_list): - for mode in modes_list: - for rect in rect_list: - for nowrite_first in nowrite_first_list: - can_sort = mix_samples_cpu is not None - cell_list_active = bool(getattr(args, "_cell_list_keys", ())) - effective_hsort_list = ( - hsort_list if (can_sort or cell_list_active) else [False] - ) - for hardcode_sort in effective_hsort_list: - _bench_config( - args, - batch, - mtp_len, - prev_ks, - state_dtype, - act_dtype, - baseline_fn, - sr_mode=sr_mode, - rectangle_for_nowrite=rect, - mode=mode, - mix_samples_cpu=mix_samples_cpu, - mix_label=mix_label, - hardcode_sort=hardcode_sort, - nowrite_first=nowrite_first, - mix_samples_sorted_cpu=mix_samples_sorted_cpu, - mix_write_frac=mix_write_frac, - ) - - _drain_pending_results(args, force=True) - - if args.profile: - torch.cuda.cudart().cudaProfilerStop() - - # JSONL is the canonical artifact (written incrementally per cell with - # host stamps). No clean-exit `.json` write — use `jsonl_to_json.py` to - # materialize a snapshot when an analyzer wants one. - if args.json_output and args._jsonl_path is not None: - print(f"\nJSONL results: {args._jsonl_path} (meta sidecar: {args.json_output}.meta.json)") - - # Write the skipped-cells sidecar. Caller can convert this list to a - # --cell-list JSON (one dict per skipped cell) to drive a retry pass in - # a fresh process. - skipped_path = getattr(args, "skipped_output", None) - if skipped_path is None and args.json_output: - # Derive default: foo.json -> foo.skipped.json - skipped_path = args.json_output.rsplit(".", 1)[0] + ".skipped.json" - if skipped_path is not None and args._skipped_cells: - payload = { - "metadata": { - "timestamp": datetime.now().isoformat(), - "cmd": " ".join(sys.argv), - "skipped_count": len(args._skipped_cells), - }, - "skipped": args._skipped_cells, - } - tmp = skipped_path + ".tmp" - with open(tmp, "w") as f: - json.dump(payload, f, indent=2) - os.replace(tmp, skipped_path) - print( - f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " - f"tags written to: {skipped_path}", - file=sys.stderr, - ) - elif args._skipped_cells: - # No output path but there are skipped cells — emit a stderr summary. - print( - f"Skipped {len(args._skipped_cells)} cells (CUPTI mismatch); " - f"first 5: {args._skipped_cells[:5]}", - file=sys.stderr, - ) - - -# CLI - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Benchmark replay_selective_state_update Triton kernel", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - parser.add_argument( - "--nheads", - type=int, - default=NHEADS, - help="Full-model nheads (divided by --tp-size for per-GPU slice)", - ) - parser.add_argument( - "--ngroups", - type=int, - default=NGROUPS, - help="Full-model ngroups (divided by --tp-size for per-GPU slice)", - ) - parser.add_argument( - "--head-dim", type=int, default=HEAD_DIM, help="Head dimension (not TP-split)" - ) - parser.add_argument( - "--d-state", type=int, default=D_STATE, help="SSM state dimension (not TP-split)" - ) - parser.add_argument( - "--tp-size", - type=int, - default=TP_SIZE, - help="Tensor parallel size; divides nheads and ngroups", - ) - parser.add_argument( - "--batch-sizes", default="1,2,4,8", help="Comma-separated decode batch sizes" - ) - parser.add_argument( - "--mtp-lengths", - default=str(DEFAULT_PMIX_T), - help="Comma-separated per-request sequence lengths (num_draft_tokens + 1 target)", - ) - parser.add_argument( - "--state-dtypes", - default="fp32", - help="Comma-separated state dtypes: fp16,bf16,fp32,int8,int16,fp8. " - "FlashInfer PR3324 baseline supports fp32/fp16/int8/fp8 only.", - ) - parser.add_argument( - "--act-dtypes", - default="bf16", - help="Comma-separated activation dtypes for x/B/C/dt: fp32,bf16", - ) - parser.add_argument( - "--warmup", - type=int, - default=4, - help="Number of warmup iterations. Default aligns with " - "the graph group-iters (default 4 for mix scenarios) so " - "warmup + iters / mix-iters lands on a clean multiple " - "without per-args rounding overhead. Earlier default of " - "20 was overkill for steady-state warming.", - ) - parser.add_argument("--iters", type=int, default=100, help="Number of timed iterations") - parser.add_argument( - "--compile-threads", - type=int, - default=64, - help="Number of THREADS used in the compile-warmup phase (one call " - "per (batch, mtp_len, prev_k, dtype, sweep) cell, parallelized over " - "N threads). Triton compile releases the GIL, so threads compile " - "in parallel and populate the persistent cache for free hits during " - "the sequential timed phase. 0 disables the phase. Default 64.", - ) - parser.add_argument( - "--mp-start-method", - choices=("spawn", "forkserver"), - default="spawn", - help="multiprocessing start method for compile-warmup workers AND " - "the CUPTI parser child process. 'spawn' (default) is robust but " - "each child re-imports the bench module (~15s torch+triton import " - "cost). 'forkserver' starts a server once, preloads the bench " - "module ONCE, then forks children cheaply (~1s each). When 4 " - "benches run concurrently with --compile-threads 26 each, spawn " - "still incurs 4*26=104 imports per round; forkserver cuts this to " - "4 (one per server).", - ) - parser.add_argument( - "--profile", - action="store_true", - help="Wrap timed region in cudaProfilerStart/Stop (for ncu --target-processes all)", - ) - parser.add_argument( - "--l2-flush", - action=argparse.BooleanOptionalAction, - default=True, - help="L2 eviction between iterations", - ) - parser.add_argument( - "--cuda-graph", - action=argparse.BooleanOptionalAction, - default=True, - help="Capture all warmup + timed iterations in a " - "single CUDA graph with per-iteration events " - "inside the graph, eliminating all host overhead.", - ) - parser.add_argument( - "--cuda-graph-group-iters", - type=int, - default=None, - help="Capture this many logical benchmark iterations per graph " - "replay when warmup + iters is divisible by this value. Default " - f"auto-selects {_DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE} for pure " - f"cells and {_DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX} for mix cells. " - "Mix cells use a per-replay device window so they can group " - "iterations too.", - ) - parser.add_argument( - "--cupti", - action=argparse.BooleanOptionalAction, - default=True, - help="Time kernels via CUPTI Activity API (1 ns from the GPU " - "profiling fabric); per-iter span = max(kernel_end) - " - "min(kernel_start). Default ON. --no-cupti disables in-bench " - "timing entirely (kernels still run, but median/p95/p99 are zero) " - "— use when wrapping the bench in nsys/ncu, where the external " - "profiler provides timings and our CUPTI subscriber would conflict.", - ) - parser.add_argument( - "--cupti-flush-period-ms", - type=int, - default=0, - help="If >0, ask CUPTI to periodically flush activity buffers during " - "the timed CUDA-graph region. This can overlap raw-buffer parsing with " - "long timed cells; 0 leaves flushing explicit at the end of each cell.", - ) - parser.add_argument( - "--cupti-defer-depth", - type=int, - default=4, - help="Maximum number of CUDA-graph CUPTI timing results that may be " - "left for the parser process while the main process starts later cells. " - "1 preserves synchronous per-cell parsing and inline retry behavior.", - ) - parser.add_argument( - "--json-output", - default=None, - help="If set, write per-cell results to this JSON file in the " - "shape consumed by collect.py / report.py. See the 'JSON output " - "schema' section at the top of this file.", - ) - parser.add_argument( - "--json-detailed", - action=argparse.BooleanOptionalAction, - default=False, - help="When --json-output is set, also include per_kernel " - "(per-iter relative start/end timestamps for each kernel). Compact " - "JSON always includes iters_us, and mix rows include " - "n_writes_per_iter. Default off keeps records compact.", - ) - parser.add_argument( - "--host-timing", - action=argparse.BooleanOptionalAction, - default=False, - help="Attach benchmark host-side phase timings to JSON/JSONL results. " - "Useful for diagnosing benchmark overhead, but it adds roughly 1 KB " - "per compact JSONL row and several perf_counter calls per cell.", - ) - parser.add_argument( - "--cupti-retry", - type=int, - default=1, - help="On CUPTI capture mismatch (kernel record count != expected), " - "retry the cell this many times in-process before giving up. CUPTI " - "gets racy after thousands of cells in one process (PDL + small " - "kernels occasionally lose records); a single retry usually catches " - "transient cases. Set 0 to disable and skip on first mismatch.", - ) - parser.add_argument( - "--skipped-output", - default=None, - help="Path to write the list of cells that failed CUPTI capture even " - "after --cupti-retry retries (JSON list of sweep_tag strings). " - "Default: derived from --json-output by replacing .json with " - ".skipped.json.", - ) - parser.add_argument( - "--cell-list", - default=None, - help="Path to a JSON list of cell dicts (one per cell to time). " - "Each dict has canonical knob keys → values: Mw, Mnw, Ww, Wnw, Sw, " - "Snw, pW, pS, H, R, CT, CPSw, CPSnw, LSw, LSnw, FL, WS, TMARL, " - "TMAWL, TMANL, TMAWS, SR, RECT, MODE; optional diagnostic HSORT is " - "also accepted (tied forms M / W / S / CPS / LS are also accepted " - "and auto-expanded). " - "When set, bench's CLI knob ranges are auto-overridden to the " - "per-knob union across all cells, and the inner-loop filter skips " - "any iteration whose knob-value tuple isn't in the list. All cells " - "must share the same key set (uniform schema).", - ) - parser.add_argument( - "--prev-tokens-fracs", - default="0,0.5,1.0", - type=lambda s: [float(x) for x in s.split(",")], - help="Fractions of mtp_len to use as prev_num_accepted_tokens " - "for the replay kernel sweep. Values are rounded " - "and clamped to [0, mtp_len].", - ) - parser.add_argument( - "--baseline", - default=None, - nargs="?", - const="flashinfer_pr3324", - choices=["flashinfer_pr3324"], - help="Baseline to benchmark alongside the replay kernel. " - "'flashinfer_pr3324': FlashInfer PR 3324 checkpointing_ssu, " - "benchmarked per prev_k for fp32/fp16/int8/fp8 state. " - "Pass --baseline alone for flashinfer_pr3324. Default: no baseline.", - ) - parser.add_argument( - "--output", - default=None, - help="Path to save results (file or directory). " - "If a directory, writes benchmark_replay_.txt inside it.", - ) - parser.add_argument( - "--block-size-m", - type=str, - default=None, - help="Override BLOCK_SIZE_M: single value or comma-separated sweep (e.g. '4,8,16,32').", - ) - parser.add_argument( - "--num-warps", - type=str, - default=None, - help="Override num_warps: single value or comma-separated sweep (e.g. '1,2,4').", - ) - parser.add_argument( - "--internal-pdl", - action=argparse.BooleanOptionalAction, - default=True, - help="Internal PDL between precompute and main kernels (default: on).", - ) - parser.add_argument( - "--num-stages", - type=str, - default=None, - help="Override num_stages for the main kernel (comma-separated sweep).", - ) - parser.add_argument( - "--block-size-m-write", - type=str, - default=None, - help="Sweep BLOCK_SIZE_M for the WRITE main only (overrides --block-size-m " - "for the write half). Tied to --block-size-m if unset.", - ) - parser.add_argument( - "--block-size-m-nowrite", - type=str, - default=None, - help="Sweep BLOCK_SIZE_M for the NOWRITE main only. Tied to --block-size-m if unset.", - ) - parser.add_argument( - "--num-warps-write", - type=str, - default=None, - help="Sweep num_warps for the WRITE main only. Tied to --num-warps if unset.", - ) - parser.add_argument( - "--num-warps-nowrite", - type=str, - default=None, - help="Sweep num_warps for the NOWRITE main only. Tied to --num-warps if unset.", - ) - parser.add_argument( - "--num-stages-write", - type=str, - default=None, - help="Sweep num_stages for the WRITE main only. Tied to --num-stages if unset.", - ) - parser.add_argument( - "--num-stages-nowrite", - type=str, - default=None, - help="Sweep num_stages for the NOWRITE main only. Tied to --num-stages if unset.", - ) - parser.add_argument( - "--cta-per-sm-write", - type=str, - default=None, - help="Sweep cta_per_sm for the WRITE persistent_main only. Tied to --cta-per-sm if unset.", - ) - parser.add_argument( - "--cta-per-sm-nowrite", - type=str, - default=None, - help="Sweep cta_per_sm for the NOWRITE persistent_main only. Tied to --cta-per-sm if unset.", - ) - parser.add_argument( - "--num-loop-stages-write", - type=str, - default=None, - help="Sweep num_loop_stages for the WRITE persistent_main only. Tied to --num-loop-stages if unset.", - ) - parser.add_argument( - "--num-loop-stages-nowrite", - type=str, - default=None, - help="Sweep num_loop_stages for the NOWRITE persistent_main only. Tied to --num-loop-stages if unset.", - ) - parser.add_argument( - "--skip-diagonal", - action=argparse.BooleanOptionalAction, - default=False, - help="When sweeping any per-main *_write / *_nowrite knobs, skip cells " - "where ALL splittable knobs satisfy write_value == nowrite_value (i.e. " - "the 'diagonal' that's already covered by a prior shared-knob sweep). " - "Useful for incremental sweeps that extend earlier results without redoing " - "the tied-knob cells.", - ) - parser.add_argument( - "--precompute-num-warps", - type=str, - default=None, - help="Override num_warps for precompute kernel (comma-separated sweep).", - ) - parser.add_argument( - "--max-window", - type=int, - default=16, - help="Cache T-axis capacity (max replay buffer length). Default 16 " - "matches Nemotron-3-Super-120B production. Pass 0 to fall back to " - "mtp_len (degenerate every-step-checkpoint case, mostly unused).", - ) - parser.add_argument( - "--prev-tokens-int", - type=lambda s: [int(x) for x in s.split(",")] if s else None, - default=None, - help="Absolute prev_num_accepted_tokens values to test, comma-separated " - "(e.g. '0,10,11,16'). Clamped to [0, max_window]. When set, " - "overrides --prev-tokens-fracs.", - ) - parser.add_argument( - "--with-conv1d", - "--with-conv-1d", - action=argparse.BooleanOptionalAction, - default=True, - dest="with_conv1d", - help="Include conv1d kernel before replay SSM. Default on; pass " - "--no-with-conv1d to time state update only. Uses realistic L2 flush: " - "cold caches flushed, hot in_proj output kept warm. Measures " - "conv1d → precompute → main span.", - ) - parser.add_argument( - "--external-pdl", - action=argparse.BooleanOptionalAction, - default=True, - help="External PDL: conv1d launches dependents, precompute waits. " - "Only relevant with --with-conv1d. --no-external-pdl disables.", - ) - parser.add_argument( - "--heads-per-block", - type=str, - default=None, - help="Override HEADS_PER_BLOCK for precompute kernel (comma-separated sweep).", - ) - parser.add_argument( - "--cta-per-sm", - type=str, - default=None, - help="CTAs per SM in the 1D persistent grid for mode=persistent_main " - "(comma-separated sweep). num_persistent = cta_per_sm × NUM_SMS. " - "Default = 1 (one CTA per SM). Replaces the old --num-persistent. " - "Ignored for non-persistent_main modes.", - ) - parser.add_argument( - "--num-loop-stages", - type=str, - default=None, - help="num_stages on the inner tl.range(...) persistent loop for " - "mode=persistent_main (comma-separated sweep). Default = 2. Note: " - "this is loop-level, NOT the kernel-arg num_stages (which only " - "pipelines dot-feeding loads). Watch Triton issue #8259 — " - "num_stages>1 + flatten=True can corrupt stores in non-dot kernels. " - "Ignored for non-persistent_main modes.", - ) - parser.add_argument( - "--flatten", - type=str, - default=None, - help="`flatten` arg on tl.range(...) for mode=persistent_main " - "(comma-separated 0/1 sweep). Default = 1. Ignored for " - "non-persistent_main modes.", - ) - parser.add_argument( - "--warp-specialize", - type=str, - default=None, - help="`warp_specialize` arg on tl.range(...) for mode=persistent_main " - "(comma-separated 0/1 sweep). Default = 0. Triton 3.6 only " - "supports it on simple matmul loops; our scan loop probably won't " - "pattern-match — exposed as a knob for sweep experiments. Requires " - "num_warps >= 4 if 1. Ignored for non-persistent_main modes.", - ) - parser.add_argument( - "--sr-modes", - type=str, - default="RN", - help="Comma-separated rounding modes to sweep: any combination of " - "{RN, SR}. SR (stochastic rounding) is silently skipped for state " - "dtypes that don't support it (bf16, fp32). Default 'RN' matches " - "legacy --philox-rounding=False behavior.", - ) - parser.add_argument( - "--rectangle-for-nowrite", - type=str, - default=None, - help="Comma-separated 0/1 values: 0 = replay-style nowrite kernel, " - "1 = dedicated rectangle nowrite kernel. Sweep both with '0,1' to " - "compare in one invocation. Silently no-op for write cells (the " - "write path always uses replay-style). When unset (default), the " - "wrapper resolves from the _DEFAULT_TUNING lookup per (batch, dtype, " - "sr) cell.", - ) - parser.add_argument( - "--nowrite-first", - type=str, - default=None, - help="Comma-separated 0/1 values for mode=persistent_main launch order. " - "0 launches write before nowrite; 1 launches nowrite before write. " - "When unset (default), the wrapper resolves from _DEFAULT_TUNING. " - "Ignored for persistent_dynamic.", - ) - parser.add_argument( - "--use-tma-rect-load", - type=str, - default=None, - help="Comma-separated 0/1 sweep. Use TMA (host-built tensor " - "descriptor) for state load in the rectangle nowrite path. " - "Cells where the rect path isn't reachable (rectangle_for_nowrite=False) " - "skip the value=1 case as a dupe.", - ) - parser.add_argument( - "--use-tma-replay-write-load", - type=str, - default=None, - help="Comma-separated 0/1 sweep. TMA state LOAD in replay main " - "for the checkpoint/write half. Independent from nowrite-load and " - "rect TMA — see CHECKPOINTING_DESIGN.md item #17 for measured perf.", - ) - parser.add_argument( - "--use-tma-replay-nowrite-load", - type=str, - default=None, - help="Comma-separated 0/1 sweep. TMA state LOAD in replay main " - "for the replay/nowrite half. Design doc reports the largest win " - "on this path (int8 b>=64: -8 to -12%%).", - ) - parser.add_argument( - "--use-tma-replay-write-store", - type=str, - default=None, - help="Comma-separated 0/1 sweep. TMA state STORE in replay main " - "for the checkpoint/write half. Independent from all load TMA flags.", - ) - parser.add_argument( - "--strided-state-cache", - action=argparse.BooleanOptionalAction, - default=False, - help="Allocate SSM state as a view with block-reuse-like gaps between " - "cache slots. Used to validate TMA descriptor handling for recurrent " - "state pools that pack SSM and conv state together.", - ) - parser.add_argument( - "--require-tma-state-layout", - action=argparse.BooleanOptionalAction, - default=False, - help="Raise instead of falling back if a TMA state path is requested " - "but the state layout cannot be represented by the 2D descriptor.", - ) - parser.add_argument( - "--modes", - type=str, - default=None, - help="Comma-separated dispatch modes to sweep, any of " - "{persistent_dynamic, persistent_main}. " - "persistent_dynamic = single persistent-CTA kernel that dispatches " - "per-slot at runtime based on PNAT. " - "persistent_main = persistent-CTA kernel with two halves (write + " - "nowrite), using caller-provided n_writes and write-first " - "replay_work_items. " - "When unset (default), the wrapper resolves from the _DEFAULT_TUNING " - "lookup per (batch, dtype, sr) cell.", - ) - parser.add_argument( - "--mix-csv", - type=str, - default=None, - help="Path to AL histogram CSV (cols: AL, count). When set, an " - "additional 'mix' cell is emitted per (batch, mtp, dtype, sr, " - "mode, RECT, M, W, ...) combo where prev_tokens varies per iter, " - "drawn from the steady-state PNAT distribution induced by the " - "AL histogram. Both persistent modes support mix scenarios. " - "Each iteration of the captured CUDA graph has a different " - "pre-baked prev_tokens vector; warmup iters use distinct samples " - "from the timed iters so nsys-included warmup leaks don't bias. " - "Mutually exclusive with --pmix.", - ) - parser.add_argument( - "--pmix", - action=argparse.BooleanOptionalAction, - default=False, - help=f"Use the built-in production-like T{DEFAULT_PMIX_T} accepted-length " - f"histogram from the replay-count column. Requires --mtp-lengths {DEFAULT_PMIX_T} " - "and is mutually exclusive with --mix-csv.", - ) - parser.add_argument( - "--mix-csv-column", - type=int, - default=1, - help="Column index (0-based) in the AL histogram CSV for the " - "count/probability column. Default 1 (second column).", - ) - parser.add_argument( - "--mix-seed", - type=int, - default=42, - help="RNG seed for the steady-state PNAT sampler. Same seed " - "across runs => same per-slot samples for reproducible " - "comparisons.", - ) - parser.add_argument( - "--hardcode-sort", - type=str, - default=None, - help="Comma-separated 0/1, default 0. Diagnostic only: when 1, " - "per-iter PNAT samples are preclustered write-first before " - "replay_work_items are built. Production-like mixed runs leave PNAT " - "unsorted and sort only replay_work_items.", - ) - parser.add_argument( - "--mix-iters", - type=int, - default=None, - help="Iteration count override for mix scenarios (each iter is a " - "different per-slot prev_tokens draw). Default (None) uses " - "--iters. Mix scenarios benefit from more iters since each " - "iter samples a different mix; pure scenarios don't.", - ) - parser.add_argument( - "--mix-only", - action=argparse.BooleanOptionalAction, - default=False, - help="When --pmix or --mix-csv is set, emit only mix scenarios and skip the " - "pure prev_k sibling scenarios. Default: false.", - ) - parser.add_argument( - "--philox-rounding", - action="store_true", - help="DEPRECATED — equivalent to --sr-modes SR. Retained for " - "backward compatibility; use --sr-modes for new scripts. fp16 SR " - "and fp8 SR require sm_100a (Blackwell B200+).", - ) - parser.add_argument( - "--philox-rounds", - type=int, - default=5, - help="Number of Philox PRNG rounds. Default 5 matches the " - "Nemotron-3-Super-120B production config (mamba_ssm_philox_rounds=5 " - "in examples/configs and tests/integration/perf configs). The " - "wrapper's generic fallback default is 10; callers without explicit " - "config see 10. Only consulted when --philox-rounding is enabled.", - ) - parser.add_argument( - "--full-import", - action="store_true", - help="Use standard tensorrt_llm import path instead of fast direct " - "module loading. Slower (~40s startup) but guaranteed correct " - "if the fast path breaks due to package changes.", - ) - args = parser.parse_args() - if args.pmix and args.mix_csv is not None: - parser.error( - "--pmix and --mix-csv are mutually exclusive; use --pmix for the " - f"built-in T{DEFAULT_PMIX_T} distribution or --mix-csv for a custom histogram." - ) - if args.pmix: - mtp_lengths_for_pmix = [int(x) for x in args.mtp_lengths.split(",") if x.strip()] - if any(t != DEFAULT_PMIX_T for t in mtp_lengths_for_pmix): - parser.error( - f"--pmix uses the built-in T{DEFAULT_PMIX_T} histogram; " - f"use --mtp-lengths {DEFAULT_PMIX_T} or pass --mix-csv for another T." - ) - if args.mix_only and args.mix_csv is None and not args.pmix: - parser.error("--mix-only requires --pmix or --mix-csv") - - # Round iter counts up so warmup + iters (and warmup + mix_iters) are clean - # multiples of the graph group-iters used downstream. Default mix group is - # 4, default pure group is 2. An explicit --cuda-graph-group-iters can - # request a larger group. We round to the max of the two so all scenarios - # in a single run (pure + mix) share a clean total_iters. The cost is at - # most (group-1) extra iters per scenario — negligible — and the win is - # that graph_group_iters never falls back to 1 (which caused ~5x slowdown - # in observed benchmark walls). - _group_for_rounding = max( - _DEFAULT_CUDA_GRAPH_GROUP_ITERS_MIX, - _DEFAULT_CUDA_GRAPH_GROUP_ITERS_PURE, - getattr(args, "cuda_graph_group_iters", None) or 0, - ) - - def _round_iters_to_group(name, val): - total = args.warmup + val - if total % _group_for_rounding == 0: - return val - new_total = ((total + _group_for_rounding - 1) // _group_for_rounding) * _group_for_rounding - new_val = new_total - args.warmup - print( - f"[bench] rounding --{name} {val} → {new_val} so warmup+{name} " - f"({new_total}) is a multiple of graph group_iters={_group_for_rounding}", - file=sys.stderr, - ) - return new_val - - args.iters = _round_iters_to_group("iters", args.iters) - if getattr(args, "mix_iters", None): - args.mix_iters = _round_iters_to_group("mix-iters", args.mix_iters) - - # Cell-list (if any) must be applied BEFORE the post-argparse string→list - # derivations below — those build args.*_list from args.* strings, so a - # cell-list override of e.g. args.modes='persistent_main' needs to land - # before args.modes_list is computed. The function populates args._cell_list_keys - # and args._cell_list_set, plus overrides args.* knob strings to the - # per-knob union of values across the listed cells. - if getattr(args, "cell_list", None): - _load_cell_list_into_args(args) - - # Backward-compat: --philox-rounding implies --sr-modes SR if --sr-modes - # was left at the default. If both are set explicitly, error. - sr_modes_default = args.sr_modes == "RN" - if args.philox_rounding: - if not sr_modes_default and args.sr_modes != "SR": - parser.error( - "--philox-rounding (deprecated) is incompatible with explicit " - f"--sr-modes={args.sr_modes!r}. Use --sr-modes SR (or " - "RN,SR) instead and drop --philox-rounding." - ) - args.sr_modes = "SR" - - sr_modes = [m.strip() for m in args.sr_modes.split(",") if m.strip()] - for m in sr_modes: - if m not in ("RN", "SR"): - parser.error(f"--sr-modes value must be RN or SR, got {m!r}") - args.sr_modes_list = sr_modes - - # rectangle_for_nowrite=None means "let the wrapper resolve from - # _DEFAULT_TUNING". Empty/unset argparse default produces [None] in the - # sweep list; the kernel call passes None and the wrapper picks per-cell. - if args.rectangle_for_nowrite is None: - rect_list = [None] - else: - rect_modes = [v.strip() for v in args.rectangle_for_nowrite.split(",") if v.strip()] - rect_list = [] - for v in rect_modes: - if v not in ("0", "1"): - parser.error(f"--rectangle-for-nowrite value must be 0 or 1, got {v!r}") - rect_list.append(v == "1") - if not rect_list: - rect_list = [None] - args.rectangle_for_nowrite_list = rect_list - - if args.nowrite_first is None: - nowrite_first_list = [None] - else: - nowrite_first_modes = [v.strip() for v in args.nowrite_first.split(",") if v.strip()] - nowrite_first_list = [] - for v in nowrite_first_modes: - if v not in ("0", "1"): - parser.error(f"--nowrite-first value must be 0 or 1, got {v!r}") - nowrite_first_list.append(v == "1") - if not nowrite_first_list: - nowrite_first_list = [None] - args.nowrite_first_list = nowrite_first_list - - hsort_modes = [ - v.strip() - for v in (args.hardcode_sort if args.hardcode_sort is not None else "0").split(",") - if v.strip() - ] - hsort_list = [] - for v in hsort_modes: - if v not in ("0", "1"): - parser.error(f"--hardcode-sort value must be 0 or 1, got {v!r}") - hsort_list.append(v == "1") - if not hsort_list: - hsort_list = [False] - args.hardcode_sort_list = hsort_list - - # mode=None means "let the wrapper resolve from _DEFAULT_TUNING". Same - # convention as --rectangle-for-nowrite. - if args.modes is None: - args.modes_list = [None] - else: - modes_raw = [v.strip() for v in args.modes.split(",") if v.strip()] - valid_modes = { - "persistent_main", - "persistent_dynamic", - } - for m in modes_raw: - if m not in valid_modes: - parser.error(f"--modes value must be one of {sorted(valid_modes)}, got {m!r}") - args.modes_list = modes_raw if modes_raw else [None] - return args - - -class _Tee: - """Write to both stdout and a file simultaneously.""" - - def __init__(self, path: str): - parent = os.path.dirname(path) - if parent: - os.makedirs(parent, exist_ok=True) - self._file = open(path, "w") # noqa: SIM115 - self._stdout = sys.stdout - - def write(self, data): - self._stdout.write(data) - self._file.write(data) - - def flush(self): - self._stdout.flush() - self._file.flush() - - def close(self): - self._file.close() - - -if __name__ == "__main__": - _args = _parse_args() - - # Configure multiprocessing start method early — must be before any - # mp.get_context() that uses the chosen method. For forkserver, also - # add this file's dir to sys.path so the forkserver can import this - # module by basename for preload (otherwise it tries to import - # __main__, which is a different beast across processes). - if _args.mp_start_method == "forkserver": - sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - mp.set_start_method("forkserver", force=True) - try: - mp.set_forkserver_preload( - [ - "benchmark_replay_selective_state_update", - ] - ) - except Exception as _e: - print( - f"[warn] set_forkserver_preload failed: {_e!r}; " - f"forks will still work but pay full import cost", - file=sys.stderr, - ) - _MP_START_METHOD = _args.mp_start_method - - _out_path = None - if _args.output != "-": - _ts = datetime.now().strftime("%Y%m%d_%H%M%S") - _fname = f"benchmark_replay_{_ts}.txt" - if _args.output is None: - _out_path = os.path.expanduser(f"~/nemo_logs/{_fname}") - elif os.path.isdir(_args.output) or _args.output.endswith("/"): - _out_path = os.path.join(_args.output, _fname) - else: - _out_path = _args.output - - if _out_path: - _tee = _Tee(_out_path) - sys.stdout = _tee - print(f"# benchmark_replay_selective_state_update {datetime.now().isoformat()}") - print(f"# cmd: {' '.join(sys.argv)}") - - try: - _run_benchmark(_args) - finally: - if _out_path: - sys.stdout = _tee._stdout - _tee.close() - print(f"\nResults saved to: {_out_path}") diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py deleted file mode 100644 index e03eb6151eea..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324.py +++ /dev/null @@ -1,325 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# 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. - -"""Benchmark-only wrapper for FlashInfer PR 3324 checkpointing SSU. - -The vendored CUDA/C++ sources in ``flashinfer_checkpointing_ssu_pr3324/`` -come from https://github.com/flashinfer-ai/flashinfer/pull/3324. -""" - -from __future__ import annotations - -import functools -import os -from pathlib import Path -from typing import Optional - -import jinja2 -import torch -from flashinfer.compilation_context import CompilationContext -from flashinfer.jit import env as jit_env -from flashinfer.jit.core import JitSpec, gen_jit_spec -from flashinfer.jit.utils import write_if_different - -_ROOT = Path(__file__).resolve().parent / "flashinfer_checkpointing_ssu_pr3324" -_CSRC_DIR = _ROOT / "csrc" -_INCLUDE_DIR = _ROOT / "include" -_SUPPORTED_MAJORS = [8, 9, 10, 11, 12] - -_DTYPE_MAP = { - torch.float16: "half", - torch.bfloat16: "nv_bfloat16", - torch.float32: "float", - torch.int8: "int8_t", - torch.int16: "int16_t", - torch.int32: "int32_t", - torch.int64: "int64_t", - torch.float8_e4m3fn: "__nv_fp8_e4m3", -} - -_FILENAME_SAFE_DTYPE_MAP = { - torch.float16: "f16", - torch.bfloat16: "bf16", - torch.float32: "f32", - torch.int8: "i8", - torch.int16: "i16", - torch.int32: "i32", - torch.int64: "i64", - torch.float8_e4m3fn: "e4m3", -} - - -def _arch_flags() -> list[str]: - context = CompilationContext() - return context.get_nvcc_flags_list(supported_major_versions=_SUPPORTED_MAJORS) - - -def _uri( - state_dtype: torch.dtype, - input_dtype: torch.dtype, - dt_dtype: torch.dtype, - weight_dtype: torch.dtype, - matrix_a_dtype: torch.dtype, - state_index_dtype: torch.dtype, - state_scale_dtype: Optional[torch.dtype], - dim: int, - dstate: int, - npredicted: int, - max_window: int, - heads_per_group: int, - philox_rounds: int, - enable_pdl: bool, -) -> str: - dtype = _FILENAME_SAFE_DTYPE_MAP - uri = ( - "trtllm_pr3324_checkpointing_ssu_v4_" - f"s_{dtype[state_dtype]}_i_{dtype[input_dtype]}_dt_{dtype[dt_dtype]}_" - f"w_{dtype[weight_dtype]}_a_{dtype[matrix_a_dtype]}_" - f"si_{dtype[state_index_dtype]}_d_{dim}_ds_{dstate}_" - f"np_{npredicted}_mw_{max_window}_hpg_{heads_per_group}" - ) - if state_scale_dtype is not None: - uri += f"_sc_{dtype[state_scale_dtype]}" - if philox_rounds > 0: - uri += f"_pr_{philox_rounds}" - if enable_pdl: - uri += "_pdl" - return uri - - -def _gen_module( - state_dtype: torch.dtype, - input_dtype: torch.dtype, - dt_dtype: torch.dtype, - weight_dtype: torch.dtype, - matrix_a_dtype: torch.dtype, - state_index_dtype: torch.dtype, - state_scale_dtype: Optional[torch.dtype], - dim: int, - dstate: int, - npredicted: int, - max_window: int, - heads_per_group: int, - philox_rounds: int, - enable_pdl: bool, -) -> JitSpec: - uri = _uri( - state_dtype, - input_dtype, - dt_dtype, - weight_dtype, - matrix_a_dtype, - state_index_dtype, - state_scale_dtype, - dim, - dstate, - npredicted, - max_window, - heads_per_group, - philox_rounds, - enable_pdl, - ) - gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri - os.makedirs(gen_directory, exist_ok=True) - - with open(_CSRC_DIR / "checkpointing_ssu_customize_config.jinja") as file: - config_template = jinja2.Template(file.read()) - - state_scale_type = _DTYPE_MAP[state_scale_dtype] if state_scale_dtype is not None else "void" - config = config_template.render( - state_dtype=_DTYPE_MAP[state_dtype], - input_dtype=_DTYPE_MAP[input_dtype], - dt_dtype=_DTYPE_MAP[dt_dtype], - weight_dtype=_DTYPE_MAP[weight_dtype], - matrixA_dtype=_DTYPE_MAP[matrix_a_dtype], - stateIndex_dtype=_DTYPE_MAP[state_index_dtype], - state_scale_type=state_scale_type, - dim=dim, - dstate=dstate, - npredicted=npredicted, - max_window=max_window, - heads_per_group=heads_per_group, - philox_rounds=philox_rounds, - enable_pdl="true" if enable_pdl else "false", - ) - write_if_different(gen_directory / "checkpointing_ssu_config.inc", config) - - source_paths = [] - for filename in ( - "checkpointing_ssu.cu", - "checkpointing_ssu_kernel_inst.cu", - "checkpointing_ssu_jit_binding.cu", - ): - source_path = _CSRC_DIR / filename - dest_path = gen_directory / filename - source_paths.append(dest_path) - with open(source_path) as file: - write_if_different(dest_path, file.read()) - - return gen_jit_spec( - uri, - source_paths, - extra_cuda_cflags=_arch_flags(), - extra_include_paths=[_INCLUDE_DIR], - ) - - -@functools.cache -def _get_module( - state_dtype: torch.dtype, - input_dtype: torch.dtype, - dt_dtype: torch.dtype, - weight_dtype: torch.dtype, - matrix_a_dtype: torch.dtype, - state_index_dtype: torch.dtype, - state_scale_dtype: Optional[torch.dtype], - dim: int, - dstate: int, - npredicted: int, - max_window: int, - heads_per_group: int, - philox_rounds: int, - enable_pdl: bool, -): - return _gen_module( - state_dtype, - input_dtype, - dt_dtype, - weight_dtype, - matrix_a_dtype, - state_index_dtype, - state_scale_dtype, - dim, - dstate, - npredicted, - max_window, - heads_per_group, - philox_rounds, - enable_pdl, - ).build_and_load() - - -def checkpointing_ssu( - state: torch.Tensor, - old_x: torch.Tensor, - old_B: torch.Tensor, - old_dt: torch.Tensor, - old_cumAdt: torch.Tensor, - cache_buf_idx: torch.Tensor, - prev_num_accepted_tokens: torch.Tensor, - x: torch.Tensor, - dt: torch.Tensor, - A: torch.Tensor, - B: torch.Tensor, - C: torch.Tensor, - out: torch.Tensor, - D: Optional[torch.Tensor] = None, - z: Optional[torch.Tensor] = None, - dt_bias: Optional[torch.Tensor] = None, - dt_softplus: bool = False, - state_batch_indices: Optional[torch.Tensor] = None, - pad_slot_id: int = -1, - state_scale: Optional[torch.Tensor] = None, - rand_seed: Optional[torch.Tensor] = None, - philox_rounds: int = 10, - d_split: Optional[int] = None, - cu_seqlens: Optional[torch.Tensor] = None, - max_seqlen: Optional[int] = None, - enable_pdl: bool = False, -) -> torch.Tensor: - quantized_state_dtypes = (torch.int8, torch.float8_e4m3fn) - if state.dtype in quantized_state_dtypes: - if state_scale is None: - raise ValueError(f"state dtype {state.dtype} requires state_scale") - elif state_scale is not None: - raise ValueError(f"state_scale must be None for non-quantized state dtype {state.dtype}") - if cu_seqlens is not None: - npredicted = max_seqlen if max_seqlen is not None else old_x.size(1) - else: - if max_seqlen is not None: - raise ValueError("max_seqlen is only valid with cu_seqlens") - npredicted = x.size(1) - - max_window = old_x.size(1) - if max_window > 16: - raise ValueError(f"PR3324 checkpointing SSU supports max_window <= 16, got {max_window}") - if npredicted > max_window: - raise ValueError(f"npredicted ({npredicted}) must be <= max_window ({max_window})") - - if d_split is None: - d_split = 1 - if state.dtype in quantized_state_dtypes and d_split != 1: - raise ValueError(f"8-bit state requires d_split=1, got {d_split}") - - state_index_dtype = ( - state_batch_indices.dtype if state_batch_indices is not None else torch.int32 - ) - nheads = state.size(1) - ngroups = B.size(-2) - if nheads % ngroups != 0: - raise ValueError(f"nheads ({nheads}) must be divisible by ngroups ({ngroups})") - heads_per_group = nheads // ngroups - - if rand_seed is None: - philox_rounds = 0 - elif philox_rounds <= 0: - raise ValueError(f"philox_rounds must be > 0 with rand_seed, got {philox_rounds}") - - weight_dtype = ( - D.dtype if D is not None else (dt_bias.dtype if dt_bias is not None else dt.dtype) - ) - - module = _get_module( - state.dtype, - x.dtype, - dt.dtype, - weight_dtype, - A.dtype, - state_index_dtype, - state_scale.dtype if state_scale is not None else None, - state.size(2), - state.size(3), - npredicted, - max_window, - heads_per_group, - philox_rounds, - enable_pdl, - ) - module.checkpointing_ssu( - state, - x, - dt, - A, - B, - C, - out, - old_x, - old_B, - old_dt, - old_cumAdt, - cache_buf_idx, - prev_num_accepted_tokens, - D, - z, - dt_bias, - dt_softplus, - state_batch_indices, - pad_slot_id, - state_scale, - rand_seed, - d_split, - cu_seqlens, - ) - return out diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu deleted file mode 100644 index e9ff3f98c188..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu.cu +++ /dev/null @@ -1,554 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -// clang-format off -// config.inc MUST come before the header: it defines DIM, DSTATE, NPREDICTED, -// MAX_WINDOW constexprs that the header's function templates rely on. -#include "checkpointing_ssu_config.inc" -#include -#include -// clang-format on -#include "tvm_ffi_utils.h" - -using namespace flashinfer; -using tvm::ffi::Optional; - -namespace flashinfer::mamba::checkpointing -{ - -void checkpointing_ssu(TensorView state, // (state_cache_size, nheads, dim, dstate) - TensorView x, // (batch, NPREDICTED, nheads, dim) / (1, total_tokens, nheads, dim) under varlen - TensorView dt, // (batch, NPREDICTED, nheads, dim) tie_hdim / (1, total_tokens, nheads, dim) - TensorView A, // (nheads, dim, dstate) tie_hdim - TensorView B, // (batch, NPREDICTED, ngroups, dstate) / (1, total_tokens, ngroups, dstate) - TensorView C, // same as B - TensorView output, // same layout as x - // Cache tensors - TensorView old_x, // (state_cache_size, MAX_WINDOW, nheads, dim) - TensorView old_B, // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) - TensorView old_dt, // (state_cache_size, 2, nheads, MAX_WINDOW) f32 - TensorView old_cumAdt, // (state_cache_size, 2, nheads, MAX_WINDOW) f32 - TensorView cache_buf_idx, // (state_cache_size,) int32 - TensorView prev_num_accepted, // (state_cache_size,) int32 - // Optional tensors - Optional D, // (nheads, dim) - Optional z, // same layout as x - Optional dt_bias, // (nheads, dim) tie_hdim - bool dt_softplus, - Optional state_batch_indices, // (batch,) int32 - int64_t pad_slot_id, - Optional state_scale, // (state_cache_size, nheads, dim) f32 - Optional rand_seed, // single int64 - int64_t d_split, // v12 §59: per-head DIM split factor (1, 2, or 4) - Optional cu_seqlens) -{ // (batch+1,) int32, varlen mode - - bool const is_varlen = cu_seqlens.has_value(); - - // ── Extract dimensions ── - auto const state_cache_size = state.size(0); - auto const nheads = state.size(1); - auto const dim = state.size(2); - auto const dstate = state.size(3); - auto const max_window = old_x.size(1); - auto const ngroups = B.size(2); - - // In non-varlen mode, batch = x.size(0) and npredicted = x.size(1) (the - // 4D batched layout). In varlen, the JIT compile-time NPREDICTED is the - // max seq_len the caller commits to; the wrapper stamped it into the JIT - // URI from `max_seqlen` and we read it back as the constexpr `NPREDICTED` - // (validated against runtime cu_seqlens on the host side below). `batch` - // = number of sequences = `cu_seqlens.size(0) - 1`. - int64_t batch; - int64_t npredicted; - if (is_varlen) - { - auto const& cs = cu_seqlens.value(); - CHECK_CUDA(cs); - CHECK_DIM(1, cs); - FLASHINFER_CHECK( - cs.size(0) >= 2, "cu_seqlens must have shape (batch+1,) with batch >= 1, got size(0)=", cs.size(0)); - FLASHINFER_CHECK(cs.dtype().code == kDLInt && cs.dtype().bits == 32, "cu_seqlens must be int32"); - CHECK_CONTIGUOUS(cs); - batch = cs.size(0) - 1; - npredicted = NPREDICTED; // JIT-stamped — wrapper ensures max(seq_lens) <= NPREDICTED. - } - else - { - batch = x.size(0); - npredicted = x.size(1); - } - - // ── JIT compile-time / runtime cross-check ── - // NPREDICTED and MAX_WINDOW are JIT compile-time constants stamped by the - // wrapper. In non-varlen NPREDICTED = x.shape[1]; in varlen NPREDICTED = - // user-supplied `max_seqlen` (upper bound on every cu_seqlens diff). - FLASHINFER_CHECK(npredicted == NPREDICTED, is_varlen ? "max_seqlen=" : "x.size(1)=", npredicted, - " must equal JIT NPREDICTED=", NPREDICTED); - FLASHINFER_CHECK(max_window == MAX_WINDOW, "old_x.size(1)=", max_window, " must equal JIT MAX_WINDOW=", MAX_WINDOW); - FLASHINFER_CHECK(npredicted <= max_window, "npredicted=", npredicted, " must be <= max_window=", max_window); - - // ── Validate state ── - CHECK_CUDA(state); - CHECK_DIM(4, state); - { - auto s = state.strides(); - auto sz = state.sizes(); - FLASHINFER_CHECK(s[3] == 1, "state dim 3 (dstate) must have stride 1, got ", s[3]); - FLASHINFER_CHECK( - s[2] == sz[3], "state dim 2 (dim) must be contiguous with dim 3, got stride ", s[2], " expected ", sz[3]); - FLASHINFER_CHECK(s[1] == sz[2] * sz[3], "state dim 1 (nheads) must be contiguous with dim 2, got stride ", s[1], - " expected ", sz[2] * sz[3]); - } - - // ── Validate x ── - // Non-varlen: shape (batch, NPREDICTED, nheads, dim). - // Varlen : shape (1, total_tokens, nheads, dim) — batch axis collapsed, - // token axis is the outer iteration. The kernel reads x via - // `bos * x_stride_token + …` so x_stride_token is the per-token - // stride in either layout (= nheads*dim for contig). - CHECK_CUDA(x); - CHECK_DIM(4, x); - if (is_varlen) - { - FLASHINFER_CHECK(x.size(0) == 1, "varlen: x.size(0)=", x.size(0), " must be 1"); - } - else - { - FLASHINFER_CHECK(x.size(0) == batch, "x.size(0)=", x.size(0), " must equal batch=", batch); - FLASHINFER_CHECK(x.size(1) == npredicted, "x.size(1)=", x.size(1), " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(x.size(2) == nheads, "x.size(2)=", x.size(2), " must equal nheads=", nheads); - FLASHINFER_CHECK(x.size(3) == dim, "x.size(3)=", x.size(3), " must equal dim=", dim); - CHECK_LAST_DIM_CONTIGUOUS(x); - FLASHINFER_CHECK(x.stride(2) == dim, "x.stride(2)=", x.stride(2), " must equal dim=", dim, - " ((nheads, dim) must be contiguous)"); - - // In varlen, all per-token tensors share the flattened token axis — use - // x.size(1) as the canonical total_tokens and cross-check the others below. - int64_t const total_tokens = is_varlen ? x.size(1) : 0; - - // ── Validate dt ── - CHECK_CUDA(dt); - CHECK_DIM(4, dt); - if (is_varlen) - { - FLASHINFER_CHECK(dt.size(0) == 1, "varlen: dt.size(0)=", dt.size(0), " must be 1"); - FLASHINFER_CHECK( - dt.size(1) == total_tokens, "varlen: dt.size(1)=", dt.size(1), " must equal x.size(1)=", total_tokens); - } - else - { - FLASHINFER_CHECK(dt.size(0) == batch, "dt.size(0)=", dt.size(0), " must equal batch=", batch); - FLASHINFER_CHECK(dt.size(1) == npredicted, "dt.size(1)=", dt.size(1), " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(dt.size(2) == nheads, "dt.size(2)=", dt.size(2), " must equal nheads=", nheads); - FLASHINFER_CHECK(dt.size(3) == dim, "dt.size(3)=", dt.size(3), " must equal dim=", dim); - FLASHINFER_CHECK(dt.stride(2) == 1, "dt.stride(2) must be 1 (tie_hdim), got ", dt.stride(2)); - FLASHINFER_CHECK(dt.stride(3) == 0, "dt.stride(3) must be 0 (tie_hdim), got ", dt.stride(3)); - - // ── Validate A: (nheads, dim, dstate) tie_hdim ── - CHECK_CUDA(A); - CHECK_DIM(3, A); - FLASHINFER_CHECK(A.size(0) == nheads, "A.size(0)=", A.size(0), " must equal nheads=", nheads); - FLASHINFER_CHECK(A.size(1) == dim, "A.size(1)=", A.size(1), " must equal dim=", dim); - FLASHINFER_CHECK(A.size(2) == dstate, "A.size(2)=", A.size(2), " must equal dstate=", dstate); - FLASHINFER_CHECK(A.stride(0) == 1, "A.stride(0) must be 1, got ", A.stride(0)); - FLASHINFER_CHECK(A.stride(1) == 0, "A.stride(1) must be 0 (tie_hdim), got ", A.stride(1)); - FLASHINFER_CHECK(A.stride(2) == 0, "A.stride(2) must be 0 (tie_hdim), got ", A.stride(2)); - - // ── Validate B ── - CHECK_CUDA(B); - CHECK_DIM(4, B); - if (is_varlen) - { - FLASHINFER_CHECK(B.size(0) == 1, "varlen: B.size(0)=", B.size(0), " must be 1"); - FLASHINFER_CHECK( - B.size(1) == total_tokens, "varlen: B.size(1)=", B.size(1), " must equal x.size(1)=", total_tokens); - } - else - { - FLASHINFER_CHECK(B.size(0) == batch, "B.size(0)=", B.size(0), " must equal batch=", batch); - FLASHINFER_CHECK(B.size(1) == npredicted, "B.size(1)=", B.size(1), " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(B.size(3) == dstate, "B.size(3)=", B.size(3), " must equal dstate=", dstate); - CHECK_LAST_DIM_CONTIGUOUS(B); - FLASHINFER_CHECK(B.stride(2) == dstate, "B.stride(2)=", B.stride(2), " must equal dstate=", dstate, - " ((ngroups, dstate) must be contiguous)"); - FLASHINFER_CHECK(nheads % ngroups == 0, "nheads=", nheads, " must be divisible by ngroups=", ngroups); - - // ── Validate C ── - CHECK_CUDA(C); - CHECK_DIM(4, C); - if (is_varlen) - { - FLASHINFER_CHECK(C.size(0) == 1, "varlen: C.size(0)=", C.size(0), " must be 1"); - FLASHINFER_CHECK( - C.size(1) == total_tokens, "varlen: C.size(1)=", C.size(1), " must equal x.size(1)=", total_tokens); - } - else - { - FLASHINFER_CHECK(C.size(0) == batch, "C.size(0)=", C.size(0), " must equal batch=", batch); - FLASHINFER_CHECK(C.size(1) == npredicted, "C.size(1)=", C.size(1), " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(C.size(2) == ngroups, "C.size(2)=", C.size(2), " must equal ngroups=", ngroups); - FLASHINFER_CHECK(C.size(3) == dstate, "C.size(3)=", C.size(3), " must equal dstate=", dstate); - CHECK_LAST_DIM_CONTIGUOUS(C); - FLASHINFER_CHECK(C.stride(2) == dstate, "C.stride(2)=", C.stride(2), " must equal dstate=", dstate, - " ((ngroups, dstate) must be contiguous)"); - - // ── Validate output ── - CHECK_CUDA(output); - CHECK_DIM(4, output); - if (is_varlen) - { - FLASHINFER_CHECK(output.size(0) == 1, "varlen: output.size(0)=", output.size(0), " must be 1"); - FLASHINFER_CHECK(output.size(1) == total_tokens, "varlen: output.size(1)=", output.size(1), - " must equal x.size(1)=", total_tokens); - } - else - { - FLASHINFER_CHECK(output.size(0) == batch, "output.size(0)=", output.size(0), " must equal batch=", batch); - FLASHINFER_CHECK( - output.size(1) == npredicted, "output.size(1)=", output.size(1), " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(output.size(2) == nheads, "output.size(2)=", output.size(2), " must equal nheads=", nheads); - FLASHINFER_CHECK(output.size(3) == dim, "output.size(3)=", output.size(3), " must equal dim=", dim); - CHECK_LAST_DIM_CONTIGUOUS(output); - FLASHINFER_CHECK(output.stride(2) == dim, "output.stride(2)=", output.stride(2), " must equal dim=", dim, - " ((nheads, dim) must be contiguous)"); - - // ── Validate cache tensors ── - // old_x: kernel uses `head * DIM + d_tile_off` → (nheads, dim) contig. - CHECK_CUDA(old_x); - CHECK_DIM(4, old_x); // (state_cache_size, MAX_WINDOW, nheads, dim) - FLASHINFER_CHECK(old_x.size(0) == state_cache_size, "old_x.size(0)=", old_x.size(0), - " must equal state_cache_size=", state_cache_size); - FLASHINFER_CHECK( - old_x.size(1) == max_window, "old_x.size(1)=", old_x.size(1), " must equal max_window=", max_window); - FLASHINFER_CHECK(old_x.size(2) == nheads, "old_x.size(2)=", old_x.size(2), " must equal nheads=", nheads); - FLASHINFER_CHECK(old_x.size(3) == dim, "old_x.size(3)=", old_x.size(3), " must equal dim=", dim); - CHECK_LAST_DIM_CONTIGUOUS(old_x); - FLASHINFER_CHECK(old_x.stride(2) == dim, "old_x.stride(2)=", old_x.stride(2), " must equal dim=", dim, - " ((nheads, dim) must be contiguous)"); - - // old_B: kernel uses `group_idx * DSTATE` → (ngroups, dstate) contig. - CHECK_CUDA(old_B); - CHECK_DIM(5, old_B); // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) - FLASHINFER_CHECK(old_B.size(0) == state_cache_size, "old_B.size(0)=", old_B.size(0), - " must equal state_cache_size=", state_cache_size); - FLASHINFER_CHECK(old_B.size(1) == 2, "old_B.size(1) must be 2 (double-buffered), got ", old_B.size(1)); - FLASHINFER_CHECK( - old_B.size(2) == max_window, "old_B.size(2)=", old_B.size(2), " must equal max_window=", max_window); - FLASHINFER_CHECK(old_B.size(3) == ngroups, "old_B.size(3)=", old_B.size(3), " must equal ngroups=", ngroups); - FLASHINFER_CHECK(old_B.size(4) == dstate, "old_B.size(4)=", old_B.size(4), " must equal dstate=", dstate); - CHECK_LAST_DIM_CONTIGUOUS(old_B); - FLASHINFER_CHECK(old_B.stride(3) == dstate, "old_B.stride(3)=", old_B.stride(3), " must equal dstate=", dstate, - " ((ngroups, dstate) must be contiguous)"); - - // old_dt: kernel only assumes last-dim contig (head row). - CHECK_CUDA(old_dt); - CHECK_DIM(4, old_dt); // (state_cache_size, 2, nheads, MAX_WINDOW) - FLASHINFER_CHECK(old_dt.size(0) == state_cache_size, "old_dt.size(0)=", old_dt.size(0), - " must equal state_cache_size=", state_cache_size); - FLASHINFER_CHECK(old_dt.size(1) == 2, "old_dt.size(1) must be 2, got ", old_dt.size(1)); - FLASHINFER_CHECK(old_dt.size(2) == nheads, "old_dt.size(2)=", old_dt.size(2), " must equal nheads=", nheads); - FLASHINFER_CHECK( - old_dt.size(3) == max_window, "old_dt.size(3)=", old_dt.size(3), " must equal max_window=", max_window); - CHECK_LAST_DIM_CONTIGUOUS(old_dt); - - // old_cumAdt: same as old_dt. - CHECK_CUDA(old_cumAdt); - CHECK_DIM(4, old_cumAdt); // (state_cache_size, 2, nheads, MAX_WINDOW) - FLASHINFER_CHECK(old_cumAdt.size(0) == state_cache_size, "old_cumAdt.size(0)=", old_cumAdt.size(0), - " must equal state_cache_size=", state_cache_size); - FLASHINFER_CHECK(old_cumAdt.size(1) == 2, "old_cumAdt.size(1) must be 2, got ", old_cumAdt.size(1)); - FLASHINFER_CHECK( - old_cumAdt.size(2) == nheads, "old_cumAdt.size(2)=", old_cumAdt.size(2), " must equal nheads=", nheads); - FLASHINFER_CHECK(old_cumAdt.size(3) == max_window, "old_cumAdt.size(3)=", old_cumAdt.size(3), - " must equal max_window=", max_window); - CHECK_LAST_DIM_CONTIGUOUS(old_cumAdt); - - CHECK_CUDA(cache_buf_idx); - CHECK_DIM(1, cache_buf_idx); - FLASHINFER_CHECK(cache_buf_idx.size(0) == state_cache_size, "cache_buf_idx.size(0)=", cache_buf_idx.size(0), - " must equal state_cache_size=", state_cache_size); - CHECK_CONTIGUOUS(cache_buf_idx); - - CHECK_CUDA(prev_num_accepted); - CHECK_DIM(1, prev_num_accepted); - FLASHINFER_CHECK(prev_num_accepted.size(0) == state_cache_size, - "prev_num_accepted.size(0)=", prev_num_accepted.size(0), " must equal state_cache_size=", state_cache_size); - CHECK_CONTIGUOUS(prev_num_accepted); - - // ── Validate optional D ── - if (D.has_value()) - { - auto& Dv = D.value(); - CHECK_CUDA(Dv); - CHECK_DIM(2, Dv); - FLASHINFER_CHECK(Dv.size(0) == nheads, "D.size(0)=", Dv.size(0), " must equal nheads=", nheads); - FLASHINFER_CHECK(Dv.size(1) == dim, "D.size(1)=", Dv.size(1), " must equal dim=", dim); - FLASHINFER_CHECK(Dv.stride(0) == 1, "D.stride(0) must be 1 (tie_hdim), got ", Dv.stride(0)); - FLASHINFER_CHECK(Dv.stride(1) == 0, "D.stride(1) must be 0 (tie_hdim), got ", Dv.stride(1)); - } - - // ── Validate optional dt_bias ── - if (dt_bias.has_value()) - { - auto& db = dt_bias.value(); - CHECK_CUDA(db); - CHECK_DIM(2, db); - FLASHINFER_CHECK(db.size(0) == nheads, "dt_bias.size(0)=", db.size(0), " must equal nheads=", nheads); - FLASHINFER_CHECK(db.size(1) == dim, "dt_bias.size(1)=", db.size(1), " must equal dim=", dim); - FLASHINFER_CHECK(db.stride(0) == 1, "dt_bias.stride(0) must be 1 (tie_hdim), got ", db.stride(0)); - FLASHINFER_CHECK(db.stride(1) == 0, "dt_bias.stride(1) must be 0 (tie_hdim), got ", db.stride(1)); - } - - // ── Validate optional z: same layout/contig rules as x ── - if (z.has_value()) - { - auto& zv = z.value(); - CHECK_CUDA(zv); - CHECK_DIM(4, zv); - if (is_varlen) - { - FLASHINFER_CHECK(zv.size(0) == 1, "varlen: z.size(0)=", zv.size(0), " must be 1"); - FLASHINFER_CHECK( - zv.size(1) == total_tokens, "varlen: z.size(1)=", zv.size(1), " must equal x.size(1)=", total_tokens); - } - else - { - FLASHINFER_CHECK(zv.size(0) == batch, "z.size(0)=", zv.size(0), " must equal batch=", batch); - FLASHINFER_CHECK(zv.size(1) == npredicted, "z.size(1)=", zv.size(1), " must equal npredicted=", npredicted); - } - FLASHINFER_CHECK(zv.size(2) == nheads, "z.size(2)=", zv.size(2), " must equal nheads=", nheads); - FLASHINFER_CHECK(zv.size(3) == dim, "z.size(3)=", zv.size(3), " must equal dim=", dim); - CHECK_LAST_DIM_CONTIGUOUS(zv); - FLASHINFER_CHECK(zv.stride(2) == dim, "z.stride(2)=", zv.stride(2), " must equal dim=", dim, - " ((nheads, dim) must be contiguous)"); - } - - // ── Validate optional state_batch_indices ── - if (state_batch_indices.has_value()) - { - auto& sbi = state_batch_indices.value(); - CHECK_CUDA(sbi); - CHECK_DIM(1, sbi); - FLASHINFER_CHECK( - sbi.size(0) == batch, "state_batch_indices.size(0)=", sbi.size(0), " must equal batch=", batch); - CHECK_CONTIGUOUS(sbi); - } - - // ── Validate optional state_scale: (state_cache_size, nheads, dim) ── - // Inner two dims (nheads, dim) must be contiguous; only batch stride is - // parameterized in the params struct. - if (state_scale.has_value()) - { - auto const& ss = state_scale.value(); - CHECK_CUDA(ss); - CHECK_DIM(3, ss); - FLASHINFER_CHECK(ss.size(0) == state_cache_size, "state_scale.size(0)=", ss.size(0), - " must equal state_cache_size=", state_cache_size); - FLASHINFER_CHECK(ss.size(1) == nheads, "state_scale.size(1)=", ss.size(1), " must equal nheads=", nheads); - FLASHINFER_CHECK(ss.size(2) == dim, "state_scale.size(2)=", ss.size(2), " must equal dim=", dim); - FLASHINFER_CHECK(ss.stride(2) == 1, "state_scale.stride(2) must be 1, got ", ss.stride(2)); - FLASHINFER_CHECK(ss.stride(1) == dim, "state_scale.stride(1)=", ss.stride(1), " must equal dim=", dim, - " ((nheads, dim) must be contiguous)"); - } - - // ── Dtype consistency ── - // input_dtype = x.dtype; all activation tensors (B, C, output, z, old_x, - // old_B) and the state cache's "input-side" mirrors must match it. - // weight_dtype = D.dtype = dt_bias.dtype (kernel template sees one - // weight_t for both). - // Cache scalar tensors have fixed dtypes hardcoded in the kernel. - { - auto input_dtype = x.dtype(); - FLASHINFER_CHECK(B.dtype() == input_dtype, "B.dtype must match x.dtype"); - FLASHINFER_CHECK(C.dtype() == input_dtype, "C.dtype must match x.dtype"); - FLASHINFER_CHECK(output.dtype() == input_dtype, "output.dtype must match x.dtype"); - FLASHINFER_CHECK(old_x.dtype() == input_dtype, "old_x.dtype must match x.dtype"); - FLASHINFER_CHECK(old_B.dtype() == input_dtype, "old_B.dtype must match x.dtype"); - if (z.has_value()) - { - FLASHINFER_CHECK(z.value().dtype() == input_dtype, "z.dtype must match x.dtype"); - } - if (D.has_value() && dt_bias.has_value()) - { - FLASHINFER_CHECK(D.value().dtype() == dt_bias.value().dtype(), - "D.dtype must equal dt_bias.dtype (kernel uses a single weight_t)"); - } - // old_dt / old_cumAdt are produced by this same kernel in f32 and - // consumed back in f32 on the next call. - FLASHINFER_CHECK(old_dt.dtype().code == kDLFloat && old_dt.dtype().bits == 32, "old_dt must be float32"); - FLASHINFER_CHECK( - old_cumAdt.dtype().code == kDLFloat && old_cumAdt.dtype().bits == 32, "old_cumAdt must be float32"); - // Index tensors used by the kernel as int32 scalars. - FLASHINFER_CHECK( - cache_buf_idx.dtype().code == kDLInt && cache_buf_idx.dtype().bits == 32, "cache_buf_idx must be int32"); - FLASHINFER_CHECK(prev_num_accepted.dtype().code == kDLInt && prev_num_accepted.dtype().bits == 32, - "prev_num_accepted must be int32"); - if (state_batch_indices.has_value()) - { - auto sbi_dt = state_batch_indices.value().dtype(); - FLASHINFER_CHECK(sbi_dt.code == kDLInt && (sbi_dt.bits == 32 || sbi_dt.bits == 64), - "state_batch_indices must be int32 or int64"); - } - if (state_scale.has_value()) - { - auto ss_dt = state_scale.value().dtype(); - FLASHINFER_CHECK(ss_dt.code == kDLFloat && ss_dt.bits == 32, "state_scale must be float32"); - } - // Quantized state dtypes (int8, fp8_e4m3fn, ...) require a state_scale - // tensor; non-quantized dtypes must not pass one. Mirrors the Python - // wrapper assertion and matches the kernel's compile-time - // `state_scale_t == void` gating. - { - auto sd = state.dtype(); - bool const is_int8 = (sd.code == kDLInt && sd.bits == 8); - bool const is_fp8 = (sd.code == kDLFloat8_e4m3fn && sd.bits == 8); - bool const is_quantized_state = is_int8 || is_fp8; - if (is_quantized_state) - { - FLASHINFER_CHECK(state_scale.has_value(), - "Quantized state.dtype (int8/fp8_e4m3fn) requires a state_scale tensor " - "of shape (state_cache_size, nheads, dim) and dtype float32"); - // The 8-bit replay path uses Layout<_4, _1> (M-shard per warp) which - // needs per-warp M = D_PER_CTA / 4 >= 16 (m16n8 atom M). This forces - // D_PER_CTA >= 64, i.e. d_split == 1. - FLASHINFER_CHECK(d_split == 1, - "Quantized state.dtype (int8/fp8_e4m3fn) requires d_split=1 (got d_split=", d_split, - "); the M-shard-per-warp replay layout needs D_PER_CTA / 4 >= 16."); - } - else - { - FLASHINFER_CHECK(!state_scale.has_value(), - "state_scale must be None for non-quantized state.dtype " - "(allowed quantized dtypes: {int8, fp8_e4m3fn})"); - } - } - } - - // ── Populate params ── - CheckpointingSsuParams p; - - // ── Validate d_split (v12 §59) ── - // Allowed for v12: {1, 2}. d_split=4 deferred to v12.x (needs warp-count - // restructure — output MMA `_1×4` layout requires D_PER_CTA ≥ 32). - FLASHINFER_CHECK( - d_split == 1 || d_split == 2, "d_split=", d_split, " must be one of {1, 2} (d_split=4 is deferred to v12.x)"); - FLASHINFER_CHECK(dim % d_split == 0, "dim=", dim, " must be divisible by d_split=", d_split); - FLASHINFER_CHECK(dim / d_split >= 32, "d_split=", d_split, " gives D_PER_CTA=", dim / d_split, - " < 32 (output MMA m16n8 atom floor with _1×4 warp layout)"); - - p.batch = batch; - p.nheads = nheads; - p.dim = dim; - p.dstate = dstate; - p.ngroups = ngroups; - p.state_cache_size = state_cache_size; - p.npredicted = npredicted; - p.max_window = max_window; - p.pad_slot_id = pad_slot_id; - p.d_split = static_cast(d_split); - p.dt_softplus = dt_softplus; - - // Pointers - p.state = state.data_ptr(); - p.x = const_cast(x.data_ptr()); - p.dt = const_cast(dt.data_ptr()); - p.A = const_cast(A.data_ptr()); - p.B = const_cast(B.data_ptr()); - p.C = const_cast(C.data_ptr()); - p.output = output.data_ptr(); - - p.old_x = old_x.data_ptr(); - p.old_B = const_cast(old_B.data_ptr()); - p.old_dt = const_cast(old_dt.data_ptr()); - p.old_cumAdt = const_cast(old_cumAdt.data_ptr()); - p.cache_buf_idx = const_cast(cache_buf_idx.data_ptr()); - p.prev_num_accepted = const_cast(prev_num_accepted.data_ptr()); - - if (D.has_value()) - p.D = const_cast(D.value().data_ptr()); - if (z.has_value()) - { - p.z = const_cast(z.value().data_ptr()); - // Same seq-dim selection as the rest of the batch-side tensors below. - p.z_stride_seq = z.value().stride(is_varlen ? 1 : 0); - p.z_stride_token = z.value().stride(1); - } - if (dt_bias.has_value()) - p.dt_bias = const_cast(dt_bias.value().data_ptr()); - if (state_batch_indices.has_value()) - p.state_batch_indices = const_cast(state_batch_indices.value().data_ptr()); - if (is_varlen) - { - p.cu_seqlens = const_cast(cu_seqlens.value().data_ptr()); - } - if (state_scale.has_value()) - { - p.state_scale = state_scale.value().data_ptr(); - p.state_scale_stride_seq = state_scale.value().stride(0); - } - if (rand_seed.has_value()) - { - auto const& rs = rand_seed.value(); - CHECK_CUDA(rs); - FLASHINFER_CHECK(rs.numel() == 1, "rand_seed must be single-element, got numel=", rs.numel()); - FLASHINFER_CHECK(rs.dtype().code == kDLInt && rs.dtype().bits == 64, "rand_seed must be int64"); - p.rand_seed = static_cast(rs.data_ptr()); - } - - // Strides - p.state_stride_seq = state.stride(0); - - // `*_stride_seq` is the outer iteration stride. Non-varlen iterates over - // dim 0 (per-batch), varlen iterates over dim 1 (per-token) — sequences - // are packed into a single batch in the (1, total_tokens, ...) layout. - // The kernel uses one formula `seq * *_stride_seq` for both modes. - int const seq_dim = is_varlen ? 1 : 0; - p.x_stride_seq = x.stride(seq_dim); - p.x_stride_token = x.stride(1); - p.dt_stride_seq = dt.stride(seq_dim); - p.dt_stride_token = dt.stride(1); - p.B_stride_seq = B.stride(seq_dim); - p.B_stride_token = B.stride(1); - p.C_stride_seq = C.stride(seq_dim); - p.C_stride_token = C.stride(1); - p.out_stride_seq = output.stride(seq_dim); - p.out_stride_token = output.stride(1); - - p.old_x_stride_seq = old_x.stride(0); - p.old_x_stride_token = old_x.stride(1); - p.old_B_stride_seq = old_B.stride(0); - p.old_B_stride_dbuf = old_B.stride(1); - p.old_B_stride_token = old_B.stride(2); - p.old_dt_stride_seq = old_dt.stride(0); - p.old_dt_stride_dbuf = old_dt.stride(1); - p.old_dt_stride_head = old_dt.stride(2); - p.old_cumAdt_stride_seq = old_cumAdt.stride(0); - p.old_cumAdt_stride_dbuf = old_cumAdt.stride(1); - p.old_cumAdt_stride_head = old_cumAdt.stride(2); - - // Launch - ffi::CUDADeviceGuard device_guard(state.device().device_id); - const cudaStream_t stream = get_stream(state.device()); - - launchCheckpointingSsu(p, stream); -} - -} // namespace flashinfer::mamba::checkpointing diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_customize_config.jinja b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_customize_config.jinja deleted file mode 100644 index 61bf35cf17ca..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_customize_config.jinja +++ /dev/null @@ -1,38 +0,0 @@ -#pragma once -#include -#include -#include -#include - -using state_t = {{ state_dtype }}; -using input_t = {{ input_dtype }}; -// dt accepted in its native dtype (e.g. bf16) — converted to f32 internally. -// Eliminates the need for a separate dtype conversion kernel launch on the host side. -using dt_t = {{ dt_dtype }}; -using weight_t = {{ weight_dtype }}; -using matrixA_t = {{ matrixA_dtype }}; -using stateIndex_t = {{ stateIndex_dtype }}; -// Type for block-scale decode factors (e.g. float, __half). -// void = no scaling (state_t is used as-is). -using state_scale_t = {{ state_scale_type }}; - -constexpr int DIM = {{ dim }}; -constexpr int DSTATE = {{ dstate }}; -constexpr int NPREDICTED = {{ npredicted }}; -constexpr int MAX_WINDOW = {{ max_window }}; -// nheads / ngroups — JIT-stamped so the kernel compiles only one -// HEADS_PER_GROUP specialization per .so (was 7 via `dispatchRatio`). -// The wrapper computes this from the runtime tensors and selects the -// matching JIT URI; the launcher reads HEADS_PER_GROUP directly without -// a runtime dispatch step. -constexpr int HEADS_PER_GROUP = {{ heads_per_group }}; -// Philox PRNG rounds for stochastic rounding of fp16 state stores. -// 0 = no stochastic rounding; typical value = 10. -constexpr int PHILOX_ROUNDS = {{ philox_rounds }}; -// Programmatic Dependent Launch. When true, the kernel emits the -// griddepcontrol.{wait,launch_dependents} PTX and the load is split around -// `gdc_wait` for cache-load-during-wait overlap. When false, a single-pass -// load_data path is used (no PDL barriers) — matches v21.0 register profile -// and load order. Stamped as a JIT URI key so each (enable_pdl) value -// compiles its own .so; no runtime branch in the kernel binary. -constexpr bool ENABLE_PDL = {{ enable_pdl }}; diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu deleted file mode 100644 index 508c52666633..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_jit_binding.cu +++ /dev/null @@ -1,51 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -#include "tvm_ffi_utils.h" - -using tvm::ffi::Optional; - -namespace flashinfer::mamba::checkpointing -{ - -void checkpointing_ssu(TensorView state, // (cache, nheads, dim, dstate) - TensorView x, // 4D (batch, T, nheads, dim) or 4D (1, total_tokens, nheads, dim) under varlen - TensorView dt, // (batch, T, nheads, dim) tie_hdim / (1, total_tokens, nheads, dim) under varlen - TensorView A, // (nheads, dim, dstate) tie_hdim - TensorView B, // (batch, T, ngroups, dstate) / (1, total_tokens, ngroups, dstate) under varlen - TensorView C, // same as B - TensorView output, // same layout as x - // Cache tensors - TensorView old_x, // (cache, T, nheads, dim) - TensorView old_B, // (cache, 2, T, ngroups, dstate) - TensorView old_dt, // (cache, 2, nheads, T) f32 - TensorView old_cumAdt, // (cache, 2, nheads, T) f32 - TensorView cache_buf_idx, // (cache,) int32 - TensorView prev_num_accepted, // (cache,) int32 - // Optional tensors - Optional D, // (nheads, dim) - Optional z, // same layout as x - Optional dt_bias, // (nheads, dim) tie_hdim - bool dt_softplus, - Optional state_batch_indices, // (batch,) int32 - int64_t pad_slot_id, - Optional state_scale, // (cache, nheads, dim) f32 - Optional rand_seed, // single int64 - int64_t d_split, // v12 §59: per-head DIM split factor (1, 2, or 4) - Optional cu_seqlens); // (batch+1,) int32 — varlen mode - -} // namespace flashinfer::mamba::checkpointing - -TVM_FFI_DLL_EXPORT_TYPED_FUNC(checkpointing_ssu, flashinfer::mamba::checkpointing::checkpointing_ssu); diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu deleted file mode 100644 index cbb945631093..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/csrc/checkpointing_ssu_kernel_inst.cu +++ /dev/null @@ -1,14 +0,0 @@ -// clang-format off -#include "checkpointing_ssu_config.inc" -#include -#include - -// clang-format on - -namespace flashinfer::mamba::checkpointing -{ - -template void launchCheckpointingSsu( - CheckpointingSsuParams&, cudaStream_t); - -} // namespace flashinfer::mamba::checkpointing diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h deleted file mode 100644 index 4521f68c5086..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/exception.h +++ /dev/null @@ -1,126 +0,0 @@ -/* - * Copyright (c) 2024 by FlashInfer team. - * - * 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. - */ -#ifndef FLASHINFER_EXCEPTION_H_ -#define FLASHINFER_EXCEPTION_H_ - -#include -#include -#include - -#define FLASHINFER_ERROR(message) throw flashinfer::Error(__FUNCTION__, __FILE__, __LINE__, message) - -// Base case for empty arguments -inline void write_to_stream(std::ostringstream& oss) -{ - // No-op for empty arguments -} - -template -void write_to_stream(std::ostringstream& oss, T&& val) -{ - oss << std::forward(val); -} - -template -void write_to_stream(std::ostringstream& oss, T&& val, Args&&... args) -{ - oss << std::forward(val) << " "; - write_to_stream(oss, std::forward(args)...); -} - -// Helper macro to handle empty __VA_ARGS__ -#define FLASHINFER_CHECK_IMPL(condition, message) \ - if (!(condition)) \ - { \ - FLASHINFER_ERROR(message); \ - } - -// Main macro that handles both cases -#define FLASHINFER_CHECK(condition, ...) \ - do \ - { \ - if (!(condition)) \ - { \ - std::ostringstream oss; \ - write_to_stream(oss, ##__VA_ARGS__); \ - std::string msg = oss.str(); \ - if (msg.empty()) \ - { \ - msg = "Check failed: " #condition; \ - } \ - FLASHINFER_ERROR(msg); \ - } \ - } while (0) - -// Warning macro -#define FLASHINFER_WARN(...) \ - do \ - { \ - std::ostringstream oss; \ - write_to_stream(oss, ##__VA_ARGS__); \ - std::string msg = oss.str(); \ - if (msg.empty()) \ - { \ - msg = "Warning triggered"; \ - } \ - flashinfer::Warning(__FUNCTION__, __FILE__, __LINE__, msg).emit(); \ - } while (0) - -namespace flashinfer -{ -class Error : public std::exception -{ -private: - std::string message_; - -public: - Error(std::string const& func, std::string const& file, int line, std::string const& message) - { - std::ostringstream oss; - oss << "Error in function '" << func << "' " - << "at " << file << ":" << line << ": " << message; - message_ = oss.str(); - } - - virtual char const* what() const noexcept override - { - return message_.c_str(); - } -}; - -class Warning -{ -private: - std::string message_; - -public: - Warning(std::string const& func, std::string const& file, int line, std::string const& message) - { - std::ostringstream oss; - oss << "Warning in function '" << func << "' " - << "at " << file << ":" << line << ": " << message; - message_ = oss.str(); - } - - void emit() const - { - std::cerr << message_ << std::endl; - } -}; - -} // namespace flashinfer - -#endif // FLASHINFER_EXCEPTION_H_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh deleted file mode 100644 index c77e2afcf8ab..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/checkpointing_ssu.cuh +++ /dev/null @@ -1,153 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -#ifndef FLASHINFER_MAMBA_CHECKPOINTING_SSU_CUH_ -#define FLASHINFER_MAMBA_CHECKPOINTING_SSU_CUH_ - -#include - -namespace flashinfer::mamba::checkpointing -{ - -struct CheckpointingSsuParams -{ - uint32_t batch{}, nheads{}, dim{}, dstate{}, ngroups{}; - uint32_t state_cache_size{}; - uint32_t npredicted{}; - uint32_t max_window{}; - int32_t pad_slot_id{-1}; - - // v12 §59: per-head DIM split factor. Must be one of {1, 2, 4}. The host - // launcher dispatches to a kernel template specialized on this value; the - // kernel cross-checks via assert(params.d_split == D_SPLIT). - int32_t d_split{1}; - - bool dt_softplus{false}; - - // Note: Programmatic Dependent Launch is JIT-stamped via the `ENABLE_PDL` - // constexpr (see checkpointing_ssu_customize_config.jinja). Each .so has - // its PDL mode baked in; no runtime field needed. - - // ── Tensor pointers ── - void* __restrict__ state{nullptr}; // (state_cache_size, nheads, dim, dstate) - void* __restrict__ x{nullptr}; // (batch, NPREDICTED, nheads, dim) - void* __restrict__ dt{nullptr}; // (batch, NPREDICTED, nheads, dim) tie_hdim - void* __restrict__ A{nullptr}; // (nheads, dim, dstate) tie_hdim - void* __restrict__ B{nullptr}; // (batch, NPREDICTED, ngroups, dstate) - void* __restrict__ C{nullptr}; // (batch, NPREDICTED, ngroups, dstate) - void* __restrict__ D{nullptr}; // (nheads, dim), optional - void* __restrict__ z{nullptr}; // (batch, NPREDICTED, nheads, dim), optional - void* __restrict__ dt_bias{nullptr}; // (nheads, dim) tie_hdim, optional - void* __restrict__ output{nullptr}; // (batch, NPREDICTED, nheads, dim) - - // ── Cache tensors for incremental replay ── - void* __restrict__ old_x{nullptr}; // (state_cache_size, MAX_WINDOW, nheads, dim) single-buffered - void* __restrict__ old_B{nullptr}; // (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) double-buffered - void* __restrict__ old_dt{nullptr}; // (state_cache_size, 2, nheads, MAX_WINDOW) double-buffered, f32 - void* __restrict__ old_cumAdt{nullptr}; // (state_cache_size, 2, nheads, MAX_WINDOW) double-buffered, f32 - void* __restrict__ cache_buf_idx{nullptr}; // (state_cache_size,) int32 - void* __restrict__ prev_num_accepted{nullptr}; // (state_cache_size,) int32 - - // ── Index tensors ── - void* __restrict__ state_batch_indices{nullptr}; // (batch,) optional - - // ── Varlen (v20): packed inputs ── - // When non-null, `x/dt/B/C/z/out` are laid out as - // `(1, total_tokens, nheads, dim)` / `(1, total_tokens, ngroups, dstate)` - // and `cu_seqlens[i]` gives the token-axis base of sequence i. - // `seq_len_i = cu_seqlens[i+1] - cu_seqlens[i]`. Kernel dispatch on - // `cu_seqlens != nullptr` selects a `VARLEN=true` template. - // - // The `*_stride_seq` fields below already encode the outer iteration - // stride for both modes — the wrapper sets them to: - // non-varlen: `tensor.stride(0)` (per-batch) - // varlen : `tensor.stride(1)` (per-token, since sequences are packed - // into a single batch of total_tokens) - // so the kernel uses one formula `seq * *_stride_seq` regardless of mode. - void* __restrict__ cu_seqlens{nullptr}; // (batch+1,) int32, optional - - // ── Block-scale decode factors for quantized state ── - void* __restrict__ state_scale{nullptr}; // float32: (state_cache_size, nheads, dim) - - // ── Philox PRNG seed for stochastic rounding ── - int64_t const* rand_seed{nullptr}; - - // ── Strides ── - // state: (state_cache_size, nheads, dim, dstate) — inner 3 dims contiguous - int64_t state_stride_seq{}; - - // For the six batch-side tensors (x, dt, B, C, out, z), `*_stride_seq` - // is the outer iteration stride — per-batch in non-varlen, per-token in - // varlen. `*_stride_token` is the inner per-row (T-axis) stride, same - // in both modes. - - // x: (batch, NPREDICTED, nheads, dim) [non-varlen] / (1, total_tokens, nheads, dim) [varlen] - int64_t x_stride_seq{}; - int64_t x_stride_token{}; - - // dt: (batch, NPREDICTED, nheads, dim) — tie_hdim (stride_dim=0) - int64_t dt_stride_seq{}; - int64_t dt_stride_token{}; - - // B: (batch, NPREDICTED, ngroups, dstate) - int64_t B_stride_seq{}; - int64_t B_stride_token{}; - - // C: (batch, NPREDICTED, ngroups, dstate) - int64_t C_stride_seq{}; - int64_t C_stride_token{}; - - // output: (batch, NPREDICTED, nheads, dim) - int64_t out_stride_seq{}; - int64_t out_stride_token{}; - - // z: (batch, NPREDICTED, nheads, dim) - int64_t z_stride_seq{}; - int64_t z_stride_token{}; - - // old_x: (state_cache_size, MAX_WINDOW, nheads, dim) — single-buffered - int64_t old_x_stride_seq{}; - int64_t old_x_stride_token{}; - - // old_B: (state_cache_size, 2, MAX_WINDOW, ngroups, dstate) — double-buffered - int64_t old_B_stride_seq{}; - int64_t old_B_stride_dbuf{}; - int64_t old_B_stride_token{}; - - // old_dt: (state_cache_size, 2, nheads, MAX_WINDOW) — double-buffered, MAX_WINDOW contiguous - int64_t old_dt_stride_seq{}; - int64_t old_dt_stride_dbuf{}; - int64_t old_dt_stride_head{}; - - // old_cumAdt: (state_cache_size, 2, nheads, MAX_WINDOW) — double-buffered, MAX_WINDOW contiguous - int64_t old_cumAdt_stride_seq{}; - int64_t old_cumAdt_stride_dbuf{}; - int64_t old_cumAdt_stride_head{}; - - // state_scale: (state_cache_size, nheads, dim) - int64_t state_scale_stride_seq{}; -}; - -// Forward declaration — defined in kernel_checkpointing_ssu.cuh. -// `launchCheckpointingSsu` is the public dispatcher: it reads -// `params.d_split` and routes to the matching `launchCheckpointingSsuImpl` -// specialization (v12 §59). Caller side stays single-entry. -template -void launchCheckpointingSsu(CheckpointingSsuParams& params, cudaStream_t stream); - -} // namespace flashinfer::mamba::checkpointing - -#endif // FLASHINFER_MAMBA_CHECKPOINTING_SSU_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh deleted file mode 100644 index 476458a24c5a..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/common.cuh +++ /dev/null @@ -1,234 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -#ifndef FLASHINFER_MAMBA_COMMON_CUH_ -#define FLASHINFER_MAMBA_COMMON_CUH_ - -#include - -#include -#include -#include - -#include - -namespace flashinfer::mamba -{ - -constexpr unsigned warpSize = 32; - -// ============================================================================= -// Common types and utilities -// ============================================================================= - -// Largest power of 2 that divides v (i.e. v & -v). Returns 1 when v == 0. -inline constexpr unsigned largestPow2Divisor(unsigned v) -{ - return v ? (v & (~v + 1)) : 1; -} - -// Simple packed vector type for loading N elements of type T. -// Alignment is the largest power-of-2 factor of the total byte size, -// so it is always valid even when N * sizeof(T) is not a power of 2 (e.g. 3 × 2 = 6). -template -struct alignas(largestPow2Divisor(N * sizeof(T))) PackedAligned -{ - T val[N]; - static constexpr int count = N; - using dtype = T; -}; - -template -__device__ __forceinline__ auto make_zeros() -> load_t -{ - load_t ret{}; -#pragma unroll - for (int i = 0; i < ret.count; i++) - ret.val[i] = typename load_t::dtype{}; // default initialization - return ret; -}; - -// Computes the vector load size that ensures full warp utilization. -// Avoids cases like: dstate=64, load_t = sizeof(float4)/sizeof(f16), warpsize=32 (32 * 8 > 64) -// in which case a part of the warp would be idle. -template -inline constexpr auto getVectorLoadSizeForFullUtilization() -> unsigned -{ - static_assert(sizeof(float4) >= sizeof(T)); - constexpr unsigned maxHardwareLoadSize = sizeof(float4) / sizeof(T); - constexpr unsigned maxLogicalLoadSize = (unsigned) DSTATE / warpSize; - return maxHardwareLoadSize < maxLogicalLoadSize ? maxHardwareLoadSize : maxLogicalLoadSize; -} - -__device__ __forceinline__ float warpReduceSum(float val) -{ - for (int s = warpSize / 2; s > 0; s /= 2) - { - val += __shfl_down_sync(UINT32_MAX, val, s); - } - return val; -} - -__device__ __forceinline__ float warpReduceMax(float val) -{ - for (int s = warpSize / 2; s > 0; s /= 2) - { - val = max(val, __shfl_down_sync(UINT32_MAX, val, s)); - } - return val; -} - -__forceinline__ __device__ float softplus(float x) -{ - return __logf(1.f + __expf(x)); -} - -__device__ __forceinline__ float thresholded_softplus(float dt_value) -{ - constexpr float threshold = 20.f; - return (dt_value <= threshold) ? softplus(dt_value) : dt_value; -} - -// ============================================================================= -// Dispatch helpers -// ============================================================================= - -// Format an integer_sequence as a comma-separated string for error messages -template -std::string format_sequence(std::integer_sequence) -{ - std::ostringstream oss; - bool first = true; - ((oss << (first ? (first = false, "") : ", ") << Values), ...); - return oss.str(); -} - -// Helper function to dispatch dim and dstate with a kernel launcher -template -void dispatchDimDstate(ParamsType& params, std::integer_sequence dims_seq, - std::integer_sequence dstates_seq, KernelLauncher&& launcher) -{ - auto dispatch_dstate = [&]() - { - auto try_dstate = [&]() - { - if (params.dstate == DSTATE) - { - launcher.template operator()(); - return true; - } - return false; - }; - bool dispatched = (try_dstate.template operator()() || ...); - FLASHINFER_CHECK(dispatched, "Unsupported dstate value: ", params.dstate, - ".\nSupported values: ", format_sequence(dstates_seq)); - }; - - auto try_dim = [&]() - { - if (params.dim == DIM) - { - dispatch_dstate.template operator()(); - return true; - } - return false; - }; - - bool dim_dispatched = (try_dim.template operator()() || ...); - FLASHINFER_CHECK( - dim_dispatched, "Unsupported dim value: ", params.dim, ".\nSupported values: ", format_sequence(dims_seq)); -} - -// Helper function to dispatch ratio with a kernel launcher -template -void dispatchRatio( - ParamsType& params, std::integer_sequence ratios_seq, KernelLauncher&& launcher) -{ - auto try_ratio = [&]() - { - if (params.nheads / params.ngroups == RATIO) - { - launcher.template operator()(); - return true; - } - return false; - }; - - bool ratio_dispatched = (try_ratio.template operator()() || ...); - FLASHINFER_CHECK(ratio_dispatched, "Unsupported nheads/ngroups ratio: ", params.nheads / params.ngroups, - ".\nSupported values: ", format_sequence(ratios_seq)); -} - -// Helper function to dispatch dim, dstate, and ntokens_mtp with a kernel launcher -// Reuses dispatchDimDstate by wrapping the launcher to add token dispatch -template -void dispatchDimDstateTokens(ParamsType& params, std::integer_sequence dims_seq, - std::integer_sequence dstates_seq, std::integer_sequence tokens_seq, - KernelLauncher&& launcher) -{ - // Wrap the launcher to add token dispatch as the innermost level - auto dim_dstate_launcher = [&]() - { - auto try_tokens = [&]() - { - if (params.ntokens_mtp == TOKENS_MTP) - { - launcher.template operator()(); - return true; - } - return false; - }; - bool dispatched = (try_tokens.template operator()() || ...); - FLASHINFER_CHECK(dispatched, "Unsupported ntokens_mtp value: ", params.ntokens_mtp, - ".\nSupported values: ", format_sequence(tokens_seq)); - }; - - dispatchDimDstate(params, dims_seq, dstates_seq, dim_dstate_launcher); -} - -// ============================================================================= -// Alignment checks -// ============================================================================= - -// Check alignment for common input variables (x, z, B, C) -// Works for both STP (SelectiveStateUpdateParams) and MTP (SelectiveStateMTPParams) -template -void check_ptr_alignment_input_vars(ParamsType const& params) -{ - using load_input_t = PackedAligned; - FLASHINFER_CHECK(reinterpret_cast(params.x) % sizeof(load_input_t) == 0, "x pointer must be aligned to ", - sizeof(load_input_t), " bytes"); - FLASHINFER_CHECK((params.x_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, - "x batch stride must be aligned to ", sizeof(load_input_t), " bytes"); - if (params.z) - { - FLASHINFER_CHECK(reinterpret_cast(params.z) % sizeof(load_input_t) == 0, - "z pointer must be aligned to ", sizeof(load_input_t), " bytes"); - FLASHINFER_CHECK((params.z_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, - "z batch stride must be aligned to ", sizeof(load_input_t), " bytes"); - } - FLASHINFER_CHECK(reinterpret_cast(params.B) % sizeof(load_input_t) == 0, "B pointer must be aligned to ", - sizeof(load_input_t), " bytes"); - FLASHINFER_CHECK(reinterpret_cast(params.C) % sizeof(load_input_t) == 0, "C pointer must be aligned to ", - sizeof(load_input_t), " bytes"); - FLASHINFER_CHECK((params.B_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, - "B batch stride must be aligned to ", sizeof(load_input_t), " bytes"); - FLASHINFER_CHECK((params.C_stride_batch * sizeof(input_t)) % sizeof(load_input_t) == 0, - "C batch stride must be aligned to ", sizeof(load_input_t), " bytes"); -} - -} // namespace flashinfer::mamba - -#endif // FLASHINFER_MAMBA_COMMON_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh deleted file mode 100644 index bab89e520ab3..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/conversion.cuh +++ /dev/null @@ -1,555 +0,0 @@ -#pragma once -#include -#include -#include -#ifdef FLASHINFER_ENABLE_BF16 -#include -#endif - -namespace flashinfer::mamba::conversion -{ - -inline __device__ float toFloat(float f) -{ - return f; -} - -inline __device__ float toFloat(__half h) -{ - return __half2float(h); -} - -#ifdef FLASHINFER_ENABLE_BF16 -inline __device__ float toFloat(__nv_bfloat16 val) -{ - return __bfloat162float(val); -} -#endif - -// No accuracy loss: int8_t / int16_t range fits exactly in float32 (24-bit -// mantissa represents all integers up to 2^24 = 16M exactly). -inline __device__ float toFloat(int8_t val) -{ - return static_cast(val); -} - -inline __device__ float toFloat(int16_t val) -{ - return static_cast(val); -} - -// fp8 e4m3 → fp32. Goes via __half (cuda_fp8 library has the implicit -// conversion that compiles to `cvt.rn.f16.e4m3` PTX on sm_89+), then -// __half2float for the final step. No direct fp8→fp32 PTX op exists. -inline __device__ float toFloat(__nv_fp8_e4m3 val) -{ - return __half2float(static_cast<__half>(val)); -} - -// Packed 2-element conversion: convert a packed pair to float2. -// Uses native packed intrinsics for bf16/fp16 (fewer PRMT/SHF instructions). -inline __device__ float2 toFloat2(float2 packed) -{ - return packed; -} - -inline __device__ float2 toFloat2(__half2 packed) -{ - return __half22float2(packed); -} - -// Pointer-based overloads: read two consecutive elements and convert to float2. -// Dispatches to the packed intrinsic for bf16/fp16 via the overloads above. -inline __device__ float2 toFloat2(float const* ptr) -{ - return {ptr[0], ptr[1]}; -} - -inline __device__ float2 toFloat2(__half const* ptr) -{ - return toFloat2(*reinterpret_cast<__half2 const*>(ptr)); -} - -#ifdef FLASHINFER_ENABLE_BF16 -// inline __device__ float2 toFloat2(__nv_bfloat162 packed) { return __bfloat1622float2(packed); } - -inline __device__ float2 toFloat2(__nv_bfloat162 packed) -{ - // bf16 is the upper 16 bits of f32 — shift/mask is cheaper than PRMT byte permutation. - // NOTE: this ignores denormals - uint32_t bits = reinterpret_cast(packed); - float2 out; - out.x = __uint_as_float(bits << 16); // low bf16 → upper 16 bits of f32 - out.y = __uint_as_float(bits & 0xFFFF0000u); // high bf16 already in upper 16 bits - return out; -} - -inline __device__ float2 toFloat2(__nv_bfloat16 const* ptr) -{ - return toFloat2(*reinterpret_cast<__nv_bfloat162 const*>(ptr)); -} - -// Paired f32 → bf16 conversion: pack two f32 values into __nv_bfloat162. -// Uses native cvt.rn.bf16x2.f32 — single instruction, round-to-nearest-even. -inline __device__ __nv_bfloat162 fromFloat2(float2 val) -{ - uint32_t result; - asm("cvt.rn.bf16x2.f32 %0, %1, %2;\n" : "=r"(result) : "f"(val.y), "f"(val.x)); - return reinterpret_cast<__nv_bfloat162 const&>(result); -} - -#endif - -inline __device__ float2 toFloat2(int8_t const* ptr) -{ - return {toFloat(ptr[0]), toFloat(ptr[1])}; -} - -inline __device__ float2 toFloat2(int16_t const* ptr) -{ - return {toFloat(ptr[0]), toFloat(ptr[1])}; -} - -inline __device__ void convertAndStore(float* output, float input) -{ - *output = input; -} - -inline __device__ void convertAndStore(__half* output, float input) -{ - *output = __float2half(input); -} - -#ifdef FLASHINFER_ENABLE_BF16 -inline __device__ void convertAndStore(__nv_bfloat16* output, float input) -{ - *output = __float2bfloat16(input); -} -#endif - -inline __device__ void convertAndStore(int16_t* output, float input) -{ - // Symmetric clip: [-max, max] (not [-max-1, max]) so that negation is safe. - // Matches Triton reference which clips to [-32767, 32767] before storing. - constexpr float int16_max = static_cast(std::numeric_limits::max()); - input = fminf(fmaxf(input, -int16_max), int16_max); - *output = static_cast(__float2int_rn(input)); -} - -// ============================================================================= -// Philox-4x32 PRNG (matches Triton's tl.randint) -// ============================================================================= - -// Generates four pseudorandom uint32s from (seed, offset) using the Philox-4x32 algorithm. -// Produces bit-identical output to Triton's tl.randint4x(seed, offset, n_rounds). -// The offset is int64 and split across Philox c0 (low 32 bits) and c1 (high -// 32 bits) — matches the i64 path of `randint4x` in triton/language/random.py. -// Provides 2^64 unique counter values per seed, avoiding collisions in large -// caches where `cache_slot * stride` exceeds 2^32. -// All four outputs (c0..c3) are independent and uniformly distributed. -template -__device__ __forceinline__ void philox_randint4x( - int64_t seed, int64_t offset, uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3) -{ - constexpr uint32_t PHILOX_KEY_A = 0x9E3779B9u; - constexpr uint32_t PHILOX_KEY_B = 0xBB67AE85u; - constexpr uint32_t PHILOX_ROUND_A = 0xD2511F53u; - constexpr uint32_t PHILOX_ROUND_B = 0xCD9E8D57u; - - uint32_t k0 = static_cast(static_cast(seed)); - uint32_t k1 = static_cast(static_cast(seed) >> 32); - uint64_t uoffset = static_cast(offset); - uint32_t c0 = static_cast(uoffset); - uint32_t c1 = static_cast(uoffset >> 32); - uint32_t c2 = 0, c3 = 0; - -#pragma unroll - for (int i = 0; i < n_rounds; i++) - { - uint32_t _c0 = c0, _c2 = c2; - c0 = __umulhi(PHILOX_ROUND_B, _c2) ^ c1 ^ k0; - c2 = __umulhi(PHILOX_ROUND_A, _c0) ^ c3 ^ k1; - c1 = PHILOX_ROUND_B * _c2; - c3 = PHILOX_ROUND_A * _c0; - k0 += PHILOX_KEY_A; - k1 += PHILOX_KEY_B; - } - r0 = c0; - r1 = c1; - r2 = c2; - r3 = c3; -} - -// Generates a pseudorandom uint32 from (seed, offset) using the Philox-4x32 algorithm. -// Produces bit-identical output to Triton's tl.randint(seed, offset, n_rounds). -// The offset is int64 (low/high split across Philox c0/c1) — see -// philox_randint4x for the full rationale. -// NOTE: This discards 3 of the 4 Philox outputs. For better throughput, use -// philox_randint4x to get all 4 outputs from a single Philox invocation. -template -__device__ __forceinline__ uint32_t philox_randint(int64_t seed, int64_t offset) -{ - uint32_t r0, r1, r2, r3; - philox_randint4x(seed, offset, r0, r1, r2, r3); - return r0; -} - -// ============================================================================= -// Stochastic rounding: fp32 → fp16 -// ============================================================================= - -// Software stochastic rounding: convert one fp32 value to fp16 using 13 random bits. -// Adds random noise at the sub-fp16-mantissa position, then truncates. -// rand13: 13-bit random value in bits [12:0]. -__device__ __forceinline__ uint16_t cvt_rs_f16_sw(float x, uint32_t rand13) -{ - uint32_t bits = __float_as_uint(x); - uint32_t sign = bits & 0x80000000u; - uint32_t abs_bits = bits & 0x7FFFFFFFu; - - // fp32 has 23 mantissa bits, fp16 has 10. The 13 LSBs are the remainder. - // Add 13-bit random noise at bits [12:0]. Carry into bit 13 → round up. - abs_bits += (rand13 & 0x1FFFu); - - // Convert to fp16 by truncation. - uint32_t f32_exp = (abs_bits >> 23) & 0xFFu; - uint32_t f32_mantissa = abs_bits & 0x7FFFFFu; - - uint16_t f16_bits; - if (f32_exp == 0xFF) - { - f16_bits = (f32_mantissa != 0) ? 0x7E00u : 0x7C00u; // NaN or Inf - } - else if (f32_exp > 142) - { // 127 + 15 = 142 → overflow to Inf - f16_bits = 0x7C00u; - } - else if (f32_exp < 113) - { // 127 - 14 = 113 → underflow to zero - f16_bits = 0; - } - else - { - uint16_t f16_exp = static_cast(f32_exp - 112); // rebias: 127→15 - uint16_t f16_mantissa = static_cast(f32_mantissa >> 13); - f16_bits = (f16_exp << 10) | f16_mantissa; - } - - return static_cast(sign >> 16) | f16_bits; -} - -// Forward declaration (defined below, after cvt_rs_f16x2_f32). -__device__ __forceinline__ uint32_t cvt_rs_f16x2_f32(float a, float b, uint32_t rbits); - -// Stochastic rounding: convert one fp32 value to fp16 using 13 random bits. -// On sm_100a+: uses PTX cvt.rs.f16x2.f32 with a dummy zero second input. -// On other archs: software emulation. -__device__ __forceinline__ __half cvt_rs_f16_f32(float x, uint32_t rand13) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL) - // Pack rand13 into rbits[12:0] (for PTX operand b → low half → our x). - // High half gets zero noise for the dummy input. - uint32_t rbits = rand13 & 0x1FFFu; - uint32_t packed = cvt_rs_f16x2_f32(x, 0.0f, rbits); - return __ushort_as_half(static_cast(packed & 0xFFFFu)); -#else - return __ushort_as_half(cvt_rs_f16_sw(x, rand13)); -#endif -} - -// Stochastic rounding: convert two fp32 values to packed fp16x2 using random bits. -// On sm_100a+: uses PTX cvt.rs.f16x2.f32 instruction. -// On other archs: software emulation matching the hardware behavior. -// -// rbits layout (from PTX docs): -// bits [28:16] = 13 random bits for PTX operand "a" (→ d[31:16], high half) -// bits [12:0] = 13 random bits for PTX operand "b" (→ d[15:0], low half) -// bits [31:29] and [15:13] = unused (zero) -// from: https://docs.nvidia.com/cuda/parallel-thread-execution/#cvt-rs-rbits-layout-f16 -// -// Our asm maps: %1→C++ a→PTX b→d[15:0], %2→C++ b→PTX a→d[31:16] -// So: C++ a uses rbits[12:0], C++ b uses rbits[28:16]. -__device__ __forceinline__ uint32_t cvt_rs_f16x2_f32(float a, float b, uint32_t rbits) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL) - uint32_t packed; - asm("cvt.rs.f16x2.f32 %0, %2, %1, %3;" - : "=r"(packed) - : "r"(__float_as_uint(a)), "r"(__float_as_uint(b)), "r"(rbits)); - return packed; -#else - uint32_t rand_a = rbits & 0x1FFFu; // bits [12:0] → C++ a (PTX b → low half) - uint32_t rand_b = (rbits >> 16) & 0x1FFFu; // bits [28:16] → C++ b (PTX a → high half) - uint16_t a_fp16 = __half_as_ushort(cvt_rs_f16_f32(a, rand_a)); - uint16_t b_fp16 = __half_as_ushort(cvt_rs_f16_f32(b, rand_b)); - return static_cast(a_fp16) | (static_cast(b_fp16) << 16); -#endif -} - -// Stochastic rounding store: generates Philox random bits and converts fp32 → fp16 in one call. -// PHILOX_ROUNDS: number of Philox rounds (compile-time), must be > 0. -// seed: Philox seed (from params.rand_seed). -// offset: unique per-element offset (e.g. d * DSTATE + i) for deterministic randomness. -template -inline __device__ void convertSRAndStore(__half* output, float input, int64_t seed, uint32_t offset) -{ - uint32_t rand = philox_randint(seed, offset); - *output = cvt_rs_f16_f32(input, rand & 0x1FFFu); -} - -// ============================================================================= -// Stochastic rounding: fp32 → fp8 (e4m3) -// ============================================================================= - -// Bit-reverse 16 bits. HW gives both elements of a pair 16 bits of independent -// SR randomness while consuming a single 16-bit chunk per pair: one element uses -// the chunk straight, the other uses bitrev16 of it. Two bitrev16'd halves of a -// uniformly random 16-bit value remain uniformly distributed and are statistically -// independent, so each element gets full 16-bit unbiased SR while sharing one -// 16-bit register with its pair-mate. -// -// Implementation: PTX `brev.b32` (CUDA intrinsic `__brev`) is a single-ALU-op -// 32-bit reverse; we right-shift by 16 to land the 16 LSBs of input in the 16 -// LSBs of output. Total: 2 SASS instructions (brev + shf/shr), vs the prior -// software 4-step mask/shift/OR chain (~12-16 inst). -__device__ __forceinline__ uint32_t bitrev16(uint32_t b) -{ - return __brev(b) >> 16; -} - -// Software stochastic rounding: convert one fp32 value to e4m3 (FN, satfinite) using 16 random -// bits. -// -// Algorithm: place `rand16` at the top of the discarded mantissa range, then truncate. -// shift_truncate = 20 for normal binade (unbiased >= -6) -// = 14 - unbiased for subnormal/underflow -// contribution = rand16 << (shift_truncate - 16) -// total = mant24 + contribution (in uint64 to avoid overflow) -// int_part = total >> shift_truncate -// Then re-encode int_part as e4m3, handling subnormal→normal transitions and saturation. -// -// Saturation (satfinite): -// |x| > 448 → ±448 (max finite e4m3 = 0x7E) -// ±Inf → ±448 -// NaN → canonical NaN with sign preserved (S|1111|111 = 0x7F or 0xFF) -// -// Verified bitwise against HW (cvt.rs.satfinite.e4m3x4.f32 on sm_100a) across 22528 -// inputs spanning subnormal, normal, and saturation regions during the SR -// reverse-engineering effort — see .plans/e4m3_stochastic_rounding.md. -__device__ __forceinline__ uint8_t cvt_rs_e4m3_sw(float x, uint32_t rand16) -{ - uint32_t bits = __float_as_uint(x); - uint32_t sign = (bits >> 31) & 1u; - uint32_t abs_bits = bits & 0x7FFFFFFFu; - uint32_t f32_exp = (abs_bits >> 23) & 0xFFu; - uint32_t f32_mant = abs_bits & 0x7FFFFFu; - - // NaN / Inf - if (f32_exp == 0xFFu) - { - if (f32_mant != 0) - { - return static_cast(0x7Fu | (sign << 7)); // canonical e4m3 NaN - } - else - { - return static_cast(0x7Eu | (sign << 7)); // Inf → ±max finite - } - } - - // fp32 zero / denormal → e4m3 zero (with sign). - if (f32_exp == 0u) - { - return static_cast(sign << 7); - } - - int unbiased = static_cast(f32_exp) - 127; - uint64_t mant24 = 0x800000u | f32_mant; // implicit-1 + mantissa, 24-bit - int shift_truncate = (unbiased >= -6) ? 20 : (14 - unbiased); - int rand_shift = shift_truncate - 16; - uint64_t rand_contrib; - if (rand_shift < 0) - { - // Defensive: shift_truncate < 16 shouldn't happen for valid normal/subnormal e4m3. - rand_contrib = static_cast(rand16 & 0xFFFFu) >> (-rand_shift); - } - else if (rand_shift < 56) - { - rand_contrib = static_cast(rand16 & 0xFFFFu) << rand_shift; - } - else - { - rand_contrib = 0; - } - uint64_t total = mant24 + rand_contrib; - // Guard the shift: shift_truncate reaches 140 for tiny-normal fp32 - // (unbiased < -49), which is UB for a uint64_t shift. Mathematically the - // result is 0 there (the value is far below e4m3's smallest subnormal), - // so flush to int_part = 0. The downstream subnormal branch rounds it to ±0. - uint32_t int_part = (shift_truncate >= 64) ? 0u : static_cast(total >> shift_truncate); - - if (unbiased >= -6) - { - // Started in normal binade. int_part ∈ [8, 15] normally; can overflow to 16+ if rand - // bumped the exponent. - int e4m3_exp = unbiased + 7; - while (int_part >= 16u) - { - int_part >>= 1; - e4m3_exp += 1; - } - if (e4m3_exp > 15 || (e4m3_exp == 15 && (int_part & 0x7u) == 7u)) - { - return static_cast(0x7Eu | (sign << 7)); - } - return static_cast((sign << 7) | (e4m3_exp << 3) | (int_part & 0x7u)); - } - else - { - // Started in subnormal/underflow. int_part: - // 0 → zero - // 1..7 → subnormal e4m3 - // 8..15 → smallest normal binade (e4m3_exp = 1) - // 16+ → higher normal binades (rare; only if rand pushed up multiple binades) - if (int_part == 0u) - { - return static_cast(sign << 7); - } - if (int_part <= 7u) - { - return static_cast((sign << 7) | int_part); - } - int e4m3_exp = 1; - while (int_part >= 16u) - { - int_part >>= 1; - e4m3_exp += 1; - } - if (e4m3_exp > 15 || (e4m3_exp == 15 && (int_part & 0x7u) == 7u)) - { - return static_cast(0x7Eu | (sign << 7)); - } - return static_cast((sign << 7) | (e4m3_exp << 3) | (int_part & 0x7u)); - } -} - -// Stochastic rounding: convert four fp32 values to packed fp8x4 e4m3 using random bits. -// On sm_100a+: uses PTX cvt.rs.satfinite.e4m3x4.f32 (combined stochastic-round + saturate). -// On other archs: software fallback via cvt_rs_e4m3_sw. -// -// Output layout (low byte first): -// packed[ 7: 0] = e4m3(a) -// packed[15: 8] = e4m3(b) -// packed[23:16] = e4m3(c) -// packed[31:24] = e4m3(d) -// -// rbits layout (per PTX docs + empirical HW oracle, see .plans/e4m3_stochastic_rounding.md): -// bits [31:16] = pair rbits for PTX operands a/b (high two outputs) -// bits [15: 0] = pair rbits for PTX operands e/f (low two outputs) -// Each PAIR shares its 16-bit chunk: HW uses the chunk straight for the "even" -// PTX operand (b, f → low byte of pair output) and bitrev16 of the chunk for -// the "odd" operand (a, e → high byte of pair output). Both elements get the -// full 16 bits of independent SR randomness this way. -// our `a` (PTX f, → byte 0) uses rbits[15: 0] -// our `b` (PTX e, → byte 1) uses bitrev16(rbits[15: 0]) -// our `c` (PTX b, → byte 2) uses rbits[31:16] -// our `d` (PTX a, → byte 3) uses bitrev16(rbits[31:16]) -// -// PTX syntax `cvt.rs.satfinite.e4m3x4.f32 d, {a3, a2, a1, a0}, rbits` writes -// e4m3(a_i) into byte i of d. We want byte 0 = e4m3(a), so the source-vector -// ordering is {d, c, b, a} = {%4, %3, %2, %1}. See PTX ISA: -// https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cvt -__device__ __forceinline__ uint32_t cvt_rs_e4m3x4_f32(float a, float b, float c, float d, uint32_t rbits) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL) - uint32_t packed; - asm("cvt.rs.satfinite.e4m3x4.f32 %0, {%4, %3, %2, %1}, %5;" - : "=r"(packed) - : "r"(__float_as_uint(a)), "r"(__float_as_uint(b)), "r"(__float_as_uint(c)), "r"(__float_as_uint(d)), - "r"(rbits)); - return packed; -#else - uint32_t low_chunk = rbits & 0xFFFFu; - uint32_t high_chunk = (rbits >> 16) & 0xFFFFu; - uint8_t pa = cvt_rs_e4m3_sw(a, low_chunk); // PTX f → byte 0 - uint8_t pb = cvt_rs_e4m3_sw(b, bitrev16(low_chunk)); // PTX e → byte 1 - uint8_t pc = cvt_rs_e4m3_sw(c, high_chunk); // PTX b → byte 2 - uint8_t pd = cvt_rs_e4m3_sw(d, bitrev16(high_chunk)); // PTX a → byte 3 - return static_cast(pa) | (static_cast(pb) << 8) | (static_cast(pc) << 16) - | (static_cast(pd) << 24); -#endif -} - -// ============================================================================= -// Round-to-nearest-even + saturate: fp32 → int8 -// ============================================================================= - -// cvt.rni.sat.s8.f32: single PTX instruction on sm_80+, replaces -// the F2I.S32 + VIMNMX(min 127) + VIMNMX(max -127) chain. -// Saturates to [-128, 127]. Callers using encode_scale = 127/amax -// guarantee |input| ≤ 127.0, so -128 is never produced. -__device__ __forceinline__ int8_t cvt_rni_sat_s8(float x) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - int32_t result; - asm("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(result) : "f"(x)); - return static_cast(result); -#else - return static_cast(max(-128, min(127, __float2int_rn(x)))); -#endif -} - -// ============================================================================= -// Stochastic rounding + saturate: fp32 → int8 -// ============================================================================= - -// Software SR for int8: add uniform noise in [0, 1) then floor. -// Matches Triton's `floor(scaled_value + rand01)` where -// `rand01 = (rand & 0x00FFFFFF) * (1.0 / (1 << 24))`. -// Saturates to [-127, 127] (symmetric, matching encode_scale = 127/amax). -__device__ __forceinline__ int8_t cvt_rs_sat_s8(float x, uint32_t rand_bits) -{ - float const rand01 = static_cast(rand_bits & 0x00FFFFFFu) * (1.0f / static_cast(1 << 24)); - // `__float2int_rd` (round toward -infinity) fuses `floorf` + `__float2int_rz` - // into a single `cvt.rmi.s32.f32` SASS instruction, saving one FRND per call. - int32_t const clamped = max(-127, min(127, __float2int_rd(x + rand01))); - return static_cast(clamped); -} - -// Stochastic rounding: convert four fp32 values to packed s8x4 using a single -// 32-bit random integer. Analogous to cvt_rs_e4m3x4_f32: 16-bit chunks are -// reused via bitrev16 so each output gets 16 bits of independent randomness -// while consuming only one shared 16-bit chunk per pair (two bitrev16'd halves -// of a uniform 16-bit value remain uniformly distributed and statistically -// independent). -// -// 16-bit entropy per element is far more than int8 SR requires: the rounding -// decision compares against a fractional residual with at most ~7 bits of -// meaningful precision for int8, so no quality loss vs the 24-bit scalar path. -// -// Amortization: 1 random u32 → 4 SR int8s. A single Philox call (4 u32s) -// covers 16 int8 conversions, a 4× reduction in PRNG cost vs the scalar -// cvt_rs_sat_s8 path. -__device__ __forceinline__ uint32_t cvt_rs_sat_s8x4_f32(float a, float b, float c, float d, uint32_t rbits) -{ - uint32_t const low_chunk = rbits & 0xFFFFu; - uint32_t const high_chunk = (rbits >> 16) & 0xFFFFu; - constexpr float kInv16 = 1.0f / static_cast(1u << 16); - - float const r_a = static_cast(low_chunk) * kInv16; - float const r_b = static_cast(bitrev16(low_chunk)) * kInv16; - float const r_c = static_cast(high_chunk) * kInv16; - float const r_d = static_cast(bitrev16(high_chunk)) * kInv16; - - // `__float2int_rd` (round toward -infinity) emits a single `cvt.rmi.s32.f32` - // SASS op, fusing the `floorf` + `__float2int_rz` chain into one instruction. - int32_t const pa = max(-127, min(127, __float2int_rd(a + r_a))); - int32_t const pb = max(-127, min(127, __float2int_rd(b + r_b))); - int32_t const pc = max(-127, min(127, __float2int_rd(c + r_c))); - int32_t const pd = max(-127, min(127, __float2int_rd(d + r_d))); - - return (static_cast(pa) & 0xFFu) | ((static_cast(pb) & 0xFFu) << 8) - | ((static_cast(pc) & 0xFFu) << 16) | ((static_cast(pd) & 0xFFu) << 24); -} - -} // namespace flashinfer::mamba::conversion diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh deleted file mode 100644 index 1cf1211bcd8f..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu.cuh +++ /dev/null @@ -1,1114 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -#ifndef FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_CUH_ -#define FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_CUH_ - -// Incremental SSU kernel — matmul-based, tensor-core MMA (single path). -// Single CTA per (batch, head). Grid: (batch, nheads). -// 4 warps per CTA, 128 threads total. -// -// Single __syncthreads() via per-warp data ownership. Every smem -// read before the final barrier is served by data the same warp loaded. -// No mbarriers, no cross-warp visibility for the first half of the kernel. -// -// Phase 0 (per-warp cp.async, no cross-warp sync): -// State: each warp loads own DIM slice (rows [16W : 16W+16]). -// B, C: redundant on W0, W1 (both do 2-warp CB). -// old_B: redundant on all 4 warps (each warp's replay needs full DSTATE). -// old_x: redundant on all 4 warps. -// x: W2 only (Phase-2 read — covered by the single syncthreads). -// z: W3 only (Phase-2 read — covered by the single syncthreads). -// Scalars + cumAdt: redundant on each warp's first NPREDICTED/MAX_WINDOW lanes. -// Each warp: __pipeline_commit → __pipeline_wait_prior(0) → __syncwarp. -// -// Phase 1 (runs with *no* barrier; CB ‖ replay parallelism preserved): -// - store_old_B hoisted here — W0,W1 only (they hold valid smem.B). -// - Warps 0,1: compute_CB_scaled_2warp (bf16 HMMA → swizzled smem.CB_scaled). -// - All warps: replay_state_mma (HMMA; state in smem updated in-place, -// each warp touches only its own DIM rows). -// -// __syncthreads() ← THE ONE. Provides cross-warp visibility of: -// CB_scaled (W0,W1→all), x (W2→all), z (W3→all). -// -// Phase 2: compute_and_store_output -// = (C @ state^T) * decay + CB_scaled @ x + D*x, z-gate → direct gmem STG. -// State writeback hoisted inside the orchestrator once matmul 3 has -// finished consuming smem.state. -// -// Phase 3: old_x / old_dt / old_cumAdt cache writes. - -#include "kernel_checkpointing_ssu_common.cuh" - -namespace flashinfer::mamba::checkpointing -{ - -// ============================================================================= -// Shared memory layout -// ============================================================================= -// smem holds the state in its native dtype (`state_t`). cp.async pulls the -// native dtype straight into smem (with the matching `SmemSwizzle`), -// and the conversion to `MMA_prop::operand_t` happens on the register read inside -// add_init_out / replay_state_mma. -// `D_PER_CTA` is the per-CTA D dimension after D-split. -// For D_SPLIT = 1 (default), D_PER_CTA == DIM (per-head DIM). At D_SPLIT > 1 -// the storage is sliced: each CTA owns a contiguous D_PER_CTA-row slice of -// the head's D axis. Buffers that aren't D-owned (B, C, old_B, scalars) are -// unaffected. -// -// Note on D_SMEM_COLS: the Swizzle<3,3,3> atom for bf16 is (8, 64), so -// `make_swizzled_layout_rc` requires col counts to be multiples of 64. When -// D_PER_CTA < 64 (e.g. D_SPLIT=2 → D_PER_CTA=32) we pad the D-owned buffer -// cols up to the swizzle atom width. The cp.async only fills the first -// D_PER_CTA cols; the padded tail is unused but keeps the swizzle layout -// well-formed. Cost: 1 KB per [NPREDICTED_PAD_MMA_M, 64] buffer at D_PER_CTA=32. -template -struct CheckpointingSsuStorage -{ - // Re-export the two T/W axis sizes so helpers that only see SmemT can - // recover them. - static constexpr int NPREDICTED = NPREDICTED_; - static constexpr int MAX_WINDOW = MAX_WINDOW_; - // Swizzle atom width for input_t (= 64 cols for 2-byte types). - static constexpr int D_SMEM_COLS = next_multiple_of::ATOM_COLS>(D_PER_CTA); - // M-dim of the output MMAs (C, x, z, CB_scaled): always m16-tiled (keyed - // off NPREDICTED, the new-tokens count). - static constexpr int NPREDICTED_PAD_MMA_M = next_multiple_of(NPREDICTED); - // N-dim of the precompute-CB MMA (matmul-1: C @ B^T). B's row count is - // the matmul N-axis → padded to MMA::N=8. When NPREDICTED ≤ 8 only warp - // 0 has valid B rows; warp 1 zero-fills its CB slice. - static constexpr int NPREDICTED_PAD_MMA_N = next_multiple_of(NPREDICTED); - // K-dim of the replay MMA (matmul-2: old_x^T @ dB_scaled). Padded to - // the small atom's K (== the LDSM unit for 2-byte elements). When - // MAX_WINDOW ≤ MMA::K_SMALL=8, replay picks the small atom (1 K-tile, - // smaller smem, +1 CTA/SM occupancy); otherwise the big atom. Assumes - // MAX_WINDOW ≤ MMA::K_BIG (asserted in the wrapper). - static constexpr int MAX_WINDOW_PAD_MMA_K = next_multiple_of(MAX_WINDOW); - // Row count for buffers padded only to the input-type swizzle atom's row - // extent (8 for 2-byte, 4 for 4-byte) — used by C and z, which alias the - // second m-tile back onto the first via `make_aliased_swizzled_layout_rc`. - // Keyed off NPREDICTED. - static constexpr int NPREDICTED_SWIZZLE_R = next_multiple_of::ATOM_ROWS>(NPREDICTED); - - // All 2D smem buffers below are stored as flat 1D arrays — the actual - // physical layout is determined by `make_swizzled_layout_rc<...>` at each - // access site, which scrambles (row, col) → physical offset via the - // Swizzle XOR. Declaring them as `T[ROWS][COLS]` would falsely suggest a - // row-major C-array layout that nobody ever uses; the only thing that - // matters here is total byte count and 16-byte alignment. - - // CB_scaled — logical (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE) Swizzle<3,3,3>. - // CB_ROW_STRIDE pads each row to one bank cycle (128 B = 32 banks × 4 B) - // worth of `input_t` so LDSM reads in matmul-4's A operand are - // conflict-free. Equals the swizzle atom's col extent for `input_t` - // (64 for 2-byte, 32 for 4-byte). Logical CB matrix is - // (NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M); trailing cols are padding. - static constexpr int CB_ROW_STRIDE = SmemSwizzle::ATOM_COLS; - alignas(16) input_t CB_scaled[NPREDICTED_PAD_MMA_M * CB_ROW_STRIDE]; - - // B — logical (NPREDICTED_PAD_MMA_N, DSTATE). Row count is matmul-1's - // N-axis (since matmul-1 = C @ B^T). Padding rows inside [NPREDICTED, - // NPREDICTED_PAD_MMA_N) contain garbage — valid output uses only - // [0, NPREDICTED). Warp-1 of compute_CB_scaled_2warp reads rows ≥ 8 of a - // 16-row view; those reads spill into C/old_B smem but are masked to 0 by - // the (j < NPREDICTED) CB-store predicate since j ≥ 8 ≥ NPREDICTED when - // NPREDICTED_PAD_MMA_N == 8. - alignas(16) input_t B[NPREDICTED_PAD_MMA_N * DSTATE]; - - // C — physical (next_multiple_of(NPREDICTED), DSTATE). Padded - // only to the swizzle atom's row extent (8 for 2-byte, 4 for 4-byte), not - // to MMA_prop::M=16. cp.async writes to this exact extent (CShape's first - // dim shrunk to match — see load_data). The MMA still views it as - // NPREDICTED_PAD_MMA_M=16 rows via `make_aliased_swizzled_layout_rc`, - // which aliases the second m-tile back onto the first via stride-0 - // row-tile mode. Garbage feeds output rows ≥ NPREDICTED — predicated - // out at gmem store. Saves up to 2 KB of smem at NPREDICTED ≤ ATOM_ROWS, - // no-op when NPREDICTED > ATOM_ROWS. - alignas(16) input_t C[NPREDICTED_SWIZZLE_R * DSTATE]; - - // x — logical (NPREDICTED_PAD_MMA_M, D_SMEM_COLS). Cols padded to - // D_SMEM_COLS for swizzle atom alignment; cp.async only fills cols - // [0, D_PER_CTA), the tail is unused. - alignas(16) input_t x[NPREDICTED_PAD_MMA_M * D_SMEM_COLS]; - - // z — physical (next_multiple_of(NPREDICTED), D_SMEM_COLS). - // Padded only to the swizzle atom's row extent (8 for 2-byte, 4 for - // 4-byte), not to MMA_prop::M=16. z is never an MMA operand — the - // z-gating epilogue reads it via `partition_C` of the m16n8 c-frag, so - // the MMA still views it as NPREDICTED_PAD_MMA_M=16 rows via - // `make_aliased_swizzled_layout_rc`, which aliases the second m-tile back - // onto the first via stride-0 row-tile mode. Garbage feeds output rows - // ≥ NPREDICTED — predicated out at gmem store. Saves up to 1 KB of smem - // at NPREDICTED ≤ ATOM_ROWS, no-op when NPREDICTED > ATOM_ROWS. - alignas(16) input_t z[NPREDICTED_SWIZZLE_R * D_SMEM_COLS]; - - // Old cache data loaded in Phase 0 (consumed in Phase 1 replay). - // old_x — logical (MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS); ldmatrix.trans - // feeds replay MMA A-operand (only the first D_PER_CTA cols are valid - // data). - alignas(16) input_t old_x[MAX_WINDOW_PAD_MMA_K * D_SMEM_COLS]; - - // old_B — logical (MAX_WINDOW_PAD_MMA_K, DSTATE) Swizzle<3,3,3>. Replay - // MMA reads via ldmatrix.trans (LDSM_T) + register scaling. Padding - // rows zero-filled via cp.async ZFILL. - alignas(16) input_t old_B[MAX_WINDOW_PAD_MMA_K * DSTATE]; - - float old_dt[MAX_WINDOW]; - float old_cumAdt[MAX_WINDOW]; - - // Processed dt for new tokens (Phase 1a uses this for CB_scaled + cumAdt) - float dt_proc[NPREDICTED]; - - // Cumulative A*dt — computed once by warp 0, read by all warps after sync - float cumAdt[NPREDICTED]; - - // state — logical (D_PER_CTA, DSTATE) in `state_t` (native dtype). The - // MMA path reinterprets 2-byte state as bf16 for LDSM; f32 state is loaded - // via UniversalCopy and converted to bf16 in registers inside - // add_init_out. - alignas(16) state_t state[D_PER_CTA * DSTATE]; -}; - -// ============================================================================= -// Stochastic-round one fp32 pair to a packed f16x2 u32 with amortized philox -// refresh. rand_idx[4] is mutated in place every 4th call (when pair_idx & 3 -// == 0): a single philox_randint4x feeds 4 consecutive cvt_rs calls, then -// gets refreshed. Each refresh uses a per-lane unique `philox_off` so the -// generated randints don't collide across threads. Triton bit-equality is -// intentionally given up here; unbiasedness still holds since each pair's -// cvt_rs gets its own dedicated 32-bit randint. -// ============================================================================= -template -__device__ __forceinline__ uint32_t stochastic_round_pair_with_philox_refresh( - float a, float b, int pair_idx, int64_t rand_seed, int64_t philox_off, uint32_t (&rand_idx)[4]) -{ - int const rand_pos = pair_idx & 3; - if (rand_pos == 0) - { - conversion::philox_randint4x( - rand_seed, philox_off, rand_idx[0], rand_idx[1], rand_idx[2], rand_idx[3]); - } - return conversion::cvt_rs_f16x2_f32(a, b, rand_idx[rand_pos]); -} - -// ============================================================================= -// Cross-pass shfl_xor + STG.64 state writeback. -// -// Given two passes' worth of post-cvt_rs packed u32s buffered in `my_packed` -// (pass-0 in [0][:], pass-1 in [1][:]), exchange via shfl_xor across lane^1 -// neighbors so that all 32 lanes can issue ONE STG.64 each per pair iter: -// - even lane k stores PASS n0 (cols (k%4)*2..(k%4)*2+3 of warp's n0 slice) -// - odd lane k stores PASS n1 (cols (k%4)*2-2..(k%4)*2+1 of warp's n1 slice) -// -// Halves the STG instruction count vs per-pass writeback: 1 STG.64 per pair -// iter covers BOTH passes' data via cross-lane participation. -// ============================================================================= -template -__device__ __forceinline__ void exchange_ntile_state_store_global(state_t* __restrict__ state_w_base, int np, int lane, - uint32_t const (&my_packed)[2][PAIRS_PER_PASS], IdPart const& id_part) -{ - using namespace cute; - static_assert( - sizeof(state_t) == 2, "exchange_ntile_state_store_global requires 2-byte state_t for STG.64 alignment"); - int const n_base_p0 = np * N_PER_PASS; - int const n_base_p1 = (np + 1) * N_PER_PASS; -#pragma unroll - for (int p = 0; p < PAIRS_PER_PASS; ++p) - { - int const i = p * 2; - // xor mask = 1 swaps neighbor lanes: lane 0 <-> lane 1, lane 2 <-> lane 3, ... - uint32_t const peer_p0 = __shfl_xor_sync(constants::MASK_ALL_LANES, my_packed[0][p], 1); - uint32_t const peer_p1 = __shfl_xor_sync(constants::MASK_ALL_LANES, my_packed[1][p], 1); - - int const row = get<0>(id_part(i)); - int const col_p0 = get<1>(id_part(i)) + n_base_p0; - int const col_p1 = get<1>(id_part(i)) + n_base_p1; - - uint64_t combined; - int32_t gmem_off; - if ((lane & 1) == 0) - { - // Even lane: store PASS n0 — my (lower col) in low, peer in high. - combined = static_cast(my_packed[0][p]) - | (static_cast(peer_p0) << constants::num_bits_uint32); - gmem_off = row * DSTATE + col_p0; - } - else - { - // Odd lane: store PASS n1 — peer (lower col) in low, my in high. - // STG addr = gmem[row*DSTATE + (peer's col base)] = col_p1 - 2. - combined = static_cast(peer_p1) - | (static_cast(my_packed[1][p]) << constants::num_bits_uint32); - gmem_off = row * DSTATE + (col_p1 - 2); - } - *reinterpret_cast(&state_w_base[gmem_off]) = combined; - } -} - -// ============================================================================= -// Phase 1b: Replay — tensor-core MMA path (matmul 2: state recurrence). -// state[D, dstate] = state * total_decay + old_x^T @ (coeff * old_B) -// All 128 threads cooperate. -// -// Warps along N=DSTATE: -// TiledMMA uses Layout<_1, _4> — per pass covers (M=DIM, N=4×MMA_prop::N=32). -// Each warp owns: full M (DIM/16 m-atoms) and one n-atom of 8 cols. -// Why: A is small (DIM × K), B is bigger (DSTATE × K). M-split (`_4×1`) -// redundantly loaded full B from each warp (4× × 4 KB = 16 KB). N-split -// (`_1×4`) instead redundantly loads full A (4× × 2 KB = 8 KB) and reads -// B disjointly across warps — net smem read drops 18 KB → 12 KB per replay -// (~33%) at K_BIG. Also unlocks D-split D_PER_CTA < 64. -// ============================================================================= -// state_w_base (f16+philox path): pre-offset gmem pointer to this CTA's owned -// [D_PER_CTA, DSTATE] state slice (params.state + cache_slot * -// state_stride_seq + head * DIM*DSTATE + d_tile * D_PER_CTA*DSTATE). -// Computed in the kernel preamble. Combining base + offset into one i64 -// pointer drops the cross-iter live-range cost from 4 regs (state_w ptr + -// state_gmem_off) to 2 regs (just the base), and the per-pair STG.32 uses an -// i32 element offset inside the chunk. Use this instead of separately -// holding params.state-ptr and state_gmem_off. -template -__device__ __forceinline__ void replay_state_mma(SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, - int prev_k, int d_tile, int64_t state_ptr_offset, state_t* state_w_base, int64_t rand_seed, bool must_checkpoint) -{ - using namespace cute; - static_assert(sizeof(input_t) == 2, "replay_state_mma requires 2-byte input type"); - static_assert(D_PER_CTA % 16 == 0, "D_PER_CTA must be divisible by 16 (m16n8 atom)"); - static_assert(D_PER_CTA >= 16, "D_PER_CTA must be at least 16"); - - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; // 8 or 16 - int const tid = warp * warpSize + lane; - - // Atom K matches the cache-window tile (MAX_WINDOW_PAD_MMA_K). - // K == MMA_prop::K_BIG (16) → m16n8k16 + x4/x2 ldmatrix.trans - // K == MMA_prop::K_SMALL (8) → m16n8k8 + x2/x1 ldmatrix.trans - using MmaAtomType - = std::conditional_t; - using LdsmA = std::conditional_t; - using LdsmB = std::conditional_t; - - // 4 warps along N=DSTATE; each warp covers full M (D_PER_CTA/16 m-atoms). - auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(tid); - - // Per-pass output tile is (D_PER_CTA, N_PER_PASS). N_PER_PASS = 4 warps × n8 = 32 cols. - constexpr int N_PER_PASS = 4 * MMA_prop::N; - static_assert(DSTATE % N_PER_PASS == 0, "DSTATE must be divisible by 4 * MMA_prop::N for _1x4 warp layout"); - constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; - - float total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - float total_decay = (prev_k > 0) ? __expf(total_cumAdt) : 1.f; - - // ── A operand: old_x [MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS] Swizzle<3,3,3>, transposed - // view [M=D_SMEM_COLS, K=MAX_WINDOW_PAD_MMA_K]. D_SMEM_COLS may be padded above - // D_PER_CTA when D_PER_CTA < swizzle atom; local_tile to D_PER_CTA - // restricts the LDSM to the valid sub-tile. Each warp loads the FULL M (4× - // redundant across warps). See header comment for traffic accounting. ── - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - auto layout_A_full = make_swizzled_layout_rc_transpose(); - Tensor smem_A_full - = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); - Tensor smem_A - = local_tile(smem_A_full, make_shape(Int{}, Int{}), make_coord(_0{}, _0{})); - - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); - Tensor frag_A = thr_mma.partition_fragment_A( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); - - cute::copy(s2r_A, smem_A_s2r, frag_A_view); - // old_x is input_t == MMA_prop::operand_t (bf16) — no conversion needed. - - // ── B operand: old_B [MAX_WINDOW_PAD_MMA_K, DSTATE] swizzled, transposed view - // [N=DSTATE, K=MAX_WINDOW_PAD_MMA_K]. Per pass loads N_PER_PASS=32 cols across - // 4 warps; partition_S splits — each warp gets its disjoint 8-col slice. ── - auto layout_B = make_swizzled_layout_rc_transpose(); - Tensor smem_B_full = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B); - - auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B = s2r_B.get_slice(tid); - - // ── State: per-CTA swizzle layout [D_PER_CTA, DSTATE]. ── - auto layout_state_swz = make_swizzled_layout_rc(); - state_t* state_base = reinterpret_cast(smem.state); - - // ── Per-pass identity for (row, col) coords ── - // partition_C of an identity tensor of the per-pass output shape gives this - // thread's (row, col) at every C-frag position, including warp-N offset. - // Frag size per thread = (M_atoms=D_PER_CTA/16) × (N_atoms_per_warp=1) × 4 elts. - auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_tile); - // Linear order from CuTe's column-major partition_C with m16n8 atom: - // i=0,1: same row (= row_lo of M-atom 0), adjacent cols (col_off, col_off+1) - // i=2,3: same row (= row_hi of M-atom 0), adjacent cols - // i=4,5: same row (= row_lo of M-atom 1) - // ... (V index 0..3 inside each m16n8, then M-atoms in M-major order) - // Pair load at (i, i+1) covers two consecutive bf16 elts → one 32-bit LDS. - - // Precompute dB coefficients once — depend only on K (lane), not on N. - constexpr int LANES_PER_N_COL = warpSize / MMA_prop::N; // = 4 for m16n8k_ - constexpr int DB_COEFFS_PER_LANE = MAX_WINDOW_PAD_MMA_K / LANES_PER_N_COL; - float dB_coeff[DB_COEFFS_PER_LANE]; - precompute_dB_coeff(dB_coeff, smem, total_cumAdt, prev_k, lane); - - using pair_t = Pair; - - // Philox state amortized across 4 consecutive pair conversions: each call - // returns 4 randints, all 4 get consumed before the next refresh (vs. 1-of-4 - // in the Triton-bit-equal layout — see writeback loop below). Compile-time - // pair_idx (n-loop and i-loop both unrolled) keeps `rand_idx[pair_idx & 3]` - // as a known register access — no local-memory spill. - constexpr bool kPhiloxF16 = (PHILOX_ROUNDS > 0) && std::is_same_v; - [[maybe_unused]] uint32_t rand_idx[4]; - // state_w_base is the pre-combined (params.state + state_gmem_off) base - // pointer — see the function header. No separate state_w / state_gmem_off - // alive in this scope. - - // ── Vectorized state writeback (cross-pass STG.64 fusion) ────────── - // smem always gets nearest-even f32→state_t (consumed by matmul 3 — must - // match Triton's f32→bf16 path as closely as possible). Gmem cache, when - // PHILOX_ROUNDS > 0 and state_t == __half, gets PTX cvt.rs.f16x2.f32 - // stochastic rounding direct from registers via cross-pass STG.64; the - // smem→gmem `store_state` is gated off in compute_and_store_output. - // - // Cross-pass STG fusion: do PASS n0 and PASS n1 back-to-back, buffering - // the post-cvt_rs packed u32s of n0 across n1's HMMA + cvt_rs. Then issue - // ONE STG.64 instruction per pair iter, all 32 lanes active: - // - even lane stores PASS n0 data at the warp's n0 column slice - // - odd lane stores PASS n1 data at the warp's n1 column slice - // Halves the STG instruction count vs per-pass writeback (16 STG.64/thread - // per 2 passes vs 16 + 16 = 32 STG.64/thread previously — same byte volume). - // - // Randint amortization: rand_idx[4] refreshed every 4 pairs; each pair's - // cvt_rs uses one of the 4 randints. Triton bit-equality is intentionally - // given up; unbiasedness still holds. - constexpr int PAIRS_PER_PASS = D_PER_CTA / 8; // = (D_PER_CTA/16) × 2 row-pair iters - static_assert(NUM_N_PASSES % 2 == 0, "Cross-pass STG fusion requires even NUM_N_PASSES"); - -#pragma unroll - for (int np = 0; np < NUM_N_PASSES; np += 2) - { - // Buffer of post-cvt_rs packed u32s for both passes (philox path only). - [[maybe_unused]] uint32_t my_packed[2][PAIRS_PER_PASS]; - -#pragma unroll - for (int local_n = 0; local_n < 2; ++local_n) - { - int const n = np + local_n; - int const n_base = n * N_PER_PASS; - - // ── Allocate per-pass C-frag (4 × M_atoms fp32 elts/thread) ── - Tensor frag_h = thr_mma.partition_fragment_C( - make_tensor((float*) 0x0, make_shape(Int{}, Int{}))); - - // ── Load state × total_decay into frag_h. ── -#pragma unroll - for (int i = 0; i < size(frag_h); i += 2) - { - int const row = get<0>(id_part(i)); - int const col = get<1>(id_part(i)) + n_base; - int const off = layout_state_swz(row, col); - pair_t const p = *reinterpret_cast(&state_base[off]); - frag_h(i) = toFloat(p[cute::Int<0>{}]) * total_decay; - frag_h(i + 1) = toFloat(p[cute::Int<1>{}]) * total_decay; - } - - // ── LDSM.T per-pass B (per warp = 1 atom of 8 cols of N) ── - Tensor smem_B_n = local_tile( - smem_B_full, make_tile(Int{}, Int{}), make_coord(n, _0{})); - auto smem_B_s2r_n = s2r_thr_B.partition_S(smem_B_n); - - Tensor frag_B = thr_mma.partition_fragment_B( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - auto frag_B_view = s2r_thr_B.retile_D(frag_B); - - cute::copy(s2r_B, smem_B_s2r_n, frag_B_view); - - compute_dB_scaling(frag_B, dB_coeff); - - // ── HMMA: frag_h += frag_A @ frag_B ── - cute::gemm(tiled_mma, frag_h, frag_A, frag_B, frag_h); - - // ── Smem write (always) + cvt_rs into my_packed (philox path) ── -#pragma unroll - for (int i = 0; i < size(frag_h); i += 2) - { - int const row = get<0>(id_part(i)); - int const col = get<1>(id_part(i)) + n_base; - int const off = layout_state_swz(row, col); - - // Smem write — always nearest-even (output's matmul 3 reads this). - pair_t const q = pack_float2(make_float2(frag_h(i), frag_h(i + 1))); - *reinterpret_cast(&state_base[off]) = q; - - if constexpr (kPhiloxF16) - { - static_assert(sizeof(state_t) == 2, "STG.64 cooperative path requires 2-byte state_t"); - int const pair_idx = n * PAIRS_PER_PASS + i / 2; - // Per-lane philox_off is unique per (thread, refresh group) — each - // pair gets its own randint bits. Always computed; only consumed - // by the refresh branch inside the helper. - int64_t const philox_off = state_ptr_offset + (int64_t) (d_tile * D_PER_CTA + row) * DSTATE + col; - // Buffer the SR'd packed u32 — store happens after BOTH passes. - my_packed[local_n][i / 2] = stochastic_round_pair_with_philox_refresh( - frag_h(i), frag_h(i + 1), pair_idx, rand_seed, philox_off, rand_idx); - } - } - } - - // ── Cross-pass STG.64: all 32 lanes active. ───────────────────────── - // m16n8 lane layout: lane k → row k/4, cols (k%4)*2..(k%4)*2+1. Lanes - // (2k, 2k+1) hold adjacent col-pairs of the same row. After shfl_xor, - // the even/odd lane each has a 4-col contiguous block (in different - // bit-orders). Even lane STG.64s the n0-pass block at its own col - // base; odd lane STG.64s the n1-pass block at the peer's (lower) col - // — both 8-byte aligned for state_t = f16. - // Runtime-gated on must_checkpoint: non-checkpoint steps skip the gmem - // STGs entirely (state HBM remains the prior checkpoint). The cvt_rs - // SR + philox refresh above still ran — only the STGs are elided — - // because skipping them would require routing must_checkpoint into the - // pair_idx amortization logic, which lives across the n-loop. - if constexpr (kPhiloxF16) - { - if (must_checkpoint) - { - exchange_ntile_state_store_global( - state_w_base, np, lane, my_packed, id_part); - } - } - } -} - -// ── Orchestrator: compute_and_store_output ───────────────────────────── -// out = (C @ state^T) * decay + CB_scaled @ x + D*x, then z-gate. -// All operations on register-resident frag_y — no smem round-trip. -// Result converted f32 → input_t in registers and stored directly to gmem -// via partition_C of the global output tensor (like CUTLASS sgemm_sm80 epilogue). -template -__device__ __forceinline__ void compute_and_store_output(SmemT& smem, CheckpointingSsuParams const& params, int warp, - int lane, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, bool must_checkpoint, - int seq_len) -{ - using namespace cute; - static_assert(sizeof(input_t) == 2, "compute_and_store_output requires 2-byte input type"); - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - int const tid = warp * warpSize + lane; - - // ── TiledMMA: 128 threads, covers [16, 32] output per step ── - auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(tid); - - // ── Swizzled smem views ── - // When D_PER_CTA < swizzle atom (= 64 for bf16), the underlying - // smem buffer is padded to D_SMEM_COLS so the swizzle layout is well-formed. - // Per-pass MMA loops only iterate D_PER_CTA / N_TILE tiles → never touch - // the padded tail. - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - - // x: swizzled [NPREDICTED_PAD_MMA_M, D_SMEM_COLS] - auto layout_x_swz = make_swizzled_layout_rc(); - Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_swz); - auto layout_x_trans_swz = make_swizzled_layout_rc_transpose(); - Tensor smem_x_trans - = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); - - // z: aliased swizzled [NPREDICTED_PAD_MMA_M, D_SMEM_COLS] — physical buffer - // is only next_multiple_of(NPREDICTED) rows tall; second m-tile - // aliases first. Ghost rows feed predicated-out output rows. - auto layout_z_swz = make_aliased_swizzled_layout_rc(); - Tensor smem_z = make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); - - // ── S2R copies ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B = s2r_B.get_slice(tid); - auto s2r_B_trans = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); - - // ── Load CB_scaled A operand from smem (precomputed by warps 0,1 between syncs) ── - // Row stride matches the buffer's padded width (one swizzle atom of `input_t`). - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - auto layout_cb_swz = make_swizzled_layout_rc(); - Tensor smem_CB - = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); - auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); - Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); - auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); - cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); - - // Decay broadcast: cumAdt[t] → [NPREDICTED_PAD_MMA_M, N_TILE] with stride-0 on N. - constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); - Tensor decay_bcast = make_tensor(make_smem_ptr(smem.cumAdt), - make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); - Tensor decay_part = thr_mma.partition_C(decay_bcast); - - // ── Gmem output: partition_C for direct register → gmem store ── - auto* __restrict__ output_ptr = reinterpret_cast(params.output); - // out_base lands on this CTA's D-slice within the head. - int64_t const out_base = out_seq_base + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; - - // Row predicate for padding. The epilogue store loop iterates i in steps - // of 2 and only consults pred(0) and pred(2) — m16n8k16 C-frag per thread - // has 4 elts at rows {t/4, t/4, t/4+8, t/4+8}, so there are only 2 unique - // row predicates. Compute them once and skip the 4-wide pred tensor. - auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_tile); - bool const pred_row_lo = get<0>(id_part(0)) < seq_len; - bool const pred_row_hi = get<0>(id_part(2)) < seq_len; - - // Number of output N-tiles per pass = D_PER_CTA / N_TILE. - // D_SPLIT=1, D_PER_CTA=64, N_TILE=32 → NUM_N_TILES = 2 (current behavior). - // D_SPLIT=2, D_PER_CTA=32 → NUM_N_TILES = 1 (uses _n1 variant). - constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; - static_assert( - NUM_N_TILES == 1 || NUM_N_TILES == 2, "Output epilogue supports NUM_N_TILES = D_PER_CTA / N_TILE in {1, 2}"); - - // ── Epilogue lambda (defined once; called per N-tile from each branch) ── - auto epilogue = [&](auto& frag_y, int n) - { - // Decay: frag_y *= exp(cumAdt[t]) -#pragma unroll - for (int i = 0; i < size(frag_y); ++i) - { - frag_y(i) *= __expf(decay_part(i)); - } - - // frag_y += CB_scaled @ x (CB from smem LDSM, x from smem via ldmatrix.trans) - add_cb_x( - frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); - - // frag_y += D * x[t, d] - add_D_skip(frag_y, smem_x, thr_mma, D_val, n); - - // frag_y *= z * sigmoid(z) - compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); - - // Store frag_y directly to gmem (register → gmem, no smem round-trip). - auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), - make_layout( - make_shape(Int{}, Int{}), make_stride(params.out_stride_token, _1{}))); - auto gOut_part = thr_mma.partition_C(gOut_tile); - // Vectorized pair store: elements i and i+1 are same-row, consecutive columns - // in the m16n8k16 partition_C layout, so &gOut_part(i+1) == &gOut_part(i) + 1. - // Address is naturally aligned to sizeof(Pair) since MMA column - // index = (lane%4)*2 → even. pack_float2 dispatches to the native packed - // cvt (e.g. cvt.rn.bf16x2.f32 for bf16) — one instruction for the pair. -#pragma unroll - for (int i = 0; i < size(frag_y); i += 2) - { - // Bit 1 of i toggles between the two row groups of the m16n8k16 - // C-frag: i∈{0,1} → row t/4, i∈{2,3} → row t/4+8 (repeats per M-atom). - bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; - if (pred_i) - { - *reinterpret_cast*>(&gOut_part(i)) - = pack_float2(make_float2(frag_y(i), frag_y(i + 1))); - } - } - }; - - // Skip the smem→gmem state copy when philox+f16: `replay_state_mma` - // already did the gmem store with stochastic rounding direct from registers. - constexpr bool kSkipSmemToGmemState = (PHILOX_ROUNDS > 0) && std::is_same_v; - - // ── Matmul 3 + store_state + epilogue, dispatching on NUM_N_TILES ── - // (NumNTiles is deduced from the variadic frag_y... pack in `add_init_out`.) - if constexpr (NUM_N_TILES == 2) - { - Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); - Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); - add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0, frag_y_1); - // State writeback hoisted here — after matmul 3 has finished consuming - // smem.state, before matmul 4 which reads only smem.x / smem.CB_scaled - // / smem.z. STGs fire-and-forget alongside the epilogue (matmul 4 + - // D*x + z-gate + output STG). Runtime-gated on must_checkpoint: - // non-checkpoint steps leave the prior state HBM intact (saving - // bandwidth — that's the perf win of the checkpointing design). - if constexpr (!kSkipSmemToGmemState) - { - if (must_checkpoint) - { - store_state( - smem, params, warp, lane, d_tile, head, cache_slot); - } - } - epilogue(frag_y_0, 0); - epilogue(frag_y_1, 1); - } - else - { // NUM_N_TILES == 1 (D_SPLIT = 2 path) - Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); - add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0); - // No sync needed before store_state: the post-replay __syncthreads() - // in the kernel already established cross-warp visibility of replay's - // writes to smem.state, and nothing after that point writes to it - // (add_init_out is read-only on smem.state). - if constexpr (!kSkipSmemToGmemState) - { - if (must_checkpoint) - { - store_state( - smem, params, warp, lane, d_tile, head, cache_slot); - } - } - epilogue(frag_y_0, 0); - } -} - -// ── Orchestrator: compute_no_write_output (must_checkpoint == false path) ── -// Skips the replay matmul entirely. smem.state still holds s_0 after Phase 0, -// so matmul-3 via add_init_out computes u^T = C @ s_0^T directly. -// -// y[t, d] = β(t) · u[t, d] -// + Σ_{j -__device__ __forceinline__ void compute_no_write_output(SmemT& smem, CheckpointingSsuParams const& params, int warp, - int lane, int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) -{ - using namespace cute; - static_assert(sizeof(input_t) == 2, "compute_no_write_output requires 2-byte input type"); - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - int const tid = warp * warpSize + lane; - - // ── TiledMMA for matmul-3 + matmul-4-new (K=NPREDICTED_PAD_MMA_M=16 fits K_BIG). ── - auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(tid); - - // ── TiledMMA for matmul-4-old: K = MAX_WINDOW_PAD_MMA_K ∈ {8, 16} → atom dispatch. ── - using MmaAtomOld = std::conditional_t; - using LdsmAOld = std::conditional_t; - using LdsmBOld = std::conditional_t; - auto tiled_mma_old = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_old = tiled_mma_old.get_slice(tid); - - // ── Swizzled smem views ── - auto layout_x_swz = make_swizzled_layout_rc(); - Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_swz); - auto layout_x_trans_swz = make_swizzled_layout_rc_transpose(); - Tensor smem_x_trans - = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); - - auto layout_old_x_trans_swz = make_swizzled_layout_rc_transpose(); - Tensor smem_old_x_trans - = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), layout_old_x_trans_swz); - - auto layout_z_swz = make_aliased_swizzled_layout_rc(); - Tensor smem_z = make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); - - // ── S2R copies (matmul-4-new) ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - auto s2r_B_trans = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); - - // ── S2R copies (matmul-4-old, K-dispatched atoms) ── - auto s2r_A_old = make_tiled_copy_A(Copy_Atom{}, tiled_mma_old); - auto s2r_thr_A_old = s2r_A_old.get_slice(tid); - auto s2r_B_old_trans = make_tiled_copy_B(Copy_Atom{}, tiled_mma_old); - auto s2r_thr_B_old_trans = s2r_B_old_trans.get_slice(tid); - - // ── Load CB_scaled A operand (cols [0, NPREDICTED_PAD_MMA_M)) ── - auto layout_cb_swz = make_swizzled_layout_rc(); - Tensor smem_CB - = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); - auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); - Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); - auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); - cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); - - // ── Load CB_old A operand (cols [NPREDICTED_PAD_MMA_M, +MAX_WINDOW_PAD_MMA_K)) ── - // Use the full physical (T_pad, CB_ROW_STRIDE) padded swizzle view — byte- - // compatible with both the CB_scaled (T_pad, T_pad, CB_ROW_STRIDE) write - // layout and compute_CB_old_2warp's wide write layout (inner offset - // r*CB_ROW_STRIDE + c is identical across the three views). - auto layout_cb_full = make_swizzled_layout_rc(); - Tensor smem_CB_full - = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); - Tensor smem_CB_old = local_tile(smem_CB_full, make_tile(Int{}, Int{}), - make_coord(_0{}, NPREDICTED_PAD_MMA_M / MAX_WINDOW_PAD_MMA_K)); - auto smem_CB_old_s2r = s2r_thr_A_old.partition_S(smem_CB_old); - Tensor frag_CB_old_A = thr_mma_old.partition_fragment_A(smem_CB_old); - auto frag_CB_old_A_view = s2r_thr_A_old.retile_D(frag_CB_old_A); - cute::copy(s2r_A_old, smem_CB_old_s2r, frag_CB_old_A_view); - - // ── Decay broadcast: cumAdt[t] (per-T scalar) with stride-0 on N. ── - constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); - Tensor decay_bcast = make_tensor(make_smem_ptr(smem.cumAdt), - make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); - Tensor decay_part = thr_mma.partition_C(decay_bcast); - - // ── β extra factor: exp(total_old_cumAdt) — uniform constant across (t, d) ── - float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - float const beta_extra = __expf(total_old_cumAdt); - - // ── Gmem output base ── - auto* __restrict__ output_ptr = reinterpret_cast(params.output); - int64_t const out_base = out_seq_base + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; - - // ── Row predicate (same pattern as compute_and_store_output) ── - auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_tile); - bool const pred_row_lo = get<0>(id_part(0)) < seq_len; - bool const pred_row_hi = get<0>(id_part(2)) < seq_len; - - constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; - static_assert(NUM_N_TILES == 1 || NUM_N_TILES == 2, "compute_no_write_output supports NUM_N_TILES in {1, 2}"); - - // ── Epilogue per N-tile ── - auto epilogue = [&](auto& frag_y, int n) - { - // 1. β-scale: frag_y(t, d) *= exp(total_old_cumAdt + cumAdt[t]). - // Matmul-3 produced u^T = C @ s_0^T; this scales the u term by β - // BEFORE matmul-4 adds the CB·x and CB_old·old_x contributions. -#pragma unroll - for (int i = 0; i < size(frag_y); ++i) - { - frag_y(i) *= beta_extra * __expf(decay_part(i)); - } - - // 2. frag_y += CB_scaled @ x (matmul-4 over new tokens). - add_cb_x( - frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); - - // 3. frag_y += CB_old @ old_x (matmul-4 over old tokens — NEW). - add_cb_old_x(frag_y, frag_CB_old_A, - smem_old_x_trans, s2r_B_old_trans, s2r_thr_B_old_trans, thr_mma_old, tiled_mma_old, n); - - // 4. frag_y += D · x[t, d]. - add_D_skip(frag_y, smem_x, thr_mma, D_val, n); - - // 5. frag_y *= z · sigmoid(z). - compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); - - // 6. Store frag_y → gmem via partition_C (same pattern as compute_and_store_output). - auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), - make_layout( - make_shape(Int{}, Int{}), make_stride(params.out_stride_token, _1{}))); - auto gOut_part = thr_mma.partition_C(gOut_tile); -#pragma unroll - for (int i = 0; i < size(frag_y); i += 2) - { - bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; - if (pred_i) - { - *reinterpret_cast*>(&gOut_part(i)) - = pack_float2(make_float2(frag_y(i), frag_y(i + 1))); - } - } - }; - - // ── Matmul-3: frag_y = C @ s_0^T (smem.state retains s_0 since replay skipped) ── - if constexpr (NUM_N_TILES == 2) - { - Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); - Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); - add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0, frag_y_1); - epilogue(frag_y_0, 0); - epilogue(frag_y_1, 1); - } - else - { // NUM_N_TILES == 1 (D_SPLIT = 2 path) - Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); - add_init_out(smem, tiled_mma, thr_mma, tid, frag_y_0); - epilogue(frag_y_0, 0); - } -} - -// ── Per-path dispatchers (called from checkpointing_ssu_kernel) ── -// ssu_checkpoint: replay → sync → output (today's body). -// ssu_nocheckpoint: sync → no-write output (skips replay). -template -__device__ __forceinline__ void ssu_checkpoint(SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, - int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) -{ - // ── DO NOT HOIST `rand_seed` ── see kernel preamble for the perf rationale. - int64_t const rand_seed = (PHILOX_ROUNDS > 0) ? *params.rand_seed : 0; - // `state_ptr_offset` is int64 — matches Triton's `base_rand = - // cache_batch_idx * stride_state_batch + ...` (cache_batch_idx is .to(int64)). - // Full 64 bits flow through `philox_randint4x`, which splits low/high - // across Philox c0/c1. No collision risk at large serving cache sizes. - int64_t const state_ptr_offset = cache_slot * params.state_stride_seq + (int64_t) head * DIM * DSTATE; - state_t* const state_w_base = reinterpret_cast(params.state) + cache_slot * params.state_stride_seq - + (int64_t) head * DIM * DSTATE + (int64_t) d_tile * D_PER_CTA * DSTATE; - replay_state_mma(smem, params, warp, lane, prev_k, d_tile, - state_ptr_offset, state_w_base, rand_seed, - /*must_checkpoint=*/true); - - __syncthreads(); - - compute_and_store_output( - smem, params, warp, lane, d_tile, out_seq_base, head, cache_slot, D_val, /*must_checkpoint=*/true, seq_len); -} - -template -__device__ __forceinline__ void ssu_nocheckpoint(SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, - int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) -{ - // Sync makes warps 0,1's CB_scaled writes and warps 2,3's CB_old writes - // visible to all warps before matmul-4 reads CB_scaled + CB_old. Also - // covers smem.x (warp 2-loaded) and smem.z (warp 3-loaded) for Phase 2. - __syncthreads(); - - compute_no_write_output( - smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); -} - -// ============================================================================= -// Kernel -// ============================================================================= -template -__global__ void checkpointing_ssu_kernel(CheckpointingSsuParams params) -{ - // Per-head DIM is sharded across `D_SPLIT` CTAs (D_PER_CTA each). - static_assert(DIM % D_SPLIT == 0, "DIM must be divisible by D_SPLIT"); - constexpr int D_PER_CTA = DIM / D_SPLIT; - static_assert(D_PER_CTA >= 32, - "D_PER_CTA must be >= 32 (output MMA m16n8 with _1×4 warp layout). " - "D_SPLIT=4 (D_PER_CTA=16) needs warp-count restructure."); - static_assert(NPREDICTED <= MAX_WINDOW, "NPREDICTED must be <= MAX_WINDOW (new tokens must fit in cache)"); - static_assert( - MAX_WINDOW <= MMA_prop::K_BIG, "MAX_WINDOW must be <= MMA::K_BIG=16 (single replay K-tile assumption)"); - // Cross-check: host launcher must dispatch the template specialization - // matching the runtime params.d_split it stamped into the struct. - assert(params.d_split == D_SPLIT); - using SmemT = CheckpointingSsuStorage; - extern __shared__ __align__(128) char smem_buf[]; - auto& smem = *reinterpret_cast(smem_buf); - - // Grid layout (D_SPLIT, batch, nheads). - int const d_tile = blockIdx.x; - int const seq = blockIdx.y; - int const head = blockIdx.z; - int const lane = threadIdx.x; - int const warp = threadIdx.y; - int const group_idx = head / HEADS_PER_GROUP; - - // ── Resolve cache slot ── - auto const* __restrict__ sbi = reinterpret_cast(params.state_batch_indices); - int64_t const cache_slot = sbi ? static_cast(sbi[seq]) : seq; - if (cache_slot == params.pad_slot_id) - return; - - // ── Double-buffer index ── - auto const* __restrict__ buf_idx_ptr = reinterpret_cast(params.cache_buf_idx); - int const buf_read = __ldg(&buf_idx_ptr[cache_slot]); - - // ── prev_num_accepted_tokens ── - auto const* __restrict__ prev_ptr = reinterpret_cast(params.prev_num_accepted); - int const prev_k = prev_ptr[cache_slot]; - - // ── Varlen vs non-varlen prologue. See checkpointing_ssu_kernel_8bit for - // the rationale: `seq_len` flows downstream as a constexpr-foldable - // NPREDICTED in non-varlen, runtime int in varlen. - // - // Uniform gmem-base formula: `outer * *_stride_seq` where - // non-varlen: outer = seq (= blockIdx.y), stride_seq = x.stride(0). - // varlen : outer = cu_seqlens[seq], stride_seq = x.stride(1). - // The wrapper picks the right stride_seq value; the kernel only branches - // on whether to load cu_seqlens. - int seq_len; - int64_t outer; - if constexpr (VARLEN) - { - auto const* __restrict__ cu_seqlens = reinterpret_cast(params.cu_seqlens); - // Two LDG.E.32 (not one LDG.E.64): cu_seqlens is only 4-byte aligned - // at `&cu_seqlens[seq]` when seq is odd, and PTX - // `ld.global.v2.b32` faults on a 4-byte-aligned address. ptxas emits - // the two scalar loads back-to-back; latency is hidden against the - // following ALU work. - int const bos = __ldg(&cu_seqlens[seq]); - int const eos = __ldg(&cu_seqlens[seq + 1]); - seq_len = eos - bos; - if (seq_len <= 0) - return; - outer = (int64_t) bos; - } - else - { - seq_len = NPREDICTED; - outer = (int64_t) seq; - } - // x/B/C bases are computed inside `load_post_pdl_wait_data` from `outer` - // so the products don't get pinned in registers across `gdc_wait` (asm - // volatile blocks rematerialization; cost was ~6 extra regs). dt/z bases - // are only consumed pre-wait, and out_base only post-replay — fine to - // precompute. - int64_t const dt_seq_base = outer * params.dt_stride_seq + head; - int64_t const z_seq_base = outer * params.z_stride_seq; - int64_t const out_seq_base = outer * params.out_stride_seq; - - // ── Per-CTA implicit checkpoint criterion ── - // When the new tokens would overflow the cache buffer, we must checkpoint: - // replay [0, prev_k) into state, write state to HBM, write the new tokens - // to the **staging** buffer (1 - buf_read) at offset 0. Otherwise, we - // append the new tokens to the **active** buffer (buf_read) at offset - // prev_k and skip the state HBM write entirely. Cache writes always - // happen — only their target buffer + offset depends on must_checkpoint. - bool const must_checkpoint = (prev_k + seq_len > MAX_WINDOW); - int const buf_write = must_checkpoint ? (1 - buf_read) : buf_read; - int const write_offset = must_checkpoint ? 0 : prev_k; - - // ── Load A (scalar, tie_hdim), dt_bias, and D (hoisted to hide gmem latency) ── - auto const* __restrict__ A_ptr = reinterpret_cast(params.A); - auto const* __restrict__ dt_bias_ptr = reinterpret_cast(params.dt_bias); - auto const* __restrict__ D_ptr = reinterpret_cast(params.D); - float const A_val = toFloat(A_ptr[head]); - float const dt_bias_val = dt_bias_ptr ? toFloat(dt_bias_ptr[head]) : 0.f; - float const D_val = D_ptr ? toFloat(D_ptr[head]) : 0.f; - - // ════════════════════════════════════════════════════════════════════════ - // Phase 0: Load all data into smem (per-warp ownership) - // ════════════════════════════════════════════════════════════════════════ - // Two-phase load around the PDL barrier: - // 1. Issue cp.async for cache (state, old_B, old_x) and in_proj-derived - // data (z); run scalar LDGs (old_dt, old_cumAdt, dt → dt_proc) and the - // cumAdt warp scan. None of these depend on conv1d, so they overlap - // with the upstream's tail. - // 2. `gdc_wait()` — wait for the upstream conv1d to signal (no-op when - // the kernel isn't launched with the PDL attribute). - // 3. Issue cp.async for conv1d outputs (x, B, C), then __pipeline_commit - // + __pipeline_wait_prior(0) + __syncwarp drains BOTH halves' cp.async - // (they share the per-thread async group). - // - // Each warp sees its own cp.async via __syncwarp. Cross-warp visibility - // is established by the post-replay __syncthreads below — replay reads of - // state are safe because (a) replay's frag_h initial load sees only the - // current warp's lane positions, and (b) the actual _1×4 cross-warp - // dependency is on writes that haven't happened yet at this point. - // ENABLE_PDL is JIT-stamped (see checkpointing_ssu_customize_config.jinja). - // `if constexpr` keeps only the chosen branch in the binary — no register - // pressure leak from the unused path. - if constexpr (ENABLE_PDL) - { - load_pre_pdl_wait_data(smem, - params, lane, warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, dt_seq_base, - z_seq_base, seq_len); - gdc_wait(); - load_post_pdl_wait_data( - smem, params, lane, warp, d_tile, head, group_idx, outer, seq_len); - } - else - { - load_data(smem, params, lane, - warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, outer, seq_len); - } - - // old_B writeback hoisted ahead of Phase 1. Source (smem.B) is consumed - // only by Phase 1a CB; the STGs fire-and-forget onto the memory subsystem - // and complete in parallel with all subsequent compute. Only W0, W1 hold - // valid smem.B at this point (they're the ones that cp.async'd B). Gate - // accordingly — store halves its thread count but B is small (4 KB) so - // still cheap. old_B is D-independent (per-group, full DSTATE) — only - // d_tile == 0 writes; other d_tiles would emit identical payloads. - if (d_tile == 0 && warp < 2) - { - store_old_B( - smem, params, warp, lane, head, group_idx, cache_slot, buf_write, write_offset, seq_len); - } - - // CB precompute (4-warp split): warps 0,1 compute CB_scaled (new tokens); - // warps 2,3 compute CB_old (old tokens) in the no-write path only. Both - // halves write to disjoint col ranges of the same swizzled smem.CB_scaled - // buffer. In the checkpoint path, warps 2,3 stay idle here and pick up - // work below in `ssu_checkpoint`'s replay matmul. - if (warp < 2) - { - compute_CB_scaled_2warp(smem, warp, lane, seq_len); - } - else if (!must_checkpoint) - { - compute_CB_old_2warp(smem, warp, lane, prev_k, seq_len); - } - - // ════════════════════════════════════════════════════════════════════════ - // Phase 1b + 2: Per-path dispatch - // ════════════════════════════════════════════════════════════════════════ - // Checkpoint path: replay + sync + compute_and_store_output (today's body). - // No-write path : sync + compute_no_write_output (skips replay; matmul-3 - // reads s_0 directly from smem.state, matmul-4 extends with - // the CB_old @ old_x contribution over [0, prev_k)). - // must_checkpoint is uniform across the CTA (derived from broadcast prev_k - // + compile-time NPREDICTED + MAX_WINDOW), so both branches contain a - // __syncthreads and divergence is balanced. - if (must_checkpoint) - { - ssu_checkpoint( - smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); - } - else - { - ssu_nocheckpoint( - smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); - } - - // ── PDL: signal downstream that `output` is written. The cache writes - // below target tensors that only the next SSU step reads, not the - // immediate downstream kernel, so we can signal before issuing them. - if constexpr (ENABLE_PDL) - { - gdc_launch_dependents(); - } - - // ── Phase 3: Store to global memory ── - // (old_B hoisted to pre-Phase-1; state hoisted into compute_and_store_output.) - - // Cache writes — old_x uses all warps (vectorized), dt/cumAdt one warp each. - // Each writes the new NPREDICTED tokens at gmem offset `write_offset` into - // buffer `buf_write` (computed above from must_checkpoint). - store_old_x( - smem, params, warp, lane, d_tile, head, cache_slot, write_offset, seq_len); - // dt_proc / cumAdt are D-independent — only d_tile == 0 writes. - if (d_tile == 0 && warp == 0 && lane < seq_len) - { - auto* __restrict__ old_dt_w = reinterpret_cast(params.old_dt); - int64_t const dt_w_base = cache_slot * params.old_dt_stride_seq + buf_write * params.old_dt_stride_dbuf - + head * params.old_dt_stride_head; - old_dt_w[dt_w_base + write_offset + lane] = smem.dt_proc[lane]; - } - if (d_tile == 0 && warp == 1 && lane < seq_len) - { - auto* __restrict__ old_cumAdt_w = reinterpret_cast(params.old_cumAdt); - int64_t const ca_w_base = cache_slot * params.old_cumAdt_stride_seq + buf_write * params.old_cumAdt_stride_dbuf - + head * params.old_cumAdt_stride_head; - old_cumAdt_w[ca_w_base + write_offset + lane] = smem.cumAdt[lane]; - } -} - -} // namespace flashinfer::mamba::checkpointing - -#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh deleted file mode 100644 index 3322e626ad9d..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_8bit.cuh +++ /dev/null @@ -1,1481 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -#ifndef FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_8BIT_CUH_ -#define FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_8BIT_CUH_ - -// 8-bit (int8, future e4m3) kernel path for the incremental SSU kernel: -// storage, replay, encode, output, and __global__ kernel. - -#include "kernel_checkpointing_ssu_common.cuh" - -namespace flashinfer::mamba::checkpointing -{ - -// ============================================================================= -// 8-bit chain-rewrite storage (sibling of CheckpointingSsuStorage) -// ============================================================================= -// Used by `checkpointing_ssu_kernel_8bit` for int8 and fp8 (e4m3) state. -// Differs from the generic `CheckpointingSsuStorage`: -// 1. No `new_state` staging buffer — matmul-3 chains state's fp32 C-frag -// directly into the next mma's A-operand in registers (à la -// `convert_layout_acc_Aregs`), so no smem round-trip is needed. -// 2. Adds an `output_transpose` buffer used to flip the (D, T) output frag -// back to (T, D) before the gmem STG. Matmul-3/4 in the chain path -// compute init_out^T[D, T] (M=D), so the per-warp M-shard frags must be -// transposed via smem before storing into the (T, D) gmem layout. -// -// All shared Phase 0/1 buffers (CB_scaled, B, C, x, z, old_x, old_B, scalars, -// state) are byte-for-byte identical to the generic struct — Phase 0/1 -// helpers (`compute_CB_scaled_2warp`, B/C/x/z loaders, etc.) are templated on -// `SmemT` and read these by name, so they work unchanged. -template -struct CheckpointingSsuStorage8bit -{ - using state_t = state_t_; - static_assert(sizeof(state_t) == 1, "CheckpointingSsuStorage8bit requires a 1-byte state_t"); - - static constexpr int NPREDICTED = NPREDICTED_; - static constexpr int MAX_WINDOW = MAX_WINDOW_; - static constexpr int D_SMEM_COLS = next_multiple_of::ATOM_COLS>(D_PER_CTA); - static constexpr int NPREDICTED_PAD_MMA_M = next_multiple_of(NPREDICTED); - static constexpr int NPREDICTED_PAD_MMA_N = next_multiple_of(NPREDICTED); - static constexpr int MAX_WINDOW_PAD_MMA_K = next_multiple_of(MAX_WINDOW); - static constexpr int NPREDICTED_SWIZZLE_R = next_multiple_of::ATOM_ROWS>(NPREDICTED); - static constexpr int CB_ROW_STRIDE = SmemSwizzle::ATOM_COLS; - - // Shared Phase 0/1 buffers (same shape/swizzle as `CheckpointingSsuStorage`). - alignas(16) input_t CB_scaled[NPREDICTED_PAD_MMA_M * CB_ROW_STRIDE]; - alignas(16) input_t B[NPREDICTED_PAD_MMA_N * DSTATE]; - alignas(16) input_t C[NPREDICTED_SWIZZLE_R * DSTATE]; - alignas(16) input_t x[NPREDICTED_PAD_MMA_M * D_SMEM_COLS]; - alignas(16) input_t z[NPREDICTED_SWIZZLE_R * D_SMEM_COLS]; - alignas(16) input_t old_x[MAX_WINDOW_PAD_MMA_K * D_SMEM_COLS]; - alignas(16) input_t old_B[MAX_WINDOW_PAD_MMA_K * DSTATE]; - - float old_dt[MAX_WINDOW]; - float old_cumAdt[MAX_WINDOW]; - float dt_proc[NPREDICTED]; - float cumAdt[NPREDICTED]; - // decay[t] = exp(cumAdt[t]) — precomputed at Phase 0 alongside cumAdt so the - // output decay broadcast in `compute_output_8bit` is a plain LDS instead of a - // per-element __expf. Each of the 4 warps redundantly writes the same values - // (same pattern as cumAdt); cross-warp visibility comes via the kernel's - // existing __syncthreads before compute_output_8bit. - float decay[NPREDICTED]; - - // state — int8 input, only LDS'd in the single replay pass. After replay - // completes its dequant + matmul into the C-frag, smem.state is dead — so - // `output_transpose` could in principle alias it (8 KB int8 vs 2 KB bf16 - // overlap easily), but for clarity we keep them separate; alias is a - // Phase-4 micro-optimization. - alignas(16) state_t state[D_PER_CTA * DSTATE]; - - // output_transpose — physical (NPREDICTED_PAD_MMA_M, OUTPUT_TRANSPOSE_ROW_STRIDE) - // input_t scratch buffer with PADDED row stride for bank-conflict-free per-thread - // STS + 16-byte-aligned cooperative LDS.128. Used by `compute_output_int8` to - // flip the per-warp `frag_y_DxT[D, T]` register layout into `(T, D)` gmem order. - // - // Row stride: D_PER_CTA + 8 = 72 bf16 elts = 144 bytes. The 8-elt (16-byte) pad - // gives: - // - 144 % 16 == 0 → LDS.128 / STG.128 stays 16-byte aligned across all rows. - // - 144 / 4 % 32 == 4 → adjacent t-rows shift bank assignment by 4 banks. - // For the m16n8 partition_C STS pattern (per-elt: 4 lanes write at fixed d, - // t ∈ {0, 2, 4, 6} → banks {0, 4, 8, 12} on the padded layout — all distinct, - // no conflicts), the padded layout cuts STS bank conflicts from ~63% of - // wavefronts (NCU v16.0) down to 0%. - // Volume: 16 × 72 × 2 B = 2.25 KB (vs unswizzled 2 KB; +256 B). - static constexpr int OUTPUT_TRANSPOSE_ROW_STRIDE = D_PER_CTA + 8; - alignas(16) input_t output_transpose[NPREDICTED_PAD_MMA_M * OUTPUT_TRANSPOSE_ROW_STRIDE]; -}; - -// ============================================================================= -// State-dtype dispatch helpers -// ============================================================================= -// `state_t` is one of: `int8_t` (symmetric int8, ±127) or `__nv_fp8_e4m3` (fp8 -// e4m3, ±448). Both are 1-byte storage; the kernel's smem layout is identical. -// Differences live in: (a) the RN encode primitive, (b) the QUANT_MAX clip -// bound, and (c) packing/unpacking the byte from a u16 `Pair`. - -template -__device__ __forceinline__ uint8_t state_byte_of(state_t v) -{ - if constexpr (std::is_same_v) - { - return static_cast(static_cast(v)); - } - else - { - static_assert(std::is_same_v, "8-bit state_t must be int8_t or __nv_fp8_e4m3"); - return reinterpret_cast<__nv_fp8_storage_t const&>(v); - } -} - -// fp32 → state_t with RN + saturate. Single-element scalar — the kernel's -// smem layout writes pairs as u16, so per-element conversion fits the -// per-thread fragment topology directly. -template -__device__ __forceinline__ state_t encode_rn_8bit(float x) -{ - if constexpr (std::is_same_v) - { - return conversion::cvt_rni_sat_s8(x); - } - else - { - static_assert(std::is_same_v, "8-bit state_t must be int8_t or __nv_fp8_e4m3"); - // cuda_fp8 ctor compiles to `cvt.rn.satfinite.e4m3.f32` on sm_89+. - return __nv_fp8_e4m3(x); - } -} - -// Per-state-dtype symmetric clip / encode-scale denominator. -// int8: ±127 (matches Triton reference, leaves -128 unused) -// fp8_e4m3fn: ±448 (max finite e4m3 value) -template -__device__ __forceinline__ constexpr float quant_max_8bit() -{ - if constexpr (std::is_same_v) - { - return 127.0f; - } - else - { - static_assert(std::is_same_v, "8-bit state_t must be int8_t or __nv_fp8_e4m3"); - return 448.0f; - } -} - -// SM80 m16n8k16 C-frag → A-frag layout reshape for chained mma (state → -// matmul-3 in the int8 chain rewrite). -// -// Pattern mirrors the SM90 helper at attention/hopper/utils.cuh:103, but: -// - SM80 m16n8 C-frag inner per-thread layout is rank-2 ((col_pair=2, -// row_pair=2)) — there's no inner "N/8" stride mode like SM90. -// - We instead `logical_divide` the *outer* MMA_N axis by 2: each pair of -// m16n8 N-atoms (= 16 cols of the producing mma's N) becomes one K=16 -// atom of the chained m16n8k16 mma's A operand. -// -// Lane-element mapping (verified by hand on the m16n8k16 PTX layout): -// C-frag at (cp, rp, mma_n=2k+kh) maps to: row=tid/4+rp*8, col=4*(tid%2)+cp+(2k+kh)*8 -// A-frag at (cp, rp, kh, mma_k=k) maps to: row=tid/4+rp*8, col=4*(tid%2)+cp+8*kh + 16k -// Same element: (2k+kh)*8 + cp == 8*kh + cp + 16k. ✓ -// -// Input layout: ((2, 2), MMA_M, MMA_N) — m16n8 C-frag -// Output layout: ((2, 2, 2), MMA_M, MMA_N / 2) — m16n8k16 A-frag, -// MMA_K = MMA_N / 2 -template -__forceinline__ __device__ auto convert_layout_acc_Aregs_sm80(Layout acc_layout) -{ - using namespace cute; - using X = Underscore; - static_assert(decltype(size<0, 0>(acc_layout))::value == 2, "C-frag inner mode must be (col_pair=2, row_pair=2)"); - static_assert(decltype(size<0, 1>(acc_layout))::value == 2, "C-frag inner mode must be (col_pair=2, row_pair=2)"); - static_assert(decltype(rank(acc_layout))::value == 3, "C-frag must be rank-3 ((C0,C1), MMA_M, MMA_N)"); - static_assert(decltype(rank(get<0>(acc_layout)))::value == 2, - "SM80 m16n8 C-frag inner is rank-2 (no inner stride mode like SM90)"); - // logical_divide the outer MMA_N axis by 2 → ((2, 2), MMA_M, (2, MMA_N/2)) - auto l = logical_divide(acc_layout, Shape{}); - return make_layout(make_layout(get<0, 0>(l), get<0, 1>(l), get<2, 0>(l)), // ((col_pair, row_pair, k_half)) - get<1>(l), // MMA_M - get<2, 1>(l)); // MMA_K = MMA_N / 2 -} - -// ============================================================================= -// Phase 1b: Replay for QUANTIZED state (int8) with RN encoding. -// ============================================================================= -// state[D, dstate] = dequant(state_q, decode_scale) * total_decay -// + old_x^T @ (coeff * old_B) -// -// Layout: per-warp M-shard via TiledMma `Layout<_4, _1>`. Each warp owns -// D_PER_CTA / 4 D-rows × full DSTATE. This makes amax-over-dstate fully -// warp-local (no atomic, no cross-warp __syncthreads), at the cost of -// loading full B (old_B) from smem in every warp (vs partitioning N -// across warps in the bf16/fp16 path). Constraint: per-warp M must equal -// the m16n8 atom M (=16), so D_PER_CTA must be 64 — the wrapper enforces -// d_split == 1 for int8. -// -// Pipeline: -// 1. Replay n-loop: m16n8 matmul, write fp32 frag → smem.new_state. -// 2. STG redistribution + amax + encode pass (warp-local): -// Each warp covers M_PER_WARP = 16 D-rows × 128 cols of new_state. -// Re-tile 32 lanes as 4 row-groups × 8 col-segments. Per round -// r ∈ [0, 4): each lane reads 16 fp32 (4× LDS.128) for one D-row, -// computes a lane-amax (16 fmaxf), `__shfl_xor` over the 8 -// col-lanes (mask 1, 2, 4) for the full-row amax, encodes 16 int8, -// and STG.128's them to gmem. One writer per row stores -// decode_scale = amax/QUANT_MAX to params.state_scale. -// -// matmul-3 reads new_state (fp32) on the same M-shard partition, so no -// cross-warp visibility is needed. The `__syncthreads` after this -// function returns is for dt_proc / cumAdt visibility (Phase 2), not for -// state. - -// ───────────────────────────────────────────────────────────────────────── -// replay_state_mma_8bit_chain: int8-state chain rewrite — PASS 1 only. -// -// Drops bf16 new_state smem buffer entirely; matmul-3 is fused inline with -// replay HMMA on a per-K-pair cadence (1 K-atom of A in flight at a time). -// -// Pipeline (single fused loop over K-pairs): -// - For kpair ∈ [0, NUM_K_PAIRS=8): -// - Replay 2 m16n8 N-atoms → fp32 frag_h × 2 (16 dstate cols of state). -// - Update per-thread amax (fp32, bit-exact). -// - Cast fp32 → bf16, pack into K-atom-sized A frag (`a_kpair`, -// 8 bf16/thread = 4 32-bit regs). Layout matches the m16n8k16 A -// operand directly (`partition_fragment_A` of the chain TiledMma). -// - LDS one K-atom of B from `smem.C[T_pad, kpair*16..+16]`. -// - `cute::gemm` accumulates one K-atom into `frag_y_DxT`: -// `frag_y_DxT[D, T] += new_state[D, kpair*16..+16] -// @ smem.C[T_pad, kpair*16..+16]^T`. -// - Both `a_kpair` and `b_kpair` go out of scope at iter end. -// Post-loop: -// - Warp-local amax reduce (`__shfl_xor` over 4 col-lanes per row pair). -// - Compute `decode_scale = amax/127`, `encode_scale = 127/amax` per row. -// - STG `decode_scale` to gmem (one writer per (cache, head, d_row)). -// - Return `encode_scale_per_row[2]` to the caller — needed by the -// PASS 2 helper (`encode_state_replay_8bit`) which runs *after* -// `compute_output_8bit` so that `frag_y_DxT`'s 8 fp32 regs are dead by -// the time PASS 2's replay-again runs. -// -// Math identity for the chain (why writing `frag_h(j)` to `a_kpair(local_n*4+j)` -// places the bytes in the m16n8k16 A operand's expected position): -// linear(cp, rp, kh, _, mma_k) = cp + 2*rp + 4*kh + 8*mma_k -// = cp + 2*rp + 4*(mma_n%2) + 8*(mma_n/2) -// = cp + 2*rp + 4*mma_n -// = same linear index as the C-frag for the 2 m16n8 N-atoms making up this -// K-pair. No layout helper needed. -// -// No internal __syncthreads — smem.C is redundantly loaded by all 4 warps -// so chain matmul-3 sees each warp's own data without cross-warp sync. -// The caller's single __syncthreads between all replay passes and -// compute_output_8bit provides smem.CB_scaled / smem.x / smem.z visibility. -template -__device__ __forceinline__ void replay_state_mma_8bit_chain(SmemT& smem, CheckpointingSsuParams const& params, int warp, - int lane, int prev_k, int d_tile, int64_t cache_slot, int head, bool must_checkpoint, FragYDxT& frag_y_DxT, - float (&encode_scale_per_row_out)[2], float (&total_scale_out)[2]) -{ - using namespace cute; - static_assert(sizeof(input_t) == 2, "replay_state_mma_8bit_chain requires 2-byte input_t"); - static_assert(sizeof(state_t) == 1, "replay_state_mma_8bit_chain is for 1-byte state_t (int8/fp8) only"); - static_assert(D_PER_CTA == 64, "replay_state_mma_8bit_chain requires D_PER_CTA == 64 (M-shard, per-warp M=16)."); - - constexpr int NUM_WARPS = 4; - constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; // 16 - static_assert(M_PER_WARP == MMA_prop::M, "Per-warp M must equal m16n8 atom M (=16)"); - - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - int const tid = warp * warpSize + lane; - - // Atom-K dispatch: K_BIG=16 (default), K_SMALL=8 if MAX_WINDOW ≤ 8. - using MmaAtomReplayType - = std::conditional_t; - using LdsmA = std::conditional_t; - using LdsmB = std::conditional_t; - - // Replay TiledMma: M-shard, 4 warps along M, 1 along N. Output is - // ((2,2), 1, NUM_N_PASSES) per thread of fp32 (or bf16 view for new_state). - auto tiled_mma_replay = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_replay = tiled_mma_replay.get_slice(tid); - - // Chain TiledMma: m16n8k16 (always K_BIG=16 since K=DSTATE/16 atoms ≥ 1), - // same M-shard layout as replay. M_per_warp=16 (1 m-atom), - // N=NPREDICTED_PAD_MMA_M (T_pad, ≤ 16 = up to 2 n-atoms per warp), K=DSTATE. - auto tiled_mma_chain = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_chain = tiled_mma_chain.get_slice(tid); - - constexpr int N_PER_PASS = MMA_prop::N; // 8 - constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; // 16 - constexpr int FRAG_SIZE = 4; - constexpr int D_ROWS_PER_THREAD = 2; - constexpr float QUANT_MAX = quant_max_8bit(); - - float const total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - float const total_decay = (prev_k > 0) ? __expf(total_cumAdt) : 1.f; - - int const lane_d = lane / 4; - int const warp_d_base = warp * M_PER_WARP; - - // ── Per-row decode_scale for state init. - auto const* __restrict__ state_scale_ptr = reinterpret_cast(params.state_scale); - int64_t const state_scale_base - = cache_slot * params.state_scale_stride_seq + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; - float decode_scale_in[D_ROWS_PER_THREAD]; - decode_scale_in[0] = state_scale_ptr[state_scale_base + warp_d_base + lane_d]; - decode_scale_in[1] = state_scale_ptr[state_scale_base + warp_d_base + lane_d + 8]; - float total_scale[D_ROWS_PER_THREAD]; - total_scale[0] = decode_scale_in[0] * total_decay; - total_scale[1] = decode_scale_in[1] * total_decay; - - // ── A operand (replay): old_x [MAX_WINDOW_PAD_MMA_K, D_SMEM_COLS] → LDSM_T. - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - auto layout_A_full = make_swizzled_layout_rc_transpose(); - Tensor smem_A_full - = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); - Tensor smem_A - = local_tile(smem_A_full, make_shape(Int{}, Int{}), make_coord(_0{}, _0{})); - - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma_replay); - auto s2r_thr_A = s2r_A.get_slice(tid); - Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); - Tensor frag_A_replay = thr_mma_replay.partition_fragment_A( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - Tensor frag_A_replay_view = s2r_thr_A.retile_D(frag_A_replay); - cute::copy(s2r_A, smem_A_s2r, frag_A_replay_view); - - // ── Bake dB coefficients into frag_A once (8 scale ops), replacing 16× - // per-N-pass compute_dB_scaling on frag_B (64 scale ops). - // dB coefficients c[k] baked into frag_A once, replacing per-N-pass B scaling. - apply_dA_coeff(frag_A_replay, smem, total_cumAdt, prev_k, lane); - - // ── B operand (replay): old_B per-pass. - auto layout_B_replay = make_swizzled_layout_rc_transpose(); - Tensor smem_B_full - = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B_replay); - auto s2r_B_replay = make_tiled_copy_B(Copy_Atom{}, tiled_mma_replay); - auto s2r_thr_B_replay = s2r_B_replay.get_slice(tid); - - // ── State: 1-byte input pointer + manual swizzle offsets (read in BOTH passes). - // Drop bf16 new_state staging — replay's fp32 frag flows directly into the - // register-resident `new_state` tensor below. - state_t* state_base = reinterpret_cast(smem.state); - - // Manual swizzle offsets for m16n8 C-fragment layout (1-byte Swizzle<3,4,3>). - // off = row * 128 + (col ^ ((row & 7) << 4)). - // row_hi = row_lo + 8; (row+8)&7 == row&7 ⇒ off_hi = off_lo + 1024. - // Fragment col within each N_PER_PASS=8 tile: (lane % 4) * 2. - int const row_lo = warp_d_base + lane_d; - int const frag_col_base = (lane & 3) << 1; - int const state_base_lo = row_lo << 7; // row_lo * DSTATE - int const state_xor = (row_lo & 7) << 4; - - float per_thread_amax[D_ROWS_PER_THREAD] = {0.f, 0.f}; - - // No __syncthreads here — smem.C is redundantly loaded by all 4 warps - // (each warp sees its own cp.async via __syncwarp in load_data). Cross-warp - // visibility for smem.CB_scaled / smem.x / smem.z is established by the - // caller's __syncthreads between this function and compute_output_8bit. - - // ── smem.C view + B-operand TiledCopy for chain matmul-3 (hoisted before - // the loop; same view per K-pair, B sliced per K-atom inside the loop). - auto layout_C_swz = make_aliased_swizzled_layout_rc(); - Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C_swz); - auto s2r_B_chain = make_tiled_copy_B(Copy_Atom{}, tiled_mma_chain); - auto s2r_thr_B_chain = s2r_B_chain.get_slice(tid); - - constexpr int NUM_K_PAIRS = NUM_N_PASSES / 2; // 8 for DSTATE=128 - static_assert(NUM_N_PASSES % 2 == 0, "Per-K-pair fusion requires NUM_N_PASSES to be even"); - static_assert(MMA_prop::K_BIG == 16, "Chain mma assumes m16n8k16 K-atom = 16"); - - // ════════════════════════════════════════════════════════════════════════ - // PASS 1 — fused replay + chain matmul-3 (per-K-pair): - // For each kpair ∈ [0, NUM_K_PAIRS): - // - Run 2 replay HMMAs (N-passes 2*kpair, 2*kpair+1) → fp32 frag_h × 2. - // - Update per-thread amax (bit-exact fp32). - // - Pack each pair's 4 fp32 → 4 bf16 into a tiny K-atom-sized A frag - // (`a_kpair` shape ((2,2,2), 1, 1) of bf16 = 8 elts/thread = 4 - // 32-bit regs). Linear positions [local_n*4 .. local_n*4+3] - // within `a_kpair` map to the m16n8k16 A operand's (kh=local_n) - // slice — proven by the linear-index identity in the deleted - // `new_state`-tensor comment above. - // - LDS one K-atom of B (smem.C[T_pad, kpair*16..+16]) into a - // similarly small `b_kpair` frag (4 32-bit regs / thread). - // - `cute::gemm` accumulates one K-atom into `frag_y_DxT`. - // - Both `a_kpair` and `b_kpair` go out of scope at iter end → the - // compiler frees those ~8 32-bit regs/thread for the next iter. - // Net: register footprint drops from the 32 regs of the old register- - // resident `new_state` array (held across the whole loop) to ~8 regs in - // flight. Frees ~24 regs/thread → potentially +1-2 blocks/SM occupancy. - // ════════════════════════════════════════════════════════════════════════ -#pragma unroll - for (int kpair = 0; kpair < NUM_K_PAIRS; ++kpair) - { - // K-atom-sized A frag for chain matmul-3 (filled across the 2 N-passes). - Tensor a_kpair = thr_mma_chain.partition_fragment_A( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - static_assert(decltype(size(a_kpair))::value == 8, "a_kpair must hold 1 m16n8k16 K-atom of A = 8 bf16/thread"); - -#pragma unroll - for (int local_n = 0; local_n < 2; ++local_n) - { - int const n = kpair * 2 + local_n; - int const n_base = n * N_PER_PASS; - - Tensor frag_h = thr_mma_replay.partition_fragment_C( - make_tensor((float*) 0x0, make_shape(Int{}, Int{}))); - static_assert( - decltype(size(frag_h))::value == FRAG_SIZE, "FRAG_SIZE must match the partitioned C-fragment size"); - - // Zero-init accumulator — MMA from scratch, state added after. - clear(frag_h); - - // Replay B operand load. - Tensor smem_B_n = local_tile( - smem_B_full, make_tile(Int{}, Int{}), make_coord(n, _0{})); - auto smem_B_s2r_n = s2r_thr_B_replay.partition_S(smem_B_n); - Tensor frag_B_replay = thr_mma_replay.partition_fragment_B( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - auto frag_B_replay_view = s2r_thr_B_replay.retile_D(frag_B_replay); - cute::copy(s2r_B_replay, smem_B_s2r_n, frag_B_replay_view); - - // Replay HMMA: frag_h = frag_A_scaled @ frag_B (c[k] baked into A). - cute::gemm(tiled_mma_replay, frag_h, frag_A_replay, frag_B_replay, frag_h); - - { - int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); - Pair const p0 = *reinterpret_cast const*>(&state_base[off_lo]); - frag_h(0) += toFloat(p0[Int<0>{}]) * total_scale[0]; - frag_h(1) += toFloat(p0[Int<1>{}]) * total_scale[0]; - Pair const p1 = *reinterpret_cast const*>(&state_base[off_lo + 1024]); - frag_h(2) += toFloat(p1[Int<0>{}]) * total_scale[1]; - frag_h(3) += toFloat(p1[Int<1>{}]) * total_scale[1]; - } - - // Update amax (fp32, bit-exact) AND pack 4 fp32 → 4 bf16 into a_kpair - // at offset local_n*4 (matches A-frag's (kh=local_n) slice). -#pragma unroll - for (int i = 0; i < FRAG_SIZE; i += 2) - { - int const d_idx = i / 2; - float const a0 = fabsf(frag_h(i)); - float const a1 = fabsf(frag_h(i + 1)); - per_thread_amax[d_idx] = fmaxf(per_thread_amax[d_idx], fmaxf(a0, a1)); - - Pair const q - = pack_float2(make_float2(frag_h(i), frag_h(i + 1))); - *reinterpret_cast*>(&a_kpair(local_n * FRAG_SIZE + i)) = q; - } - } - - // ── B operand for chain matmul-3 K-atom: smem.C[T_pad, kpair*16..+16] ── - Tensor smem_C_k = local_tile( - smem_C, make_tile(Int{}, Int{}), make_coord(_0{}, kpair)); - auto smem_C_k_s2r = s2r_thr_B_chain.partition_S(smem_C_k); - Tensor b_kpair = thr_mma_chain.partition_fragment_B( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - auto b_kpair_view = s2r_thr_B_chain.retile_D(b_kpair); - cute::copy(s2r_B_chain, smem_C_k_s2r, b_kpair_view); - - // Single K-atom chain matmul-3: frag_y_DxT += a_kpair @ b_kpair - // frag_y_DxT (pre-zeroed by caller) accumulates across all 8 K-atoms. - cute::gemm(tiled_mma_chain, frag_y_DxT, a_kpair, b_kpair, frag_y_DxT); - } - - // ── Warp-local amax reduce (Layout<_4,_1> → fully warp-local; no atomics). -#pragma unroll - for (int i = 0; i < D_ROWS_PER_THREAD; ++i) - { - per_thread_amax[i] - = fmaxf(per_thread_amax[i], __shfl_xor_sync(constants::MASK_ALL_LANES, per_thread_amax[i], 1)); - per_thread_amax[i] - = fmaxf(per_thread_amax[i], __shfl_xor_sync(constants::MASK_ALL_LANES, per_thread_amax[i], 2)); - } - - // ── encode scale (Triton fall-through for amax==0). - // decode_scale = 1 / encode_scale (mathematically: decode = amax/QUANT_MAX, - // encode = QUANT_MAX/amax, so decode = 1/encode; and when amax==0 both fall - // through to 1.f → 1/1 == 1). Computed inline at the STG below — keeping - // only `encode_scale_per_row` in regs saves 2 fp32 regs across PASS 2. - float encode_scale_per_row[D_ROWS_PER_THREAD]; -#pragma unroll - for (int i = 0; i < D_ROWS_PER_THREAD; ++i) - { - float const a = per_thread_amax[i]; - encode_scale_per_row[i] = (a == 0.f) ? 1.f : (QUANT_MAX / a); - } - - // ── STG decode_scale (one writer per (cache, head, d_row)). - if (must_checkpoint && (lane & 3) == 0) - { - auto* __restrict__ state_scale_w = reinterpret_cast(params.state_scale); -#pragma unroll - for (int i = 0; i < D_ROWS_PER_THREAD; ++i) - { - int const d_row_in_atom = lane_d + (i & 1) * 8; - int const d_row = warp_d_base + d_row_in_atom; - state_scale_w[state_scale_base + d_row] = 1.f / encode_scale_per_row[i]; - } - } - - // Hand `encode_scale_per_row` AND `total_scale` (= OLD decode_scale_in × - // total_decay) to the caller so PASS 2 (encode replay-again) can: - // - dequantize the OLD int8 state with the OLD decode_scale (NOT the NEW - // one we just STG'd above — re-reading params.state_scale in PASS 2 - // would pick up the new value and corrupt the encode), and - // - encode the NEW state with the right encode_scale = 127 / amax. - encode_scale_per_row_out[0] = encode_scale_per_row[0]; - encode_scale_per_row_out[1] = encode_scale_per_row[1]; - total_scale_out[0] = total_scale[0]; - total_scale_out[1] = total_scale[1]; -} - -// ───────────────────────────────────────────────────────────────────────── -// encode_state_replay_8bit: PASS 2 of the int8 chain rewrite. -// -// Re-runs the replay matmul fresh (replay-again), encodes the post-replay -// state fp32 → int8 using `encode_scale_per_row[]` from PASS 1, and STG.16's -// the int8 pairs to gmem. Bit-exact with Triton's fp32-encode path. -// -// Called *after* `compute_output_8bit` so that: -// - `frag_y_DxT`'s 8 fp32 regs are dead (chain matmul-3's accumulator -// was consumed by the output STG). -// - PASS 2's gmem STGs fire alongside `store_old_x` / dt_proc / cumAdt -// writes — all gmem traffic at the kernel tail where there's nothing -// else to do. -// -// The setup (TiledMma, frag_A_replay, smem layouts) is duplicated -// from `replay_state_mma_8bit_chain` — separate stack frame keeps register -// allocation simple and avoids cross-function lifetime tracking. -template -__device__ __forceinline__ void encode_state_replay_8bit(SmemT& smem, CheckpointingSsuParams const& params, int warp, - int lane, int prev_k, int d_tile, int64_t cache_slot, int head, float const (&encode_scale_per_row)[2], - float const (&total_scale)[2], int64_t rand_seed, int64_t state_ptr_offset) -{ - using namespace cute; - static_assert(sizeof(input_t) == 2, "encode_state_replay_8bit requires 2-byte input_t"); - static_assert(sizeof(state_t) == 1, "encode_state_replay_8bit is for 1-byte state_t (int8/fp8) only"); - static_assert(D_PER_CTA == 64, "encode_state_replay_8bit requires D_PER_CTA == 64 (M-shard, per-warp M=16)."); - - constexpr int NUM_WARPS = 4; - constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; - static_assert(M_PER_WARP == MMA_prop::M, "Per-warp M must equal m16n8 atom M (=16)"); - - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - int const tid = warp * warpSize + lane; - - using MmaAtomReplayType - = std::conditional_t; - using LdsmA = std::conditional_t; - using LdsmB = std::conditional_t; - - auto tiled_mma_replay = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_replay = tiled_mma_replay.get_slice(tid); - - constexpr int N_PER_PASS = MMA_prop::N; - constexpr int NUM_N_PASSES = DSTATE / N_PER_PASS; - constexpr int FRAG_SIZE = 4; - constexpr int D_ROWS_PER_THREAD = 2; - - float const total_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - - // total_scale (= OLD decode_scale_in × total_decay) was computed in PASS 1 - // and is passed in by reference. We MUST NOT re-load decode_scale_in from - // params.state_scale here — by the time PASS 2 runs, PASS 1 has already - // STG'd the NEW decode_scale to that same gmem location. - - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - auto layout_A_full = make_swizzled_layout_rc_transpose(); - Tensor smem_A_full - = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), layout_A_full); - Tensor smem_A - = local_tile(smem_A_full, make_shape(Int{}, Int{}), make_coord(_0{}, _0{})); - - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma_replay); - auto s2r_thr_A = s2r_A.get_slice(tid); - Tensor smem_A_s2r = s2r_thr_A.partition_S(smem_A); - Tensor frag_A_replay = thr_mma_replay.partition_fragment_A( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - Tensor frag_A_replay_view = s2r_thr_A.retile_D(frag_A_replay); - cute::copy(s2r_A, smem_A_s2r, frag_A_replay_view); - - // dB coefficients baked into frag_A (same identity as PASS 1). - apply_dA_coeff(frag_A_replay, smem, total_cumAdt, prev_k, lane); - - auto layout_B_replay = make_swizzled_layout_rc_transpose(); - Tensor smem_B_full - = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_B)), layout_B_replay); - auto s2r_B_replay = make_tiled_copy_B(Copy_Atom{}, tiled_mma_replay); - auto s2r_thr_B_replay = s2r_B_replay.get_slice(tid); - - state_t* state_base = reinterpret_cast(smem.state); - - // Manual swizzle offsets (same derivation as replay_state_mma_8bit_chain). - int const lane_d = lane / 4; - int const warp_d_base = warp * M_PER_WARP; - int const row_lo = warp_d_base + lane_d; - int const frag_col_base = (lane & 3) << 1; - int const state_base_lo = row_lo << 7; - int const state_xor = (row_lo & 7) << 4; - - // Philox state for SR — one refresh every 4 n-passes (cvt_rs_sat_s8x4_f32 - // packs 4 int8s per u32 of randomness, so 1 Philox call covers 16 int8s). - [[maybe_unused]] uint32_t rand_idx[4]; - -#pragma unroll - for (int n = 0; n < NUM_N_PASSES; ++n) - { - int const n_base = n * N_PER_PASS; - - Tensor frag_h = thr_mma_replay.partition_fragment_C( - make_tensor((float*) 0x0, make_shape(Int{}, Int{}))); - - // Zero-init accumulator — MMA from scratch, state added after. - clear(frag_h); - - Tensor smem_B_n - = local_tile(smem_B_full, make_tile(Int{}, Int{}), make_coord(n, _0{})); - auto smem_B_s2r_n = s2r_thr_B_replay.partition_S(smem_B_n); - Tensor frag_B_replay = thr_mma_replay.partition_fragment_B( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - auto frag_B_replay_view = s2r_thr_B_replay.retile_D(frag_B_replay); - cute::copy(s2r_B_replay, smem_B_s2r_n, frag_B_replay_view); - - // HMMA: frag_h = frag_A_scaled @ frag_B (c[k] baked into A). - cute::gemm(tiled_mma_replay, frag_h, frag_A_replay, frag_B_replay, frag_h); - - { - int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); - Pair const p0 = *reinterpret_cast const*>(&state_base[off_lo]); - frag_h(0) += toFloat(p0[Int<0>{}]) * total_scale[0]; - frag_h(1) += toFloat(p0[Int<1>{}]) * total_scale[0]; - Pair const p1 = *reinterpret_cast const*>(&state_base[off_lo + 1024]); - frag_h(2) += toFloat(p1[Int<0>{}]) * total_scale[1]; - frag_h(3) += toFloat(p1[Int<1>{}]) * total_scale[1]; - } - - // ── Encode + in-place STS to smem.state at cols [n*8, n*8+8) ── - // Overwrites the OLD 8-bit input *for this n-pass's cols only*. The next - // n-pass dequants from a DIFFERENT col band [(n+1)*8, +8) — still OLD — - // so no read-after-write hazard. Each warp writes only to its own - // M-shard rows; cross-warp visibility is established by the - // caller's __syncthreads before the cooperative store_state call. - { - int const off_lo = state_base_lo + ((frag_col_base + n_base) ^ state_xor); - float const e0 = encode_scale_per_row[0]; - float const e1 = encode_scale_per_row[1]; - - if constexpr (PHILOX_ROUNDS > 0) - { - // One Philox4x call yields 4 independent u32s — enough for 4 n-passes - // when each pass packs all 4 of its 8-bit outputs into one u32 via the - // dtype-specific x4 cvt_rs (int8: `cvt_rs_sat_s8x4_f32`, 16-bit - // randomness/elt via bitrev16 trick; fp8 e4m3: `cvt_rs_e4m3x4_f32`, - // native PTX `cvt.rs.satfinite.e4m3x4.f32` on sm_100a+ with SW fallback). - int const rand_pos = n & 3; - if (rand_pos == 0) - { - int64_t const philox_off = state_ptr_offset + (int64_t) row_lo * DSTATE + (frag_col_base + n_base); - conversion::philox_randint4x( - rand_seed, philox_off, rand_idx[0], rand_idx[1], rand_idx[2], rand_idx[3]); - } - // Packed layout: byte 0 = q0_lo, byte 1 = q1_lo (→ row_lo store at off_lo) - // byte 2 = q0_hi, byte 3 = q1_hi (→ row_hi store at off_lo + 1024) - uint32_t packed; - if constexpr (std::is_same_v) - { - packed = conversion::cvt_rs_sat_s8x4_f32( - frag_h(0) * e0, frag_h(1) * e0, frag_h(2) * e1, frag_h(3) * e1, rand_idx[rand_pos]); - } - else - { - static_assert( - std::is_same_v, "8-bit SR supports state_t in {int8_t, __nv_fp8_e4m3}"); - packed = conversion::cvt_rs_e4m3x4_f32( - frag_h(0) * e0, frag_h(1) * e0, frag_h(2) * e1, frag_h(3) * e1, rand_idx[rand_pos]); - } - Pair q_lo, q_hi; - q_lo.raw = static_cast(packed & 0xFFFFu); - q_hi.raw = static_cast(packed >> 16); - *reinterpret_cast*>(&state_base[off_lo]) = q_lo; - *reinterpret_cast*>(&state_base[off_lo + 1024]) = q_hi; - } - else - { - // d_idx=0: row_lo - state_t const q0_lo = encode_rn_8bit(frag_h(0) * e0); - state_t const q1_lo = encode_rn_8bit(frag_h(1) * e0); - Pair q_lo; - q_lo.raw - = static_cast(state_byte_of(q0_lo)) | (static_cast(state_byte_of(q1_lo)) << 8); - *reinterpret_cast*>(&state_base[off_lo]) = q_lo; - // d_idx=1: row_hi = row_lo + 8, off_hi = off_lo + 1024 - state_t const q0_hi = encode_rn_8bit(frag_h(2) * e1); - state_t const q1_hi = encode_rn_8bit(frag_h(3) * e1); - Pair q_hi; - q_hi.raw - = static_cast(state_byte_of(q0_hi)) | (static_cast(state_byte_of(q1_hi)) << 8); - *reinterpret_cast*>(&state_base[off_lo + 1024]) = q_hi; - } - } - } - - // No __syncthreads or cooperative STG here — the caller's single sync - // provides cross-warp smem.state visibility, then calls store_state. -} - -// ──────────────────────────────────────────────────────────────────────── -// compute_output_8bit: transposed matmul-4 + epilogue + smem-transpose STG -// ──────────────────────────────────────────────────────────────────────── -// Companion to `replay_state_mma_int8_chain`. Consumes the per-warp -// `frag_y_DxT` (shape ((2,2), 1, T_pad/8) of fp32 per thread; M=D-shard, -// N=T_pad) — pre-loaded with init_out^T from chain matmul-3 — and: -// 1. Decay broadcast: frag_y_DxT *= exp(cumAdt[t]) (per T-col, scalar LDS). -// 2. Chain matmul-4 transposed: frag_y_DxT += x^T[D, T] @ CB_scaled^T[T, T] -// A operand: smem.x viewed via x_trans (D, T) → LDSM_N feeds A(M=D, K=T). -// B operand: smem.CB_scaled (T, T) → LDSM_T feeds B(K=T, N=T). -// 3. D*x skip: frag_y_DxT(d, t) += D_val * x[t, d] (scalar LDS per element; -// consecutive frag elts at fixed D, varying T → not pair-loadable). -// 4. z-gate: frag_y_DxT *= z * sigmoid(z) (scalar LDS per element). -// 5. fp32 → input_t pack (in-place register cvt via pack_float2). -// 6. Per-thread STS to smem.output_transpose at (T, D) layout. -// 7. __syncthreads. -// 8. Cooperative STG.128 from smem.output_transpose (T, D) to gmem (T, D). -// -// Cross-warp dependencies (smem.x, smem.z, smem.CB_scaled) are already -// visible because the caller's __syncthreads fires between all replay -// passes and this function. -template -__device__ __forceinline__ void compute_output_8bit(SmemT& smem, CheckpointingSsuParams const& params, int warp, - int lane, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len, - FragYDxT& frag_y_DxT) -{ - using namespace cute; - static_assert(sizeof(input_t) == 2, "compute_output_8bit requires 2-byte input_t"); - static_assert(D_PER_CTA == 64, "compute_output_8bit requires D_PER_CTA == 64"); - static_assert(NUM_WARPS == 4, "compute_output_8bit requires 4 warps"); - - int const tid = warp * warpSize + lane; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - constexpr int M_PER_WARP = D_PER_CTA / NUM_WARPS; // 16 - - // Same TiledMma as replay_state_mma_int8_chain (M-shard, m16n8k16). - auto tiled_mma_chain = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_chain = tiled_mma_chain.get_slice(tid); - - // ── Smem views ── - // x_trans: x physically stored at (T, D); transposed view at (D, T). - // Used as the A operand of the chain matmul-4. - auto layout_x_trans = make_swizzled_layout_rc_transpose(); - Tensor smem_x_trans - = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans); - Tensor smem_x_trans_tile - = local_tile(smem_x_trans, make_shape(Int{}, Int{}), make_coord(_0{}, _0{})); - - // x natural (T, D) view — for D-skip + z-gate per-element scalar LDS. - auto layout_x = make_swizzled_layout_rc(); - - // z natural (T, D) view (aliased so padded rows alias valid rows). - auto layout_z = make_aliased_swizzled_layout_rc(); - - // CB_scaled (T, T_pad) within (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE). - auto layout_cb = make_swizzled_layout_rc(); - Tensor smem_CB - = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb); - - // ── Per-thread (d, t) coord lookup for epilogue scalar reads + smem-transpose write ── - auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma_chain.partition_C(id_tile); - - // ── 1. Decay broadcast: frag_y(i) *= exp(cumAdt[t]) ── - // Reads `smem.decay[t]` (= exp(cumAdt[t])) precomputed in Phase 0 by - // `compute_cumAdt` — fused EX2 with the cumsum write, replacing ~512 per-CTA - // __expf calls in this inner loop with a single LDS per element. - // For padded T-cols (t >= NPREDICTED), the read returns garbage but the STG - // at the end is predicated on t < NPREDICTED, so the garbage never reaches gmem. -#pragma unroll - for (int i = 0; i < size(frag_y_DxT); ++i) - { - int const t = get<1>(id_part(i)); - if (t < seq_len) - { - frag_y_DxT(i) *= smem.decay[t]; - } - } - - // ── 2. Chain matmul-4: frag_y_DxT += x^T @ CB^T ── - // A operand: smem.x physically (T, D); transposed view (D, T) used as - // A(M=D, K=T). The transposed view has D-stride=1, T-stride=D — same - // pattern as replay's A from old_x — so use LDSM_T to produce row-major - // A from this column-wise smem source. - auto s2r_A_x = make_tiled_copy_A(Copy_Atom{}, tiled_mma_chain); - auto s2r_thr_A_x = s2r_A_x.get_slice(tid); - auto smem_x_s2r = s2r_thr_A_x.partition_S(smem_x_trans_tile); - Tensor frag_A_x = thr_mma_chain.partition_fragment_A( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - auto frag_A_x_view = s2r_thr_A_x.retile_D(frag_A_x); - cute::copy(s2r_A_x, smem_x_s2r, frag_A_x_view); - - // B operand for chain matmul-4 = CB^T. smem.CB natural view shape (T, T) - // already has T_inner stride 1 = K-major. Use LDSM_N (no transpose). - auto s2r_B_CB = make_tiled_copy_B(Copy_Atom{}, tiled_mma_chain); - auto s2r_thr_B_CB = s2r_B_CB.get_slice(tid); - auto smem_CB_s2r = s2r_thr_B_CB.partition_S(smem_CB); - Tensor frag_B_CB = thr_mma_chain.partition_fragment_B( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - auto frag_B_CB_view = s2r_thr_B_CB.retile_D(frag_B_CB); - cute::copy(s2r_B_CB, smem_CB_s2r, frag_B_CB_view); - - cute::gemm(tiled_mma_chain, frag_y_DxT, frag_A_x, frag_B_CB, frag_y_DxT); - - // ── 3. D*x skip: frag_y(d, t) += D_val * x[t, d] (scalar LDS per element) ── - if (D_val != 0.f) - { - auto* __restrict__ smem_x_base = reinterpret_cast(smem.x); -#pragma unroll - for (int i = 0; i < size(frag_y_DxT); ++i) - { - int const d = get<0>(id_part(i)); - int const t = get<1>(id_part(i)); - if (t < seq_len) - { - int const off = layout_x(t, d); - frag_y_DxT(i) += D_val * toFloat(smem_x_base[off]); - } - } - } - - // ── 4. z-gate: frag_y *= z * sigmoid(z) (scalar LDS per element) ── - if (params.z != nullptr) - { - auto* __restrict__ smem_z_base = reinterpret_cast(smem.z); -#pragma unroll - for (int i = 0; i < size(frag_y_DxT); ++i) - { - int const d = get<0>(id_part(i)); - int const t = get<1>(id_part(i)); - if (t < seq_len) - { - int const off = layout_z(t, d); - float const z = toFloat(smem_z_base[off]); - frag_y_DxT(i) *= z * __fdividef(1.f, (1.f + __expf(-z))); - } - } - } - - // ── 5. Pack fp32 → input_t per element + 6. STS to smem.output_transpose (T, D) ── - // Padded row stride (D_PER_CTA + 8 = 72 bf16 = 144 bytes) gives: - // - 16-byte-aligned LDS.128 / STG.128 across all rows. - // - 4-bank shift per row → m16n8 STS pattern hits {bank 0, 4, 8, 12} for the - // 4 t-rows of an elt → bank-conflict-free (vs 4-way conflict at stride 64). - // See CheckpointingSsuStorage8bit::OUTPUT_TRANSPOSE_ROW_STRIDE for derivation. - constexpr int kSmemRowStride = SmemT::OUTPUT_TRANSPOSE_ROW_STRIDE; // 72 bf16 elts - auto* __restrict__ smem_out_base = reinterpret_cast(smem.output_transpose); -#pragma unroll - for (int i = 0; i < size(frag_y_DxT); ++i) - { - int const d = get<0>(id_part(i)); - int const t = get<1>(id_part(i)); - if (t < seq_len) - { - // Pack via pack_float2(f, 0.f) and take low elt — emits a single cvt - // (compiler folds the dummy into a no-op for the discarded high half). - smem_out_base[t * kSmemRowStride + d] = pack_float2(make_float2(frag_y_DxT(i), 0.f))[Int<0>{}]; - } - } - - // ── 7. Warp sync for cross-lane STS→LDS ordering ── - __syncwarp(); - - // ── 8. Warp-local cooperative STG.128: 32 lanes → one warp's 16 D-rows ── - // Each warp's data: 16 D-rows × T_pad=16 cols × 2 B = 512 B. - // Re-tile 32 lanes: (t = lane%16, d_group = lane/16 ∈ {0, 1}) → covers - // T_pad × 2 D-groups = 32 slots, each STG.128 = 8 D-cols × 2 B = 16 B. - // No cross-warp coordination → no __syncthreads. - constexpr int kElsPerSTG = 16 / sizeof(input_t); // 8 bf16 elts per STG.128 - constexpr int kDGroupsPerWarp = M_PER_WARP / kElsPerSTG; // = 16 / 8 = 2 - static_assert(NPREDICTED_PAD_MMA_M * kDGroupsPerWarp == 32, - "warp-local STG re-tile: T_pad × dGroupsPerWarp must equal warpSize"); - - int const stg_t = lane % NPREDICTED_PAD_MMA_M; - int const stg_d_group = lane / NPREDICTED_PAD_MMA_M; - int const warp_d_base = warp * M_PER_WARP; - int const stg_d = warp_d_base + stg_d_group * kElsPerSTG; - - if (stg_t < seq_len) - { - int const smem_off = stg_t * kSmemRowStride + stg_d; - - auto* __restrict__ output_ptr = reinterpret_cast(params.output); - int64_t const out_base = out_seq_base + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; - int64_t const gmem_off = out_base + (int64_t) stg_t * params.out_stride_token + stg_d; - - // 128-bit copy. smem_off * 2 B = (t * 144 + d * 2) is 16-byte aligned - // for any t when d % 8 == 0 (here d_offset_within_warp = 0 or 8). - using Vec = uint4; - *reinterpret_cast(&output_ptr[gmem_off]) = *reinterpret_cast(&smem_out_base[smem_off]); - } -} - -// ============================================================================= -// add_init_out_8bit: matmul-3 for the no-checkpoint path (N-shard). -// ============================================================================= -// Computes `frag_y[T, D] = C @ dequant(smem.state)^T` in N-shard `Layout<_1,_4>`, -// reading int8/fp8 state directly from smem and dequanting per-element to bf16 -// in registers before the HMMA. Mirrors the bf16 path's `add_init_out` but with -// a custom B-operand loader since the 1-byte state can't use LDSM directly. -// -// Per-thread B-frag layout (m16n8k16, PTX ISA): -// Per-lane 4 bf16 elts at (K, N) = {(2t, gID), (2t+1, gID), (2t+8, gID), (2t+9, gID)} -// where t = lane%4, gID = lane/4. -// In our matmul-3: B = state^T = (DSTATE = K-axis, D_PER_CTA = N-axis). N-shard -// gives each warp `MMA::N = 8` D-cols per N-tile, so this lane's d_row = -// n_tile*N_TILE + warp*8 + lane/4. -// -// Per K-tile: load 4 int8 bytes per lane (2 byte-pair LDS into Pair), -// CAST to bf16 (no per-row scale), pack into B-frag, HMMA into the n-th frag_y. -// `decode_scale[d]` is per-output-col (constant across the K reduction), so -// the caller pulls it OUT of the inner product and applies it post-matmul in -// the β-scale loop: -// y[t, d] = decode_scale[d] · Σ_n C[t, n] · state_byte[d, n] -// This eliminates the per-cell `... * scale` FMUL chain (was the -// long_scoreboard hotspot at line 1066) at the cost of 2 extra FMUL/elt in -// the post-matmul C-frag scale (net 1792× fewer FMUL per warp). -template -__device__ __forceinline__ void add_init_out_8bit( - SmemT const& smem, int warp, int lane, TiledMma const& tiled_mma, ThrMma const& thr_mma, int tid, FragY&... frag_y) -{ - using namespace cute; - static_assert(sizeof(state_t) == 1, "add_init_out_8bit requires 1-byte state"); - static_assert(D_PER_CTA == 64, "add_init_out_8bit requires D_PER_CTA == 64 (8-bit D_SPLIT=1)"); - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int K_TILE = MMA_prop::K_BIG; // 16 - constexpr int NUM_K_TILES = DSTATE / K_TILE; // 8 - constexpr int N_TILE = cute::tile_size<1>(TiledMma{}); // 32 (4 warps × MMA::N=8) - constexpr int NUM_N_TILES = sizeof...(FragY); // 2 (D_PER_CTA / N_TILE) - static_assert(NUM_N_TILES * N_TILE == D_PER_CTA, "FragY count must match D_PER_CTA / N_TILE"); - - // ── Per-thread coords ── - int const t = lane & 3; // K-pair index within K-atom - int const lane_d = lane >> 2; // gID = lane/4; selects N-col within atom - int const warp_d_base = warp * MMA_prop::N; // warp's 8-col offset within an N-tile - - // ── A operand (C): swizzled (T_pad, DSTATE), K-tiled per K-loop iter ── - auto layout_C_swz = make_aliased_swizzled_layout_rc(); - Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C_swz); - Tensor smem_C_ktiled - = local_tile(smem_C, make_tile(Int{}, Int{}), make_coord(_0{}, _)); - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - auto smem_A_s2r = s2r_thr_A.partition_S(smem_C_ktiled); - Tensor frag_A = thr_mma.partition_fragment_A(smem_C_ktiled(_, _, _0{})); - auto frag_A_view = s2r_thr_A.retile_D(frag_A); - - // ── B-frag (one m16n8k16 atom of B; 4 bf16 elts/lane) ── - Tensor frag_B = thr_mma.partition_fragment_B( - make_tensor((MMA_prop::operand_t*) 0x0, make_shape(Int{}, Int{}))); - static_assert(decltype(size(frag_B))::value == 4, "B-frag must be 4 elts/lane for m16n8k16"); - - // ── Smem state base (1-byte) + Swizzle<3,4,3> XOR formula: - // off = d_row * DSTATE + (K XOR ((d_row & 7) << 4)) - // K within the same swizzle row-group {0..7 mod 8} shares the same XOR mask. ── - state_t const* state_base = reinterpret_cast(smem.state); - - // Pre-clear accumulators (caller doesn't pre-zero — matches bf16 add_init_out). - (clear(frag_y), ...); - - // Parameter-pack indexing via pointer array (same pattern as pipelined_kloop_gemm). - using FragY0 = std::tuple_element_t<0, std::tuple>; - FragY0* frag_y_p[NUM_N_TILES] = {(&frag_y)...}; - - // ── K-loop ── -#pragma unroll - for (int k = 0; k < NUM_K_TILES; ++k) - { - int const K_base = k * K_TILE; - - // Load A K-tile via LDSM (shared across all N-tiles within this K-tile). - cute::copy(s2r_A, smem_A_s2r(_, _, _, k), frag_A_view); - - CUTE_UNROLL - for (int n = 0; n < NUM_N_TILES; ++n) - { - int const d_row = n * N_TILE + warp_d_base + lane_d; - int const state_base_lo = d_row << 7; // d_row * DSTATE (DSTATE=128 → <<7) - int const state_xor = (d_row & 7) << 4; // Swizzle<3,4,3> - int const off_lo = state_base_lo + ((K_base + (t << 1)) ^ state_xor); - int const off_hi = state_base_lo + ((K_base + (t << 1) + 8) ^ state_xor); - - Pair const p_lo = *reinterpret_cast const*>(&state_base[off_lo]); - Pair const p_hi = *reinterpret_cast const*>(&state_base[off_hi]); - - // Pure int8/fp8 → bf16 cast. decode_scale is applied post-matmul in - // the caller's β-scale loop. - Pair const b_lo - = pack_float2(make_float2(toFloat(p_lo[Int<0>{}]), toFloat(p_lo[Int<1>{}]))); - Pair const b_hi - = pack_float2(make_float2(toFloat(p_hi[Int<0>{}]), toFloat(p_hi[Int<1>{}]))); - - // frag_B(0,1) = K-pair at {K_base+2t, K_base+2t+1}; (2,3) = at {+8, +9}. - *reinterpret_cast*>(&frag_B(0)) = b_lo; - *reinterpret_cast*>(&frag_B(2)) = b_hi; - - cute::gemm(tiled_mma, *frag_y_p[n], frag_A, frag_B, *frag_y_p[n]); - } - } -} - -// ============================================================================= -// compute_no_write_output_8bit — N-shard output for the no-checkpoint path. -// ============================================================================= -// Mirror of the bf16 path's `compute_no_write_output`, but with int8/fp8 state. -// Uses N-shard `Layout<_1,_4>` (best smem traffic) instead of the M-shard chain -// — the M-shard exists only for amax reduction in the checkpoint path, which -// doesn't run here. -// -// y[t, d] = β(t) · u[t, d] (matmul-3 via -// add_init_out_8bit) -// + Σ_{j -__device__ __forceinline__ void compute_no_write_output_8bit(SmemT& smem, CheckpointingSsuParams const& params, - int warp, int lane, int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, - int seq_len) -{ - using namespace cute; - static_assert(sizeof(input_t) == 2, "compute_no_write_output_8bit requires 2-byte input_t"); - static_assert(sizeof(state_t) == 1, "compute_no_write_output_8bit is for 1-byte state"); - static_assert(D_PER_CTA == 64, "compute_no_write_output_8bit requires D_PER_CTA == 64"); - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - int const tid = warp * warpSize + lane; - - // ── TiledMMA for matmul-3 + matmul-4-new (m16n8k16) ── - auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(tid); - - // ── TiledMMA for matmul-4-old: K = MAX_WINDOW_PAD_MMA_K ∈ {8, 16} → atom dispatch ── - using MmaAtomOld = std::conditional_t; - using LdsmAOld = std::conditional_t; - using LdsmBOld = std::conditional_t; - auto tiled_mma_old = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_old = tiled_mma_old.get_slice(tid); - - // ── Swizzled smem views ── - auto layout_x_swz = make_swizzled_layout_rc(); - Tensor smem_x = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_swz); - auto layout_x_trans_swz = make_swizzled_layout_rc_transpose(); - Tensor smem_x_trans - = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_trans_swz); - - auto layout_old_x_trans_swz = make_swizzled_layout_rc_transpose(); - Tensor smem_old_x_trans - = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_x)), layout_old_x_trans_swz); - - auto layout_z_swz = make_aliased_swizzled_layout_rc(); - Tensor smem_z = make_tensor(make_smem_ptr(reinterpret_cast(smem.z)), layout_z_swz); - - // ── S2R copies (matmul-4-new) ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - auto s2r_B_trans = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B_trans = s2r_B_trans.get_slice(tid); - - // ── S2R copies (matmul-4-old) ── - auto s2r_A_old = make_tiled_copy_A(Copy_Atom{}, tiled_mma_old); - auto s2r_thr_A_old = s2r_A_old.get_slice(tid); - auto s2r_B_old_trans = make_tiled_copy_B(Copy_Atom{}, tiled_mma_old); - auto s2r_thr_B_old_trans = s2r_B_old_trans.get_slice(tid); - - // ── Load CB_scaled A operand (cols [0, NPREDICTED_PAD_MMA_M)) ── - auto layout_cb_swz = make_swizzled_layout_rc(); - Tensor smem_CB - = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); - auto smem_CB_s2r = s2r_thr_A.partition_S(smem_CB); - Tensor frag_CB_A = thr_mma.partition_fragment_A(smem_CB); - auto frag_CB_A_view = s2r_thr_A.retile_D(frag_CB_A); - cute::copy(s2r_A, smem_CB_s2r, frag_CB_A_view); - - // ── Load CB_old A operand (cols [NPREDICTED_PAD_MMA_M, +MAX_WINDOW_PAD_MMA_K)) ── - auto layout_cb_full = make_swizzled_layout_rc(); - Tensor smem_CB_full - = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); - Tensor smem_CB_old = local_tile(smem_CB_full, make_tile(Int{}, Int{}), - make_coord(_0{}, NPREDICTED_PAD_MMA_M / MAX_WINDOW_PAD_MMA_K)); - auto smem_CB_old_s2r = s2r_thr_A_old.partition_S(smem_CB_old); - Tensor frag_CB_old_A = thr_mma_old.partition_fragment_A(smem_CB_old); - auto frag_CB_old_A_view = s2r_thr_A_old.retile_D(frag_CB_old_A); - cute::copy(s2r_A_old, smem_CB_old_s2r, frag_CB_old_A_view); - - // ── Decay broadcast (per-T scalar, stride-0 on N) ── - constexpr int N_TILE = cute::tile_size<1>(decltype(tiled_mma){}); - Tensor decay_bcast = make_tensor(make_smem_ptr(smem.cumAdt), - make_layout(make_shape(Int{}, Int{}), make_stride(_1{}, _0{}))); - Tensor decay_part = thr_mma.partition_C(decay_bcast); - - float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - float const beta_extra = __expf(total_old_cumAdt); - - // ── Post-matmul-3 state decode_scale (factored OUT of the inner K-product). ── - // y[t, d] = decode_scale[d] · raw_y[t, d] where raw_y = C @ (bf16)(state_byte)^T. - // Per lane, the m16n8 C-frag holds 4 elts spanning 2 unique d-cols (d_lo = 2t, - // d_hi = 2t+1) per N-tile. Indexed by `i & 1` in the epilogue scale loop. - auto const* __restrict__ state_scale_ptr = reinterpret_cast(params.state_scale); - int64_t const state_scale_base - = cache_slot * params.state_scale_stride_seq + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; - constexpr int NUM_N_TILES = D_PER_CTA / N_TILE; - static_assert(NUM_N_TILES == 2, "compute_no_write_output_8bit assumes NUM_N_TILES == 2 (D_PER_CTA=64, N_TILE=32)"); - int const t_col = lane & 3; // 2t and 2t+1 are this lane's two C-frag d-cols - float decode_scale[NUM_N_TILES][2]; - CUTE_UNROLL - for (int n = 0; n < NUM_N_TILES; ++n) - { - int const d_lo = n * N_TILE + warp * MMA_prop::N + (t_col << 1); - decode_scale[n][0] = state_scale_ptr[state_scale_base + d_lo]; - decode_scale[n][1] = state_scale_ptr[state_scale_base + d_lo + 1]; - } - - // ── Gmem output base ── - auto* __restrict__ output_ptr = reinterpret_cast(params.output); - int64_t const out_base = out_seq_base + (int64_t) head * DIM + (int64_t) d_tile * D_PER_CTA; - - // ── Row predicate ── - auto id_tile = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_tile); - bool const pred_row_lo = get<0>(id_part(0)) < seq_len; - bool const pred_row_hi = get<0>(id_part(2)) < seq_len; - - auto epilogue = [&](auto& frag_y, int n) - { - // β-scale + state decode_scale fused: frag_y(i) *= β · exp(cumAdt[t]) · decode_scale[d]. - // The decode_scale[d] absorbs the per-row state quant factor (was previously - // multiplied into each B-element during dequant). -#pragma unroll - for (int i = 0; i < size(frag_y); ++i) - { - int const d_idx = i & 1; // i=0,2 → d_lo; i=1,3 → d_hi - frag_y(i) *= beta_extra * __expf(decay_part(i)) * decode_scale[n][d_idx]; - } - - // matmul-4-new (CB_scaled @ x). - add_cb_x( - frag_y, frag_CB_A, smem_x_trans, s2r_B_trans, s2r_thr_B_trans, thr_mma, tiled_mma, n); - - // matmul-4-old (CB_old @ old_x). - add_cb_old_x(frag_y, frag_CB_old_A, - smem_old_x_trans, s2r_B_old_trans, s2r_thr_B_old_trans, thr_mma_old, tiled_mma_old, n); - - // D·x. - add_D_skip(frag_y, smem_x, thr_mma, D_val, n); - - // z-gate. - compute_z_gating(frag_y, smem_z, thr_mma, params.z, n); - - // Direct partition_C STG. - auto gOut_tile = make_tensor(make_gmem_ptr(output_ptr + out_base + n * N_TILE), - make_layout( - make_shape(Int{}, Int{}), make_stride(params.out_stride_token, _1{}))); - auto gOut_part = thr_mma.partition_C(gOut_tile); -#pragma unroll - for (int i = 0; i < size(frag_y); i += 2) - { - bool const pred_i = (i & 2) ? pred_row_hi : pred_row_lo; - if (pred_i) - { - *reinterpret_cast*>(&gOut_part(i)) - = pack_float2(make_float2(frag_y(i), frag_y(i + 1))); - } - } - }; - - // ── Matmul-3: frag_y = C @ (bf16)(state_byte)^T (smem.state retains s_0 since - // replay skipped; decode_scale[d] applied post-matmul in the epilogue). ── - Tensor frag_y_0 = thr_mma.partition_fragment_C(id_tile); - Tensor frag_y_1 = thr_mma.partition_fragment_C(id_tile); - add_init_out_8bit( - smem, warp, lane, tiled_mma, thr_mma, tid, frag_y_0, frag_y_1); - epilogue(frag_y_0, 0); - epilogue(frag_y_1, 1); -} - -// ============================================================================= -// Per-path dispatcher: no-checkpoint branch (must_checkpoint == false). -// ============================================================================= -// Sync makes warps 0,1's CB_scaled writes AND warps 2,3's CB_old writes -// visible to all warps before matmul-3 and matmul-4 read smem.{CB_scaled, -// CB_old, x, z}. Matches the bf16 path's `ssu_nocheckpoint`. -template -__device__ __forceinline__ void ssu_nocheckpoint_8bit(SmemT& smem, CheckpointingSsuParams const& params, int warp, - int lane, int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) -{ - __syncthreads(); - compute_no_write_output_8bit( - smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); -} - -// ============================================================================= -// Per-path dispatcher: checkpoint branch (must_checkpoint == true). -// ============================================================================= -// Encapsulates the existing M-shard chain: PASS 1 replay+matmul-3 → frag_y_DxT, -// PASS 2 re-replay+encode (always runs since must_checkpoint==true here), the -// single __syncthreads, cooperative state STG, and the transposed matmul-4 + -// transpose-STG output. -// -// Pulled out of `checkpointing_ssu_kernel_8bit` to mirror the bf16 path's -// `ssu_checkpoint` and make the kernel-body dispatch on -// must_checkpoint readable. -template -__device__ __forceinline__ void ssu_checkpoint_8bit(SmemT& smem, CheckpointingSsuParams const& params, int warp, - int lane, int prev_k, int d_tile, int64_t out_seq_base, int head, int64_t cache_slot, float D_val, int seq_len) -{ - using namespace cute; - int const tid = warp * warpSize + lane; - - // ── Allocate per-warp frag_y_DxT (chain mma C-frag, fp32) ── - // Layout ((2, 2), MMA_M=1, MMA_N=NPREDICTED_PAD_MMA_M/8) per thread. - // Caller must zero before chain matmul-3 accumulates. - auto tiled_mma_chain = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma_chain = tiled_mma_chain.get_slice(tid); - auto id_DxT = make_identity_tensor(make_shape(Int{}, Int{})); - Tensor frag_y_DxT = thr_mma_chain.partition_fragment_C(id_DxT); - cute::clear(frag_y_DxT); - - // ── Phase 1b: replay + amax + chain matmul-3 → frag_y_DxT (init_out^T). - // `encode_scale_per_row[]` is computed at the end of PASS 1 (from the warp- - // reduced amax) and consumed by `encode_state_replay_8bit` further below — - // *after* `compute_output_8bit` consumes `frag_y_DxT` and STGs the output. - float encode_scale_per_row[2]; - float total_scale[2]; - replay_state_mma_8bit_chain(smem, params, warp, lane, prev_k, d_tile, - cache_slot, head, /*must_checkpoint=*/true, frag_y_DxT, encode_scale_per_row, total_scale); - - // ── Philox seed for stochastic rounding (deferred to reduce register pressure) ── - [[maybe_unused]] int64_t const rand_seed = (PHILOX_ROUNDS > 0) ? *params.rand_seed : 0; - // `state_ptr_offset` is int64 — matches Triton's `base_rand = - // cache_batch_idx * stride_state_batch + ...` (cache_batch_idx is .to(int64)). - // Full 64 bits flow through `philox_randint4x`, which splits low/high - // across Philox c0/c1. No collision risk at large serving cache sizes. - int64_t const state_ptr_offset = cache_slot * params.state_stride_seq + (int64_t) head * DIM * DSTATE; - - // ── PASS 2 (replay-again): re-run replay HMMA, encode fp32 → int8 to - // smem.state. Runs BEFORE the sync so both replay passes overlap with - // warps 0,1's CB precompute — one fewer __syncthreads in the kernel. - // frag_y_DxT stays live through PASS 2 (extra register pressure accepted). ── - encode_state_replay_8bit(smem, params, warp, lane, prev_k, - d_tile, cache_slot, head, encode_scale_per_row, total_scale, rand_seed, state_ptr_offset); - - // ── Single sync: cross-warp visibility for smem.CB_scaled (warps 0,1) / - // smem.x (warp 2) / smem.z (warp 3) / smem.state (all warps' M-shards). ── - __syncthreads(); - - // ── Cooperative STG.128 for encoded state (after sync for cross-warp - // smem.state visibility). Fire-and-forget before compute_output_8bit. ── - store_state(smem, params, warp, lane, d_tile, head, cache_slot); - - // ── Phase 2: transposed matmul-4 + epilogue + smem-transpose STG ── - compute_output_8bit( - smem, params, warp, lane, d_tile, out_seq_base, head, cache_slot, D_val, seq_len, frag_y_DxT); -} - -// ============================================================================= -// Kernel — int8 chain rewrite (separate kernel from the generic path) -// ============================================================================= -// The int8 path uses a fundamentally different output computation: -// 1. M-shard replay (Layout<_4, _1>) — same as v15.4. -// 2. Chained matmul-3: replay's fp32 C-frag → bf16 A-frag in registers via -// `convert_layout_acc_Aregs_sm80` (no smem.new_state staging). -// 3. Transposed matmul-4: x as A (M=D), CB^T as B → output^T(D, T) in regs. -// 4. Smem-transpose + cooperative STG.128 to (T, D) gmem. -// To keep the generic kernel uncluttered (no `if constexpr (sizeof(state_t) == 1)` -// branches), the int8 kernel is a standalone function that calls the new -// helpers (`replay_state_mma_8bit_chain`, `compute_output_8bit`) and uses -// `CheckpointingSsuStorage8bit` for smem. Phase 0/1 helpers (`load_data`, -// `store_old_B`, `compute_CB_scaled_2warp`) are reused verbatim — they only -// touch shared smem fields that both storage structs expose by name. -// -template -__global__ void checkpointing_ssu_kernel_8bit(CheckpointingSsuParams params) -{ - using namespace cute; - static_assert(sizeof(state_t) == 1, "checkpointing_ssu_kernel_8bit requires 1-byte state_t (int8 or fp8 e4m3)"); - static_assert(NPREDICTED <= MAX_WINDOW); - static_assert(MAX_WINDOW <= MMA_prop::K_BIG); - // int8 path uses M-shard layout (Layout<_4,_1>): per-warp M = 16 = m16n8 - // atom M. D_PER_CTA must equal DIM (D_SPLIT=1) to give 4×16=64 D-rows/CTA. - // The wrapper enforces d_split == 1 for int8. - constexpr int D_PER_CTA = DIM; - static_assert(D_PER_CTA == 64, "int8 chain kernel requires DIM == 64"); - assert(params.d_split == 1); - - using SmemT = CheckpointingSsuStorage8bit; - extern __shared__ __align__(128) char smem_buf[]; - auto& smem = *reinterpret_cast(smem_buf); - - // Grid: (1, batch, nheads). D-tile is always 0 for int8 (D_SPLIT=1). - int const d_tile = blockIdx.x; - int const seq = blockIdx.y; - int const head = blockIdx.z; - int const lane = threadIdx.x; - int const warp = threadIdx.y; - int const group_idx = head / HEADS_PER_GROUP; - - // ── Resolve cache slot ── - auto const* __restrict__ sbi = reinterpret_cast(params.state_batch_indices); - int64_t const cache_slot = sbi ? static_cast(sbi[seq]) : seq; - if (cache_slot == params.pad_slot_id) - return; - - auto const* __restrict__ buf_idx_ptr = reinterpret_cast(params.cache_buf_idx); - int const buf_read = __ldg(&buf_idx_ptr[cache_slot]); - - auto const* __restrict__ prev_ptr = reinterpret_cast(params.prev_num_accepted); - int const prev_k = prev_ptr[cache_slot]; - - // ── Varlen vs non-varlen prologue. The kernel branches once on the - // VARLEN template; downstream helpers receive `seq_len` (constexpr-foldable - // NPREDICTED in non-varlen, runtime in varlen) and pre-computed per-sequence - // gmem base offsets (`x_seq_base` etc.) — they're varlen-agnostic. - // - // Uniform gmem-base formula: `outer * *_stride_seq` where - // non-varlen: outer = seq (= blockIdx.y), stride_seq = x.stride(0). - // varlen : outer = cu_seqlens[seq], stride_seq = x.stride(1). - // The wrapper picks the right stride_seq value; the kernel only branches - // on whether to load cu_seqlens. - int seq_len; - int64_t outer; - if constexpr (VARLEN) - { - auto const* __restrict__ cu_seqlens = reinterpret_cast(params.cu_seqlens); - // Two LDG.E.32 (not one LDG.E.64): cu_seqlens is only 4-byte aligned - // at `&cu_seqlens[seq]` when seq is odd, and PTX - // `ld.global.v2.b32` faults on a 4-byte-aligned address. ptxas emits - // the two scalar loads back-to-back; latency is hidden against the - // following ALU work. - int const bos = __ldg(&cu_seqlens[seq]); - int const eos = __ldg(&cu_seqlens[seq + 1]); - seq_len = eos - bos; - if (seq_len <= 0) - return; - outer = (int64_t) bos; - } - else - { - seq_len = NPREDICTED; - outer = (int64_t) seq; - } - // x/B/C bases computed inside `load_post_pdl_wait_data` from `outer` — - // see generic kernel for rationale (avoid pinning 6 regs across gdc_wait). - int64_t const dt_seq_base = outer * params.dt_stride_seq + head; - int64_t const z_seq_base = outer * params.z_stride_seq; - int64_t const out_seq_base = outer * params.out_stride_seq; - - bool const must_checkpoint = (prev_k + seq_len > MAX_WINDOW); - int const buf_write = must_checkpoint ? (1 - buf_read) : buf_read; - int const write_offset = must_checkpoint ? 0 : prev_k; - - // ── Load scalars (A, dt_bias, D) ── - auto const* __restrict__ A_ptr = reinterpret_cast(params.A); - auto const* __restrict__ dt_bias_ptr = reinterpret_cast(params.dt_bias); - auto const* __restrict__ D_ptr = reinterpret_cast(params.D); - float const A_val = toFloat(A_ptr[head]); - float const dt_bias_val = dt_bias_ptr ? toFloat(dt_bias_ptr[head]) : 0.f; - float const D_val = D_ptr ? toFloat(D_ptr[head]) : 0.f; - - // ── Phase 0: two-phase load around the PDL barrier (see generic kernel - // for the full rationale). Pre-wait: state + old_* cache + in_proj - // outputs (dt, z) + scalar scans. Post-wait: x/B/C from conv1d. ── - // ENABLE_PDL is JIT-stamped; `if constexpr` keeps only one load path in - // the binary (no register pressure leak from the unused path). - if constexpr (ENABLE_PDL) - { - load_pre_pdl_wait_data(smem, - params, lane, warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, dt_seq_base, - z_seq_base, seq_len); - gdc_wait(); - load_post_pdl_wait_data( - smem, params, lane, warp, d_tile, head, group_idx, outer, seq_len); - } - else - { - load_data(smem, params, lane, - warp, d_tile, head, group_idx, cache_slot, buf_read, A_val, dt_bias_val, outer, seq_len); - } - - // ── store_old_B hoist (warps 0,1 only, d_tile == 0) ── - if (d_tile == 0 && warp < 2) - { - store_old_B( - smem, params, warp, lane, head, group_idx, cache_slot, buf_write, write_offset, seq_len); - } - - // ── CB precompute (4-warp split): warps 0,1 compute CB_scaled (new tokens); - // warps 2,3 compute CB_old (old tokens) in the no-write path only. Mirrors - // the bf16 path's dispatch — warps 2,3 stay idle in checkpoint mode and - // pick up work below inside `ssu_checkpoint_8bit`'s replay. ── - if (warp < 2) - { - compute_CB_scaled_2warp(smem, warp, lane, seq_len); - } - else if (!must_checkpoint) - { - compute_CB_old_2warp(smem, warp, lane, prev_k, seq_len); - } - - // ── Phase 1b + 2: per-path dispatch ── - // Checkpoint: M-shard chain (PASS 1 + PASS 2 + sync + state STG + transposed - // matmul-4 with smem-transpose STG). - // No-write : N-shard matmul-3 from int8/fp8 state + matmul-4-new + matmul-4-old - // + direct partition_C STG (mirrors the bf16 no-write path). - // must_checkpoint is uniform across the CTA — both branches contain a - // __syncthreads so divergence is balanced. - if (must_checkpoint) - { - ssu_checkpoint_8bit( - smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); - } - else - { - ssu_nocheckpoint_8bit( - smem, params, warp, lane, prev_k, d_tile, out_seq_base, head, cache_slot, D_val, seq_len); - } - - // ── PDL: signal downstream that `output` is written. Cache writes below - // target tensors only the next SSU step reads, not the immediate - // downstream kernel — safe to signal first. ── - if constexpr (ENABLE_PDL) - { - gdc_launch_dependents(); - } - - // ── Phase 3: cache writes (old_x, dt_proc, cumAdt) ── - store_old_x( - smem, params, warp, lane, d_tile, head, cache_slot, write_offset, seq_len); - if (d_tile == 0 && warp == 0 && lane < seq_len) - { - auto* __restrict__ old_dt_w = reinterpret_cast(params.old_dt); - int64_t const dt_w_base = cache_slot * params.old_dt_stride_seq + buf_write * params.old_dt_stride_dbuf - + head * params.old_dt_stride_head; - old_dt_w[dt_w_base + write_offset + lane] = smem.dt_proc[lane]; - } - if (d_tile == 0 && warp == 1 && lane < seq_len) - { - auto* __restrict__ old_cumAdt_w = reinterpret_cast(params.old_cumAdt); - int64_t const ca_w_base = cache_slot * params.old_cumAdt_stride_seq + buf_write * params.old_cumAdt_stride_dbuf - + head * params.old_cumAdt_stride_head; - old_cumAdt_w[ca_w_base + write_offset + lane] = smem.cumAdt[lane]; - } -} - -} // namespace flashinfer::mamba::checkpointing - -#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_8BIT_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh deleted file mode 100644 index c22e9c62c304..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/kernel_checkpointing_ssu_common.cuh +++ /dev/null @@ -1,1733 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -#ifndef FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_COMMON_CUH_ -#define FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_COMMON_CUH_ - -// Shared infrastructure for the incremental SSU kernel: utilities, loaders, -// stores, MMA helpers, and functions used by both the 2/4-byte (bf16/fp16/fp32) -// and 8-bit (int8, future e4m3) kernel paths. - -#include -#include - -#include -#include - -#include "../utils.cuh" -#include "../vec_dtypes.cuh" -#include "checkpointing_ssu.cuh" -#include "common.cuh" -#include "conversion.cuh" -#include "cute/tensor.hpp" -#include "ssu_mtp_common.cuh" - -namespace flashinfer::mamba::checkpointing -{ - -using namespace conversion; - -// ldmatrix.b8 (SM100_U8x16_LDSM_T) was tried as a replacement for per-lane -// LDS.16 int8 state loads. It's 5-18% slower (bench v16.7b vs v16.8) because: -// (1) inherent 2-way bank conflicts (16 threads × 16B vs 128B banks), -// (2) state is the accumulator (C-frag), not an A/B operand — layout -// remapping costs 8 shuffles + byte extractions, -// (3) dynamic byte selection via SHF adds 15%+ short_scoreboard stalls. -namespace constants -{ -constexpr unsigned int MASK_ALL_LANES = 0xFFFFFFFFu; -constexpr unsigned int num_bits_uint32 = 32u; -} // namespace constants - -// ── Programmatic Dependent Launch (PDL) helpers ──────────────────────────── -// `gdc_wait` enforces no gmem access before the upstream PDL-paired kernel -// has signaled. `gdc_launch_dependents` hints the downstream PDL-paired -// kernel to launch early. Both are no-ops on SM<90 and harmless without -// the launch-time `cudaLaunchAttributeProgrammaticStreamSerialization` -// attribute, so the kernel can always emit them; the host-side `enable_pdl` -// toggle is what flips the launch attribute. -__forceinline__ __device__ void gdc_wait() -{ -#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); -#endif -} - -__forceinline__ __device__ void gdc_launch_dependents() -{ -#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); -#endif -} - -// Round x up to the next multiple of Y (Y must be a power of 2). -template -constexpr int next_multiple_of(int x) -{ - static_assert(Y > 0 && (Y & (Y - 1)) == 0, "Y must be a power of 2"); - return (x + Y - 1) & ~(Y - 1); -} - -// NativeOf::type: scalar T → 2-wide native CUDA vector type. -template -struct NativeOf; - -template <> -struct NativeOf -{ - using type = float2; -}; - -template <> -struct NativeOf<__half> -{ - using type = __half2; -}; - -template <> -struct NativeOf<__nv_bfloat16> -{ - using type = __nv_bfloat162; -}; - -// Pair: thin wrapper over the native 2-wide type, adding compile-time -// `[cute::Int{}]` indexing so index-driven loops stay branchless. Same -// layout as the native type — `=` compiles to one LDS.U32 / STS.U32 and the -// pair stays in one register. -template -struct Pair -{ - typename NativeOf::type raw; - - template - __device__ __forceinline__ auto operator[](cute::Int) const - { - static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); - if constexpr (I == 0) - return raw.x; - else - return raw.y; - } -}; - -// Pair: explicit 16-bit packed specialization. A struct of two -// `int8_t` fields would let the compiler split the pair into two 32-bit -// registers (CUDA registers are 32-bit; sub-word fields are zero-extended -// per access). By backing the pair with a single 16-bit word and -// extracting via shift+cast, we keep the two elements in one register -// throughout the load → unpack → cast pipeline. -template <> -struct Pair -{ - uint16_t raw; // [bits 7:0] = element 0, [bits 15:8] = element 1 - - template - __device__ __forceinline__ int8_t operator[](cute::Int) const - { - static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); - if constexpr (I == 0) - return static_cast(raw & 0xFFu); - else - return static_cast(raw >> 8); - } -}; - -// Pair<__nv_fp8_e4m3>: same single-u16 backing as `Pair` — fp8 e4m3 -// is also a 1-byte storage type, so the load → unpack → cast pipeline runs -// through one 16-bit register. -template <> -struct Pair<__nv_fp8_e4m3> -{ - uint16_t raw; - - template - __device__ __forceinline__ __nv_fp8_e4m3 operator[](cute::Int) const - { - static_assert(I == 0 || I == 1, "Pair index must be 0 or 1"); - __nv_fp8_storage_t const byte - = (I == 0) ? static_cast<__nv_fp8_storage_t>(raw & 0xFFu) : static_cast<__nv_fp8_storage_t>(raw >> 8); - return reinterpret_cast<__nv_fp8_e4m3 const&>(byte); - } -}; - -// pack_float2: float2 → Pair, using packed hardware cvt when available. -template -__device__ __forceinline__ Pair pack_float2(float2 val); - -template <> -__device__ __forceinline__ Pair pack_float2(float2 val) -{ - return {val}; -} - -template <> -__device__ __forceinline__ Pair<__half> pack_float2<__half>(float2 val) -{ - return {__float22half2_rn(val)}; -} - -template <> -__device__ __forceinline__ Pair<__nv_bfloat16> pack_float2<__nv_bfloat16>(float2 val) -{ - return {conversion::fromFloat2(val)}; -} - -// ============================================================================= -// cp.async copy atoms -// ============================================================================= -// 128-bit vector loads, shared by every gmem→smem copy in the kernel. The -// ldmatrix unit has the same vector width so `vec_bytes` is derived from the -// atom's source-register type and reused as the LDSM vector width. -struct Copy_prop -{ - using Atom = cute::SM80_CP_ASYNC_CACHEALWAYS; - using AtomZFill = cute::SM80_CP_ASYNC_CACHEALWAYS_ZFILL; - static constexpr int vec_bytes = sizeof(std::remove_extent_t); -}; - -// ============================================================================= -// MMA constants -// ============================================================================= -// All MMA-related atom types, dtype, and dims for this kernel, grouped so the -// header doesn't sprinkle loose aliases. The replay step chooses between the -// k=8 and k=16 atoms at compile time (MAX_WINDOW ≤ 8 picks K8 for smaller smem, -// +1 CTA/SM); dims are pulled from MMA_Traits so they stay in sync with the -// atom choice (e.g. m16n8k32 for int8 would just need AtomK16/K8 swapped). -struct MMA_prop -{ - using AtomK16 = cute::SM80_16x8x16_F32BF16BF16F32_TN; - using AtomK8 = cute::SM80_16x8x8_F32BF16BF16F32_TN; - // Operand dtype — matches the bf16 input of the atoms above. - using operand_t = __nv_bfloat16; - - static constexpr int M = cute::size<0>(typename cute::MMA_Traits::Shape_MNK{}); - static constexpr int N = cute::size<1>(typename cute::MMA_Traits::Shape_MNK{}); - static constexpr int K_BIG = cute::size<2>(typename cute::MMA_Traits::Shape_MNK{}); - static constexpr int K_SMALL = cute::size<2>(typename cute::MMA_Traits::Shape_MNK{}); -}; - -// ============================================================================= -// Swizzled smem layout for mma.sync operands (row-major). -// ============================================================================= -// The swizzle picks the `M` parameter to make each ldmatrix / cp.async atom -// exactly 16 bytes of contiguous element data (one 128-bit vector), and keeps -// B = S = 3 so that each 8-row block XORs row↔column bits to stay -// bank-conflict-free on the 128-byte bank cycle. -// -// sizeof(T) Swizzle atom rows × cols row bytes -// 2B Swizzle<3, 3, 3> 8 × 64 128 -// 4B Swizzle<3, 2, 3> 8 × 32 128 -// 1B Swizzle<3, 4, 3> 8 × 128 128 -// -// The MMA operand element type dictates the smem buffer element type, which in -// turn dictates the swizzle — so every call site passes its own element type. -constexpr int log2_pow2(int x) -{ - int r = 0; - while (x > 1) - { - x >>= 1; - ++r; - } - return r; -} - -template -struct SmemSwizzle -{ - static_assert(Copy_prop::vec_bytes % sizeof(Elem) == 0, "element size must divide LDSM atom (16 bytes)"); - static constexpr int ELEMS_PER_ATOM = Copy_prop::vec_bytes / sizeof(Elem); - using type = cute::Swizzle<3, log2_pow2(ELEMS_PER_ATOM), 3>; - static constexpr int ATOM_ROWS = 1 << type::num_bits; - static constexpr int ATOM_COLS = 1 << (type::num_base + type::num_shft); -}; - -// Default (ROW_STRIDE == COLS): tile the swizzle atom into a (ROWS, COLS) -// physical extent — the canonical CuTe pattern. -// Padded (ROW_STRIDE > COLS): logical (ROWS, COLS) view with the row stride -// inflated to ROW_STRIDE. Used when COLS doesn't tile cleanly with the -// swizzle atom's col extent (e.g. CB_scaled: logical 16 cols, atom 64) but -// we want the atom-aligned bank pattern. The extra cols-per-row are not -// "wasted padding" — the swizzle XOR scatters logical cells across the full -// ROW_STRIDE, so the physical extent is what the bijection actually needs. -template -__device__ __forceinline__ auto make_swizzled_layout_rc() -{ - using namespace cute; - using S = SmemSwizzle; - static_assert(ROWS % S::ATOM_ROWS == 0, "ROWS must be a multiple of the swizzle atom rows"); - static_assert(ROW_STRIDE % S::ATOM_COLS == 0, "ROW_STRIDE must be a multiple of the swizzle atom cols"); - static_assert(ROW_STRIDE >= COLS, "ROW_STRIDE must be at least COLS"); - if constexpr (ROW_STRIDE == COLS) - { - auto atom = composition(typename S::type{}, - make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); - return tile_to_shape(atom, make_shape(Int{}, Int{})); - } - else - { - return composition(typename S::type{}, - make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); - } -} - -// Aliased-row swizzled smem layout: logical (LOGICAL_ROWS, COLS) view over a -// physical buffer sized to next_multiple_of(VALID_ROWS) rows. -// When VALID_ROWS ≤ ATOM_ROWS the physical buffer is just one row-atom tall -// (e.g. 8 rows for bf16) but the MMA still wants to address LOGICAL_ROWS=16 -// rows (m16n8k16's M). We achieve the alias with a stride-0 outer mode on -// the row-tile axis: row r ∈ [0, LOGICAL_ROWS) maps to physical row -// (r mod PHYS_ROWS), col c maps unchanged. The first m-tile carries the -// real C data; the second m-tile reads the same bytes, feeds garbage into -// MMA accumulator rows ≥ VALID_ROWS, predicated out at gmem store. -// -// When VALID_ROWS > ATOM_ROWS (VALID_ROWS > 8 for 2-byte), PHYS_ROWS == LOGICAL_ROWS -// and the alias factor collapses to 1 — this then degenerates to the same -// layout that `make_swizzled_layout_rc` produces. -template -__device__ __forceinline__ auto make_aliased_swizzled_layout_rc() -{ - using namespace cute; - using S = SmemSwizzle; - static_assert(LOGICAL_ROWS % S::ATOM_ROWS == 0, "LOGICAL_ROWS must be a multiple of the swizzle atom rows"); - static_assert(COLS % S::ATOM_COLS == 0, "COLS must be a multiple of the swizzle atom cols"); - constexpr int PHYS_ROWS = next_multiple_of(VALID_ROWS); - constexpr int LOG_M_TILES = LOGICAL_ROWS / S::ATOM_ROWS; - constexpr int PHYS_M_TILES = PHYS_ROWS / S::ATOM_ROWS; - static_assert(LOG_M_TILES % PHYS_M_TILES == 0, "LOGICAL_ROWS must be a multiple of PHYS_ROWS for clean alias"); - constexpr int ALIAS = LOG_M_TILES / PHYS_M_TILES; - constexpr int N_TILES = COLS / S::ATOM_COLS; - auto atom = composition(typename S::type{}, - make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{}))); - // Outer layout (in atom-units): row-tile mode = (PHYS_M_TILES, ALIAS) strides - // (1, 0); col-tile mode = N_TILES stride PHYS_M_TILES. blocked_product - // scales these by the atom cosize (= ATOM_ROWS * ATOM_COLS). - auto outer = make_layout(make_shape(make_shape(Int{}, Int{}), Int{}), - make_stride(make_stride(_1{}, _0{}), Int{})); - return blocked_product(atom, outer); -} - -// Transposed swizzled smem layout: maps (col, row) → same physical offset as -// make_swizzled_layout_rc maps (row, col). Enables bank-conflict-free ldmatrix.trans -// reads on data stored with make_swizzled_layout_rc. -// Built by swapping modes of the original inner layout (before swizzle), which -// guarantees correct cross-atom offsets when both dimensions have multiple atoms. -template -__device__ __forceinline__ auto make_swizzled_layout_rc_transpose() -{ - using namespace cute; - using S = SmemSwizzle; - static_assert(ROWS % S::ATOM_ROWS == 0, "ROWS must be a multiple of the swizzle atom rows"); - static_assert(COLS % S::ATOM_COLS == 0, "COLS must be a multiple of the swizzle atom cols"); - // Build the inner (un-swizzled) tiled layout for the original (ROWS, COLS) layout - auto inner = tile_to_shape( - make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, _1{})), - make_shape(Int{}, Int{})); - // Swap modes to get true transpose: result(c, r) == original(r, c) - auto inner_T = make_layout(get<1>(inner), get<0>(inner)); - return composition(typename S::type{}, inner_T); -} - -// ============================================================================= -// B/C/x/z load helper: single-warp cp.async into swizzled smem. -// ============================================================================= - -// Generic swizzled cp.async with ZFILL for padding rows — the six [ROWS_PAD, -// COLS] single-warp loaders (B, old_B, x, z, old_x, C) all collapse into this. -// Gmem tile is [ROWS_PAD, COLS] with runtime row stride; rows >= VALID_ROWS -// are zero-filled in smem without touching gmem (cp.async.ca.ZFILL). Thread -// layout Shape<_4,_8>×val Shape<_1,_8> = 32 threads × 16B each = one warp -// covers 4 rows × 64 cols per step. -// -// Template args are all compile-time so the ZFILL predicate constant-folds; -// the caller pre-offsets `gmem_src` by the tile base, keeping 64-bit pointer -// math at the callsite. `SmemShape` is a CuTe static shape, e.g. -// `cute::Shape, cute::Int<128>>` — (rows_pad, cols_pad). The -// shape travels as a single type so later we can pad cols too (e.g. DSTATE=96 -// rounded up to a full bank cycle) without growing the parameter list. -// `valid_rows_rt` is a runtime bound that overrides the compile-time -// `VALID_ROWS` template parameter — used by the varlen (v20) path to tighten -// the predicate from `< NPREDICTED` to `< seq_len`. When the caller omits it, -// the default is the constexpr `VALID_ROWS` so non-varlen call sites fold to -// the same SASS as before. -template -__device__ __forceinline__ void load_tile_async(input_t* __restrict__ smem_dst, input_t const* __restrict__ gmem_src, - int gmem_row_stride, int lane, int valid_rows_rt = VALID_ROWS) -{ - using namespace cute; - constexpr int ROWS_PAD = size<0>(SmemShape{}); - constexpr int VALID_COLS = size<1>(SmemShape{}); - // Smem cols are padded up to the swizzle atom width. We always use the - // wide thread layout (4 thread-rows × 8 thread-cols × 1×8 val = 4 rows × - // 64 cols/pass for bf16) so each thread-row covers all 8 vec-cols of the - // Swizzle atom — the design contract that makes cp.async writes - // bank-conflict-free (each row consumes one full bank cycle, rows - // serialize across cycles). A "narrow" layout (½-atom-width per row) - // would force adjacent rows to compete for the same 16 banks, costing - // ~3-4× replay (observed as 12-way LDGSTS conflicts in d_split=2 ncu). - // When VALID_COLS < SMEM_COLS (D_SPLIT > 1 path), cp.async ZFILL drops - // the predicated-out cols as zeros without touching gmem — same mechanism - // as ZFILL'ing rows ≥ VALID_ROWS. The padded smem cells are unused. - constexpr int SMEM_COLS = next_multiple_of::ATOM_COLS>(VALID_COLS); - Tensor s_full = make_tensor(make_smem_ptr(smem_dst), make_swizzled_layout_rc()); - Tensor g_full = make_tensor(make_gmem_ptr(gmem_src), - make_layout(make_shape(Int{}, Int{}), make_stride(gmem_row_stride, Int<1>{}))); - - constexpr int VAL_COLS_PER_THREAD = Copy_prop::vec_bytes / sizeof(input_t); - static_assert(SMEM_COLS % VAL_COLS_PER_THREAD == 0, "SMEM_COLS must be divisible by VAL_COLS_PER_THREAD"); - using ThrLayout = Layout, Stride<_8, _1>>; - static_assert(size<1>(ThrLayout{}) * VAL_COLS_PER_THREAD == SmemSwizzle::ATOM_COLS, - "wide thread layout must cover one full swizzle atom width per row"); - auto g2s = make_tiled_copy( - Copy_Atom{}, ThrLayout{}, Layout>>{}); - auto thr = g2s.get_slice(lane); - - auto id = make_identity_tensor(make_shape(Int{}, Int{})); - auto thr_id = thr.partition_S(id); - auto pred = make_tensor(shape(thr_id)); - CUTE_UNROLL - for (int i = 0; i < size(pred); ++i) - { - pred(i) = (get<0>(thr_id(i)) < valid_rows_rt) && (get<1>(thr_id(i)) < VALID_COLS); - } - copy_if(g2s, pred, thr.partition_S(g_full), thr.partition_D(s_full)); -} - -// State load — D_SPLIT-conditional dispatch: -// -// D_SPLIT == 1: per-warp partition (warp W loads rows -// [W*DIM/4 : (W+1)*DIM/4)). Large coalesced gmem -// reads per warp. Tests pass without an extra CTA-wide barrier -// because the post-replay __syncthreads covers the eventual -// cross-warp state reads. -// -// D_SPLIT >= 2: 128-thread cooperative load. Required because at -// D_PER_CTA = 16 / 4 = 4 D-rows per warp the per-warp layout no -// longer divides cleanly into the (4, 8) thread-tile atom that -// `Copy_prop::Atom` expects. Cooperative load works for any -// D_PER_CTA ∈ {DIM, DIM/2, DIM/4} that's a multiple of 16. -// -// Both variants write through `make_swizzled_layout_rc` -// followed by a `local_tile` to the (D_PER_CTA, DSTATE) slice this CTA -// owns — the swizzle outer-stride is invariant across D_SPLIT. -template -__device__ __forceinline__ void load_state_per_warp( - SmemT& smem, state_t const* __restrict__ state_ptr, int64_t state_base, int warp, int lane) -{ - using namespace cute; - static_assert(NUM_WARPS == 4, "Expected 4 warps"); - static_assert(D_PER_CTA % NUM_WARPS == 0, "D_PER_CTA must be divisible by NUM_WARPS"); - constexpr int DIM_PER_WARP = D_PER_CTA / NUM_WARPS; - - // Single-local_tile path — swizzle layout sized to this CTA's - // D_PER_CTA slice; one local_tile splits it directly per-warp. - Tensor sState_full = make_tensor( - make_smem_ptr(reinterpret_cast(smem.state)), make_swizzled_layout_rc()); - Tensor gState_full = make_tensor(make_gmem_ptr(state_ptr + state_base), - make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{}))); - - Tensor sState = local_tile(sState_full, make_shape(Int{}, Int{}), make_coord(warp, _0{})); - Tensor gState = local_tile(gState_full, make_shape(Int{}, Int{}), make_coord(warp, _0{})); - - constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); - auto g2s = make_tiled_copy(Copy_Atom{}, Layout, Stride<_8, _1>>{}, - Layout>>{}); - auto thr = g2s.get_slice(lane); - copy(g2s, thr.partition_S(gState), thr.partition_D(sState)); -} - -template -__device__ __forceinline__ void load_state_cta( - SmemT& smem, state_t const* __restrict__ state_ptr, int64_t state_base, int tid) -{ - using namespace cute; - static_assert(NUM_WARPS == 4, "Expected 4 warps"); - - Tensor sState = make_tensor( - make_smem_ptr(reinterpret_cast(smem.state)), make_swizzled_layout_rc()); - Tensor gState = make_tensor(make_gmem_ptr(state_ptr + state_base), - make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{}))); - - constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); - using ThrLayout = Layout, Stride<_8, _1>>; - constexpr int THR_ROWS = decltype(size<0>(ThrLayout{}))::value; - static_assert(D_PER_CTA % THR_ROWS == 0, "D_PER_CTA must be divisible by the thread layout's row count"); - auto g2s = make_tiled_copy(Copy_Atom{}, ThrLayout{}, Layout>>{}); - auto thr = g2s.get_slice(tid); - copy(g2s, thr.partition_S(gState), thr.partition_D(sState)); -} - -// ============================================================================= -// Phase 0: cooperative data load into smem (all warps). -// Compute cumAdt[T] = cumsum(A * dt_proc) → smem. -// Warp-level inclusive prefix sum using Hillis-Steele shuffles. -// Only the first NPREDICTED lanes participate; the rest are idle. -// -// If the storage struct exposes a `decay` field, this also writes -// `decay[lane] = exp(val)` — fuses the EX2 with the cumsum write so the output -// decay broadcast in `compute_output_8bit` becomes a plain LDS (no per-element -// __expf). Detected via SFINAE so the 2/4-byte storage (no decay field) is -// unaffected. -namespace detail -{ -template -struct has_decay : std::false_type -{ -}; - -template -struct has_decay().decay[0])>> : std::true_type -{ -}; -} // namespace detail - -template -__device__ __forceinline__ void compute_cumAdt(SmemT& smem, int lane, float A_val) -{ - float val = (lane < NPREDICTED) ? A_val * smem.dt_proc[lane] : 0.f; - // Inclusive prefix sum (Hillis-Steele) - for (int offset = 1; offset < NPREDICTED; offset *= 2) - { - float other = __shfl_up_sync(constants::MASK_ALL_LANES, val, offset); - if (lane >= offset) - val += other; - } - if (lane < NPREDICTED) - { - smem.cumAdt[lane] = val; - if constexpr (detail::has_decay::value) - { - smem.decay[lane] = __expf(val); - } - } -} - -// Load phase. Split into two halves around the PDL barrier (`gdc_wait`): -// -// load_pre_pdl_wait_data: data NOT produced by the immediate upstream -// kernel (conv1d) — state and old_* are cache from the previous SSU -// step; dt and z are in_proj outputs (in_proj fully completed before -// conv1d began, so they are visible by the time we hit `gdc_wait`). -// Issues cp.async for cache + z, runs the scalar LDGs (old_dt, -// old_cumAdt, dt→dt_proc) and the cumAdt warp scan. No commit/wait — -// the cp.async stays in flight while we `gdc_wait` on conv1d. -// -// load_post_pdl_wait_data: x/B/C cp.async (conv1d outputs — must wait) -// and the single `__pipeline_commit + __pipeline_wait_prior(0) + -// __syncwarp` that drains both halves. Cache cp.async issued in the -// pre-wait half share the per-thread async group with these, so one -// wait_prior(0) covers them all. -// -// Per-warp data ownership (unchanged from the pre-split version): -// state: per-warp contiguous DIM slice (warp W owns rows [16W : 16W+16]). -// B, C: redundant on W0, W1 (both compute 2-warp CB, both need full). -// old_B: redundant on all 4 warps (each warp's replay reads full DSTATE). -// old_x: redundant on all 4 warps (small, ~2 KB — partitioning not worth -// the complication). -// x: W2 only (Phase-2 read, covered by final __syncthreads). -// z: W3 only (Phase-2 read, covered by final __syncthreads). -// scalars (old_dt, old_cumAdt, dt→dt_proc) + cumAdt cumsum: -// redundant on each warp's first NPREDICTED/MAX_WINDOW lanes. Writes -// are idempotent across warps (identical payloads to same slots). -// ============================================================================= -// Per-sequence gmem base offsets (`x_seq_base`, etc.) are computed once in -// the kernel prologue — they encode the "start of this sequence" along the -// outer axis (batch in non-varlen, packed-token in varlen). Helper indexing -// is then uniform `seq_base + inner`. -// -// `seq_len` is the per-sequence new-token count (== NPREDICTED constexpr in -// non-varlen, runtime int in varlen). Used as the cp.async row predicate -// and the dt/scalar lane predicate so trailing rows past `seq_len` ZFILL to -// zero in smem. -template -__device__ __forceinline__ void load_pre_pdl_wait_data(SmemT& smem, CheckpointingSsuParams const& params, int lane, - int warp, int d_tile, int head, int group_idx, int64_t cache_slot, int buf_read, float A_val, float dt_bias_val, - int64_t dt_seq_base, int64_t z_seq_base, int seq_len) -{ - constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 - static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); - static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); - - int const d_tile_off = d_tile * D_PER_CTA; - - auto const* __restrict__ z_ptr = reinterpret_cast(params.z); - auto const* __restrict__ old_x_ptr = reinterpret_cast(params.old_x); - auto const* __restrict__ old_B_ptr = reinterpret_cast(params.old_B); - auto const* __restrict__ old_dt_ptr = reinterpret_cast(params.old_dt); - auto const* __restrict__ old_cumAdt_ptr = reinterpret_cast(params.old_cumAdt); - auto const* __restrict__ dt_ptr = reinterpret_cast(params.dt); - - int64_t const ox_base = cache_slot * params.old_x_stride_seq + head * DIM + d_tile_off; - int64_t const oB_base - = cache_slot * params.old_B_stride_seq + buf_read * params.old_B_stride_dbuf + group_idx * DSTATE; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - // ZShape: shrunk to the swizzle atom's row extent (z is read via - // partition_C alias, the physical buffer only needs to be one swizzle - // row-atom tall). - using ZShape = cute::Shape, cute::Int>; - // old_B / old_x's smem row count = replay matmul K-axis = MAX_WINDOW_PAD_MMA_K. - using OldBShape = cute::Shape, cute::Int>; - using OxShape = cute::Shape, cute::Int>; - - // ── State: per-CTA D-slice ([D_PER_CTA, DSTATE]). Dispatch on D_SPLIT - // (= DIM == D_PER_CTA): per-warp coalesced load when one CTA owns the - // full head's D, cooperative 128-thread load when D is sharded. ── - { - auto const* __restrict__ state_ptr = reinterpret_cast(params.state); - int64_t const state_base - = cache_slot * params.state_stride_seq + (int64_t) head * DIM * DSTATE + (int64_t) d_tile_off * DSTATE; - if constexpr (DIM == D_PER_CTA) - { - // D_SPLIT=1: per-warp partition (warp w loads contiguous DIM/4 D-rows). - load_state_per_warp(smem, state_ptr, state_base, warp, lane); - } - else - { - // D_SPLIT>=2: 128-thread cooperative load (per-warp doesn't divide - // cleanly when D_PER_CTA/4 is too small for the (4,8) thread atom). - int const tid = warp * warpSize + lane; - load_state_cta(smem, state_ptr, state_base, tid); - } - } - - // ── old_B: redundant on all 4 warps (each warp's replay consumes full - // DSTATE). Identical payloads to same smem dest — final bytes - // deterministic. VALID_ROWS = MAX_WINDOW (cache rows). ── - load_tile_async(smem.old_B, old_B_ptr + oB_base, params.old_B_stride_token, lane); - - // ── old_x: redundant on all 4 warps (small, simpler than partitioning). - // VALID_ROWS = MAX_WINDOW (cache rows). ── - load_tile_async(smem.old_x, old_x_ptr + ox_base, params.old_x_stride_token, lane); - - // ── z: W3 only (Phase-2 read, final __syncthreads makes it visible). - // Sourced from in_proj — not from conv1d — so safe to issue pre-wait. ── - if (warp == 3 && z_ptr) - { - int64_t const z_base = z_seq_base + head * DIM + d_tile_off; - load_tile_async(smem.z, z_ptr + z_base, params.z_stride_token, lane, seq_len); - } - - // Commit the cache cp.async group BEFORE the caller's `gdc_wait()` so the - // hardware actually issues the gmem→smem transfers while the wait is in - // flight (without commit, the operations sit pending and only kick off - // once the post-wait commit fires — no overlap). Placed immediately after - // the last cp.async (z); the synchronous LDGs + cumAdt scan below are not - // part of any pipeline group and run in parallel with the in-flight - // transfers. The post half issues a second group; `__pipeline_wait_prior(0)` - // there drains both. - __pipeline_commit(); - - // ── Scalar loads + cumAdt cumsum: redundant per warp. - // old_dt / old_cumAdt: load up to MAX_WINDOW lanes (cache scalars). - // dt_proc: load up to NPREDICTED lanes (new-token scalars from in_proj). - // Synchronous LDG + plain smem stores — no cp.async. Writes from 4 - // warps to the same slots are idempotent (same payloads). ── - static_assert(MAX_WINDOW <= warpSize, "MAX_WINDOW must fit in a single warp"); - if (lane < MAX_WINDOW) - { - int64_t const dt_rd_base = cache_slot * params.old_dt_stride_seq + buf_read * params.old_dt_stride_dbuf - + head * params.old_dt_stride_head; - smem.old_dt[lane] = old_dt_ptr[dt_rd_base + lane]; - - int64_t const ca_rd_base = cache_slot * params.old_cumAdt_stride_seq + buf_read * params.old_cumAdt_stride_dbuf - + head * params.old_cumAdt_stride_head; - smem.old_cumAdt[lane] = old_cumAdt_ptr[ca_rd_base + lane]; - } - // dt → softplus → smem.dt_proc. Under varlen the active lane range is - // `[0, seq_len)`; lanes `[seq_len, NPREDICTED)` are left uninitialized — - // `compute_cumAdt` will scan over them and produce garbage in the - // `cumAdt[seq_len:NPREDICTED]` tail, but every downstream consumer - // (`compute_CB_scaled_2warp` mask, output STG, dt_proc/cumAdt tape writes) - // is gated on `seq_len`, so the garbage never reaches gmem or contaminates - // valid rows. - // - // Per-lane stride along the T-axis is `dt_stride_token` in both layouts - // (4D batch and 1D packed varlen); the caller bakes `head` into - // `dt_seq_base` so the inner indexing is `dt_seq_base + lane * - // dt_stride_token`. - if (lane < seq_len) - { - float dt_val = toFloat(dt_ptr[dt_seq_base + (int64_t) lane * params.dt_stride_token]); - dt_val += dt_bias_val; - if (params.dt_softplus) - dt_val = thresholded_softplus(dt_val); - smem.dt_proc[lane] = dt_val; - } - // cumAdt = cumsum(A * dt_proc) — warp-local Hillis-Steele shuffle. Each - // of the 4 warps runs the same reduction on identical inputs (dt_proc - // just written above) and writes the same smem.cumAdt slots. - compute_cumAdt(smem, lane, A_val); -} - -// Post-wait half. Issues cp.async for conv1d outputs (x, B, C) and drains -// the per-thread async group (which includes both the cache cp.async issued -// in `load_pre_pdl_wait_data` and these conv1d cp.async). Caller must have -// called `gdc_wait()` between the two halves; otherwise this reads stale -// conv1d data. -// -// Takes `outer` (the per-sequence outer index) rather than pre-multiplied -// `*_seq_base` scalars. Computing `outer * stride_seq` inside this function -// keeps the multipliers transient instead of pinning them across the -// `gdc_wait()` asm-volatile barrier (which the compiler can't reorder around -// and thus can't rematerialize through). Saves ~6 registers vs. pre-computed -// bases. -template -__device__ __forceinline__ void load_post_pdl_wait_data(SmemT& smem, CheckpointingSsuParams const& params, int lane, - int warp, int d_tile, int head, int group_idx, int64_t outer, int seq_len) -{ - constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 - static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); - static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); - - int const d_tile_off = d_tile * D_PER_CTA; - - auto const* __restrict__ B_ptr = reinterpret_cast(params.B); - auto const* __restrict__ C_ptr = reinterpret_cast(params.C); - auto const* __restrict__ x_ptr = reinterpret_cast(params.x); - - int64_t const B_base = outer * params.B_stride_seq + (int64_t) group_idx * DSTATE; - int64_t const C_base = outer * params.C_stride_seq + (int64_t) group_idx * DSTATE; - int64_t const x_base = outer * params.x_stride_seq + head * DIM + d_tile_off; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; - // CShape: first dim shrunk to the swizzle atom's row extent so cp.async - // writes don't spill past the (also shrunk) C smem buffer. When NPREDICTED - // > ATOM_ROWS this falls back to NPREDICTED_PAD_MMA_M. - using CShape = cute::Shape, cute::Int>; - // B's smem row count = matmul-1 N-axis = NPREDICTED_PAD_MMA_N. - using BShape = cute::Shape, cute::Int>; - using XShape = cute::Shape, cute::Int>; - - // ── B: redundant on W0, W1 (both do 2-warp CB compute) ── - if (warp < 2) - { - load_tile_async(smem.B, B_ptr + B_base, params.B_stride_token, lane, seq_len); - } - // ── C: redundant on all 4 warps — chain matmul-3 reads smem.C from every - // warp, so each warp must see its own cp.async without a cross-warp sync. - // Identical payloads to same smem dest (same pattern as old_B / old_x). ── - load_tile_async(smem.C, C_ptr + C_base, params.C_stride_token, lane, seq_len); - - // ── x: W2 only (Phase-2 read, final __syncthreads makes it visible) ── - if (warp == 2) - { - load_tile_async(smem.x, x_ptr + x_base, params.x_stride_token, lane, seq_len); - } - - // Commit the conv1d cp.async group and drain BOTH groups: the cache - // group committed in `load_pre_pdl_wait_data` (pre-`gdc_wait`) and this - // conv1d group. `__pipeline_wait_prior(0)` waits for ≤0 pending groups. - // __syncwarp() provides acquire semantics across the 32 lanes of each - // warp. No cross-warp sync here; the only __syncthreads is after - // CB + replay. - __pipeline_commit(); - __pipeline_wait_prior(0); - __syncwarp(); -} - -// Single-pass load — used when `params.enable_pdl == false`. All cp.async -// (state, B, C, old_B, old_x, x, z) issue together into one async group; -// scalars + cumAdt scan run while the cp.async are in flight; one commit + -// wait_prior(0) + syncwarp drains the whole thing. This restores the v21.0 -// load order: the synchronous LDG-then-STS for old_dt/old_cumAdt/dt benefits -// from overlap with the conv1d cp.async (B/C/x) — which the split (pre + -// gdc_wait + post) form sacrifices since conv1d cp.async only issue after -// the wait. When PDL is paired with an upstream conv1d, the split's -// cache-load-during-wait overlap dominates; when not paired, the split is -// pure overhead (gdc_wait is a no-op, but the cp.async are delayed). -template -__device__ __forceinline__ void load_data(SmemT& smem, CheckpointingSsuParams const& params, int lane, int warp, - int d_tile, int head, int group_idx, int64_t cache_slot, int buf_read, float A_val, float dt_bias_val, - int64_t outer, int seq_len) -{ - constexpr int INPUT_PACK = 16 / sizeof(input_t); // 8 for bf16 - static_assert(DSTATE % INPUT_PACK == 0, "DSTATE must be divisible by input pack size"); - static_assert(D_PER_CTA % INPUT_PACK == 0, "D_PER_CTA must be divisible by input pack size"); - - int const d_tile_off = d_tile * D_PER_CTA; - - auto const* __restrict__ B_ptr = reinterpret_cast(params.B); - auto const* __restrict__ C_ptr = reinterpret_cast(params.C); - auto const* __restrict__ x_ptr = reinterpret_cast(params.x); - auto const* __restrict__ z_ptr = reinterpret_cast(params.z); - auto const* __restrict__ old_x_ptr = reinterpret_cast(params.old_x); - auto const* __restrict__ old_B_ptr = reinterpret_cast(params.old_B); - auto const* __restrict__ old_dt_ptr = reinterpret_cast(params.old_dt); - auto const* __restrict__ old_cumAdt_ptr = reinterpret_cast(params.old_cumAdt); - auto const* __restrict__ dt_ptr = reinterpret_cast(params.dt); - - int64_t const B_base = outer * params.B_stride_seq + (int64_t) group_idx * DSTATE; - int64_t const C_base = outer * params.C_stride_seq + (int64_t) group_idx * DSTATE; - int64_t const x_base = outer * params.x_stride_seq + head * DIM + d_tile_off; - int64_t const ox_base = cache_slot * params.old_x_stride_seq + head * DIM + d_tile_off; - int64_t const oB_base - = cache_slot * params.old_B_stride_seq + buf_read * params.old_B_stride_dbuf + group_idx * DSTATE; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - using CShape = cute::Shape, cute::Int>; - using BShape = cute::Shape, cute::Int>; - using XShape = cute::Shape, cute::Int>; - using ZShape = cute::Shape, cute::Int>; - using OldBShape = cute::Shape, cute::Int>; - using OxShape = cute::Shape, cute::Int>; - - // ── State: per-CTA D-slice ([D_PER_CTA, DSTATE]) ── - { - auto const* __restrict__ state_ptr = reinterpret_cast(params.state); - int64_t const state_base - = cache_slot * params.state_stride_seq + (int64_t) head * DIM * DSTATE + (int64_t) d_tile_off * DSTATE; - if constexpr (DIM == D_PER_CTA) - { - load_state_per_warp(smem, state_ptr, state_base, warp, lane); - } - else - { - int const tid = warp * warpSize + lane; - load_state_cta(smem, state_ptr, state_base, tid); - } - } - - if (warp < 2) - { - load_tile_async(smem.B, B_ptr + B_base, params.B_stride_token, lane, seq_len); - } - load_tile_async(smem.C, C_ptr + C_base, params.C_stride_token, lane, seq_len); - - load_tile_async(smem.old_B, old_B_ptr + oB_base, params.old_B_stride_token, lane); - load_tile_async(smem.old_x, old_x_ptr + ox_base, params.old_x_stride_token, lane); - - if (warp == 2) - { - load_tile_async(smem.x, x_ptr + x_base, params.x_stride_token, lane, seq_len); - } - if (warp == 3 && z_ptr) - { - int64_t const z_base = outer * params.z_stride_seq + head * DIM + d_tile_off; - load_tile_async(smem.z, z_ptr + z_base, params.z_stride_token, lane, seq_len); - } - - // ── Scalar loads (overlap with cp.async) + cumAdt cumsum ── - static_assert(MAX_WINDOW <= warpSize, "MAX_WINDOW must fit in a single warp"); - if (lane < MAX_WINDOW) - { - int64_t const dt_rd_base = cache_slot * params.old_dt_stride_seq + buf_read * params.old_dt_stride_dbuf - + head * params.old_dt_stride_head; - smem.old_dt[lane] = old_dt_ptr[dt_rd_base + lane]; - - int64_t const ca_rd_base = cache_slot * params.old_cumAdt_stride_seq + buf_read * params.old_cumAdt_stride_dbuf - + head * params.old_cumAdt_stride_head; - smem.old_cumAdt[lane] = old_cumAdt_ptr[ca_rd_base + lane]; - } - int64_t const dt_seq_base_local = outer * params.dt_stride_seq + head; - if (lane < seq_len) - { - float dt_val = toFloat(dt_ptr[dt_seq_base_local + (int64_t) lane * params.dt_stride_token]); - dt_val += dt_bias_val; - if (params.dt_softplus) - dt_val = thresholded_softplus(dt_val); - smem.dt_proc[lane] = dt_val; - } - compute_cumAdt(smem, lane, A_val); - - __pipeline_commit(); - __pipeline_wait_prior(0); - __syncwarp(); -} - -// (compute_cumAdt moved above load_pre_pdl_wait_data so it can be called from there) - -// Compute CB_scaled[T,T] = (C @ B^T) * decay * dt_proc * causal_mask. -// Split across 2 warps: warp 0 computes columns 0:8, warp 1 computes columns 8:16. -// Result stored to swizzled smem.CB_scaled (input_t, row stride 64, Swizzle<3,3,3>). -// Called between the two __syncthreads by warps 0 and 1 only. -// `seq_len` is the runtime row/col bound on the (T, T) CB_scaled tile. -// Caller passes `NPREDICTED` (constexpr) for non-varlen — the mask -// `j <= t && t < seq_len && j < seq_len` then folds to today's SASS. -// Varlen passes the per-sequence `seq_len ≤ NPREDICTED`; rows/cols past it -// get zeroed so downstream matmul-4 / chain matmul-3 see zeros there. -template -__device__ __forceinline__ void compute_CB_scaled_2warp(SmemT& smem, int warp, int lane, int seq_len) -{ - using namespace cute; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; - // 2-warp output split: each warp owns NPREDICTED_PAD_MMA_M / 2 cols of - // the (NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M) CB tile. Must be a - // multiple of the MMA atom's N for the partition to be atom-aligned - // (currently 8 == MMA_prop::N for NPREDICTED_PAD_MMA_M=16; if M-pad ever - // grows, this still holds as long as M-pad is a multiple of 2 * MMA_prop::N). - constexpr int N_HALF = NPREDICTED_PAD_MMA_M / 2; - static_assert( - N_HALF % MMA_prop::N == 0, "compute_CB_scaled_2warp: NPREDICTED_PAD_MMA_M / 2 must be a multiple of MMA::N"); - - // CB_scaled output tile layout (used by both warp 0 compute and warp 1 - // zero-fill when smem.B has only 8 rows). Row stride matches the buffer's - // padded width (one swizzle atom of `input_t`). - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - auto layout_cb_swz = make_swizzled_layout_rc(); - - // ── NPREDICTED_PAD_MMA_N == 8: warp 1 has no valid B rows to read. But - // CB_scaled[:, 8:16] must still be zero so matmul-4's K-reduction sees - // zeros for j ≥ NPREDICTED. Do a simple 32-thread zero-fill and return. - if constexpr (NPREDICTED_PAD_MMA_N == 8) - { - if (warp == 1) - { - auto* __restrict__ cb = reinterpret_cast(smem.CB_scaled); - constexpr int COLS_TO_CLEAR = NPREDICTED_PAD_MMA_M - N_HALF; // 8 -#pragma unroll - for (int i = lane; i < NPREDICTED_PAD_MMA_M * COLS_TO_CLEAR; i += warpSize) - { - int const r = i / COLS_TO_CLEAR; - int const c = N_HALF + (i % COLS_TO_CLEAR); - cb[layout_cb_swz(r, c)] = MMA_prop::operand_t(0.f); - } - return; - } - } - - // ── Swizzled smem views ── - // C is padded to NPREDICTED_PAD_MMA_M; B has NPREDICTED_PAD_MMA_N rows. - // Use NPREDICTED_PAD_MMA_N for smem_B so the physical layout matches the - // write layout from load_tile_async — `tile_to_shape` produces different - // outer strides for (8, 128) vs (16, 128). - // Aliased C view: physical buffer is just next_multiple_of(NPREDICTED) - // rows tall but the MMA atom needs M=16; second m-tile aliases first m-tile - // (predicated rows discarded at output store). - auto layout_C = make_aliased_swizzled_layout_rc(); - auto layout_B = make_swizzled_layout_rc(); - Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C); - Tensor smem_B = make_tensor(make_smem_ptr(reinterpret_cast(smem.B)), layout_B); - - // ── TiledMMA: _1x_1 = 32 threads, one [16, 8] atom ── - auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(lane); - - // ── K-tile A operand (C): full [NPREDICTED_PAD_MMA_M, K_TILE], shared by both warps ── - constexpr int K_TILE = MMA_prop::K_BIG; - Tensor smem_C_tiled - = local_tile(smem_C, make_tile(Int{}, Int{}), make_coord(_0{}, _)); - - // ── K-tile B operand ── - // NPREDICTED_PAD_MMA_N == 16: warp 0 → N=[0,8), warp 1 → N=[8,16). - // NPREDICTED_PAD_MMA_N == 8 : only warp 0 runs (warp 1 took the early - // exit above), tile at (_0, _). - Tensor smem_B_half = local_tile(smem_B, make_tile(Int{}, Int{}), make_coord(warp, _)); - - // ── Register fragments ── - Tensor frag_A = thr_mma.partition_fragment_A(smem_C_tiled(_, _, _0{})); - Tensor frag_B = thr_mma.partition_fragment_B(smem_B_half(_, _, _0{})); - - // ── Output accumulator: [NPREDICTED_PAD_MMA_M, N_HALF] f32 ── - auto layout_cb_half = make_layout(make_shape(Int{}, Int{})); - Tensor frag_acc = thr_mma.partition_fragment_C(make_tensor((float*) nullptr, layout_cb_half)); - clear(frag_acc); - - // ── S2R copies ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(lane); - Tensor smem_C_s2r = s2r_thr_A.partition_S(smem_C_tiled); - Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); - - auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B = s2r_B.get_slice(lane); - Tensor smem_B_s2r = s2r_thr_B.partition_S(smem_B_half); - Tensor frag_B_view = s2r_thr_B.retile_D(frag_B); - - // ── Gemm: 8 K-tiles, 1 HMMA each ── - constexpr int NUM_K_TILES = DSTATE / K_TILE; -#pragma unroll - for (int k = 0; k < NUM_K_TILES; ++k) - { - cute::copy(s2r_A, smem_C_s2r(_, _, _, k), frag_A_view); - cute::copy(s2r_B, smem_B_s2r(_, _, _, k), frag_B_view); - cute::gemm(tiled_mma, frag_acc, frag_A, frag_B, frag_acc); - } - - // ── Elementwise: decay * dt_proc * causal mask, convert f32 → MMA_prop::operand_t ── - auto id_half = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_half); - - // ── Store to swizzled smem.CB_scaled ── - Tensor smem_CB = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_swz); - // Tile into [NPREDICTED_PAD_MMA_M, N_HALF] halves; warp selects its half - Tensor smem_CB_half - = local_tile(smem_CB, make_tile(Int{}, Int{}), make_coord(_0{}, warp)); - Tensor smem_CB_part = thr_mma.partition_C(smem_CB_half); - -#pragma unroll - for (int i = 0; i < size(frag_acc); ++i) - { - int t = get<0>(id_part(i)); - int j = warp * N_HALF + get<1>(id_part(i)); - float val; - if (j <= t && t < seq_len && j < seq_len) - { - val = frag_acc(i) * __expf(smem.cumAdt[t] - smem.cumAdt[j]) * smem.dt_proc[j]; - } - else - { - val = 0.f; - } - smem_CB_part(i) = MMA_prop::operand_t(val); - } -} - -// Compute CB_old[t, i] = (C @ old_B^T)[t, i] * exp(cumAdt[t]) * dB_old(i) for -// i ∈ [0, prev_k); 0 otherwise. -// dB_old(i) = exp(total_old_cumAdt − smem.old_cumAdt[i]) * smem.old_dt[i]. -// The per-t factor exp(cumAdt[t]) is baked in here (vs at matmul-4 time) so the -// epilogue's β-scale = exp(total_old_cumAdt) * exp(cumAdt[t]) on init_out -// composes with a single CB_old @ old_x add — no extra elementwise pass. -// Identity (matches Triton's combined-sequence SSU): -// y_old_contrib[t, d] = exp(cumAdt[t]) * Σ_i dB_old(i) * x_old[i, d] * (C[t] · B_old[i]) -// = Σ_i CB_old[t, i] * old_x[i, d]. -// Written into smem.CB_scaled at cols [NPREDICTED_PAD_MMA_M, NPREDICTED_PAD_MMA_M + -// MAX_WINDOW_PAD_MMA_K). Sibling of compute_CB_scaled_2warp — runs on warps 2, 3 in parallel with -// warps 0, 1 writing the new-token half at cols [0, NPREDICTED_PAD_MMA_M). Uses the no-write -// path's CB_old region of the same swizzled buffer (32 cols total ≤ CB_ROW_STRIDE=64). -// -// 2-warp N-split: each warp owns one m16n8 N-atom (MMA::N=8 cols). -// MAX_WINDOW_PAD_MMA_K == 16: warp 2 → cols [0, 8); warp 3 → cols [8, 16). -// MAX_WINDOW_PAD_MMA_K == 8 : warp 2 covers all 8 cols; warp 3 returns early. -template -__device__ __forceinline__ void compute_CB_old_2warp(SmemT& smem, int warp, int lane, int prev_k, int seq_len) -{ - using namespace cute; - - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int MAX_WINDOW_PAD_MMA_K = SmemT::MAX_WINDOW_PAD_MMA_K; - constexpr int CB_ROW_STRIDE = SmemT::CB_ROW_STRIDE; - constexpr int N_HALF = MMA_prop::N; // 8 — one m16n8 atom per warp - constexpr int NUM_N_ATOMS = MAX_WINDOW_PAD_MMA_K / N_HALF; - static_assert( - MAX_WINDOW_PAD_MMA_K % N_HALF == 0, "compute_CB_old_2warp: MAX_WINDOW_PAD_MMA_K must be a multiple of MMA::N"); - static_assert(NPREDICTED_PAD_MMA_M + MAX_WINDOW_PAD_MMA_K <= CB_ROW_STRIDE, - "CB_scaled buffer must fit both CB_new (cols [0,T_pad)) and CB_old " - "(cols [T_pad, T_pad+K_old)) within its physical row stride"); - - int const sub_warp = warp - 2; // ∈ {0, 1} - if (sub_warp >= NUM_N_ATOMS) - return; - - float const total_old_cumAdt = (prev_k > 0) ? smem.old_cumAdt[prev_k - 1] : 0.f; - - // ── Swizzled smem views (A = C, B = old_B; same shapes as the replay path's - // C/old_B reads, so we get cache locality with no extra cp.async). ── - auto layout_C = make_aliased_swizzled_layout_rc(); - auto layout_old_B = make_swizzled_layout_rc(); - Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C); - Tensor smem_old_B = make_tensor(make_smem_ptr(reinterpret_cast(smem.old_B)), layout_old_B); - - // ── TiledMMA: 32 threads, single m16n8k16 atom (K-loops over DSTATE) ── - auto tiled_mma = make_tiled_mma(MMA_Atom>{}, Layout>{}); - auto thr_mma = tiled_mma.get_slice(lane); - - // ── K-tile A operand (C): full M, K-loop dim ── - constexpr int K_TILE = MMA_prop::K_BIG; - Tensor smem_C_tiled - = local_tile(smem_C, make_tile(Int{}, Int{}), make_coord(_0{}, _)); - - // ── K-tile B operand (old_B): warp picks its 8-col N-atom slice ── - Tensor smem_old_B_half = local_tile(smem_old_B, make_tile(Int{}, Int{}), make_coord(sub_warp, _)); - - // ── Register fragments ── - Tensor frag_A = thr_mma.partition_fragment_A(smem_C_tiled(_, _, _0{})); - Tensor frag_B = thr_mma.partition_fragment_B(smem_old_B_half(_, _, _0{})); - - auto layout_cb_half = make_layout(make_shape(Int{}, Int{})); - Tensor frag_acc = thr_mma.partition_fragment_C(make_tensor((float*) nullptr, layout_cb_half)); - clear(frag_acc); - - // ── S2R copies ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(lane); - Tensor smem_C_s2r = s2r_thr_A.partition_S(smem_C_tiled); - Tensor frag_A_view = s2r_thr_A.retile_D(frag_A); - - auto s2r_B = make_tiled_copy_B(Copy_Atom{}, tiled_mma); - auto s2r_thr_B = s2r_B.get_slice(lane); - Tensor smem_old_B_s2r = s2r_thr_B.partition_S(smem_old_B_half); - Tensor frag_B_view = s2r_thr_B.retile_D(frag_B); - - // ── GEMM: DSTATE / K_BIG = 8 K-tiles ── - constexpr int NUM_K_TILES = DSTATE / K_TILE; -#pragma unroll - for (int k = 0; k < NUM_K_TILES; ++k) - { - cute::copy(s2r_A, smem_C_s2r(_, _, _, k), frag_A_view); - cute::copy(s2r_B, smem_old_B_s2r(_, _, _, k), frag_B_view); - cute::gemm(tiled_mma, frag_acc, frag_A, frag_B, frag_acc); - } - - // ── Identity coords for elementwise / store ── - auto id_half = make_identity_tensor(make_shape(Int{}, Int{})); - auto id_part = thr_mma.partition_C(id_half); - - // ── Store to swizzled smem.CB_scaled at the CB_old region (cols [T_pad, T_pad+K_old)). - // Use the full physical (NPREDICTED_PAD_MMA_M, CB_ROW_STRIDE) padded swizzle view. - // Byte-compatible with compute_CB_scaled_2warp's (T_pad, T_pad, CB_ROW_STRIDE) - // padded view: both produce inner offset r*CB_ROW_STRIDE + c, same Swizzle. ── - auto layout_cb_full = make_swizzled_layout_rc(); - Tensor smem_CB = make_tensor(make_smem_ptr(reinterpret_cast(smem.CB_scaled)), layout_cb_full); - // (16, CB_ROW_STRIDE) tiled by (16, N_HALF) → (1, CB_ROW_STRIDE/N_HALF) tiles. - // Coord (0, T_pad/N_HALF + sub_warp) lands inside the CB_old region. - constexpr int CB_OLD_TILE_BASE = NPREDICTED_PAD_MMA_M / N_HALF; - Tensor smem_CB_half = local_tile( - smem_CB, make_tile(Int{}, Int{}), make_coord(_0{}, CB_OLD_TILE_BASE + sub_warp)); - Tensor smem_CB_part = thr_mma.partition_C(smem_CB_half); - -#pragma unroll - for (int i = 0; i < size(frag_acc); ++i) - { - int t = get<0>(id_part(i)); - int j = sub_warp * N_HALF + get<1>(id_part(i)); - float val; - if (j < prev_k && t < seq_len) - { - val = frag_acc(i) * __expf(smem.cumAdt[t] + total_old_cumAdt - smem.old_cumAdt[j]) * smem.old_dt[j]; - } - else - { - val = 0.f; - } - smem_CB_part(i) = MMA_prop::operand_t(val); - } -} - -// ============================================================================= -// Precompute dB scaling coefficients (called once before the N-tile loop). -// Returns DB_COEFFS_PER_LANE floats in coeff[], one per B-fragment element of -// the replay MMA — equal to (replay K) / 4 (each m16n8k* B-frag holds K*8/32 -// elts per lane). -// DB_COEFFS_PER_LANE = 4 for m16n8k16 B fragment -// DB_COEFFS_PER_LANE = 2 for m16n8k8 B fragment -// coeff[i] = 0 when k >= prev_k, embedding the causal mask so the inner -// loop needs no branch. -// -// K-index derivation (row-major TN MMA, lane = tid % 32): -// K_base = (lane % 4) * 2 -// m16n8k16 B frag (4 elts): 0→K_base, 1→K_base+1, 2→K_base+8, 3→K_base+9 -// m16n8k8 B frag (2 elts): 0→K_base, 1→K_base+1 -// ============================================================================= -template -__device__ __forceinline__ void precompute_dB_coeff( - float coeff[DB_COEFFS_PER_LANE], SmemT const& smem, float total_cumAdt, int prev_k, int lane) -{ - static_assert(DB_COEFFS_PER_LANE == 2 || DB_COEFFS_PER_LANE == 4, "DB_COEFFS_PER_LANE must be 2 (k8) or 4 (k16)"); - int const K_base = (lane % 4) * 2; -#pragma unroll - for (int i = 0; i < DB_COEFFS_PER_LANE; ++i) - { - // m16n8k_ V-index → K-offset: (V & 1) is the col-pair offset; (V & 2) ? 8 : 0 - // covers the second K-tile inside the K_BIG (k16) atom. - int const k = K_base + (i & 1) + ((i & 2) << 2); - coeff[i] = (k < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[k]) * smem.old_dt[k] : 0.f; - } -} - -// Apply precomputed dB coefficients to frag_B in-place. -// Scales frag_B in-place by per-coefficient multiplier; frag dtype inferred -// from FragB::value_type. -// coeff[i] = 0 encodes both causal mask and zero-fill for k >= prev_k. -// ============================================================================= -template -__device__ __forceinline__ void compute_dB_scaling(FragB& frag_B, float const coeff[DB_COEFFS_PER_LANE]) -{ - using namespace cute; - static_assert(size(FragB{}) == DB_COEFFS_PER_LANE, "frag_B size must match DB_COEFFS_PER_LANE"); - using frag_t = typename FragB::value_type; -#pragma unroll - for (int i = 0; i < DB_COEFFS_PER_LANE; ++i) - { - frag_B(i) = frag_t(toFloat(frag_B(i)) * coeff[i]); - } -} - -// Scale frag_A by dB coefficients ONCE before the N-pass loop, replacing -// 16 per-N-pass compute_dB_scaling calls on frag_B (64 scale ops → 8 or 4). -// Identity: sum_k A[m,k]*(c[k]*B[k,n]) == sum_k (c[k]*A[m,k])*B[k,n]. -// -// K-index derivation (PTX ISA, m16n8k{8,16} mma.sync, A operand row-major): -// groupID = lane / 4, threadID_in_group = lane % 4 -// K_base = (lane % 4) * 2 (same formula as B operand) -// m16n8k8 A regs: a0 = A[groupID, K_base:K_base+2] -// a1 = A[groupID+8, K_base:K_base+2] -// frag (4 elts): {0,2}→K_base, {1,3}→K_base+1 (2 unique K) -// m16n8k16 A regs: a0..a1 as above, -// a2 = A[groupID, K_base+8:K_base+10] -// a3 = A[groupID+8, K_base+8:K_base+10] -// frag (8 elts): {0,2}→K_base, {1,3}→K_base+1, -// {4,6}→K_base+8, {5,7}→K_base+9 (4 unique K) -// ============================================================================= -template -__device__ __forceinline__ void apply_dA_coeff( - FragA& frag_A, SmemT const& smem, float total_cumAdt, int prev_k, int lane) -{ - using namespace cute; - constexpr int FRAG_A_SIZE = size(FragA{}); - static_assert((MAX_WINDOW_PAD_MMA_K == 16 && FRAG_A_SIZE == 8) || (MAX_WINDOW_PAD_MMA_K == 8 && FRAG_A_SIZE == 4), - "apply_dA_coeff: unsupported MMA K / frag_A size combination"); - using frag_t = typename FragA::value_type; - - int const K_base = (lane % 4) * 2; - - if constexpr (MAX_WINDOW_PAD_MMA_K == 8) - { - float const c0 = (K_base < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[K_base]) * smem.old_dt[K_base] : 0.f; - float const c1 = (K_base + 1 < prev_k) - ? __expf(total_cumAdt - smem.old_cumAdt[K_base + 1]) * smem.old_dt[K_base + 1] - : 0.f; -#pragma unroll - for (int i = 0; i < 4; ++i) - { - frag_A(i) = frag_t(toFloat(frag_A(i)) * ((i & 1) ? c1 : c0)); - } - } - else - { - float c[4]; -#pragma unroll - for (int j = 0; j < 4; ++j) - { - int const k = K_base + (j & 1) + ((j & 2) ? 8 : 0); - c[j] = (k < prev_k) ? __expf(total_cumAdt - smem.old_cumAdt[k]) * smem.old_dt[k] : 0.f; - } -#pragma unroll - for (int i = 0; i < 8; ++i) - { - int const ci = (i & 1) | ((i & 4) >> 1); - frag_A(i) = frag_t(toFloat(frag_A(i)) * c[ci]); - } - } -} - -// ── CuTe mma.sync output sub-functions ────────────────────────────────────── -// Each operates on a register-resident frag_y accumulator (f32). -// Called from compute_and_store_output's N-tile loop. - -// Convert fragment elements from src_t to MmaT in-place. -// No-op when src_t == MmaT. For the cross-dtype case: reads a src_t pair, -// converts via f32 intermediate, writes an MmaT pair. `pack_float2` -// dispatches to the native packed cvt for the destination type (e.g. -// cvt.rn.bf16x2.f32 for bf16). -template -__device__ __forceinline__ void convert_frag(Frag& frag) -{ - if constexpr (!std::is_same_v) - { -#pragma unroll - for (int i = 0; i < cute::size(frag); i += 2) - { - float2 const vals = toFloat2(reinterpret_cast(&frag(i))); - *reinterpret_cast*>(&frag(i)) = pack_float2(vals); - } - } -} - -// State → MMA B operand: dtype-aware TiledCopy. -// 2-byte smem: LDSM (SM75_U32x2_LDSM_N) — vectorized 16-bit ldmatrix. -// 4-byte smem: scalar UniversalCopy; pairs are converted to -// bf16 in registers by `convert_frag` after the load. -template -__device__ __forceinline__ auto make_state_b_s2r(TiledMma const& tm) -{ - using namespace cute; - if constexpr (sizeof(state_t) == 2) - { - return make_tiled_copy_B(Copy_Atom{}, tm); - } - else - { - static_assert(sizeof(state_t) == 4, "wide state path expects 4-byte smem"); - return make_tiled_copy_B(Copy_Atom, state_t>{}, tm); - } -} - -// Src → dst fragment conversion — a strict superset of the in-place overload -// above: supports narrowing (e.g. f32 → bf16) via a separate src fragment. -// Three paths: -// (1) src_t == dst_t: bit copy via Pair (sidesteps cutlass-wrapper vs -// native dtype mismatches like cutlass::bfloat16_t vs __nv_bfloat16). -// (2) Same width, different dtype (e.g. fp16 → bf16): paired cvt through f32. -// Works in-place when `src` aliases `dst`. -// (3) Different width (e.g. f32 → bf16): paired element load + pack_float2. -template -__device__ __forceinline__ void convert_frag(SrcFrag const& src, DstFrag& dst) -{ - using namespace cute; - if constexpr (std::is_same_v) - { -#pragma unroll - for (int i = 0; i < size(src); i += 2) - { - *reinterpret_cast*>(&dst(i)) = *reinterpret_cast const*>(&src(i)); - } - } - else if constexpr (sizeof(src_t) == sizeof(dst_t)) - { -#pragma unroll - for (int i = 0; i < size(src); i += 2) - { - float2 const vals = toFloat2(reinterpret_cast(&src(i))); - *reinterpret_cast*>(&dst(i)) = pack_float2(vals); - } - } - else - { - static_assert(sizeof(dst_t) == 2, "only narrowing to 2-byte dst supported"); -#pragma unroll - for (int i = 0; i < size(src); i += 2) - { - *reinterpret_cast*>(&dst(i)) = pack_float2(make_float2(src(i), src(i + 1))); - } - } -} - -// 2b. frag_y += CB_scaled @ x (matmul 4, single K-tile) -// CB_scaled A operand loaded from swizzled smem via LDSM (precomputed by warps 0,1). -// x B operand loaded from smem via ldmatrix.trans. -template -__device__ __forceinline__ void add_cb_x(FragY& frag_y, FragCB const& frag_CB, SmemXTrans const& smem_x_trans, - S2RBTrans const& s2r_B_trans, S2RThrBTrans const& s2r_thr_B_trans, ThrMma const& thr_mma, TiledMma const& tiled_mma, - int n) -{ - using namespace cute; - Tensor smem_x_trans_ntile - = local_tile(smem_x_trans, make_tile(Int{}, Int{}), make_coord(n, _0{})); - auto smem_x_trans_s2r = s2r_thr_B_trans.partition_S(smem_x_trans_ntile); - auto frag_B_x = thr_mma.partition_fragment_B( - make_tensor((MmaT*) 0x0, make_shape(Int{}, Int{}))); - auto frag_B_x_view = s2r_thr_B_trans.retile_D(frag_B_x); - - cute::copy(s2r_B_trans, smem_x_trans_s2r, frag_B_x_view); - cute::gemm(tiled_mma, frag_y, frag_CB, frag_B_x, frag_y); -} - -// 2c. frag_y += CB_old @ old_x (matmul-4 over old tokens; sibling of add_cb_x). -// CB_old A-operand: pre-loaded by caller (m16n8k_old A-frag). -// old_x B-operand: ldmatrix.trans from smem.old_x viewed transposed. -// K_OLD = MAX_WINDOW_PAD_MMA_K ∈ {8, 16}. Caller's tiled_mma_old uses the -// matching m16n8k_OLD atom (K_BIG=16 or K_SMALL=8). frag_y partitioned by a -// different (K_BIG) tiled_mma is layout-compatible — the m16n8 C-frag shape is -// the same regardless of K. -template -__device__ __forceinline__ void add_cb_old_x(FragY& frag_y, FragCBOld const& frag_CB_old, - SmemOldXTrans const& smem_old_x_trans, S2RBTransOld const& s2r_B_trans_old, - S2RThrBTransOld const& s2r_thr_B_trans_old, ThrMmaOld const& thr_mma_old, TiledMmaOld const& tiled_mma_old, int n) -{ - using namespace cute; - Tensor smem_old_x_ntile - = local_tile(smem_old_x_trans, make_tile(Int{}, Int{}), make_coord(n, _0{})); - auto smem_old_x_s2r = s2r_thr_B_trans_old.partition_S(smem_old_x_ntile); - auto frag_B_old_x = thr_mma_old.partition_fragment_B( - make_tensor((MmaT*) 0x0, make_shape(Int{}, Int{}))); - auto frag_B_old_x_view = s2r_thr_B_trans_old.retile_D(frag_B_old_x); - - cute::copy(s2r_B_trans_old, smem_old_x_s2r, frag_B_old_x_view); - cute::gemm(tiled_mma_old, frag_y, frag_CB_old, frag_B_old_x, frag_y); -} - -// 3b. frag_y += D * x[t, d] (per-thread skip connection via partition_C) -template -__device__ __forceinline__ void add_D_skip( - FragY& frag_y, SmemX const& smem_x, ThrMma const& thr_mma, float D_val, int n) -{ - using namespace cute; - if (D_val == 0.f) - return; - Tensor smem_x_tile = local_tile(smem_x, make_tile(Int{}, Int{}), make_coord(_0{}, n)); - Tensor x_part = thr_mma.partition_C(smem_x_tile); - // Load pairs of consecutive bf16 elements and convert via paired toFloat2. - // m16n8k16 partition_C places consecutive N-column pairs adjacent in smem. - static_assert(sizeof(input_t) == 2, "vectorized D_skip requires 2-byte input_t"); -#pragma unroll - for (int i = 0; i < size(frag_y); i += 2) - { - float2 vals = toFloat2(reinterpret_cast(&x_part(i))); - frag_y(i) += D_val * vals.x; - frag_y(i + 1) += D_val * vals.y; - } -} - -// 4b. frag_y *= z * sigmoid(z) (z-gating via partition_C) -template -__device__ __forceinline__ void compute_z_gating( - FragY& frag_y, SmemZ const& smem_z, ThrMma const& thr_mma, void const* z_ptr, int n) -{ - using namespace cute; - if (!z_ptr) - return; - Tensor smem_z_tile = local_tile(smem_z, make_tile(Int{}, Int{}), make_coord(_0{}, n)); - Tensor z_part = thr_mma.partition_C(smem_z_tile); -#pragma unroll - for (int i = 0; i < size(frag_y); i += 2) - { - float2 const z = toFloat2(reinterpret_cast(&z_part(i))); - frag_y(i) *= z.x * __fdividef(1.f, (1.f + __expf(-z.x))); - frag_y(i + 1) *= z.y * __fdividef(1.f, (1.f + __expf(-z.y))); - } -} - -// ============================================================================= -// Pipelined K-loop GEMM -// ============================================================================= -// Computes frag_y[n] += A @ B[n] for n ∈ [0, NumNTiles), where A is shared -// across N-tiles and B[n] is the n-th N-tile of `smem_B` (sliced inside). -// -// NumStages-deep register pipeline hides LDSM → HMMA latency: at steady state -// slot (k+NumStages-1) is loading while HMMA consumes slot k. ATypeIn → MmaT -// and BTypeIn → MmaT conversions happen in registers between load and consume -// (in-place when widths match — see `convert_frag`). -// -// Used by matmul 3 (init_out += C @ state^T): A = C (shared), B = state. -// NumNTiles = sizeof...(FragY) = D_PER_CTA / N_TILE (1 for D_SPLIT=2, 2 for -// D_SPLIT=1). -template -__device__ __forceinline__ void pipelined_kloop_gemm(TiledMma const& tiled_mma, ThrMma const& thr_mma, int tid, - SmemAKtiled const& smem_A_ktiled, SmemB const& smem_B, FragY&... frag_y) -{ - using namespace cute; - constexpr int NumNTiles = sizeof...(FragY); - static_assert(NumStages >= 2, "NumStages must be >= 2 for pipelining"); - static_assert(NumKTiles >= NumStages - 1, "NumKTiles must be >= NumStages - 1 for full prologue"); - static_assert(NumNTiles >= 1, "NumNTiles must be >= 1"); - - constexpr int N_TILE = cute::tile_size<1>(TiledMma{}); - constexpr int K_TILE = cute::tile_size<2>(TiledMma{}); - - // ── S2R copies ── - auto s2r_A = make_tiled_copy_A(Copy_Atom{}, tiled_mma); - auto s2r_thr_A = s2r_A.get_slice(tid); - auto s2r_B = make_state_b_s2r(tiled_mma); - auto s2r_thr_B = s2r_B.get_slice(tid); - - // ── Tile B by (N, K): shape (N_TILE, K_TILE, N_OUTER, NumKTiles) ── - auto smem_B_tiled = local_tile(smem_B, make_tile(Int{}, Int{}), make_coord(_, _)); - - // ── Partitioned smem (A shared, B per-N-tile) ── - auto smem_A_s2r = s2r_thr_A.partition_S(smem_A_ktiled); - auto sample_smem_B_n = smem_B_tiled(_, _, _0{}, _); - using SmemBS2RType = decltype(s2r_thr_B.partition_S(sample_smem_B_n)); - SmemBS2RType smem_B_s2r[NumNTiles]; - CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) - { - smem_B_s2r[n] = s2r_thr_B.partition_S(smem_B_tiled(_, _, n, _)); - } - - // ── Fragment / view types ── - using FragA = decltype(thr_mma.partition_fragment_A(smem_A_ktiled(_, _, _0{}))); - using FragB = decltype(thr_mma.partition_fragment_B(sample_smem_B_n(_, _, _0{}))); - using b_view_t = std::conditional_t; - using FragBStg = decltype(make_fragment_like(std::declval())); - using FragAView = decltype(s2r_thr_A.retile_D(std::declval())); - using FragBStgView = decltype(s2r_thr_B.retile_D(std::declval())); - - // ── Multi-stage register fragments ── - // Storage type matches the MMA fragment for A; for B the staging buffer is - // BTypeIn-typed (when narrowing) or MmaT-typed (when widths match — the two - // alias the same registers and `convert_frag` collapses to a bit-copy / - // in-place reinterpret). - FragA frag_A[NumStages]; - FragB frag_B[NumNTiles][NumStages]; - FragBStg frag_B_stg[NumNTiles][NumStages]; - FragAView frag_A_view[NumStages]; - FragBStgView frag_B_stg_view[NumNTiles][NumStages]; - CUTE_UNROLL - for (int s = 0; s < NumStages; ++s) - { - frag_A_view[s] = s2r_thr_A.retile_D(frag_A[s]); - CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) - { - frag_B_stg_view[n][s] = s2r_thr_B.retile_D(frag_B_stg[n][s]); - } - } - - // Pack frag_y into a pointer array for indexed access (replay kernel pattern). - using FragY0 = std::tuple_element_t<0, std::tuple>; - static_assert((std::is_same_v && ...), "all FragY parameters must be the same type"); - FragY0* frag_y_p[NumNTiles] = {(&frag_y)...}; - - // ── Per-stage operations (slot is constant after #pragma unroll) ── - auto load_one = [&](int k_src, int slot) - { - cute::copy(s2r_A, smem_A_s2r(_, _, _, k_src), frag_A_view[slot]); - CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) - { - cute::copy(s2r_B, smem_B_s2r[n](_, _, _, k_src), frag_B_stg_view[n][slot]); - } - }; - auto convert_one = [&](int slot) - { - convert_frag(frag_A[slot]); - CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) - { - convert_frag(frag_B_stg[n][slot], frag_B[n][slot]); - } - }; - auto compute_one = [&](int slot) - { - CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) - { - cute::gemm(tiled_mma, *frag_y_p[n], frag_A[slot], frag_B[n][slot], *frag_y_p[n]); - } - }; - - // ── Clear accumulators ── - CUTE_UNROLL - for (int n = 0; n < NumNTiles; ++n) - clear(*frag_y_p[n]); - - // ── Prologue: load + convert stages 0..NumStages-2 ── - CUTE_UNROLL - for (int s = 0; s < NumStages - 1; ++s) - { - load_one(s, s); - convert_one(s); - } - - // ── Main K-loop: load slot (k+NumStages-1) % NumStages, compute slot k % NumStages ── -#pragma unroll - for (int k = 0; k < NumKTiles; ++k) - { - int const k_load = k + NumStages - 1; - int const slot_load = k_load % NumStages; - int const slot_compute = k % NumStages; - if (k_load < NumKTiles) - load_one(k_load, slot_load); - compute_one(slot_compute); - if (k_load < NumKTiles) - convert_one(slot_load); - } -} - -// ── Matmul 3: init_out = C @ state^T ──────────────────────────────────────── -// Thin wrapper: builds the swizzled smem views for C and state, then dispatches -// to `pipelined_kloop_gemm`. NumNTiles = sizeof...(FragY) = D_PER_CTA / N_TILE -// (1 for D_SPLIT=2, 2 for D_SPLIT=1). -// -// On C: aliased view (see compute_CB_scaled_2warp). On state: 2-byte smem is -// reinterpret-cast to MMA_prop::operand_t so the 16-bit LDSM atom matches the view -// (actual element type recovered inside `convert_frag`); ≥4-byte smem keeps -// the native dtype and uses scalar UniversalCopy + register conversion. -template -__device__ __forceinline__ void add_init_out( - SmemT const& smem, TiledMma const& tiled_mma, ThrMma const& thr_mma, int tid, FragY&... frag_y) -{ - using namespace cute; - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - constexpr int K_TILE = cute::tile_size<2>(TiledMma{}); - constexpr int NUM_K_TILES = DSTATE / K_TILE; - // Smem source dtype for matmul-3 (generic kernel only; 8-bit state goes - // through the dedicated `checkpointing_ssu_kernel_8bit` path): - // - sizeof(state_t) == 2 (fp16/bf16): LDSM the native 16-bit, view as bf16. - // - sizeof(state_t) == 4 (fp32): scalar UniversalCopy + on-the-fly convert. - static_assert(sizeof(state_t) != 1, - "add_init_out is the 2/4-byte path; 1-byte state goes through " - "compute_output_8bit"); - constexpr bool is_2byte_smem = (sizeof(state_t) == 2); - using state_view_t = std::conditional_t; - using BTypeIn = state_t; - - auto layout_C_swz = make_aliased_swizzled_layout_rc(); - Tensor smem_C = make_tensor(make_smem_ptr(reinterpret_cast(smem.C)), layout_C_swz); - Tensor smem_C_ktiled - = local_tile(smem_C, make_tile(Int{}, Int{}), make_coord(_0{}, _)); - - // Swizzle layout matches the dtype of the buffer being viewed. - auto const layout_state_swz = make_swizzled_layout_rc(); - state_view_t const* smem_state_ptr = reinterpret_cast(smem.state); - Tensor smem_state = make_tensor(make_smem_ptr(smem_state_ptr), layout_state_swz); - - pipelined_kloop_gemm<3, NUM_K_TILES, input_t, BTypeIn, MMA_prop::operand_t>( - tiled_mma, thr_mma, tid, smem_C_ktiled, smem_state, frag_y...); -} - -// store_state: vectorized smem → gmem state writeback (128 threads). -// Defined here (rather than alongside the other Phase 3 store helpers -// below) because compute_and_store_output calls it inline — issued right -// after matmul 3 so the STGs fire-and-forget in parallel with matmul 4 + -// epilogue. smem and gmem hold the same dtype now (no on-egress -// conversion) so this is always a direct 128-bit copy. -template -__device__ __forceinline__ void store_state( - SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, int d_tile, int head, int64_t cache_slot) -{ - using namespace cute; - int const flat_tid = warp * warpSize + lane; - auto* __restrict__ state_w = reinterpret_cast(params.state); - // gmem dest = head's full state base + d_tile's row slice. - int64_t const state_base - = cache_slot * params.state_stride_seq + (int64_t) head * DIM * DSTATE + (int64_t) d_tile * D_PER_CTA * DSTATE; - - // ── Per-CTA smem swizzle layout [D_PER_CTA, DSTATE]. ── - auto layout_smem_swz = make_swizzled_layout_rc(); - state_t const* smem_state_base = reinterpret_cast(smem.state); - - Tensor sState = make_tensor(make_smem_ptr(smem_state_base), layout_smem_swz); - Tensor gState = make_tensor(make_gmem_ptr(state_w + state_base), - make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{}))); - // Each store is 16 bytes — adjust val cols to the dtype. - constexpr int VAL_COLS = Copy_prop::vec_bytes / sizeof(state_t); - auto s2g = make_tiled_copy(Copy_Atom, state_t>{}, Layout, Stride<_8, _1>>{}, - Layout>>{}); - auto thr = s2g.get_slice(flat_tid); - copy(s2g, thr.partition_S(sState), thr.partition_D(gState)); -} - -// ── Store functions (called from kernel after compute_y + sync) ── -// (store_state moved above compute_and_store_output — used there for -// the state-writeback hoist.) - -template -__device__ __forceinline__ void store_old_x(SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, - int d_tile, int head, int64_t cache_slot, int write_offset, int seq_len) -{ - using namespace cute; - constexpr int NPREDICTED_PAD_MMA_M = SmemT::NPREDICTED_PAD_MMA_M; - int const flat_tid = warp * warpSize + lane; - - auto* __restrict__ old_x_w = reinterpret_cast(params.old_x); - // gmem dest = head's full slot + d_tile's D-slice offset, shifted by - // `write_offset` along the T-axis (must_checkpoint ? 0 : prev_k). - int64_t const ox_w_base = cache_slot * params.old_x_stride_seq + (int64_t) write_offset * params.old_x_stride_token - + head * DIM + (int64_t) d_tile * D_PER_CTA; - - // Smem and gmem are both viewed at the full atom-padded width D_SMEM_COLS. - // The wide thread layout (16 row × 8 col × 1×8 val = 16 rows × 64 cols/pass - // for bf16) covers one full atom width per thread-row, which is the - // swizzle's bank-conflict-free contract on the LDS side (load-from-smem). - // A narrow layout would (a) waste 64 threads (warps 2, 3 idle) and - // (b) cause LDS bank conflicts on the smem-read side (observed as - // 4-way LDS conflict in d_split=2 ncu). Cols ≥ D_PER_CTA are predicated - // off via copy_if so STG never fires for them — no OOB write into the - // next d_tile / next head's gmem region. - constexpr int D_SMEM_COLS = SmemT::D_SMEM_COLS; - auto layout_x_swz = make_swizzled_layout_rc(); - Tensor sX = make_tensor(make_smem_ptr(reinterpret_cast(smem.x)), layout_x_swz); - Tensor gX = make_tensor(make_gmem_ptr(old_x_w + ox_w_base), - make_layout(make_shape(Int{}, Int{}), - make_stride(params.old_x_stride_token, Int<1>{}))); - - using ThrLayoutX = Layout, Stride<_8, _1>>; - auto s2g = make_tiled_copy(Copy_Atom, input_t>{}, ThrLayoutX{}, Layout>{}); - auto thr_s2g = s2g.get_slice(flat_tid); - - auto tSsX = thr_s2g.partition_S(sX); - auto tSgX = thr_s2g.partition_D(gX); - - // Per-(row, col) predicate: skip rows ≥ NPREDICTED (m-padding) and cols ≥ - // D_PER_CTA (atom-padding past the d_tile's data). - auto cX = make_identity_tensor(make_shape(Int{}, Int{})); - auto tScX = thr_s2g.partition_D(cX); - auto pred = make_tensor(shape(tScX)); - CUTE_UNROLL - for (int i = 0; i < size(pred); ++i) - { - pred(i) = (get<0>(tScX(i)) < seq_len) && (get<1>(tScX(i)) < D_PER_CTA); - } - copy_if(s2g, pred, tSsX, tSgX); -} - -// store_old_B runs on W0, W1 only (64 threads). Caller must gate -// with `if (warp < 2)` — these are the warps that hold valid smem.B -// after their own cp.async + wait. Halving the thread count keeps the -// overlap (writeback fires before CB+replay consume smem.B). -// -// Source: smem.B with NPREDICTED_PAD_MMA_N rows. -// Destination: gmem old_B[buf_write][write_offset:write_offset+NPREDICTED, :]. -// The `write_offset` argument shifts the gmem T-axis base — it's added to the -// base pointer below; the per-element predicate masks rows ≥ NPREDICTED. -// -// Thread layout `(8, 8) × (1, 8)` — **atom-aligned** with the Swizzle<3,3,3> -// (8, 64) atom for conflict-free smem reads. Per-tile 8 × 64 covers one -// full atom. For NPREDICTED_PAD_MMA_N=16: iters (2, 2) = 4 tiles, each -// thread owns 2 rows (t/8 and t/8+8) → per-iteration row predicate. For -// =8: iters (1, 2) = 2 tiles, each thread owns 1 row. The per-element -// predicate works for both. -template -__device__ __forceinline__ void store_old_B(SmemT& smem, CheckpointingSsuParams const& params, int warp, int lane, - int head, int group_idx, int64_t cache_slot, int buf_write, int write_offset, int seq_len) -{ - using namespace cute; - if (head % HEADS_PER_GROUP != 0) - return; - constexpr int NPREDICTED_PAD_MMA_N = SmemT::NPREDICTED_PAD_MMA_N; // matches smem.B row count - // Called only from warps 0, 1 — flat_tid ∈ [0, 64). - int const flat_tid = warp * warpSize + lane; - - auto* __restrict__ old_B_w = reinterpret_cast(params.old_B); - int64_t const oB_base = cache_slot * params.old_B_stride_seq + buf_write * params.old_B_stride_dbuf - + (int64_t) write_offset * params.old_B_stride_token + group_idx * DSTATE; - - auto layout_B_swz = make_swizzled_layout_rc(); - Tensor sB = make_tensor(make_smem_ptr(reinterpret_cast(smem.B)), layout_B_swz); - Tensor gB = make_tensor(make_gmem_ptr(old_B_w + oB_base), - make_layout( - make_shape(Int{}, Int{}), make_stride(params.old_B_stride_token, Int<1>{}))); - - // 64 threads, (8, 8) × (1, 8) = atom-aligned per-tile (8, 64). - auto s2g = make_tiled_copy(Copy_Atom, input_t>{}, Layout, Stride<_8, _1>>{}, - Layout>{}); - auto thr_s2g = s2g.get_slice(flat_tid); - auto tSsB = thr_s2g.partition_S(sB); - auto tSgB = thr_s2g.partition_D(gB); - - // Fast path: no smem-side row padding AND no varlen-side truncation. - // The runtime `seq_len == NPREDICTED` is a constexpr-foldable compare in - // the non-varlen path (kernel prologue assigns `seq_len = NPREDICTED`), - // so it eliminates at -O3. In varlen with `seq_len == NPREDICTED` it's - // a runtime check that picks the cheaper unpredicated STG. - if constexpr (NPREDICTED == NPREDICTED_PAD_MMA_N) - { - if (seq_len == NPREDICTED) - { - copy(s2g, tSsB, tSgB); - return; - } - } - // Predicated: either smem rows > NPREDICTED (m-padding) OR varlen with - // seq_len < NPREDICTED. Mask each iter against `seq_len`. - auto cB = make_identity_tensor(make_shape(Int{}, Int{})); - auto tScB = thr_s2g.partition_D(cB); - auto pred = make_tensor(shape(tScB)); - CUTE_UNROLL - for (int i = 0; i < size(pred); ++i) - { - pred(i) = get<0>(tScB(i)) < seq_len; - } - copy_if(s2g, pred, tSsB, tSgB); -} - -} // namespace flashinfer::mamba::checkpointing - -#endif // FLASHINFER_MAMBA_KERNEL_CHECKPOINTING_SSU_COMMON_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh deleted file mode 100644 index c338e2a9eb90..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/launch_checkpointing_ssu.cuh +++ /dev/null @@ -1,192 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ -#ifndef FLASHINFER_MAMBA_LAUNCH_CHECKPOINTING_SSU_CUH_ -#define FLASHINFER_MAMBA_LAUNCH_CHECKPOINTING_SSU_CUH_ - -// Launcher functions for the incremental SSU kernel. -// Includes both the bf16/fp16/fp32 and 8-bit kernel headers. - -#include "kernel_checkpointing_ssu.cuh" -#include "kernel_checkpointing_ssu_8bit.cuh" - -namespace flashinfer::mamba::checkpointing -{ - -// ── Dispatcher ───────────────────────────────────────────────────────────── -// `D_SPLIT` splits each head's DIM axis across `D_SPLIT` CTAs. -// `VARLEN` selects the packed-token gmem layout (cu_seqlens-driven). -// `launchCheckpointingSsuImpl` is the per-(D_SPLIT, VARLEN) specialization; -// `launchCheckpointingSsu` (below) is the runtime dispatcher. -template -void launchCheckpointingSsuImpl(CheckpointingSsuParams& params, cudaStream_t stream) -{ - constexpr int NUM_WARPS = 4; - - FLASHINFER_CHECK(params.nheads % params.ngroups == 0, "nheads (", params.nheads, ") must be divisible by ngroups (", - params.ngroups, ")"); - - // cp.async.ca with .L2::128B requires 16B-aligned pointers (128-bit / sizeof element). - // The .L2::128B hint further requires the base address to be 128B-aligned for full - // cache line utilization, but the hardware only faults on < 16B alignment. - // All cp.async-loaded operands need 16B alignment; output is also vectorized - // (Pair stores partitioned by m16n8k16 partition_C — base must be at - // least 16B-aligned for the stride math to keep per-thread stores aligned). - FLASHINFER_CHECK_ALIGNMENT(params.B, 16); - FLASHINFER_CHECK_ALIGNMENT(params.C, 16); - FLASHINFER_CHECK_ALIGNMENT(params.x, 16); - FLASHINFER_CHECK_ALIGNMENT(params.state, 16); - FLASHINFER_CHECK_ALIGNMENT(params.old_x, 16); - FLASHINFER_CHECK_ALIGNMENT(params.old_B, 16); - FLASHINFER_CHECK_ALIGNMENT(params.output, 16); - if (params.z != nullptr) - { - FLASHINFER_CHECK_ALIGNMENT(params.z, 16); - } - - // Per-CTA D = DIM / D_SPLIT. Smem footprint shrinks for D-owned - // buffers (state, x, z, old_x); non-D buffers (B, C, old_B, scalars) unchanged. - constexpr int D_PER_CTA = DIM / D_SPLIT; - - // HEADS_PER_GROUP is JIT-stamped via the customize_config jinja, so only - // one (nheads / ngroups) specialization gets baked into this .so. The - // wrapper has already validated `nheads / ngroups == HEADS_PER_GROUP` - // before reaching us — the kernel cross-checks with an assert below. - FLASHINFER_CHECK(params.nheads / params.ngroups == HEADS_PER_GROUP, - "nheads/ngroups (=", params.nheads / params.ngroups, ") must match JIT HEADS_PER_GROUP=", HEADS_PER_GROUP); - // PDL launch attribute. ENABLE_PDL is JIT-stamped (see - // checkpointing_ssu_customize_config.jinja); the kernel's body has its - // PDL PTX gated on the same constexpr via `if constexpr (ENABLE_PDL)`, so - // the .so contains exactly one load path. When ENABLE_PDL is false the - // attribute is set to 0 (effectively no PDL) — cudaLaunchKernelEx is - // used either way per FlashInfer convention (see norm.cuh:135). - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = ENABLE_PDL ? 1 : 0; - - auto launch_kernel = [&]() - { - if constexpr (sizeof(state_t) == 1) - { - // int8 chain rewrite — uses checkpointing_ssu_kernel_8bit + - // CheckpointingSsuStorage8bit. Only D_SPLIT == 1 is valid (the wrapper - // asserts this); D_SPLIT == 2 still gets template-instantiated by the - // public dispatcher's switch but is unreachable at runtime — gate the - // body with `if constexpr (D_SPLIT == 1)` so that path doesn't launch. - if constexpr (D_SPLIT == 1) - { - auto func = checkpointing_ssu_kernel_8bit; - constexpr size_t smem_size - = sizeof(CheckpointingSsuStorage8bit); - - if constexpr (smem_size > 0) - { - FLASHINFER_CUDA_CHECK( - cudaFuncSetAttribute(func, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - } - - cudaLaunchConfig_t config; - config.gridDim = dim3(D_SPLIT, params.batch, params.nheads); - config.blockDim = dim3(warpSize, NUM_WARPS); - config.dynamicSmemBytes = smem_size; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = 1; - FLASHINFER_CUDA_CHECK(cudaLaunchKernelEx(&config, func, params)); - } - else - { - FLASHINFER_CHECK(false, - "checkpointing_ssu_kernel_8bit: unsupported D_SPLIT != 1 for 8-bit " - "state_t (got D_SPLIT=", - D_SPLIT, ")"); - } - } - else - { - // Generic kernel: bf16 / fp16 / fp32 state, supports D_SPLIT ∈ {1, 2}. - auto func - = checkpointing_ssu_kernel; - - constexpr size_t smem_size - = sizeof(CheckpointingSsuStorage); - - if constexpr (smem_size > 0) - { - FLASHINFER_CUDA_CHECK( - cudaFuncSetAttribute(func, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - } - - // Grid is (D_SPLIT, batch, nheads). D-tile is the fastest axis so the - // `D_SPLIT` CTAs of the same head land on adjacent SMs and share L2 - // lines for the redundantly-loaded inputs (C, B, dt, ...). - cudaLaunchConfig_t config; - config.gridDim = dim3(D_SPLIT, params.batch, params.nheads); - config.blockDim = dim3(warpSize, NUM_WARPS); - config.dynamicSmemBytes = smem_size; - config.stream = stream; - config.attrs = attrs; - config.numAttrs = 1; - FLASHINFER_CUDA_CHECK(cudaLaunchKernelEx(&config, func, params)); - } - }; - - launch_kernel(); -} - -// Public dispatcher: routes on `params.d_split` ({1, 2}) and varlen -// (`params.cu_seqlens != nullptr` → VARLEN=true). Each (D_SPLIT, VARLEN) -// pair gets its own template specialization — the JIT URI distinguishes them -// only via `d_split` today, so the same compiled `.so` will hold all four -// specializations after this commit. -template -void launchCheckpointingSsu(CheckpointingSsuParams& params, cudaStream_t stream) -{ - bool const is_varlen = (params.cu_seqlens != nullptr); - auto launch = [&]() - { - launchCheckpointingSsuImpl(params, stream); - }; - auto launch_d_split = [&]() - { - if (is_varlen) - { - launch.template operator()(); - } - else - { - launch.template operator()(); - } - }; - switch (params.d_split) - { - case 1: launch_d_split.template operator()<1>(); break; - case 2: launch_d_split.template operator()<2>(); break; - default: - FLASHINFER_CHECK(false, "Unsupported d_split: ", params.d_split, - ". Allowed values: {1, 2}. d_split=4 needs " - "warp-count restructure."); - } -} - -} // namespace flashinfer::mamba::checkpointing - -#endif // FLASHINFER_MAMBA_LAUNCH_CHECKPOINTING_SSU_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh deleted file mode 100644 index 8c6d8e93dd5b..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/mamba/ssu_mtp_common.cuh +++ /dev/null @@ -1,160 +0,0 @@ -/* - * Copyright (c) 2025 by FlashInfer team. - * - * 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. - */ - -// Shared definitions for the vertical and horizontal MTP kernels. - -#pragma once - -#include - -#include "conversion.cuh" - -namespace flashinfer::mamba::mtp -{ - -// Round up to next power of 2 (compile-time). -constexpr int nextPow2(int v) -{ - v--; - v |= v >> 1; - v |= v >> 2; - v |= v >> 4; - v |= v >> 8; - v |= v >> 16; - return v + 1; -} - -using barrier_t = cuda::barrier; - -enum class WarpRole -{ - kCompute, - kTMALoad, - kEpilogue -}; - -__device__ __forceinline__ WarpRole get_warp_role(int warp) -{ - if (warp < 12) - return WarpRole::kCompute; - if (warp < 15) - return WarpRole::kTMALoad; - return WarpRole::kEpilogue; -} - -// XOR-based bank-conflict-free swizzle for horizontal state traversal. -// Operates on flat byte addresses: XORs the bank index with the row (cycle) index. -// cycle_length = row stride in bytes, bank_size = sizeof(uint32_t). -template -__device__ __forceinline__ int xor_swizzle(int address) -{ - int const cycle = address / cycle_length; - int const delta = address % cycle_length; - int const bank_idx = delta / bank_size; - int const intra_bank = delta % bank_size; - int const new_bank_idx = bank_idx ^ cycle; - return cycle * cycle_length + new_bank_idx * bank_size + intra_bank; -} - -// ── Parity-based barrier helpers (tight spin, no NANOSLEEP) ───────────────── -// More efficient than cuda::barrier::wait() for latency-sensitive pipelines. -// The standard cuda::barrier::wait() adds a NANOSLEEP backoff loop between -// try_wait attempts, which can overshoot and waste cycles. The raw -// mbarrier.try_wait.parity instruction does a tight spin instead. -// See CUDA Programming Guide §4.9.3 "Explicit Phase Tracking". - -__device__ __forceinline__ void arrive_and_wait_parity(barrier_t& bar, uint32_t& parity) -{ - uint32_t const smem_addr - = static_cast(__cvta_generic_to_shared(cuda::device::barrier_native_handle(bar))); - asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" ::"r"(smem_addr) : "memory"); - uint32_t ready = 0; - while (!ready) - { - asm volatile( - "{\n" - ".reg .pred p;\n" - "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n" - "selp.b32 %0, 1, 0, p;\n" - "}\n" - : "=r"(ready) - : "r"(smem_addr), "r"(parity)); - } - parity ^= 1; -} - -// ── SM100 f32x2 packed SIMD helpers ────────────────────────────────────────── -// On Blackwell (SM100+), {mul,fma}.f32x2 pack two fp32 operations into one -// instruction and issue on the dedicated FMUL2 pipeline, which runs in parallel -// with the regular FMA pipe. This halves instruction count for element-wise -// fp32 math on independent pairs (e.g. adjacent state-vector components). -// On older architectures the fallback is two scalar ops — zero overhead. -// See: https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/simd_sm100.hpp - -__device__ __forceinline__ void mul_f32x2(float2& c, float2 const& a, float2 const& b) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 - asm("mul.f32x2 %0, %1, %2;\n" - : "=l"(reinterpret_cast(c)) - : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b))); -#else - c.x = a.x * b.x; - c.y = a.y * b.y; -#endif -} - -__device__ __forceinline__ void fma_f32x2(float2& d, float2 const& a, float2 const& b, float2 const& c) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 - asm("fma.rn.f32x2 %0, %1, %2, %3;\n" - : "=l"(reinterpret_cast(d)) - : "l"(reinterpret_cast(a)), "l"(reinterpret_cast(b)), - "l"(reinterpret_cast(c))); -#else - d.x = a.x * b.x + c.x; - d.y = a.y * b.y + c.y; -#endif -} - -// ============================================================================= -// convertAndStoreSRHorizontal — convert a pair of f32 state values to half. -// When PHILOX_ROUNDS > 0: stochastic rounding via f16x2. -// When PHILOX_ROUNDS == 0: plain nearest-even conversion. -// e is the pair-aligned index within the tile (must be even). -// ============================================================================= - -template -__device__ __forceinline__ void convertAndStoreSRHorizontal(state_t& out0, state_t& out1, float s0, float s1, - int64_t rand_seed, int state_ptr_offset, int dd, int col0, int e, uint32_t (&rand_ints)[4]) -{ - using namespace conversion; - if constexpr (PHILOX_ROUNDS > 0) - { - if (e % 4 == 0) - philox_randint4x(rand_seed, state_ptr_offset + dd * DSTATE + col0 + e, rand_ints[0], - rand_ints[1], rand_ints[2], rand_ints[3]); - uint32_t packed = cvt_rs_f16x2_f32(s0, s1, rand_ints[e / 2 % 2]); - out0 = __ushort_as_half(static_cast(packed & 0xFFFFu)); - out1 = __ushort_as_half(static_cast(packed >> 16)); - } - else - { - convertAndStore(&out0, s0); - convertAndStore(&out1, s1); - } -} - -} // namespace flashinfer::mamba::mtp diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh deleted file mode 100644 index 334cd3c5f2ee..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/utils.cuh +++ /dev/null @@ -1,648 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer team. - * - * 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. - */ -#ifndef FLASHINFER_UTILS_CUH_ -#define FLASHINFER_UTILS_CUH_ -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include - -#include "exception.h" - -#define STR_HELPER(x) #x -#define STR(x) STR_HELPER(x) - -// macro to turn off fp16 qk reduction to reduce binary -#ifndef FLASHINFER_ALWAYS_DISUSE_FP16_QK_REDUCTION -#define FLASHINFER_ALWAYS_DISUSE_FP16_QK_REDUCTION 0 -#endif - -#ifndef NDEBUG -#define FLASHINFER_CUDA_CALL(func, ...) \ - { \ - cudaError_t e = (func); \ - if (e != cudaSuccess) \ - { \ - std::cerr << "CUDA Error: " << cudaGetErrorString(e) << " (" << e << ") " << __FILE__ << ": line " \ - << __LINE__ << " at function " << STR(func) << std::endl; \ - return e; \ - } \ - } -#else -#define FLASHINFER_CUDA_CALL(func, ...) \ - { \ - cudaError_t e = (func); \ - if (e != cudaSuccess) \ - { \ - return e; \ - } \ - } -#endif - -#define FLASHINFER_CUDA_CHECK(func) \ - do \ - { \ - cudaError_t e = (func); \ - FLASHINFER_CHECK(e == cudaSuccess, "CUDA Error: ", cudaGetErrorString(e), " (", int(e), ") at ", __FILE__, \ - ":", __LINE__, " in ", STR(func)); \ - } while (0) - -#define FLASHINFER_CHECK_ALIGNMENT(ptr, size_bytes) \ - FLASHINFER_CHECK(reinterpret_cast(ptr) % (size_bytes) == 0, #ptr, " must be aligned to ", (size_bytes), \ - " bytes, got address ", (uintptr_t) (ptr)) - -#define FLASHINFER_CHECK_TMA_ALIGNED(ptr) FLASHINFER_CHECK_ALIGNMENT(ptr, 128) - -#define DISPATCH_USE_FP16_QK_REDUCTION(use_fp16_qk_reduction, USE_FP16_QK_REDUCTION, ...) \ - if (use_fp16_qk_reduction) \ - { \ - FLASHINFER_ERROR("FP16_QK_REDUCTION disabled at compile time"); \ - } \ - else \ - { \ - constexpr bool USE_FP16_QK_REDUCTION = false; \ - __VA_ARGS__ \ - } - -#define DISPATCH_NUM_MMA_Q(num_mma_q, NUM_MMA_Q, ...) \ - if (num_mma_q == 1) \ - { \ - constexpr size_t NUM_MMA_Q = 1; \ - __VA_ARGS__ \ - } \ - else if (num_mma_q == 2) \ - { \ - constexpr size_t NUM_MMA_Q = 2; \ - __VA_ARGS__ \ - } \ - else \ - { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported num_mma_q: " << num_mma_q; \ - FLASHINFER_ERROR(err_msg.str()); \ - } - -#define DISPATCH_NUM_MMA_KV(max_mma_kv, NUM_MMA_KV, ...) \ - if (max_mma_kv >= 8) \ - { \ - constexpr size_t NUM_MMA_KV = 8; \ - __VA_ARGS__ \ - } \ - else if (max_mma_kv >= 4) \ - { \ - constexpr size_t NUM_MMA_KV = 4; \ - __VA_ARGS__ \ - } \ - else if (max_mma_kv >= 2) \ - { \ - constexpr size_t NUM_MMA_KV = 2; \ - __VA_ARGS__ \ - } \ - else if (max_mma_kv >= 1) \ - { \ - constexpr size_t NUM_MMA_KV = 1; \ - __VA_ARGS__ \ - } \ - else \ - { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported max_mma_kv: " << max_mma_kv; \ - FLASHINFER_ERROR(err_msg.str()); \ - } - -#define DISPATCH_CTA_TILE_Q(cta_tile_q, CTA_TILE_Q, ...) \ - switch (cta_tile_q) \ - { \ - case 128: \ - { \ - constexpr uint32_t CTA_TILE_Q = 128; \ - __VA_ARGS__ \ - break; \ - } \ - case 64: \ - { \ - constexpr uint32_t CTA_TILE_Q = 64; \ - __VA_ARGS__ \ - break; \ - } \ - case 16: \ - { \ - constexpr uint32_t CTA_TILE_Q = 16; \ - __VA_ARGS__ \ - break; \ - } \ - default: \ - { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported cta_tile_q: " << cta_tile_q; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_GQA_GROUP_SIZE(group_size, GROUP_SIZE, ...) \ - if (group_size == 1) \ - { \ - constexpr size_t GROUP_SIZE = 1; \ - __VA_ARGS__ \ - } \ - else if (group_size == 2) \ - { \ - constexpr size_t GROUP_SIZE = 2; \ - __VA_ARGS__ \ - } \ - else if (group_size == 3) \ - { \ - constexpr size_t GROUP_SIZE = 3; \ - __VA_ARGS__ \ - } \ - else if (group_size == 4) \ - { \ - constexpr size_t GROUP_SIZE = 4; \ - __VA_ARGS__ \ - } \ - else if (group_size == 8) \ - { \ - constexpr size_t GROUP_SIZE = 8; \ - __VA_ARGS__ \ - } \ - else \ - { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported group_size: " << group_size; \ - FLASHINFER_ERROR(err_msg.str()); \ - } - -#define DISPATCH_MASK_MODE(mask_mode, MASK_MODE, ...) \ - switch (mask_mode) \ - { \ - case MaskMode::kNone: \ - { \ - constexpr MaskMode MASK_MODE = MaskMode::kNone; \ - __VA_ARGS__ \ - break; \ - } \ - case MaskMode::kCausal: \ - { \ - constexpr MaskMode MASK_MODE = MaskMode::kCausal; \ - __VA_ARGS__ \ - break; \ - } \ - case MaskMode::kCustom: \ - { \ - constexpr MaskMode MASK_MODE = MaskMode::kCustom; \ - __VA_ARGS__ \ - break; \ - } \ - case MaskMode::kMultiItemScoring: \ - { \ - constexpr MaskMode MASK_MODE = MaskMode::kMultiItemScoring; \ - __VA_ARGS__ \ - break; \ - } \ - default: \ - { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported mask_mode: " << int(mask_mode); \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -// convert head_dim to compile-time constant -#define DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, ...) \ - switch (head_dim) \ - { \ - case 64: \ - { \ - constexpr size_t HEAD_DIM = 64; \ - __VA_ARGS__ \ - break; \ - } \ - case 128: \ - { \ - constexpr size_t HEAD_DIM = 128; \ - __VA_ARGS__ \ - break; \ - } \ - case 256: \ - { \ - constexpr size_t HEAD_DIM = 256; \ - __VA_ARGS__ \ - break; \ - } \ - case 512: \ - { \ - constexpr size_t HEAD_DIM = 512; \ - __VA_ARGS__ \ - break; \ - } \ - default: \ - { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported head_dim: " << head_dim; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -// convert interleave to compile-time constant -#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \ - if (interleave) \ - { \ - constexpr bool INTERLEAVE = true; \ - __VA_ARGS__ \ - } \ - else \ - { \ - constexpr bool INTERLEAVE = false; \ - __VA_ARGS__ \ - } - -#define DISPATCH_ROPE_DIM(rope_dim, ROPE_DIM, ...) \ - switch (rope_dim) \ - { \ - case 16: \ - { \ - constexpr uint32_t ROPE_DIM = 16; \ - __VA_ARGS__ \ - break; \ - } \ - case 32: \ - { \ - constexpr uint32_t ROPE_DIM = 32; \ - __VA_ARGS__ \ - break; \ - } \ - case 64: \ - { \ - constexpr uint32_t ROPE_DIM = 64; \ - __VA_ARGS__ \ - break; \ - } \ - case 128: \ - { \ - constexpr uint32_t ROPE_DIM = 128; \ - __VA_ARGS__ \ - break; \ - } \ - case 256: \ - { \ - constexpr uint32_t ROPE_DIM = 256; \ - __VA_ARGS__ \ - break; \ - } \ - default: \ - { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported ROPE_DIM: " << rope_dim; \ - err_msg << ". Supported values: 16, 32, 64, 128, 256"; \ - err_msg << " in DISPATCH_ROPE_DIM"; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_POS_ENCODING_MODE(pos_encoding_mode, POS_ENCODING_MODE, ...) \ - switch (pos_encoding_mode) \ - { \ - case PosEncodingMode::kNone: \ - { \ - constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kNone; \ - __VA_ARGS__ \ - break; \ - } \ - case PosEncodingMode::kRoPELlama: \ - { \ - constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kRoPELlama; \ - __VA_ARGS__ \ - break; \ - } \ - case PosEncodingMode::kALiBi: \ - { \ - constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kALiBi; \ - __VA_ARGS__ \ - break; \ - } \ - default: \ - { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported pos_encoding_mode: " << int(pos_encoding_mode); \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_ALIGNED_VEC_SIZE(aligned_vec_size, ALIGNED_VEC_SIZE, ...) \ - switch (aligned_vec_size) \ - { \ - case 16: \ - { \ - constexpr size_t ALIGNED_VEC_SIZE = 16; \ - __VA_ARGS__ \ - break; \ - } \ - case 8: \ - { \ - constexpr size_t ALIGNED_VEC_SIZE = 8; \ - __VA_ARGS__ \ - break; \ - } \ - case 4: \ - { \ - constexpr size_t ALIGNED_VEC_SIZE = 4; \ - __VA_ARGS__ \ - break; \ - } \ - case 2: \ - { \ - constexpr size_t ALIGNED_VEC_SIZE = 2; \ - __VA_ARGS__ \ - break; \ - } \ - case 1: \ - { \ - constexpr size_t ALIGNED_VEC_SIZE = 1; \ - __VA_ARGS__ \ - break; \ - } \ - default: \ - { \ - std::ostringstream err_msg; \ - err_msg << "Unsupported aligned_vec_size: " << aligned_vec_size; \ - FLASHINFER_ERROR(err_msg.str()); \ - } \ - } - -#define DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, ...) \ - if (compute_capacity.first >= 8) \ - { \ - constexpr uint32_t NUM_STAGES_SMEM = 2; \ - __VA_ARGS__ \ - } \ - else \ - { \ - constexpr uint32_t NUM_STAGES_SMEM = 1; \ - __VA_ARGS__ \ - } - -namespace flashinfer -{ - -template -__forceinline__ __device__ __host__ constexpr T1 ceil_div(const T1 x, const T2 y) noexcept -{ - return (x + y - 1) / y; -} - -template -__forceinline__ __device__ __host__ constexpr T1 round_up(const T1 x, const T2 y) noexcept -{ - return ceil_div(x, y) * y; -} - -template -__forceinline__ __device__ __host__ constexpr T1 round_down(const T1 x, const T2 y) noexcept -{ - return (x / y) * y; -} - -inline std::pair GetCudaComputeCapability() -{ - int device_id = 0; - cudaGetDevice(&device_id); - int major = 0, minor = 0; - cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_id); - cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device_id); - return std::make_pair(major, minor); -} - -// This function is thread-safe and cached the sm_count. -// But it will only check the current CUDA device, thus assuming each process handles single GPU. -inline int GetCudaMultiProcessorCount() -{ - static std::atomic sm_count{0}; - int cached = sm_count.load(std::memory_order_relaxed); - if (cached == 0) - { - int device_id; - cudaGetDevice(&device_id); - cudaDeviceProp device_prop; - cudaGetDeviceProperties(&device_prop, device_id); - cached = device_prop.multiProcessorCount; - sm_count.store(cached, std::memory_order_relaxed); - } - return cached; -} - -template -inline void DebugPrintCUDAArray(T* device_ptr, size_t size, std::string prefix = "") -{ - std::vector host_array(size); - std::cout << prefix; - cudaMemcpy(host_array.data(), device_ptr, size * sizeof(T), cudaMemcpyDeviceToHost); - for (size_t i = 0; i < size; ++i) - { - std::cout << host_array[i] << " "; - } - std::cout << std::endl; -} - -inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_dim) -{ - if (avg_packed_qo_len > 64 && head_dim < 256) - { - return 128; - } - else - { - auto compute_capacity = GetCudaComputeCapability(); - if (compute_capacity.first >= 8) - { - // Ampere or newer - if (avg_packed_qo_len > 16) - { - // avg_packed_qo_len <= 64 - return 64; - } - else - { - // avg_packed_qo_len <= 16 - return 16; - } - } - else - { - // NOTE(Zihao): not enough shared memory on Turing for 1x4 warp layout - return 64; - } - } -} - -inline int UpPowerOfTwo(int x) -{ - // Returns the smallest power of two greater than or equal to x - if (x <= 0) - return 1; - --x; - x |= x >> 1; - x |= x >> 2; - x |= x >> 4; - x |= x >> 8; - x |= x >> 16; - return x + 1; -} - -#define LOOP_SPLIT_MASK(iter, COND1, COND2, ...) \ - { \ - _Pragma("unroll 1") for (; (COND1); (iter) -= 1) \ - { \ - constexpr bool WITH_MASK = true; \ - __VA_ARGS__ \ - } \ - _Pragma("unroll 1") for (; (COND2); (iter) -= 1) \ - { \ - constexpr bool WITH_MASK = false; \ - __VA_ARGS__ \ - } \ - } - -/*! - * \brief Return x - y if x > y, otherwise return 0. - */ -__device__ __forceinline__ uint32_t sub_if_greater_or_zero(uint32_t x, uint32_t y) -{ - return (x > y) ? x - y : 0U; -} - -// ======================= PTX Memory Utility Functions ======================= -// Non-atomic global memory access with cache streaming hint (cs) -// These are useful for streaming memory access patterns where data is used once - -/*! - * \brief Get the lane ID within a warp (0-31) - */ -__forceinline__ __device__ int get_lane_id() -{ - int lane_id; - asm("mov.u32 %0, %%laneid;" : "=r"(lane_id)); - return lane_id; -} - -/*! - * \brief Non-atomic global load for short (2 bytes) with cache streaming hint - */ -__forceinline__ __device__ short ld_na_global_s16(short const* addr) -{ - short val; - asm volatile("ld.global.cs.b16 %0, [%1];" : "=h"(val) : "l"(addr)); - return val; -} - -/*! - * \brief Non-atomic global store for short (2 bytes) with cache streaming hint - */ -__forceinline__ __device__ void st_na_global_s16(short* addr, short val) -{ - asm volatile("st.global.cs.b16 [%0], %1;" ::"l"(addr), "h"(val)); -} - -/*! - * \brief Non-atomic global load for int (4 bytes) with cache streaming hint - */ -__forceinline__ __device__ int ld_na_global_v1(int const* addr) -{ - int val; - asm volatile("ld.global.cs.b32 %0, [%1];" : "=r"(val) : "l"(addr)); - return val; -} - -/*! - * \brief Non-atomic global load for int2 (8 bytes) with cache streaming hint - */ -__forceinline__ __device__ int2 ld_na_global_v2(int2 const* addr) -{ - int2 val; - asm volatile("ld.global.cs.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr)); - return val; -} - -/*! - * \brief Non-atomic global store for int (4 bytes) with cache streaming hint - */ -__forceinline__ __device__ void st_na_global_v1(int* addr, int val) -{ - asm volatile("st.global.cs.b32 [%0], %1;" ::"l"(addr), "r"(val)); -} - -/*! - * \brief Non-atomic global store for int2 (8 bytes) with cache streaming hint - */ -__forceinline__ __device__ void st_na_global_v2(int2* addr, int2 val) -{ - asm volatile("st.global.cs.v2.b32 [%0], {%1, %2};" ::"l"(addr), "r"(val.x), "r"(val.y)); -} - -/*! - * \brief Prefetch data to L2 cache - */ -template -__forceinline__ __device__ void prefetch_L2(T const* addr) -{ - asm volatile("prefetch.global.L2 [%0];" ::"l"(addr)); -} - -__device__ __forceinline__ void swap(uint32_t& a, uint32_t& b) -{ - uint32_t tmp = a; - a = b; - b = tmp; -} - -__device__ __forceinline__ uint32_t dim2_offset(uint32_t const& dim_a, uint32_t const& idx_b, uint32_t const& idx_a) -{ - return idx_b * dim_a + idx_a; -} - -__device__ __forceinline__ uint32_t dim3_offset( - uint32_t const& dim_b, uint32_t const& dim_a, uint32_t const& idx_c, uint32_t const& idx_b, uint32_t const& idx_a) -{ - return (idx_c * dim_b + idx_b) * dim_a + idx_a; -} - -__device__ __forceinline__ uint32_t dim4_offset(uint32_t const& dim_c, uint32_t const& dim_b, uint32_t const& dim_a, - uint32_t const& idx_d, uint32_t const& idx_c, uint32_t const& idx_b, uint32_t const& idx_a) -{ - return ((idx_d * dim_c + idx_c) * dim_b + idx_b) * dim_a + idx_a; -} - -#define DEFINE_HAS_MEMBER(member) \ - template \ - struct has_##member : std::false_type \ - { \ - }; \ - template \ - struct has_##member().member)>> : std::true_type \ - { \ - }; \ - template \ - inline constexpr bool has_##member##_v = has_##member::value; - -} // namespace flashinfer - -#endif // FLASHINFER_UTILS_CUH_ diff --git a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh b/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh deleted file mode 100644 index 54ebaebb6e7d..000000000000 --- a/tests/unittest/_torch/modules/mamba/flashinfer_checkpointing_ssu_pr3324/include/flashinfer/vec_dtypes.cuh +++ /dev/null @@ -1,3201 +0,0 @@ -/* - * Copyright (c) 2023 by FlashInfer team. - * - * 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. - */ -#ifndef VEC_DTYPES_CUH_ -#define VEC_DTYPES_CUH_ - -#include -#include -#include -#include -#if CUDA_VERSION >= 12080 -#include -#endif -#include - -#include - -namespace flashinfer -{ - -#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 900)) -#define FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED -#endif - -#define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__ - -__device__ __forceinline__ void st_global_release(int4 const& val, int4* addr) -{ - asm volatile("st.release.global.sys.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), "r"(val.z), - "r"(val.w), "l"(addr)); -} - -__device__ __forceinline__ int4 ld_global_acquire(int4* addr) -{ - int4 val; - asm volatile("ld.acquire.global.sys.v4.b32 {%0, %1, %2, %3}, [%4];" - : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) - : "l"(addr)); - return val; -} - -__device__ __forceinline__ void st_global_volatile(int4 const& val, int4* addr) -{ - asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), - "l"(addr)); -} - -__device__ __forceinline__ int4 ld_global_volatile(int4* addr) -{ - int4 val; - asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" - : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) - : "l"(addr)); - return val; -} - -#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 < 120200) \ - && (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)) -// CUDA version < 12.2 and GPU architecture < 80 -FLASHINFER_INLINE __nv_bfloat162 make_bfloat162(const __nv_bfloat16 x, const __nv_bfloat16 y) -{ - __nv_bfloat162 t; - t.x = x; - t.y = y; - return t; -} - -FLASHINFER_INLINE __nv_bfloat16 __hmul(const __nv_bfloat16 a, const __nv_bfloat16 b) -{ - __nv_bfloat16 val; - float const fa = __bfloat162float(a); - float const fb = __bfloat162float(b); - // avoid ftz in device code - val = __float2bfloat16(__fmaf_ieee_rn(fa, fb, -0.0f)); - return val; -} - -FLASHINFER_INLINE __nv_bfloat162 __hmul2(const __nv_bfloat162 a, const __nv_bfloat162 b) -{ - __nv_bfloat162 val; - val.x = __hmul(a.x, b.x); - val.y = __hmul(a.y, b.y); - return val; -} - -FLASHINFER_INLINE __nv_bfloat162 __floats2bfloat162_rn(float const a, float const b) -{ - __nv_bfloat162 val; - val = __nv_bfloat162(__float2bfloat16_rn(a), __float2bfloat16_rn(b)); - return val; -} - -FLASHINFER_INLINE __nv_bfloat162 __float22bfloat162_rn(const float2 a) -{ - __nv_bfloat162 val = __floats2bfloat162_rn(a.x, a.y); - return val; -} - -FLASHINFER_INLINE float2 __bfloat1622float2(const __nv_bfloat162 a) -{ - float hi_float; - float lo_float; - lo_float = __internal_bfloat162float(((__nv_bfloat162_raw) a).x); - hi_float = __internal_bfloat162float(((__nv_bfloat162_raw) a).y); - return make_float2(lo_float, hi_float); -} -#endif - -/******************* vec_t type cast *******************/ - -template -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(dst_t* dst, src_t const* src) - { -#pragma unroll - for (size_t i = 0; i < vec_size; ++i) - { - dst[i] = (dst_t) src[i]; - } - } -}; - -template <> -struct vec_cast<__nv_fp8_e4m3, float> -{ - template - FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, float const* src) - { - if constexpr (vec_size == 1) - { - dst[0] = __nv_fp8_e4m3(src[0]); - } - else - { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - ((__nv_fp8x2_storage_t*) dst)[i] - = __nv_cvt_float2_to_fp8x2(((float2*) src)[i], __NV_SATFINITE, __NV_E4M3); - } - } - } -}; - -template <> -struct vec_cast<__nv_fp8_e5m2, float> -{ - template - FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, float const* src) - { - if constexpr (vec_size == 1) - { - dst[0] = __nv_fp8_e5m2(src[0]); - } - else - { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - ((__nv_fp8x2_storage_t*) dst)[i] - = __nv_cvt_float2_to_fp8x2(((float2*) src)[i], __NV_SATFINITE, __NV_E5M2); - } - } - } -}; - -template <> -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(float* dst, half const* src) - { - if constexpr (vec_size == 1) - { - dst[0] = (float) src[0]; - } - else - { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - ((float2*) dst)[i] = __half22float2(((half2*) src)[i]); - } - } - } -}; - -template <> -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(half* dst, float const* src) - { - if constexpr (vec_size == 1) - { - dst[0] = __float2half(src[0]); - } - else - { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - ((half2*) dst)[i] = __float22half2_rn(((float2*) src)[i]); - } - } - } -}; - -template -constexpr FLASHINFER_INLINE int get_exponent_bits() -{ - if constexpr (std::is_same_v) - { - return 4; - } - else if constexpr (std::is_same_v) - { - return 5; - } - else if constexpr (std::is_same_v) - { - return 5; - } - else if constexpr (std::is_same_v) - { - return 8; - } -} - -template -constexpr FLASHINFER_INLINE int get_mantissa_bits() -{ - if constexpr (std::is_same_v) - { - return 3; - } - else if constexpr (std::is_same_v) - { - return 2; - } - else if constexpr (std::is_same_v) - { - return 11; - } - else if constexpr (std::is_same_v) - { - return 7; - } -} - -/*! - * \brief Fallback to software fast dequant implementation if hardware dequantization is not - * available. - * \note Inspired by Marlin's fast dequantization, but here we don't have to permute - * weights order. - * \ref - * https://github.com/vllm-project/vllm/blob/6dffa4b0a6120159ef2fe44d695a46817aff65bc/csrc/quantization/fp8/fp8_marlin.cu#L120 - */ -template -__device__ void fast_dequant_f8f16x4(uint32_t* input, uint2* output) -{ - uint32_t q = *input; - if constexpr (std::is_same_v && std::is_same_v) - { - output->x = __byte_perm(0U, q, 0x5140); - output->y = __byte_perm(0U, q, 0x7362); - } - else - { - constexpr int FP8_EXPONENT = get_exponent_bits(); - constexpr int FP8_MANTISSA = get_mantissa_bits(); - constexpr int FP16_EXPONENT = get_exponent_bits(); - - constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT; - // Calculate MASK for extracting mantissa and exponent - constexpr int MASK1 = 0x80000000; - constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA); - constexpr int MASK3 = MASK2 & 0x7fffffff; - constexpr int MASK = MASK3 | (MASK3 >> 16); - q = __byte_perm(q, q, 0x1302); - - // Extract and shift FP8 values to FP16 format - uint32_t Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - uint32_t Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT); - - constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); - // Construct and apply exponent bias - if constexpr (std::is_same_v) - { - const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); - - // Convert to half2 and apply bias - *(half2*) &(output->x) = __hmul2(*reinterpret_cast(&Out1), bias_reg); - *(half2*) &(output->y) = __hmul2(*reinterpret_cast(&Out2), bias_reg); - } - else - { - constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; - const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); - // Convert to bfloat162 and apply bias - *(nv_bfloat162*) &(output->x) = __hmul2(*reinterpret_cast(&Out1), bias_reg); - *(nv_bfloat162*) &(output->y) = __hmul2(*reinterpret_cast(&Out2), bias_reg); - } - } -} - -template <> -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(nv_bfloat16* dst, __nv_fp8_e4m3 const* src) - { - if constexpr (vec_size == 1) - { - dst[0] = nv_bfloat16(src[0]); - } - else if constexpr (vec_size == 2) - { - dst[0] = nv_bfloat16(src[0]); - dst[1] = nv_bfloat16(src[1]); - } - else - { - static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); -#pragma unroll - for (uint32_t i = 0; i < vec_size / 4; ++i) - { - fast_dequant_f8f16x4<__nv_fp8_e4m3, nv_bfloat16>((uint32_t*) &src[i * 4], (uint2*) &dst[i * 4]); - } - } - } -}; - -template <> -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(nv_bfloat16* dst, __nv_fp8_e5m2 const* src) - { - if constexpr (vec_size == 1) - { - dst[0] = nv_bfloat16(src[0]); - } - else if constexpr (vec_size == 2) - { - dst[0] = nv_bfloat16(src[0]); - dst[1] = nv_bfloat16(src[1]); - } - else - { - static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); -#pragma unroll - for (uint32_t i = 0; i < vec_size / 4; ++i) - { - fast_dequant_f8f16x4<__nv_fp8_e5m2, nv_bfloat16>((uint32_t*) &src[i * 4], (uint2*) &dst[i * 4]); - } - } - } -}; - -template <> -struct vec_cast<__nv_fp8_e4m3, half> -{ - template - FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, half const* src) - { -#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - if constexpr (vec_size == 1) - { - dst[0] = __nv_fp8_e4m3(src[0]); - } - else - { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - uint16_t y; - uint32_t x = *(uint32_t*) &src[i * 2]; - asm volatile("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(y) : "r"(x)); - *(uint16_t*) &dst[i * 2] = y; - } - } -#else -#pragma unroll - for (size_t i = 0; i < vec_size; ++i) - { - dst[i] = __nv_fp8_e4m3(src[i]); - } -#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - } -}; - -template <> -struct vec_cast<__nv_fp8_e5m2, half> -{ - template - FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, half const* src) - { -#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - if constexpr (vec_size == 1) - { - dst[0] = __nv_fp8_e5m2(src[0]); - } - else - { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - uint16_t y; - uint32_t x = *(uint32_t*) &src[i * 2]; - asm volatile("cvt.rn.satfinite.e5m2x2.f16x2 %0, %1;" : "=h"(y) : "r"(x)); - *(uint16_t*) &dst[i * 2] = y; - } - } -#else -#pragma unroll - for (size_t i = 0; i < vec_size; ++i) - { - dst[i] = __nv_fp8_e5m2(src[i]); - } -#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - } -}; - -template <> -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(half* dst, __nv_fp8_e4m3 const* src) - { -#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - if constexpr (vec_size == 1) - { - dst[0] = half(src[0]); - } - else - { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - uint32_t y; - uint16_t x = *(uint16_t*) &src[i * 2]; - asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(y) : "h"(x)); - *(uint32_t*) &dst[i * 2] = y; - } - } -#else - if constexpr (vec_size == 1) - { - dst[0] = half(src[0]); - } - else if constexpr (vec_size == 2) - { - dst[0] = half(src[0]); - dst[1] = half(src[1]); - } - else - { - static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); -#pragma unroll - for (uint32_t i = 0; i < vec_size / 4; ++i) - { - fast_dequant_f8f16x4<__nv_fp8_e4m3, half>((uint32_t*) &src[i * 4], (uint2*) &dst[i * 4]); - } - } -#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - } -}; - -template <> -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(half* dst, __nv_fp8_e5m2 const* src) - { -#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - if constexpr (vec_size == 1) - { - dst[0] = half(src[0]); - } - else - { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - uint32_t y; - uint16_t x = *(uint16_t*) &src[i * 2]; - asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;" : "=r"(y) : "h"(x)); - *(uint32_t*) &dst[i * 2] = y; - } - } -#else - if constexpr (vec_size == 1) - { - dst[0] = half(src[0]); - } - else if constexpr (vec_size == 2) - { - dst[0] = half(src[0]); - dst[1] = half(src[1]); - } - else - { - static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); -#pragma unroll - for (uint32_t i = 0; i < vec_size / 4; ++i) - { - fast_dequant_f8f16x4<__nv_fp8_e5m2, half>((uint32_t*) &src[i * 4], (uint2*) &dst[i * 4]); - } - } -#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED - } -}; - -#if defined(FLASHINFER_ENABLE_FP4_E2M1) && CUDA_VERSION >= 12080 -// Convert __nv_fp4x2_e2m1 (2 fp4 values per byte) to fp16. -// vec_size counts fp16 output elements; src has stride-2 layout: -// src[0] holds x0,x1 src[1] is padding -// src[2] holds x2,x3 src[3] is padding ... etc. -// Each valid byte encodes 2 fp4 values -> 2 fp16 via cvt.rn.f16x2.e2m1x2. -template <> -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(half* dst, __nv_fp4x2_e2m1 const* src) - { - static_assert(vec_size % 2 == 0, "vec_size must be even for fp4x2 dequantization"); -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - uint32_t y; - // Valid fp4x2 bytes are at even positions (stride 2); odd positions are padding. - uint32_t b = reinterpret_cast(src)[i * 2]; - asm volatile( - "{\n" - ".reg .b8 fp4_byte;\n" - "mov.b32 {fp4_byte, _, _, _}, %1;\n" - "cvt.rn.f16x2.e2m1x2 %0, fp4_byte;\n" - "}" - : "=r"(y) - : "r"(b)); - reinterpret_cast(dst)[i] = y; - } -#else - // Software LUT fallback for arch < SM100. - // e2m1 encoding: bit[3]=sign, bit[2:0]=magnitude index in {0,0.5,1,1.5,2,3,4,6}. - // Each packed byte holds two fp4 values: bits[3:0]=first, bits[7:4]=second. - constexpr uint16_t lut[16] = { - 0x0000, // +0.0 - 0x3800, // +0.5 - 0x3C00, // +1.0 - 0x3E00, // +1.5 - 0x4000, // +2.0 - 0x4200, // +3.0 - 0x4400, // +4.0 - 0x4600, // +6.0 - 0x8000, // -0.0 - 0xB800, // -0.5 - 0xBC00, // -1.0 - 0xBE00, // -1.5 - 0xC000, // -2.0 - 0xC200, // -3.0 - 0xC400, // -4.0 - 0xC600, // -6.0 - }; -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - uint8_t b = reinterpret_cast(src)[i * 2]; - reinterpret_cast(dst)[i * 2 + 0] = lut[b & 0x0F]; - reinterpret_cast(dst)[i * 2 + 1] = lut[(b >> 4) & 0x0F]; - } -#endif - } -}; - -template <> -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(nv_bfloat16* dst, __nv_fp4x2_e2m1 const* src) - { - static_assert(vec_size % 2 == 0, "vec_size must be even for fp4x2 dequantization"); -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - uint32_t y; - // Valid fp4x2 bytes are at even positions (stride 2); odd positions are padding. - uint32_t b = reinterpret_cast(src)[i * 2]; -#if (defined __CUDACC_VER_MAJOR__) && (defined __CUDACC_VER_MINOR__) \ - && ((__CUDACC_VER_MAJOR__ > 13) || ((__CUDACC_VER_MAJOR__ == 13) && (__CUDACC_VER_MINOR__ >= 2))) - // cvt.rn.bf16x2.e2m1x2 requires CUDA Toolkit >= 13.2 - asm volatile( - "{\n" - ".reg .b8 fp4_byte;\n" - "mov.b32 {fp4_byte, _, _, _}, %1;\n" - "cvt.rn.bf16x2.e2m1x2 %0, fp4_byte;\n" - "}" - : "=r"(y) - : "r"(b)); -#else - // Fallback: convert e2m1 -> fp16 -> bf16 when cvt.rn.bf16x2.e2m1x2 is unavailable - uint32_t fp16x2; - asm volatile( - "{\n" - ".reg .b8 fp4_byte;\n" - "mov.b32 {fp4_byte, _, _, _}, %1;\n" - "cvt.rn.f16x2.e2m1x2 %0, fp4_byte;\n" - "}" - : "=r"(fp16x2) - : "r"(b)); - __half2 h2 = reinterpret_cast<__half2&>(fp16x2); - __nv_bfloat162 bf16x2 = __float22bfloat162_rn(__half22float2(h2)); - y = reinterpret_cast(bf16x2); -#endif - reinterpret_cast(dst)[i] = y; - } -#else - // Software LUT fallback for arch < SM100. - // e2m1 encoding: bit[3]=sign, bit[2:0]=magnitude index in {0,0.5,1,1.5,2,3,4,6}. - // Each packed byte holds two fp4 values: bits[3:0]=first, bits[7:4]=second. - constexpr uint16_t lut[16] = { - 0x0000, // +0.0 - 0x3F00, // +0.5 - 0x3F80, // +1.0 - 0x3FC0, // +1.5 - 0x4000, // +2.0 - 0x4040, // +3.0 - 0x4080, // +4.0 - 0x40C0, // +6.0 - 0x8000, // -0.0 - 0xBF00, // -0.5 - 0xBF80, // -1.0 - 0xBFC0, // -1.5 - 0xC000, // -2.0 - 0xC040, // -3.0 - 0xC080, // -4.0 - 0xC0C0, // -6.0 - }; -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - uint8_t b = reinterpret_cast(src)[i * 2]; - reinterpret_cast(dst)[i * 2 + 0] = lut[b & 0x0F]; - reinterpret_cast(dst)[i * 2 + 1] = lut[(b >> 4) & 0x0F]; - } -#endif - } -}; - -#endif // FLASHINFER_ENABLE_FP4_E2M1 && CUDA_VERSION >= 12080 - -template <> -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(float* dst, nv_bfloat16 const* src) - { - if constexpr (vec_size == 1) - { - dst[0] = (float) src[0]; - } - else - { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - ((float2*) dst)[i] = __bfloat1622float2(((nv_bfloat162*) src)[i]); - } - } - } -}; - -template <> -struct vec_cast -{ - template - FLASHINFER_INLINE static void cast(nv_bfloat16* dst, float const* src) - { - if constexpr (vec_size == 1) - { - dst[0] = nv_bfloat16(src[0]); - } - else - { -#pragma unroll - for (size_t i = 0; i < vec_size / 2; ++i) - { - ((nv_bfloat162*) dst)[i] = __float22bfloat162_rn(((float2*) src)[i]); - } - } - } -}; - -template -struct vec_t -{ - FLASHINFER_INLINE float_t& operator[](size_t i); - FLASHINFER_INLINE float_t const& operator[](size_t i) const; - FLASHINFER_INLINE void fill(float_t val); - FLASHINFER_INLINE void load(float_t const* ptr); - FLASHINFER_INLINE void store(float_t* ptr) const; - FLASHINFER_INLINE void load_global_acquire(float* addr); - FLASHINFER_INLINE void store_global_release(float* addr) const; - FLASHINFER_INLINE void load_global_volatile(float* addr); - FLASHINFER_INLINE void store_global_volatile(float* addr) const; - template - FLASHINFER_INLINE void cast_from(vec_t const& src); - template - FLASHINFER_INLINE void cast_load(T const* ptr); - template - FLASHINFER_INLINE void cast_store(T* ptr) const; - FLASHINFER_INLINE static void memcpy(float_t* dst, float_t const* src); - FLASHINFER_INLINE float_t* ptr(); -}; - -template -FLASHINFER_INLINE void cast_from_impl(vec_t& dst, vec_t const& src) -{ - vec_cast::cast( - dst.ptr(), const_cast*>(&src)->ptr()); -} - -template -FLASHINFER_INLINE void cast_load_impl(vec_t& dst, src_float_t const* src_ptr) -{ - if constexpr (std::is_same_v) - { - dst.load(src_ptr); - } - else - { - vec_t tmp; - tmp.load(src_ptr); - dst.cast_from(tmp); - } -} - -template -FLASHINFER_INLINE void cast_store_impl(tgt_float_t* dst_ptr, vec_t const& src) -{ - if constexpr (std::is_same_v) - { - src.store(dst_ptr); - } - else - { - vec_t tmp; - tmp.cast_from(src); - tmp.store(dst_ptr); - } -} - -/******************* vec_t<__nv_fp8_e4m3> *******************/ - -// __nv_fp8_e4m3 x 1 -template <> -struct vec_t<__nv_fp8_e4m3, 1> -{ - __nv_fp8_e4m3 data; - - FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) - { - return ((__nv_fp8_e4m3*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e4m3 const& operator[](size_t i) const - { - return ((__nv_fp8_e4m3 const*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e4m3* ptr() - { - return reinterpret_cast<__nv_fp8_e4m3*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); - FLASHINFER_INLINE void load(__nv_fp8_e4m3 const* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src); -}; - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::fill(__nv_fp8_e4m3 val) -{ - data = val; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::load(__nv_fp8_e4m3 const* ptr) -{ - data = *ptr; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::store(__nv_fp8_e4m3* ptr) const -{ - *ptr = data; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src) -{ - *dst = *src; -} - -// __nv_fp8_e4m3 x 2 -template <> -struct vec_t<__nv_fp8_e4m3, 2> -{ - __nv_fp8x2_e4m3 data; - - FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) - { - return ((__nv_fp8_e4m3*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e4m3 const& operator[](size_t i) const - { - return ((__nv_fp8_e4m3 const*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e4m3* ptr() - { - return reinterpret_cast<__nv_fp8_e4m3*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); - FLASHINFER_INLINE void load(__nv_fp8_e4m3 const* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src); -}; - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::fill(__nv_fp8_e4m3 val) -{ - data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::load(__nv_fp8_e4m3 const* ptr) -{ - data = *((__nv_fp8x2_e4m3*) ptr); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::store(__nv_fp8_e4m3* ptr) const -{ - *((__nv_fp8x2_e4m3*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src) -{ - *((__nv_fp8x2_e4m3*) dst) = *((__nv_fp8x2_e4m3*) src); -} - -// __nv_fp8_e4m3 x 4 - -template <> -struct vec_t<__nv_fp8_e4m3, 4> -{ - __nv_fp8x4_e4m3 data; - - FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) - { - return ((__nv_fp8_e4m3*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e4m3 const& operator[](size_t i) const - { - return ((__nv_fp8_e4m3 const*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e4m3* ptr() - { - return reinterpret_cast<__nv_fp8_e4m3*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); - FLASHINFER_INLINE void load(__nv_fp8_e4m3 const* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src); -}; - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::fill(__nv_fp8_e4m3 val) -{ - data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) - | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::load(__nv_fp8_e4m3 const* ptr) -{ - data = *((__nv_fp8x4_e4m3*) ptr); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::store(__nv_fp8_e4m3* ptr) const -{ - *((__nv_fp8x4_e4m3*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src) -{ - *((__nv_fp8x4_e4m3*) dst) = *((__nv_fp8x4_e4m3*) src); -} - -// __nv_fp8_e4m3 x 8 - -template <> -struct vec_t<__nv_fp8_e4m3, 8> -{ - uint2 data; - - FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) - { - return ((__nv_fp8_e4m3*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e4m3 const& operator[](size_t i) const - { - return ((__nv_fp8_e4m3 const*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e4m3* ptr() - { - return reinterpret_cast<__nv_fp8_e4m3*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); - FLASHINFER_INLINE void load(__nv_fp8_e4m3 const* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src); -}; - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::fill(__nv_fp8_e4m3 val) -{ - ((__nv_fp8x4_e4m3*) (&data.x))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) - | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e4m3*) (&data.y))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) - | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::load(__nv_fp8_e4m3 const* ptr) -{ - data = *((uint2*) ptr); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::store(__nv_fp8_e4m3* ptr) const -{ - *((uint2*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src) -{ - *((uint2*) dst) = *((uint2*) src); -} - -// __nv_fp8_e4m3 x 16 or more -template -struct vec_t<__nv_fp8_e4m3, vec_size> -{ - static_assert(vec_size % 16 == 0, "Invalid vector size"); - int4 data[vec_size / 16]; - - FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) - { - return ((__nv_fp8_e4m3*) data)[i]; - } - - FLASHINFER_INLINE __nv_fp8_e4m3 const& operator[](size_t i) const - { - return ((__nv_fp8_e4m3 const*) data)[i]; - } - - FLASHINFER_INLINE __nv_fp8_e4m3* ptr() - { - return reinterpret_cast<__nv_fp8_e4m3*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - ((__nv_fp8x4_e4m3*) (&(data[i].x)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) - | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) - | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e4m3*) (&(data[i].y)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) - | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) - | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e4m3*) (&(data[i].z)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) - | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) - | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e4m3*) (&(data[i].w)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) - | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) - | __nv_fp8x4_storage_t(val.__x); - } - } - - FLASHINFER_INLINE void load(__nv_fp8_e4m3 const* ptr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - data[i] = ((int4*) ptr)[i]; - } - } - - FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - ((int4*) ptr)[i] = data[i]; - } - } - - FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e4m3* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - *((int4*) (data + i)) = ld_global_acquire((int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void store_global_release(__nv_fp8_e4m3* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - st_global_release(data[i], (int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e4m3* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - data[i] = ld_global_volatile((int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e4m3* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - st_global_volatile(data[i], (int4*) (addr + i * 16)); - } - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, __nv_fp8_e4m3 const* src) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - ((int4*) dst)[i] = ((int4*) src)[i]; - } - } -}; - -/******************* vec_t<__nv_fp8_e5m2> *******************/ - -// __nv_fp8_e5m2 x 1 -template <> -struct vec_t<__nv_fp8_e5m2, 1> -{ - __nv_fp8_e5m2 data; - - FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) - { - return ((__nv_fp8_e5m2*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e5m2 const& operator[](size_t i) const - { - return ((__nv_fp8_e5m2 const*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e5m2* ptr() - { - return reinterpret_cast<__nv_fp8_e5m2*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); - FLASHINFER_INLINE void load(__nv_fp8_e5m2 const* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src); -}; - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::fill(__nv_fp8_e5m2 val) -{ - data = val; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::load(__nv_fp8_e5m2 const* ptr) -{ - data = *ptr; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::store(__nv_fp8_e5m2* ptr) const -{ - *ptr = data; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src) -{ - *dst = *src; -} - -// __nv_fp8_e5m2 x 2 -template <> -struct vec_t<__nv_fp8_e5m2, 2> -{ - __nv_fp8x2_e5m2 data; - - FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) - { - return ((__nv_fp8_e5m2*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e5m2 const& operator[](size_t i) const - { - return ((__nv_fp8_e5m2 const*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e5m2* ptr() - { - return reinterpret_cast<__nv_fp8_e5m2*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); - FLASHINFER_INLINE void load(__nv_fp8_e5m2 const* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src); -}; - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::fill(__nv_fp8_e5m2 val) -{ - data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::load(__nv_fp8_e5m2 const* ptr) -{ - data = *((__nv_fp8x2_e5m2*) ptr); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::store(__nv_fp8_e5m2* ptr) const -{ - *((__nv_fp8x2_e5m2*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src) -{ - *((__nv_fp8x2_e5m2*) dst) = *((__nv_fp8x2_e5m2*) src); -} - -// __nv_fp8_e5m2 x 4 - -template <> -struct vec_t<__nv_fp8_e5m2, 4> -{ - __nv_fp8x4_e5m2 data; - - FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) - { - return ((__nv_fp8_e5m2*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e5m2 const& operator[](size_t i) const - { - return ((__nv_fp8_e5m2 const*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e5m2* ptr() - { - return reinterpret_cast<__nv_fp8_e5m2*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); - FLASHINFER_INLINE void load(__nv_fp8_e5m2 const* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src); -}; - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::fill(__nv_fp8_e5m2 val) -{ - data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) - | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::load(__nv_fp8_e5m2 const* ptr) -{ - data = *((__nv_fp8x4_e5m2*) ptr); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::store(__nv_fp8_e5m2* ptr) const -{ - *((__nv_fp8x4_e5m2*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src) -{ - *((__nv_fp8x4_e5m2*) dst) = *((__nv_fp8x4_e5m2*) src); -} - -// __nv_fp8_e5m2 x 8 - -template <> -struct vec_t<__nv_fp8_e5m2, 8> -{ - uint2 data; - - FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) - { - return ((__nv_fp8_e5m2*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e5m2 const& operator[](size_t i) const - { - return ((__nv_fp8_e5m2 const*) (&data))[i]; - } - - FLASHINFER_INLINE __nv_fp8_e5m2* ptr() - { - return reinterpret_cast<__nv_fp8_e5m2*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); - FLASHINFER_INLINE void load(__nv_fp8_e5m2 const* ptr); - FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src); -}; - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::fill(__nv_fp8_e5m2 val) -{ - ((__nv_fp8x4_e5m2*) (&data.x))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) - | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e5m2*) (&data.y))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) - | (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::load(__nv_fp8_e5m2 const* ptr) -{ - data = *((uint2*) ptr); -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::store(__nv_fp8_e5m2* ptr) const -{ - *((uint2*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src) -{ - *((uint2*) dst) = *((uint2*) src); -} - -// __nv_fp8_e5m2 x 16 or more - -template -struct vec_t<__nv_fp8_e5m2, vec_size> -{ - static_assert(vec_size % 16 == 0, "Invalid vector size"); - int4 data[vec_size / 16]; - - FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) - { - return ((__nv_fp8_e5m2*) data)[i]; - } - - FLASHINFER_INLINE __nv_fp8_e5m2 const& operator[](size_t i) const - { - return ((__nv_fp8_e5m2 const*) data)[i]; - } - - FLASHINFER_INLINE __nv_fp8_e5m2* ptr() - { - return reinterpret_cast<__nv_fp8_e5m2*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - ((__nv_fp8x4_e5m2*) (&(data[i].x)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) - | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) - | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e5m2*) (&(data[i].y)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) - | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) - | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e5m2*) (&(data[i].z)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) - | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) - | __nv_fp8x4_storage_t(val.__x); - ((__nv_fp8x4_e5m2*) (&(data[i].w)))->__x = (__nv_fp8x4_storage_t(val.__x) << 24) - | (__nv_fp8x4_storage_t(val.__x) << 16) | (__nv_fp8x4_storage_t(val.__x) << 8) - | __nv_fp8x4_storage_t(val.__x); - } - } - - FLASHINFER_INLINE void load(__nv_fp8_e5m2 const* ptr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - data[i] = ((int4*) ptr)[i]; - } - } - - FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - ((int4*) ptr)[i] = data[i]; - } - } - - FLASHINFER_INLINE void store_global_release(__nv_fp8_e5m2* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - st_global_release(data[i], (int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e5m2* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - data[i] = ld_global_acquire((int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e5m2* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - st_global_volatile(data[i], (int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e5m2* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - data[i] = ld_global_volatile((int4*) (addr + i * 16)); - } - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, __nv_fp8_e5m2 const* src) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - ((int4*) dst)[i] = ((int4*) src)[i]; - } - } -}; - -#if defined(FLASHINFER_ENABLE_FP4_E2M1) && CUDA_VERSION >= 12080 -/******************* vec_t<__nv_fp4_e2m1> *******************/ - -// __nv_fp4_e2m1 x 2 -template <> -struct vec_t<__nv_fp4_e2m1, 2> -{ - uint8_t data; - - // index access is not supported for sub-byte data type - FLASHINFER_INLINE __nv_fp4_e2m1* ptr() - { - return reinterpret_cast<__nv_fp4_e2m1*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) - { - data = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); - } - - FLASHINFER_INLINE void load(__nv_fp4_e2m1 const* ptr) - { - data = *((uint8_t*) ptr); - } - - FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const - { - *((uint8_t*) ptr) = data; - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, __nv_fp4_e2m1 const* src) - { - *((uint8_t*) dst) = *((uint8_t*) src); - } -}; - -// __nv_fp4_e2m1 x 4 -template <> -struct vec_t<__nv_fp4_e2m1, 4> -{ - uint16_t data; - - FLASHINFER_INLINE __nv_fp4_e2m1* ptr() - { - return reinterpret_cast<__nv_fp4_e2m1*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) - { - __nv_fp4x2_storage_t val8 = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); - data = (uint16_t(val8) << 8) | uint16_t(val8); - } - - FLASHINFER_INLINE void load(__nv_fp4_e2m1 const* ptr) - { - data = *((uint16_t*) ptr); - } - - FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const - { - *((uint16_t*) ptr) = data; - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, __nv_fp4_e2m1 const* src) - { - *((uint16_t*) dst) = *((uint16_t*) src); - } -}; - -// __nv_fp4_e2m1 x 8 -template <> -struct vec_t<__nv_fp4_e2m1, 8> -{ - uint32_t data; - - FLASHINFER_INLINE __nv_fp4_e2m1* ptr() - { - return reinterpret_cast<__nv_fp4_e2m1*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) - { - __nv_fp4x2_storage_t val8 = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); - uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); - data = (uint32_t(val16) << 16) | uint32_t(val16); - } - - FLASHINFER_INLINE void load(__nv_fp4_e2m1 const* ptr) - { - data = *((uint32_t*) ptr); - } - - FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const - { - *((uint32_t*) ptr) = data; - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, __nv_fp4_e2m1 const* src) - { - *((uint32_t*) dst) = *((uint32_t*) src); - } -}; - -template <> -struct vec_t<__nv_fp4_e2m1, 16> -{ - uint2 data; - - FLASHINFER_INLINE __nv_fp4_e2m1* ptr() - { - return reinterpret_cast<__nv_fp4_e2m1*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) - { - __nv_fp4x2_storage_t val8 = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); - uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); - uint32_t val32 = (uint32_t(val16) << 16) | uint32_t(val16); - data.x = val32; - data.y = val32; - } - - FLASHINFER_INLINE void load(__nv_fp4_e2m1 const* ptr) - { - data = *((uint2*) ptr); - } - - FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const - { - *((uint2*) ptr) = data; - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, __nv_fp4_e2m1 const* src) - { - *((uint2*) dst) = *((uint2*) src); - } -}; - -// __nv_fp4_e2m1 x 32 or more -template -struct vec_t<__nv_fp4_e2m1, vec_size> -{ - static_assert(vec_size % 32 == 0, "Invalid vector size"); - int4 data[vec_size / 32]; - - FLASHINFER_INLINE __nv_fp4_e2m1* ptr() - { - return reinterpret_cast<__nv_fp4_e2m1*>(&data); - } - - FLASHINFER_INLINE void fill(__nv_fp4_e2m1 val) - { - __nv_fp4x2_storage_t val8 = (__nv_fp4x2_storage_t(val.__x) << 4) | __nv_fp4x2_storage_t(val.__x); - uint16_t val16 = (uint16_t(val8) << 8) | uint16_t(val8); - uint32_t val32 = (uint32_t(val16) << 16) | uint32_t(val16); -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) - { - data[i].x = val32; - data[i].y = val32; - data[i].z = val32; - data[i].w = val32; - } - } - - FLASHINFER_INLINE void load(__nv_fp4_e2m1 const* ptr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) - { - data[i] = ((int4*) ptr)[i]; - } - } - - FLASHINFER_INLINE void store(__nv_fp4_e2m1* ptr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) - { - ((int4*) ptr)[i] = data[i]; - } - } - - FLASHINFER_INLINE void store_global_release(__nv_fp4_e2m1* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) - { - st_global_release(*(int4*) &data[i], (int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void load_global_acquire(__nv_fp4_e2m1* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) - { - *(int4*) &data[i] = ld_global_acquire((int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void store_global_volatile(__nv_fp4_e2m1* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) - { - st_global_volatile(*(int4*) &data[i], (int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void load_global_volatile(__nv_fp4_e2m1* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) - { - *(int4*) &data[i] = ld_global_volatile((int4*) (addr + i * 16)); - } - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(__nv_fp4_e2m1* dst, __nv_fp4_e2m1 const* src) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 32; ++i) - { - ((int4*) dst)[i] = ((int4*) src)[i]; - } - } -}; - -#endif // FLASHINFER_ENABLE_FP4_E2M1 && CUDA_VERSION >= 12080 - -/******************* vec_t *******************/ - -// half x 1 -template <> -struct vec_t -{ - half data; - - FLASHINFER_INLINE half& operator[](size_t i) - { - return ((half*) (&data))[i]; - } - - FLASHINFER_INLINE half const& operator[](size_t i) const - { - return ((half const*) (&data))[i]; - } - - FLASHINFER_INLINE half* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(half val); - FLASHINFER_INLINE void load(half const* ptr); - FLASHINFER_INLINE void store(half* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(half* dst, half const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(half val) -{ - data = val; -} - -FLASHINFER_INLINE void vec_t::load(half const* ptr) -{ - data = *ptr; -} - -FLASHINFER_INLINE void vec_t::store(half* ptr) const -{ - *ptr = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(half* dst, half const* src) -{ - *dst = *src; -} - -// half x 2 -template <> -struct vec_t -{ - half2 data; - - FLASHINFER_INLINE half& operator[](size_t i) - { - return ((half*) (&data))[i]; - } - - FLASHINFER_INLINE half const& operator[](size_t i) const - { - return ((half const*) (&data))[i]; - } - - FLASHINFER_INLINE half* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(half val); - FLASHINFER_INLINE void load(half const* ptr); - FLASHINFER_INLINE void store(half* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(half* dst, half const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(half val) -{ - data = make_half2(val, val); -} - -FLASHINFER_INLINE void vec_t::load(half const* ptr) -{ - data = *((half2*) ptr); -} - -FLASHINFER_INLINE void vec_t::store(half* ptr) const -{ - *((half2*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(half* dst, half const* src) -{ - *((half2*) dst) = *((half2*) src); -} - -// half x 4 - -template <> -struct vec_t -{ - uint2 data; - - FLASHINFER_INLINE half& operator[](size_t i) - { - return ((half*) (&data))[i]; - } - - FLASHINFER_INLINE half const& operator[](size_t i) const - { - return ((half const*) (&data))[i]; - } - - FLASHINFER_INLINE half* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(half val); - FLASHINFER_INLINE void load(half const* ptr); - FLASHINFER_INLINE void store(half* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(half* dst, half const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(half val) -{ - *(half2*) (&data.x) = make_half2(val, val); - *(half2*) (&data.y) = make_half2(val, val); -} - -FLASHINFER_INLINE void vec_t::load(half const* ptr) -{ - data = *((uint2*) ptr); -} - -FLASHINFER_INLINE void vec_t::store(half* ptr) const -{ - *((uint2*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(half* dst, half const* src) -{ - *((uint2*) dst) = *((uint2*) src); -} - -// half x 8 or more - -template -struct vec_t -{ - static_assert(vec_size % 8 == 0, "Invalid vector size"); - int4 data[vec_size / 8]; - - FLASHINFER_INLINE half& operator[](size_t i) - { - return ((half*) data)[i]; - } - - FLASHINFER_INLINE half const& operator[](size_t i) const - { - return ((half const*) data)[i]; - } - - FLASHINFER_INLINE half* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(half val) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - *(half2*) (&(data[i].x)) = make_half2(val, val); - *(half2*) (&(data[i].y)) = make_half2(val, val); - *(half2*) (&(data[i].z)) = make_half2(val, val); - *(half2*) (&(data[i].w)) = make_half2(val, val); - } - } - - FLASHINFER_INLINE void load(half const* ptr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - data[i] = ((int4*) ptr)[i]; - } - } - - FLASHINFER_INLINE void store(half* ptr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - ((int4*) ptr)[i] = data[i]; - } - } - - FLASHINFER_INLINE void load_global_acquire(half* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - data[i] = ld_global_acquire((int4*) (addr + i * 8)); - } - } - - FLASHINFER_INLINE void store_global_release(half* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - st_global_release(data[i], (int4*) (addr + i * 8)); - } - } - - FLASHINFER_INLINE void store_global_volatile(half* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - st_global_volatile(data[i], (int4*) (addr + i * 8)); - } - } - - FLASHINFER_INLINE void load_global_volatile(half* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - data[i] = ld_global_volatile((int4*) (addr + i * 8)); - } - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(half* dst, half const* src) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - ((int4*) dst)[i] = ((int4*) src)[i]; - } - } -}; - -/******************* vec_t *******************/ - -// nv_bfloat16 x 1 -template <> -struct vec_t -{ - nv_bfloat16 data; - - FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) - { - return ((nv_bfloat16*) (&data))[i]; - } - - FLASHINFER_INLINE nv_bfloat16 const& operator[](size_t i) const - { - return ((nv_bfloat16 const*) (&data))[i]; - } - - FLASHINFER_INLINE nv_bfloat16* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(nv_bfloat16 val); - FLASHINFER_INLINE void load(nv_bfloat16 const* ptr); - FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, nv_bfloat16 const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) -{ - data = val; -} - -FLASHINFER_INLINE void vec_t::load(nv_bfloat16 const* ptr) -{ - data = *ptr; -} - -FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const -{ - *ptr = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, nv_bfloat16 const* src) -{ - *dst = *src; -} - -// nv_bfloat16 x 2 -template <> -struct vec_t -{ - nv_bfloat162 data; - - FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) - { - return ((nv_bfloat16*) (&data))[i]; - } - - FLASHINFER_INLINE nv_bfloat16 const& operator[](size_t i) const - { - return ((nv_bfloat16 const*) (&data))[i]; - } - - FLASHINFER_INLINE nv_bfloat16* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(nv_bfloat16 val); - FLASHINFER_INLINE void load(nv_bfloat16 const* ptr); - FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, nv_bfloat16 const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) -{ - data = make_bfloat162(val, val); -} - -FLASHINFER_INLINE void vec_t::load(nv_bfloat16 const* ptr) -{ - data = *((nv_bfloat162*) ptr); -} - -FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const -{ - *((nv_bfloat162*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, nv_bfloat16 const* src) -{ - *((nv_bfloat162*) dst) = *((nv_bfloat162*) src); -} - -// nv_bfloat16 x 4 - -template <> -struct vec_t -{ - uint2 data; - - FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) - { - return ((nv_bfloat16*) (&data))[i]; - } - - FLASHINFER_INLINE nv_bfloat16 const& operator[](size_t i) const - { - return ((nv_bfloat16 const*) (&data))[i]; - } - - FLASHINFER_INLINE nv_bfloat16* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(nv_bfloat16 val); - FLASHINFER_INLINE void load(nv_bfloat16 const* ptr); - FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, nv_bfloat16 const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) -{ - *(nv_bfloat162*) (&data.x) = make_bfloat162(val, val); - *(nv_bfloat162*) (&data.y) = make_bfloat162(val, val); -} - -FLASHINFER_INLINE void vec_t::load(nv_bfloat16 const* ptr) -{ - data = *((uint2*) ptr); -} - -FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const -{ - *((uint2*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, nv_bfloat16 const* src) -{ - *((uint2*) dst) = *((uint2*) src); -} - -// nv_bfloat16 x 8 or more - -template -struct vec_t -{ - static_assert(vec_size % 8 == 0, "Invalid vector size"); - int4 data[vec_size / 8]; - - FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) - { - return ((nv_bfloat16*) data)[i]; - } - - FLASHINFER_INLINE nv_bfloat16 const& operator[](size_t i) const - { - return ((nv_bfloat16 const*) data)[i]; - } - - FLASHINFER_INLINE nv_bfloat16* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(nv_bfloat16 val) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - *(nv_bfloat162*) (&(data[i].x)) = make_bfloat162(val, val); - *(nv_bfloat162*) (&(data[i].y)) = make_bfloat162(val, val); - *(nv_bfloat162*) (&(data[i].z)) = make_bfloat162(val, val); - *(nv_bfloat162*) (&(data[i].w)) = make_bfloat162(val, val); - } - } - - FLASHINFER_INLINE void load(nv_bfloat16 const* ptr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - data[i] = ((int4*) ptr)[i]; - } - } - - FLASHINFER_INLINE void store(nv_bfloat16* ptr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - ((int4*) ptr)[i] = data[i]; - } - } - - FLASHINFER_INLINE void store_global_release(nv_bfloat16* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - st_global_release(data[i], (int4*) (addr + i * 8)); - } - } - - FLASHINFER_INLINE void load_global_acquire(nv_bfloat16* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - data[i] = ld_global_acquire((int4*) (addr + i * 8)); - } - } - - FLASHINFER_INLINE void store_global_volatile(nv_bfloat16* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - st_global_volatile(data[i], (int4*) (addr + i * 8)); - } - } - - FLASHINFER_INLINE void load_global_volatile(nv_bfloat16* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - data[i] = ld_global_volatile((int4*) (addr + i * 8)); - } - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, nv_bfloat16 const* src) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 8; ++i) - { - ((int4*) dst)[i] = ((int4*) src)[i]; - } - } -}; - -/******************* vec_t *******************/ - -// uint8_t x 1 -template <> -struct vec_t -{ - uint8_t data; - - FLASHINFER_INLINE uint8_t& operator[](size_t i) - { - return ((uint8_t*) (&data))[i]; - } - - FLASHINFER_INLINE uint8_t const& operator[](size_t i) const - { - return ((uint8_t const*) (&data))[i]; - } - - FLASHINFER_INLINE uint8_t* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(uint8_t val); - FLASHINFER_INLINE void load(uint8_t const* ptr); - FLASHINFER_INLINE void store(uint8_t* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(uint8_t* dst, uint8_t const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(uint8_t val) -{ - data = val; -} - -FLASHINFER_INLINE void vec_t::load(uint8_t const* ptr) -{ - data = *ptr; -} - -FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const -{ - *ptr = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, uint8_t const* src) -{ - *dst = *src; -} - -// uint8_t x 2 -template <> -struct vec_t -{ - uint16_t data; - - FLASHINFER_INLINE uint8_t& operator[](size_t i) - { - return ((uint8_t*) (&data))[i]; - } - - FLASHINFER_INLINE uint8_t const& operator[](size_t i) const - { - return ((uint8_t const*) (&data))[i]; - } - - FLASHINFER_INLINE uint8_t* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(uint8_t val); - FLASHINFER_INLINE void load(uint8_t const* ptr); - FLASHINFER_INLINE void store(uint8_t* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(uint8_t* dst, uint8_t const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(uint8_t val) -{ - data = (uint16_t(val) << 8) | uint16_t(val); -} - -FLASHINFER_INLINE void vec_t::load(uint8_t const* ptr) -{ - data = *((uint16_t*) ptr); -} - -FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const -{ - *((uint16_t*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, uint8_t const* src) -{ - *((uint16_t*) dst) = *((uint16_t*) src); -} - -// uint8_t x 4 - -template <> -struct vec_t -{ - uint32_t data; - - FLASHINFER_INLINE uint8_t& operator[](size_t i) - { - return ((uint8_t*) (&data))[i]; - } - - FLASHINFER_INLINE uint8_t const& operator[](size_t i) const - { - return ((uint8_t const*) (&data))[i]; - } - - FLASHINFER_INLINE uint8_t* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(uint8_t val); - FLASHINFER_INLINE void load(uint8_t const* ptr); - FLASHINFER_INLINE void store(uint8_t* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(uint8_t* dst, uint8_t const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(uint8_t val) -{ - data = (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); -} - -FLASHINFER_INLINE void vec_t::load(uint8_t const* ptr) -{ - data = *((uint32_t*) ptr); -} - -FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const -{ - *((uint32_t*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, uint8_t const* src) -{ - *((uint32_t*) dst) = *((uint32_t*) src); -} - -// uint8_t x 8 - -template <> -struct vec_t -{ - uint2 data; - - FLASHINFER_INLINE uint8_t& operator[](size_t i) - { - return ((uint8_t*) (&data))[i]; - } - - FLASHINFER_INLINE uint8_t const& operator[](size_t i) const - { - return ((uint8_t const*) (&data))[i]; - } - - FLASHINFER_INLINE uint8_t* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(uint8_t val); - FLASHINFER_INLINE void load(uint8_t const* ptr); - FLASHINFER_INLINE void store(uint8_t* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(uint8_t* dst, uint8_t const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(uint8_t val) -{ - uint32_t val32 = (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); - data.x = val32; - data.y = val32; -} - -FLASHINFER_INLINE void vec_t::load(uint8_t const* ptr) -{ - data = *((uint2*) ptr); -} - -FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const -{ - *((uint2*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, uint8_t const* src) -{ - *((uint2*) dst) = *((uint2*) src); -} - -// uint8_t x 16 or more - -template -struct vec_t -{ - static_assert(vec_size % 16 == 0, "Invalid vector size"); - int4 data[vec_size / 16]; - - FLASHINFER_INLINE uint8_t& operator[](size_t i) - { - return ((uint8_t*) data)[i]; - } - - FLASHINFER_INLINE uint8_t const& operator[](size_t i) const - { - return ((uint8_t const*) data)[i]; - } - - FLASHINFER_INLINE uint8_t* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(uint8_t val) - { - uint32_t val32 = (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - data[i].x = val32; - data[i].y = val32; - data[i].z = val32; - data[i].w = val32; - } - } - - FLASHINFER_INLINE void load(uint8_t const* ptr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - data[i] = ((int4*) ptr)[i]; - } - } - - FLASHINFER_INLINE void store(uint8_t* ptr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - ((int4*) ptr)[i] = data[i]; - } - } - - FLASHINFER_INLINE void load_global_acquire(uint8_t* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - data[i] = ld_global_acquire((int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void store_global_release(uint8_t* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - st_global_release(data[i], (int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void load_global_volatile(uint8_t* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - data[i] = ld_global_volatile((int4*) (addr + i * 16)); - } - } - - FLASHINFER_INLINE void store_global_volatile(uint8_t* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - st_global_volatile(data[i], (int4*) (addr + i * 16)); - } - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(uint8_t* dst, uint8_t const* src) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 16; ++i) - { - ((int4*) dst)[i] = ((int4*) src)[i]; - } - } -}; - -/******************* vec_t *******************/ - -// float x 1 - -template <> -struct vec_t -{ - float data; - - FLASHINFER_INLINE float& operator[](size_t i) - { - return ((float*) (&data))[i]; - } - - FLASHINFER_INLINE float const& operator[](size_t i) const - { - return ((float const*) (&data))[i]; - } - - FLASHINFER_INLINE float* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(float val); - FLASHINFER_INLINE void load(float const* ptr); - FLASHINFER_INLINE void store(float* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(float* dst, float const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(float val) -{ - data = val; -} - -FLASHINFER_INLINE void vec_t::load(float const* ptr) -{ - data = *ptr; -} - -FLASHINFER_INLINE void vec_t::store(float* ptr) const -{ - *ptr = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(float* dst, float const* src) -{ - *dst = *src; -} - -// float x 2 - -template <> -struct vec_t -{ - float2 data; - - FLASHINFER_INLINE float& operator[](size_t i) - { - return ((float*) (&data))[i]; - } - - FLASHINFER_INLINE float const& operator[](size_t i) const - { - return ((float const*) (&data))[i]; - } - - FLASHINFER_INLINE float* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(float val); - FLASHINFER_INLINE void load(float const* ptr); - FLASHINFER_INLINE void store(float* ptr) const; - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(float* dst, float const* src); -}; - -FLASHINFER_INLINE void vec_t::fill(float val) -{ - data = make_float2(val, val); -} - -FLASHINFER_INLINE void vec_t::load(float const* ptr) -{ - data = *((float2*) ptr); -} - -FLASHINFER_INLINE void vec_t::store(float* ptr) const -{ - *((float2*) ptr) = data; -} - -FLASHINFER_INLINE void vec_t::memcpy(float* dst, float const* src) -{ - *((float2*) dst) = *((float2*) src); -} - -// float x 4 or more -template -struct vec_t -{ - static_assert(vec_size % 4 == 0, "Invalid vector size"); - float4 data[vec_size / 4]; - - FLASHINFER_INLINE float& operator[](size_t i) - { - return ((float*) (data))[i]; - } - - FLASHINFER_INLINE float const& operator[](size_t i) const - { - return ((float const*) (data))[i]; - } - - FLASHINFER_INLINE float* ptr() - { - return reinterpret_cast(&data); - } - - FLASHINFER_INLINE void fill(float val) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) - { - data[i] = make_float4(val, val, val, val); - } - } - - FLASHINFER_INLINE void load(float const* ptr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) - { - data[i] = ((float4*) ptr)[i]; - } - } - - FLASHINFER_INLINE void store(float* ptr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) - { - ((float4*) ptr)[i] = data[i]; - } - } - - FLASHINFER_INLINE void store_global_release(float* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) - { - st_global_release(*(int4*) (data + i), (int4*) (addr + i * 4)); - } - } - - FLASHINFER_INLINE void load_global_acquire(float* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) - { - *((int4*) (data + i)) = ld_global_acquire((int4*) (addr + i * 4)); - } - } - - FLASHINFER_INLINE void store_global_volatile(float* addr) const - { -#pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) - { - st_global_volatile(*(int4*) (data + i), (int4*) (addr + i * 4)); - } - } - - FLASHINFER_INLINE void load_global_volatile(float* addr) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) - { - *((int4*) (data + i)) = ld_global_volatile((int4*) (addr + i * 4)); - } - } - - template - FLASHINFER_INLINE void cast_from(vec_t const& src) - { - cast_from_impl(*this, src); - } - - template - FLASHINFER_INLINE void cast_load(T const* ptr) - { - cast_load_impl(*this, ptr); - } - - template - FLASHINFER_INLINE void cast_store(T* ptr) const - { - cast_store_impl(ptr, *this); - } - - FLASHINFER_INLINE static void memcpy(float* dst, float const* src) - { -#pragma unroll - for (size_t i = 0; i < vec_size / 4; ++i) - { - ((float4*) dst)[i] = ((float4*) src)[i]; - } - } -}; - -template -struct vec2_dtype -{ - using type = T; -}; - -template <> -struct vec2_dtype -{ - using type = half2; -}; - -template <> -struct vec2_dtype<__nv_bfloat16> -{ - using type = __nv_bfloat162; -}; - -template <> -struct vec2_dtype<__nv_fp8_e4m3> -{ - using type = __nv_fp8x2_e4m3; -}; - -template <> -struct vec2_dtype<__nv_fp8_e5m2> -{ - using type = __nv_fp8x2_e5m2; -}; - -template -using vec2_dtype_t = typename vec2_dtype::type; - -template -FLASHINFER_INLINE vec2_dtype_t get_vec2_element(vec_t& vec, int i) -{ - static_assert(VEC_SIZE % 2 == 0, "VEC_SIZE must be a multiple of 2"); - return ((vec2_dtype_t*) &(vec[0]))[i]; -} - -} // namespace flashinfer - -#endif // VEC_DTYPES_CUH_