sm121a_matmul.py

"""
Custom BF16 matmul kernel tuned for NVIDIA DGX Spark (GB10, sm_121a).

Design constraints (from Phase 4 research, Agents 1-3):
  - 48 SMs, 1536 threads/SM, 100 KB shared memory per SM (101 KB strict ceiling)
  - 273 GB/s LPDDR5X unified memory — bandwidth-bound for large MLPs
  - TMEM hardware-absent (NVIDIA staff confirmed), no tcgen05 path
  - mma.sync.aligned.m16n8k16 is the right MMA shape
  - cluster_dims TMA broken on sm_121 (Triton #8335), do NOT use
  - Target sm_120a SASS for forward compatibility (binary compat to sm_121)

Tile configs derived from NVIDIA forum 154 TFLOPS @ 256x128 result for prefill
on GB10. We chose 128x256 (same FLOPS, fits 101 KB SMEM more comfortably).

Target shapes for Qwen3.5-2B LoRA training (per packed step):
  MLP gate/up (M=5120, N=11008, K=2048)   ← HIGHEST ROI per Agent 2
  MLP down    (M=5120, N=2048,  K=11008)
  QKV proj   (M=5120, N=2048,  K=2048) and (M=5120, N=512, K=2048) for GQA

This file exposes:
  - `sm121a_matmul(a, b) -> c`           plain BF16 mm

  The kernel is dispatched via the per-shape selector in
  `lab/kernels/fused_lora/v3_hybrid.py` for MLP gate/up shapes
  (N >= 4096 and K <= 8000), where it outperforms cuBLAS by ~5%.
"""
from __future__ import annotations
import torch
import triton
import triton.language as tl


# ── Autotune configs ────────────────────────────────────────────────────────
# Keep (BLOCK_M+BLOCK_N) * BLOCK_K * 2 * num_stages under ~96 KB to leave
# headroom for register spill / accumulator under the 101 KB sm_121a ceiling.

_AUTOTUNE_CONFIGS = [
    # NVIDIA-forum-reported peak for sm_121a prefill — large MLP
    triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 64, "GROUP_M": 8},
                  num_warps=8, num_stages=3),
    # Same tile, fewer stages (lower SMEM, more SMs idle on small batches)
    triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 64, "GROUP_M": 8},
                  num_warps=8, num_stages=2),
    # Square tile — Q/K/V/O projections (2048×2048)
    triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8},
                  num_warps=4, num_stages=3),
    # Skinny N — V projection with GQA (2048×512)
    triton.Config({"BLOCK_M": 128, "BLOCK_N": 64,  "BLOCK_K": 128, "GROUP_M": 8},
                  num_warps=4, num_stages=2),
    # Deep K — MLP down (11008→2048)
    triton.Config({"BLOCK_M": 256, "BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8},
                  num_warps=8, num_stages=2),
    # Fallback for small problems
    triton.Config({"BLOCK_M": 64,  "BLOCK_N": 64,  "BLOCK_K": 64, "GROUP_M": 8},
                  num_warps=4, num_stages=2),
]


@triton.autotune(configs=_AUTOTUNE_CONFIGS, key=["M", "N", "K"])
@triton.jit
def _sm121a_matmul_kernel(
    A_ptr, B_ptr, C_ptr,
    M, N, K,
    stride_am: tl.constexpr, stride_ak: tl.constexpr,
    stride_bk: tl.constexpr, stride_bn: tl.constexpr,
    stride_cm: tl.constexpr, stride_cn: tl.constexpr,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    """BF16 × BF16 → BF16 matmul with FP32 accumulator. Group-major
    program order improves L2 reuse for tall-skinny shapes.

    NOTE: Uses Triton's `tl.dot()` which emits ldmatrix + mma.sync.aligned
    on sm_120/121. We avoid cluster_dims and TMA bulk because of broken
    sm_121a paths (Triton issue #8335).
    """
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = tl.minimum(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M
    offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N
    offs_k = tl.arange(0, BLOCK_K)

    a_ptrs = A_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = B_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        k_remaining = K - k * BLOCK_K
        a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
        accumulator = tl.dot(a, b, accumulator)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk
    c = accumulator.to(tl.bfloat16)

    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    c_ptrs = C_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)


def sm121a_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """BF16 matmul, sm_121a-tuned. Same semantics as torch.matmul but
    dispatches the kernel from this file."""
    assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16, \
        f"sm121a_matmul: only bf16 (got {a.dtype}, {b.dtype})"
    assert a.shape[-1] == b.shape[0], \
        f"shape mismatch {a.shape} × {b.shape}"
    M, K = a.shape
    K2, N = b.shape
    a = a.contiguous() if not a.is_contiguous() else a
    b = b.contiguous() if not b.is_contiguous() else b
    c = torch.empty((M, N), device=a.device, dtype=torch.bfloat16)

    def grid(META):
        return (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),)

    _sm121a_matmul_kernel[grid](
        a, b, c, M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
    )
    return c


__all__ = ["sm121a_matmul"]