composable_kernel/dispatcher/examples/fmha/python/24_vlayout_col_fmha.py

#!/usr/bin/env python3

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

"""
Example 24: Column-Major V Layout FMHA

Demonstrates column-major (vlayout="c") vs row-major (vlayout="r") for
the V tensor. In row-major, V is [batch, nhead, seqlen_k, hdim_v]; in
column-major, V is [batch, nhead, hdim_v, seqlen_k].

Column-major V can improve performance when hdim_v access patterns
benefit from the transposed layout (e.g., certain tile sizes or memory
coalescing characteristics on specific GPU architectures).

The prebuilt library uses row-major V. This example shows both layouts
with CPU reference and validates correctness.

Usage:
    python3 24_vlayout_col_fmha.py
    python3 24_vlayout_col_fmha.py --seqlen 512
    python3 24_vlayout_col_fmha.py --batch 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,
    FmhaValidator,
    cpu_attention_fwd,
    detect_gpu_arch,
    setup_fmha_dispatcher,
)


def cpu_attention_fwd_vlayout_col(
    Q: np.ndarray,
    K: np.ndarray,
    V_col: np.ndarray,
    scale: float,
) -> np.ndarray:
    """CPU reference: attention with column-major V.

    Args:
        Q:     [batch, nhead_q, seqlen_q, hdim_q]  float32 (row-major)
        K:     [batch, nhead_k, seqlen_k, hdim_q]  float32 (row-major)
        V_col: [batch, nhead_k, hdim_v, seqlen_k]  float32 (column-major)
        scale: softmax scale

    Returns:
        O: [batch, nhead_q, seqlen_q, hdim_v]  float32
    """
    V_row = V_col.transpose(0, 1, 3, 2)
    return cpu_attention_fwd(Q, K, V_row, scale)


def analyze_strides(name: str, arr: np.ndarray, dim_names: list):
    """Print stride information for a tensor."""
    strides_bytes = arr.strides
    itemsize = arr.itemsize
    strides_elems = tuple(s // itemsize for s in strides_bytes)
    print(f"  {name}:")
    print(f"    Shape:   {arr.shape}")
    print(f"    Strides: {strides_elems} (elements)")
    for i, (dname, s) in enumerate(zip(dim_names, strides_elems)):
        contiguous = "(contiguous)" if i == len(dim_names) - 1 and s == 1 else ""
        print(f"      {dname}: stride={s} {contiguous}")


def main():
    parser = argparse.ArgumentParser(
        description="Column-Major V Layout FMHA Example",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    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=128)
    parser.add_argument("--hdim", type=int, default=128)
    args = parser.parse_args()

    print("=" * 70)
    print("Example 24: Column-Major V Layout FMHA")
    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,
    )

    # Step 1: Layout comparison
    print("\nStep 1: V Tensor Layouts")

    np.random.seed(42)
    V_row = np.ascontiguousarray(
        (np.random.randn(*prob.v_shape()) * 0.3).astype(np.float32)
    )
    V_col = np.ascontiguousarray(V_row.transpose(0, 1, 3, 2))

    analyze_strides(
        "V row-major [B, H, SeqK, Hdim]",
        V_row,
        ["batch", "nhead", "seqlen_k", "hdim_v"],
    )
    analyze_strides(
        "V col-major [B, H, Hdim, SeqK]",
        V_col,
        ["batch", "nhead", "hdim_v", "seqlen_k"],
    )

    print("\n  Row-major: last dim is hdim_v -> sequential hdim access per token")
    print("  Col-major: last dim is seqlen_k -> sequential token access per hdim")

    # Step 2: CPU reference for both layouts
    print("\nStep 2: CPU Reference (both layouts)")

    Q = (np.random.randn(*prob.q_shape()) * 0.3).astype(np.float32)
    K = (np.random.randn(*prob.k_shape()) * 0.3).astype(np.float32)

    O_from_row = cpu_attention_fwd(Q, K, V_row, prob.scale)
    O_from_col = cpu_attention_fwd_vlayout_col(Q, K, V_col, prob.scale)

    validator = FmhaValidator(rtol=1e-5, atol=1e-5)
    ok, max_abs, max_rel = validator.check(O_from_row, O_from_col)

    print(
        f"  O from row-major V: shape={O_from_row.shape}  "
        f"range=[{O_from_row.min():.4f}, {O_from_row.max():.4f}]"
    )
    print(
        f"  O from col-major V: shape={O_from_col.shape}  "
        f"range=[{O_from_col.min():.4f}, {O_from_col.max():.4f}]"
    )
    print(f"  Max abs error: {max_abs:.2e}")
    print(f"  Match: {'PASS' if ok else 'FAIL'}")

    # Step 3: Memory access pattern analysis
    print("\nStep 3: Memory Access Pattern Analysis")

    tile_sizes = [(128, 128), (64, 128), (128, 64)]
    print("\n  For P @ V matmul (P: [sq, sk] x V: [sk, hdim_v]):")
    print(f"  {'Tile(M,N)':>12} {'V Row Accesses':>18} {'V Col Accesses':>18}")
    print("  " + "-" * 52)

    for tm, tn in tile_sizes:
        row_access = f"sk_stride={args.hdim}"
        col_access = "sk_stride=1"
        print(f"  {f'{tm}x{tn}':>12} {row_access:>18} {col_access:>18}")

    print("\n  Row-major V: coalesced reads when accessing hdim_v (inner loop)")
    print("  Col-major V: coalesced reads when accessing seqlen_k (inner loop)")
    print("  Optimal layout depends on tile shape and GPU memory subsystem")

    # Step 4: Shape sweep with both layouts
    print("\nStep 4: Correctness Sweep")

    shapes = [
        (1, 4, 64, 64, 64),
        (2, 8, 128, 128, 128),
        (1, 8, 256, 256, 128),
        (2, 4, 128, 128, 64),
        (1, 16, 64, 64, 128),
    ]

    print(f"\n  {'Shape':<32} {'MaxErr':>12} {'Status':>8}")
    print("  " + "-" * 55)

    all_ok = True
    for b, h, sq, sk, d in shapes:
        Q_t = (np.random.randn(b, h, sq, d) * 0.3).astype(np.float32)
        K_t = (np.random.randn(b, h, sk, d) * 0.3).astype(np.float32)
        V_r = (np.random.randn(b, h, sk, d) * 0.3).astype(np.float32)
        V_c = np.ascontiguousarray(V_r.transpose(0, 1, 3, 2))

        scale = 1.0 / (d**0.5)
        O_r = cpu_attention_fwd(Q_t, K_t, V_r, scale)
        O_c = cpu_attention_fwd_vlayout_col(Q_t, K_t, V_c, scale)

        ok_t, max_abs_t, _ = validator.check(O_r, O_c)
        all_ok = all_ok and ok_t
        shape_str = f"B{b}_H{h}_S{sq}x{sk}_D{d}"
        print(f"  {shape_str:<32} {max_abs_t:>12.2e} {'PASS' if ok_t else 'FAIL':>8}")

    # Step 5: GPU API pattern
    print("\nStep 5: GPU Kernel Configuration")
    print("  NOTE: The prebuilt library uses row-major V (vlayout='r').")
    print("  For column-major V, compile a kernel with vlayout='c':")
    print()
    print("    FmhaSignature()")
    print("        .vlayout('c')   // column-major V: [B, H, Hdim, SeqK]")
    print()
    print("    FmhaKernelConfig(vlayout='c', ...)")

    # Step 6: GPU baseline (row-major)
    print("\nStep 6: GPU Baseline (row-major V)")

    config = FmhaKernelConfig(
        data_type="fp16",
        hdim_q=args.hdim,
        hdim_v=args.hdim,
        gfx_arch=args.arch,
    )
    setup = setup_fmha_dispatcher(config)
    if not setup.success:
        print(f"  JIT build failed: {setup.error}")
    else:
        runner = setup.runner
        print(f"  JIT build: {setup.build_time_s:.1f}s")

        Q_f16 = Q.astype(np.float16)
        K_f16 = K.astype(np.float16)
        V_f16 = V_row.astype(np.float16)

        result = runner.run(Q_f16, K_f16, V_f16, prob)
        if result.success:
            ok_gpu, max_abs_gpu, _ = validator.check(result.output, O_from_row)
            print(
                f"  GPU (row-major V): time={result.time_ms:.4f}ms  TFLOPS={result.tflops:.2f}  "
                f"max_err={max_abs_gpu:.2e}  {'PASS' if ok_gpu else 'FAIL'}"
            )
        else:
            print(f"  GPU error: {result.error}")

    # Summary
    print("\n" + "=" * 70)
    print("  vlayout='r': V is [B, H, SeqK, Hdim] (default, row-major)")
    print("  vlayout='c': V is [B, H, Hdim, SeqK] (column-major)")
    print(
        f"  Both layouts produce identical results (verified: {'PASS' if all_ok else 'FAIL'})"
    )
    print("  Choice depends on upstream memory layout and GPU tile access patterns")
    print("=" * 70)

    return 0


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