composable_kernel/dispatcher/examples/fmha/python/37_bwd_deterministic_fmha.py

#!/usr/bin/env python3

# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

"""
Example 37: Backward Pass Deterministic Mode

Demonstrates deterministic vs non-deterministic backward computation.

Non-deterministic mode (default):
  - dQ is accumulated via atomicAdd across seqlen_k tiles
  - Faster but produces slightly different results each run
  - Acceptable for training where stochastic noise is tolerable

Deterministic mode:
  - Uses multi-buffer reduction instead of atomics
  - Each tile writes to a separate buffer, then a final reduction sums them
  - Bit-exact reproducible gradients across runs
  - Slower due to extra memory and reduction pass

CPU reference simulates both modes. On CPU, both modes are numerically
identical (no atomics), but this example demonstrates the API pattern
and compares GPU-style multi-buffer reduction semantics.

Usage:
    python3 37_bwd_deterministic_fmha.py
    python3 37_bwd_deterministic_fmha.py --seqlen 128
    python3 37_bwd_deterministic_fmha.py --num-tiles 4
"""

import sys
import argparse
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python"))
import numpy as np

from fmha_utils import (
    FmhaKernelConfig,
    FmhaProblem,
    setup_fmha_dispatcher,
    detect_gpu_arch,
)


def cpu_fwd_with_intermediates(
    Q: np.ndarray,
    K: np.ndarray,
    V: np.ndarray,
    scale: float,
) -> tuple:
    """Forward returning out, P, LSE."""
    S = np.matmul(Q, K.transpose(0, 1, 3, 2)) * scale
    S_max = S.max(axis=-1, keepdims=True)
    S_exp = np.exp(S - S_max)
    S_sum = S_exp.sum(axis=-1, keepdims=True)
    P = S_exp / S_sum
    out = np.matmul(P, V)
    lse = (np.log(S_sum.squeeze(-1)) + S_max.squeeze(-1)).astype(np.float32)
    return out, P, lse


def cpu_bwd_nondeterministic(
    Q: np.ndarray,
    K: np.ndarray,
    V: np.ndarray,
    out: np.ndarray,
    dO: np.ndarray,
    P: np.ndarray,
    scale: float,
) -> tuple:
    """Standard backward (single accumulation). Returns (dQ, dK, dV)."""
    D = (dO * out).sum(axis=-1, keepdims=True)
    dP = np.matmul(dO, V.transpose(0, 1, 3, 2))
    dS = P * (dP - D)
    dQ = np.matmul(dS, K) * scale
    dK = np.matmul(dS.transpose(0, 1, 3, 2), Q) * scale
    dV = np.matmul(P.transpose(0, 1, 3, 2), dO)
    return dQ, dK, dV


def cpu_bwd_deterministic(
    Q: np.ndarray,
    K: np.ndarray,
    V: np.ndarray,
    out: np.ndarray,
    dO: np.ndarray,
    P: np.ndarray,
    scale: float,
    num_tiles_k: int = 4,
) -> tuple:
    """Deterministic backward with explicit multi-buffer reduction for dQ.

    Simulates the GPU pattern where seqlen_k is split into tiles,
    each tile writes dQ to a separate buffer, then buffers are summed.

    Returns: (dQ, dK, dV, dQ_buffers)
    """
    B, Hq, Sq, Dq = Q.shape
    Sk = K.shape[2]

    D = (dO * out).sum(axis=-1, keepdims=True)

    tile_sk = max(1, Sk // num_tiles_k)
    actual_tiles = (Sk + tile_sk - 1) // tile_sk

    dQ_buffers = np.zeros((actual_tiles, B, Hq, Sq, Dq), dtype=np.float32)
    dK = np.zeros_like(K)
    dV = np.zeros_like(V)

    for t in range(actual_tiles):
        sk_start = t * tile_sk
        sk_end = min(sk_start + tile_sk, Sk)

        K_tile = K[:, :, sk_start:sk_end, :]
        V_tile = V[:, :, sk_start:sk_end, :]
        P_tile = P[:, :, :, sk_start:sk_end]

        dP_tile = np.matmul(dO, V_tile.transpose(0, 1, 3, 2))
        dS_tile = P_tile * (dP_tile - D)

        dQ_buffers[t] = np.matmul(dS_tile, K_tile) * scale
        dK[:, :, sk_start:sk_end, :] = (
            np.matmul(dS_tile.transpose(0, 1, 3, 2), Q) * scale
        )
        dV[:, :, sk_start:sk_end, :] = np.matmul(P_tile.transpose(0, 1, 3, 2), dO)

    dQ = dQ_buffers.sum(axis=0)
    return dQ, dK, dV, dQ_buffers


def main():
    parser = argparse.ArgumentParser(description="Backward Deterministic Mode")
    parser.add_argument("--arch", default=detect_gpu_arch())
    parser.add_argument("--batch", type=int, default=2)
    parser.add_argument("--nhead", type=int, default=8)
    parser.add_argument("--seqlen", type=int, default=64)
    parser.add_argument("--hdim", type=int, default=128)
    parser.add_argument(
        "--num-tiles",
        type=int,
        default=4,
        help="Number of seqlen_k tiles for deterministic mode",
    )
    args = parser.parse_args()

    print("=" * 70)
    print("Example 37: Backward Pass Deterministic Mode")
    print("=" * 70)

    prob = FmhaProblem(
        batch=args.batch,
        nhead_q=args.nhead,
        nhead_k=args.nhead,
        seqlen_q=args.seqlen,
        seqlen_k=args.seqlen,
        hdim_q=args.hdim,
        hdim_v=args.hdim,
    )

    print(
        f"\n  Problem:   B={prob.batch} H={prob.nhead_q} S={args.seqlen} D={args.hdim}"
    )
    print(f"  Tiles:     {args.num_tiles} (seqlen_k split)")
    print(f"  Tile size: {max(1, args.seqlen // args.num_tiles)}")

    # --- JIT compile a basic fp16 h128 fwd kernel ---
    print("\n--- JIT Compilation ---")
    config = FmhaKernelConfig(
        data_type="fp16",
        hdim_q=args.hdim,
        hdim_v=args.hdim,
        gfx_arch=args.arch,
    )
    setup = setup_fmha_dispatcher(config)
    if setup.success:
        print(f"  Fwd kernel compiled: {setup.build_time_s:.1f}s")
        print("  Backward deterministic kernel: separate JIT with deterministic=True")
    else:
        print(f"  JIT build: {setup.error}")
        print("  Continuing with CPU reference only")

    # --- Generate data ---
    np.random.seed(42)
    Q = (np.random.randn(*prob.q_shape()) * 0.1).astype(np.float32)
    K = (np.random.randn(*prob.k_shape()) * 0.1).astype(np.float32)
    V = (np.random.randn(*prob.v_shape()) * 0.1).astype(np.float32)
    dO = (np.random.randn(*prob.o_shape()) * 0.1).astype(np.float32)

    out, P, lse = cpu_fwd_with_intermediates(Q, K, V, prob.scale)

    # --- Non-deterministic backward ---
    print("\n--- Non-Deterministic Backward ---")
    dQ_nd, dK_nd, dV_nd = cpu_bwd_nondeterministic(Q, K, V, out, dO, P, prob.scale)

    print(f"  dQ range: [{dQ_nd.min():.4e}, {dQ_nd.max():.4e}]")
    print(f"  dK range: [{dK_nd.min():.4e}, {dK_nd.max():.4e}]")
    print(f"  dV range: [{dV_nd.min():.4e}, {dV_nd.max():.4e}]")

    # --- Deterministic backward ---
    print(f"\n--- Deterministic Backward ({args.num_tiles} tiles) ---")
    dQ_det, dK_det, dV_det, dQ_bufs = cpu_bwd_deterministic(
        Q,
        K,
        V,
        out,
        dO,
        P,
        prob.scale,
        num_tiles_k=args.num_tiles,
    )

    print(f"  dQ range: [{dQ_det.min():.4e}, {dQ_det.max():.4e}]")
    print(f"  dK range: [{dK_det.min():.4e}, {dK_det.max():.4e}]")
    print(f"  dV range: [{dV_det.min():.4e}, {dV_det.max():.4e}]")
    print(f"  dQ buffers: {dQ_bufs.shape[0]} x {dQ_bufs.shape[1:]}")

    # --- Per-buffer analysis ---
    print("\n--- Per-Tile dQ Buffer Analysis ---")
    print(f"\n  {'Tile':>6} {'|buf| mean':>12} {'|buf| max':>12} {'% of total':>12}")
    print("  " + "-" * 46)

    total_l1 = float(np.abs(dQ_det).sum())
    for t in range(dQ_bufs.shape[0]):
        buf = dQ_bufs[t]
        buf_mean = float(np.abs(buf).mean())
        buf_max = float(np.abs(buf).max())
        buf_pct = float(np.abs(buf).sum()) / (total_l1 + 1e-15) * 100
        print(f"  {t:>6} {buf_mean:>12.4e} {buf_max:>12.4e} {buf_pct:>11.1f}%")

    # --- Compare deterministic vs non-deterministic ---
    print("\n--- Deterministic vs Non-Deterministic Comparison ---")
    print(f"\n  {'Grad':<6} {'Max abs diff':>14} {'Mean abs diff':>14} {'Match':>8}")
    print("  " + "-" * 46)

    for name, g_det, g_nd in [
        ("dQ", dQ_det, dQ_nd),
        ("dK", dK_det, dK_nd),
        ("dV", dV_det, dV_nd),
    ]:
        abs_diff = np.abs(g_det - g_nd)
        max_abs = float(abs_diff.max())
        mean_abs = float(abs_diff.mean())
        match = max_abs < 1e-6
        print(
            f"  {name:<6} {max_abs:>14.2e} {mean_abs:>14.2e} {'YES' if match else 'NO':>8}"
        )

    print("\n  NOTE: On CPU, both modes produce identical results.")
    print("  On GPU, non-deterministic mode uses atomicAdd for dQ,")
    print("  causing order-dependent floating-point rounding differences.")

    # --- Reproducibility test ---
    print("\n--- Reproducibility Test (Deterministic Mode) ---")
    num_runs = 5
    dQ_runs = []
    for run in range(num_runs):
        dQ_r, _, _, _ = cpu_bwd_deterministic(
            Q,
            K,
            V,
            out,
            dO,
            P,
            prob.scale,
            num_tiles_k=args.num_tiles,
        )
        dQ_runs.append(dQ_r)

    max_variation = 0.0
    for i in range(1, num_runs):
        diff = float(np.abs(dQ_runs[i] - dQ_runs[0]).max())
        max_variation = max(max_variation, diff)

    print(f"  Runs: {num_runs}")
    print(f"  Max dQ variation across runs: {max_variation:.2e}")
    print(f"  Bit-exact reproducible: {'YES' if max_variation == 0.0 else 'NO'}")

    # --- Memory overhead analysis ---
    print("\n--- Deterministic Mode Memory Overhead ---")
    dq_size = Q.nbytes
    buf_size = dQ_bufs.nbytes
    overhead = buf_size / dq_size

    print(f"  dQ single buffer:    {dq_size:>10,} bytes")
    print(f"  dQ multi-buffer:     {buf_size:>10,} bytes ({args.num_tiles} tiles)")
    print(f"  Memory overhead:     {overhead:.1f}x")
    print(f"  Extra memory:        {buf_size - dq_size:>10,} bytes")

    # --- GPU API pattern ---
    print("\n--- GPU Deterministic API Pattern ---")
    print("  Non-deterministic (default):")
    print("    FmhaKernelConfig(family='bwd', deterministic=False)")
    print("    dQ accumulated via atomicAdd (fast, non-reproducible)")
    print()
    print("  Deterministic:")
    print("    FmhaKernelConfig(family='bwd', deterministic=True)")
    print("    dQ via multi-buffer + final reduction (reproducible)")
    print("    Requires extra workspace: num_tiles_k * sizeof(dQ)")

    # --- Summary ---
    print("\n" + "=" * 70)
    print(f"  Tiles:          {args.num_tiles}")
    print(f"  Memory overhead: {overhead:.1f}x for deterministic dQ")
    print("  Reproducible:   Deterministic mode guarantees bit-exact results")
    print("  Performance:    Deterministic ~10-20% slower on GPU (extra reduction)")
    print("  GPU:            Requires bwd-family JIT kernel")
    print("  Status:         DEMO")
    print("=" * 70)

    return 0


if __name__ == "__main__":
    sys.exit(main())