diff --git a/docs/audit/backend/nvidia/NVIDIA_AUDIT.md b/docs/audit/backend/nvidia/NVIDIA_AUDIT.md index 8cf92054e..0cb5e66c3 100644 --- a/docs/audit/backend/nvidia/NVIDIA_AUDIT.md +++ b/docs/audit/backend/nvidia/NVIDIA_AUDIT.md @@ -63,14 +63,24 @@ This document consolidates NVIDIA-specific audit material. ## Still Open The original "no execution row / not hardware-proven" gaps are **closed** -(above). What remains is breadth beyond the first proven kernel: +(above), and the **compiler-generated lane is now landed + hardware-proven** on +sm_120 (RTX 5070 Ti, PRs #290–#297): -- **Compiler-GENERATED NVIDIA lane** — the analog of ROCm's `rocm_compiled` - (a `tessera-opt` NVIDIA pipeline that generates + serializes the kernel - in-process, rather than dispatching the hand-shipped `libtessera_nvidia_gemm` - symbol). Today's `nvidia_mma` lane mirrors `rocm_wmma` (shipped symbol), not - the compiler-generated path. -- **NVFP4 block-scale matmul (#9)** — the warp `mma.sync…block_scale` +- **Compiler-GENERATED NVIDIA lane — LANDED.** `emit/nvidia_cuda.py` is a full + three-seam plugin (emitter + `nvcc` compile + ctypes runner) that synthesizes, + compiles, and launches kernels in-process for **all four `fusion_core` region + kinds** — fused matmul-epilogue, flash-attention (C4), SwiGLU gate + pointwise + DAG (C5) — each F4-gated on-device. The emit-path `mma.sync` GEMM + (`ptx_emit.py` → the shipped `tessera_nvidia_ptx_launch` bridge: driver-JIT + + `cuLaunchKernel`) executes the *emitted* PTX, distinct from the hand-shipped + `libtessera_nvidia_gemm` symbol. Both are first-class D1 arbiter candidates: + the shipped GEMM is **Tier-3 hand-tuned**, the emitted GEMM **Tier-2 emitted** + (B1), with D2 measured autotune + D3 fallback logging choosing/observing between + them. So NVIDIA now has the `rocm_compiled` analog it lacked, plus the arbiter + surface. +- **NVFP4 block-scale matmul (#9)** — **emit + ptxas-assemble landed** (#291, + `emit_nvfp4_block_scale_mma_ptx`); on-device execution + non-unit-scale numerics + stay gated on the PTX-ISA scale-distribution spec. The warp `mma.sync…block_scale` instruction already assembles + executes on `sm_120a` (see `spikes/sm120_mma_sync/`); productization is pending the PTX ISA scale-distribution spec for numerics. @@ -87,14 +97,20 @@ The original "no execution row / not hardware-proven" gaps are **closed** ## Next Work -1. Build the compiler-generated NVIDIA lane (the `rocm_compiled` analog) via a - `tessera-opt` NVIDIA pipeline; add its executable matrix row. -2. Land NVFP4 block-scale matmul once the scale-distribution numerics are - grounded; flip its manifest row when execute-and-compare passes on `sm_120a`. -3. Bring the sm_120 `mma.sync` flash-attention forward to the same - execute-and-compare bar (attention analog of the GEMM proof). -4. Author an sm_120 `mma.sync` kernel inventory (sibling to - `docs/nvidia_cuda13_kernel_inventory.md`). +Done (2026-07-07): the compiler-generated lane (#290–#297), the sm_120 `mma.sync` +flash-attention execute-compare (C4), and the sm_120 kernel-inventory doc +(`docs/nvidia_sm120_mma_sync_kernel_inventory.md`). Remaining: + +1. **NVFP4 block-scale execution + numerics** — bind the fp4 fragment packing and + flip the manifest row once execute-and-compare passes on `sm_120a` and the + scale-distribution numerics are grounded (emit + ptxas-assemble already land). +2. **mma.sync tensor-core versions** of the flash-attention + fused lanes (the + correctness-first synth lanes are proven; these are the perf follow-on) + the + **D2 fleet-shared autotune corpus** (persist `MeasureCache`, Theory §7.5). +3. **Dtypes beyond f32** for the fused / attention / gated lanes (16-bit storage + is served by the B1 matmul lane today). +4. **`wgmma` sm_90a** — complete the instruction-encoding skeleton into a real + Hopper WGMMA kernel (assemble-only until a Hopper box) — and **sm_100 tcgen05**. 5. Promote sm_80/90/100 manifest rows only when their own silicon is available and the generated dashboards agree. diff --git a/docs/nvidia_sm120_mma_sync_kernel_inventory.md b/docs/nvidia_sm120_mma_sync_kernel_inventory.md new file mode 100644 index 000000000..a5726eedd --- /dev/null +++ b/docs/nvidia_sm120_mma_sync_kernel_inventory.md @@ -0,0 +1,102 @@ +--- +status: Informative +classification: Reference / Kernel Inventory +authority: Companion to docs/nvidia_cuda13_kernel_inventory.md +last_updated: 2026-07-07 +--- + +# NVIDIA sm_120 (consumer Blackwell) mma.sync Kernel Inventory + +> Hardware-verified companion to +> [`docs/nvidia_cuda13_kernel_inventory.md`](nvidia_cuda13_kernel_inventory.md) +> (which is the SM_90+ WGMMA *planning* inventory). This doc enumerates the +> kernels Tessera actually **synthesizes, compiles, and runs on sm_120** — proven +> on an RTX 5070 Ti (compute cap 12.0, CUDA 13.3, driver ≥610.43.02). Statuses +> here are execution truth, not roofline targets; the generated dashboards +> (`docs/audit/generated/`) remain the drift-gated count/status surface. + +Three kernel families feed the NVIDIA lane, each a distinct compile path: + +1. **Compiler-emitted CUDA** — `emit/nvidia_cuda.py` synthesizes CUDA C, compiles + it with `nvcc -arch=sm_120a`, and launches via `ctypes` (the three-seam generic + plugin). Correctness-first, arch-agnostic; F4-gated on-device. +2. **Emit-path PTX** — `ptx_emit.py` emits PTX text, the shipped + `tessera_nvidia_ptx_launch` bridge driver-JITs it (`cuModuleLoadDataEx`, cached + by kernel name) and launches it (`cuLaunchKernel`). +3. **Shipped C-ABI** — `libtessera_nvidia_gemm.so` NVRTC-compiles a hand-written + CUDA-C `mma.sync` GEMM at first call (the hand-tuned lane). + +Execution status legend: **✅ proven** (execute-and-compare on sm_120) · +**🟡 assemble-only** (ptxas accepts; not launched/numerically gated) · +**⬜ skeleton** (instruction-encoding only, not assemblable). + +--- + +## 1. Compiler-emitted CUDA lanes (`emit/nvidia_cuda.py`) + +| Entry symbol | Op / region | Shape model | dtype | Kernel shape | Status | +|---|---|---|---|---|---| +| `tessera_nvidia_fused` | `FusedRegion` — matmul + prologue/epilogue/residual/reduction | runtime M/N/K | f32 | one thread per output row | ✅ | +| `tessera_nvidia_attn` | `AttentionRegion` — `O = softmax(scale·Q·Kᵀ)·V` | runtime M/Nk/D/Dv (Dv ≤ 256) | f32 | flash: one query/thread, online softmax, streaming KV, causal + transpose flags | ✅ (C4) | +| `tessera_nvidia_gated` | `GatedMatmulRegion` — SwiGLU gate `gate_act(A·Wg) ⊙ (A·Wu)` | runtime M/K/H | f32 | one output row/thread, shared A load | ✅ (C5) | +| `tessera_nvidia_pointwise` | `PointwiseGraphRegion` — same-shape pointwise DAG | runtime numel | f32 | one element/thread; DAG from `POINTWISE_OPS` C-expr table (+ NaN-safe `sign`/`clamp` shims) | ✅ (C5) | + +All four are Tier-1 **synthesized** D1 arbiter candidates +(`Nvidia{Generic,FlashAttn,Gated,Pointwise}…Candidate`). f32 is the +correctness-first floor; 16-bit storage is served by the GEMM lanes below. + +## 2. Emit-path PTX lanes (`ptx_emit.py` → `tessera_nvidia_ptx_launch`) + +| Entry symbol | Op | Shape | dtype | Instruction | Status | +|---|---|---|---|---|---| +| `tessera_mma_m16n8k16_bf16` | matmul (single tile) | 16×8×16 | bf16→f32 | `mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32` | ✅ | +| `tessera_mma_gemm_bf16` | matmul (general) | aligned M%16/N%8/K%16, index < 2³¹ | bf16→f32 | m16n8k16 tile, K-loop + grid-tiled | ✅ (C2 breadth) | +| `tessera_mma_gemm_f16` | matmul (general) | aligned M%16/N%8/K%16, index < 2³¹ | f16→f32 | m16n8k16 tile, K-loop + grid-tiled | ✅ | +| `tessera_nvfp4_mma_m16n8k64` | matmul (block-scale) | 16×8×64 | fp4 e2m1 + ue4m3 scales → f32 | `mma.sync…m16n8k64…kind::mxf4nvf4.block_scale.scale_vec::4X` | 🟡 assemble-only (numerics gated on PTX-ISA scale spec) | +| `tessera_wgmma_matmul_bf16` | matmul (Hopper) | m64n{64,128,256}k16 | bf16→f32 | `wgmma.mma_async…` | ⬜ skeleton (needs smem descriptors + TMA; sm_90a, no Hopper box) | + +The emitted general GEMM (`tessera_mma_gemm_{bf16,f16}`) is the **Tier-2 emitted** +D1 candidate (`NvidiaMmaGemmEmittedCandidate`); the launch bridge rejects +i32-index-overflow shapes (element count > 2³¹) honestly. + +## 3. Shipped C-ABI lane (`libtessera_nvidia_gemm.so`) + +| Entry symbol | Op | Shape | dtype | Status | +|---|---|---|---|---| +| `tessera_nvidia_mma_gemm_f16` | matmul (general tiled) | any M/N/K | f16→f32 | ✅ | +| `tessera_nvidia_mma_gemm_bf16` | matmul (general tiled) | any M/N/K | bf16→f32 | ✅ | +| `tessera_nvidia_mma_gemm_tf32` | matmul (general tiled) | any M/N/K | fp32/tf32-math→f32 | ✅ | + +The shipped GEMM is the **Tier-3 hand-tuned** D1 candidate +(`NvidiaMmaGemmShippedCandidate`) — the arbiter default (lead-safe, Decision #28), +displaced only when D2's measured loop proves the emitted lane faster + in budget. + +--- + +## Arbiter mapping (D1/D2/D3) + +| op | Tier-3 hand-tuned | Tier-2 emitted | Tier-1 synthesized | +|---|---|---|---| +| `matmul` | `tessera_nvidia_mma_gemm_*` (shipped) | `tessera_mma_gemm_*` (PTX bridge) | — | +| `fused_region` | — | — | `tessera_nvidia_fused` | +| `attention` | — | — | `tessera_nvidia_attn` | +| `gated_matmul` | — | — | `tessera_nvidia_gated` | +| `pointwise` | — | — | `tessera_nvidia_pointwise` | + +Selection: tier-priority by default; `emit/autotune.py` measures on-device and +caches the fastest per `(device, op, shape-bucket, dtype)` (D2); every dispatch is +recorded in the arbiter fallback log (D3, `arbiter_dispatch_histogram`). + +## Toolchain pin + +CUDA Toolkit **13.3** (PTX ISA 9.3); target `sm_120a` (FP4 +`mma.sync.block_scale`); driver ≥610.43.02; smem 100 KB/SM. The emit-path and +shipped lanes need only the host compiler + CUDA driver (libcuda) + NVRTC at load +time; the compiler-emitted CUDA lane needs `nvcc`. + +## Still open + +NVFP4 execution + non-unit-scale numerics; mma.sync tensor-core versions of the +attention + fused lanes (perf); dtypes beyond f32 for the fused/attention/gated +lanes; the Hopper `wgmma` completion (sm_90a) and sm_100 `tcgen05` (their own +silicon). See [`docs/audit/backend/nvidia/NVIDIA_AUDIT.md`](audit/backend/nvidia/NVIDIA_AUDIT.md). diff --git a/python/tessera/compiler/emit/nvidia_cuda.py b/python/tessera/compiler/emit/nvidia_cuda.py index 4840bfab5..8148e1436 100644 --- a/python/tessera/compiler/emit/nvidia_cuda.py +++ b/python/tessera/compiler/emit/nvidia_cuda.py @@ -38,6 +38,7 @@ import ctypes import os +import re import shutil import subprocess import tempfile @@ -47,7 +48,9 @@ from tessera.compiler.emit.candidate import ( OP_ATTENTION, OP_FUSED_REGION, + OP_GATED_MATMUL, OP_MATMUL, + OP_POINTWISE, Candidate, Tier, register_candidate, @@ -63,12 +66,21 @@ register_emitter, register_runner, ) -from tessera.compiler.fusion_core import AttentionRegion, FusedRegion, MatmulRegion +from tessera.compiler.fusion_core import ( + POINTWISE_OPS, + AttentionRegion, + FusedRegion, + GatedMatmulRegion, + MatmulRegion, + PointwiseGraphRegion, +) _TARGET = "nvidia" _LANG = "cuda" _ENTRY = "tessera_nvidia_fused" _ATTN_ENTRY = "tessera_nvidia_attn" +_GATED_ENTRY = "tessera_nvidia_gated" +_PW_ENTRY = "tessera_nvidia_pointwise" _REAL_TAG = "nvidia_cuda" #: Max head dim (Dv) the one-thread-per-query flash kernel holds in its per-thread #: online-softmax accumulator; larger Dv declines to the reference. @@ -181,21 +193,116 @@ def _synthesize_attention_cuda() -> str: ) +def _synthesize_gated_cuda(region: GatedMatmulRegion) -> str: + """CUDA source for the SwiGLU gate ``O = f(A @ Wg) ⊙ (A @ Wu)`` (f32) — one + thread per output row, sharing the A load across the two K-contractions, then + the gate activation + elementwise multiply. A(M,K), Wg/Wu(K,H), O(M,H).""" + from tessera.compiler.emit._fused_scalar_body import pointwise_snippet + act = pointwise_snippet(region.gate_act, "g") # e.g. g = g/(1+expf(-g)); + return ( + "#include \n" + "#include \n" + f"__global__ void {_GATED_ENTRY}_kernel(const float* A, const float* Wg,\n" + " const float* Wu, float* O, int M, int K, int H) {\n" + " int m = blockIdx.x*blockDim.x + threadIdx.x;\n" + " if (m >= M) return;\n" + " for (int h=0; h>>(dA,dWg,dWu,dO,M,K,H);\n" + " int ok = (cudaDeviceSynchronize()==cudaSuccess) ? 1 : 3;\n" + " if (ok==1) cudaMemcpy(hO,dO,szO,cudaMemcpyDeviceToHost);\n" + " cudaFree(dA); cudaFree(dWg); cudaFree(dWu); cudaFree(dO);\n" + " return ok;\n" + "}\n" + ) + + +def _pw_cvar(vid: str) -> str: + """A valid C identifier for a pointwise value-id.""" + return "v_" + re.sub(r"\W", "_", str(vid)) + + +def _synthesize_pointwise_cuda(region: PointwiseGraphRegion) -> str: + """CUDA source for a same-shape pointwise DAG (f32) — one thread per element. + The DAG is emitted from the ``POINTWISE_OPS`` C-expression table (topo order), + with device shims for the two ops the table names but CUDA doesn't define: + ``sign`` and ``clamp`` (the latter appears inside the ``gelu`` expression). + Both preserve NaN to match numpy (``np.sign`` / ``np.clip`` propagate NaN), + so a DAG on NaN-containing data agrees with the reference. One kernel per + region (the DAG + input count are baked in).""" + n = len(region.inputs) + params = ", ".join(f"const float* i{j}" for j in range(n)) + loads = "".join(f" float {_pw_cvar(v)} = i{j}[idx];\n" + for j, v in enumerate(region.inputs)) + body = "" + for key, ins, out in region.ops: + _arity, expr, _ref = POINTWISE_OPS[key] + body += f" float {_pw_cvar(out)} = {expr.format(*[_pw_cvar(i) for i in ins])};\n" + hparams = ", ".join(f"const float* hi{j}" for j in range(n)) + allocs = "".join( + f" float* d{j}=0; if (cudaMalloc(&d{j},sz)!=cudaSuccess) return 3;\n" + f" cudaMemcpy(d{j},hi{j},sz,cudaMemcpyHostToDevice);\n" for j in range(n)) + dargs = ", ".join(f"d{j}" for j in range(n)) + frees = " ".join(f"cudaFree(d{j});" for j in range(n)) + return ( + "#include \n" + "#include \n" + "__device__ __forceinline__ float sign(float x){ return isnan(x) ? x : (float)((x>0.0f)-(x<0.0f)); }\n" + "__device__ __forceinline__ float clamp(float x, float lo, float hi){ return isnan(x) ? x : fminf(fmaxf(x,lo),hi); }\n" + f"__global__ void {_PW_ENTRY}_kernel({params}, float* out, long numel) {{\n" + " long idx = (long)blockIdx.x*blockDim.x + threadIdx.x;\n" + " if (idx >= numel) return;\n" + f"{loads}{body}" + f" out[idx] = {_pw_cvar(region.output)};\n" + "}\n" + f'extern "C" int {_PW_ENTRY}({hparams}, float* hout, long numel) {{\n' + " size_t sz=(size_t)numel*4;\n" + f"{allocs}" + " float* dout=0; if (cudaMalloc(&dout,sz)!=cudaSuccess) return 3;\n" + " int t=256; long b=(numel+t-1)/t;\n" + f" {_PW_ENTRY}_kernel<<>>({dargs}, dout, numel);\n" + " int ok = (cudaDeviceSynchronize()==cudaSuccess) ? 1 : 3;\n" + " if (ok==1) cudaMemcpy(hout,dout,sz,cudaMemcpyDeviceToHost);\n" + f" {frees} cudaFree(dout);\n" + " return ok;\n" + "}\n" + ) + + class NvidiaCudaEmitter(KernelEmitter): target = _TARGET lang = _LANG def can_emit(self, region: Any) -> bool: - return isinstance(region, (FusedRegion, AttentionRegion)) + return isinstance(region, (FusedRegion, AttentionRegion, + GatedMatmulRegion, PointwiseGraphRegion)) def emit(self, region: Any, *, spec: SpecPolicy = SpecPolicy.BUCKET, dtype: str = "f32", dims: tuple[int, ...] | None = None) -> KernelSource: - if not isinstance(region, (FusedRegion, AttentionRegion)): + if not self.can_emit(region): raise EmitError( f"NvidiaCudaEmitter cannot emit a region of type " - f"{type(region).__name__} (FusedRegion / AttentionRegion; the " - "shipped mma.sync GEMM lane serves single matmuls via the jit " - "nvidia_mma executor)") + f"{type(region).__name__} (FusedRegion / AttentionRegion / " + "GatedMatmulRegion / PointwiseGraphRegion; the shipped mma.sync " + "GEMM lane serves single matmuls via the jit nvidia_mma executor)") if spec is SpecPolicy.DYNAMIC: raise EmitError("NvidiaCudaEmitter does not yet support SpecPolicy.DYNAMIC " "(bucket/static only)") @@ -203,6 +310,10 @@ def emit(self, region: Any, *, spec: SpecPolicy = SpecPolicy.BUCKET, raise EmitError(f"NvidiaCudaEmitter only supports f32 so far, got {dtype!r}") if isinstance(region, AttentionRegion): source, entry = _synthesize_attention_cuda(), _ATTN_ENTRY + elif isinstance(region, GatedMatmulRegion): + source, entry = _synthesize_gated_cuda(region), _GATED_ENTRY + elif isinstance(region, PointwiseGraphRegion): + source, entry = _synthesize_pointwise_cuda(region), _PW_ENTRY else: source, entry = _synthesize_fused_cuda(region), _ENTRY key = bucket_key(dims, spec, dim_names=getattr(region, "dim_names", None)) @@ -247,6 +358,16 @@ def _nvidia_cuda_compile_fn(source: KernelSource) -> str: _LIB_CACHE: dict[str, Any] = {} +def _load_lib(artifact: str): + """dlopen ``artifact`` (cached) and return the raw handle — callers bind the + entry symbol + argtypes for their own ABI.""" + lib = _LIB_CACHE.get(artifact) + if lib is None: + lib = ctypes.CDLL(artifact) + _LIB_CACHE[artifact] = lib + return lib + + def _load_entry(artifact: str): """dlopen ``artifact`` (cached) and return its bound entry symbol with the fixed C ABI: ``int(A, B, bias, residual, out, M, N, K)``.""" @@ -344,10 +465,50 @@ def run_fused_attention(self, region: Any, Q: Any, K: Any, V: Any, def run_gated_matmul_region(self, region: Any, A: Any, Wg: Any, Wu: Any, *a: Any, **k: Any) -> tuple[Any, str]: + # C5: the SwiGLU gate lane O = gate_act(A@Wg) * (A@Wu), one row per thread. + import numpy as np + try: + Af = np.ascontiguousarray(A, np.float32) + Wgf = np.ascontiguousarray(Wg, np.float32) + Wuf = np.ascontiguousarray(Wu, np.float32) + M, K = Af.shape + Kg, H = Wgf.shape + if Kg != K or Wuf.shape != (K, H): + return region.reference(A, Wg, Wu), "reference" + compiled = build(region, _TARGET, dtype="f32", dims=None) + fn = getattr(_load_lib(compiled.artifact), _GATED_ENTRY) + fn.restype = ctypes.c_int + fn.argtypes = [ctypes.c_void_p] * 4 + [ctypes.c_int] * 3 + out = np.zeros((M, H), np.float32) + rc = fn(_ptr(Af), _ptr(Wgf), _ptr(Wuf), _ptr(out), M, K, H) + if rc == 1: + return out, _REAL_TAG + except Exception: + pass return region.reference(A, Wg, Wu), "reference" def run_pointwise_graph(self, region: Any, arrays: Any, *a: Any, **k: Any) -> tuple[Any, str]: + # C5: the same-shape pointwise-DAG lane, one thread per element. + import numpy as np + try: + ins = [np.ascontiguousarray(x, np.float32) for x in arrays] + if len(ins) != len(region.inputs) or not ins: + return region.reference(*arrays), "reference" + shape = ins[0].shape + if any(x.shape != shape for x in ins): + return region.reference(*arrays), "reference" + numel = int(np.prod(shape)) if shape else 1 + compiled = build(region, _TARGET, dtype="f32", dims=None) + fn = getattr(_load_lib(compiled.artifact), _PW_ENTRY) + fn.restype = ctypes.c_int + fn.argtypes = [ctypes.c_void_p] * len(ins) + [ctypes.c_void_p, ctypes.c_long] + out = np.zeros(shape, np.float32) + rc = fn(*[_ptr(x) for x in ins], _ptr(out), numel) + if rc == 1: + return out, _REAL_TAG + except Exception: + pass return region.reference(*arrays), "reference" @@ -400,6 +561,33 @@ def run(self, region: Any, Q: Any, K: Any, V: Any, return _SHARED_RUNNER.run_fused_attention(region, Q, K, V) +class NvidiaGatedCandidate(Candidate): + """Tier-1 (C5): the synthesized SwiGLU-gate CUDA lane + (``O = gate_act(A·Wg) ⊙ (A·Wu)``). Serves any ``GatedMatmulRegion``.""" + + name = "nvidia_gated" + tier = Tier.SYNTHESIZED + target = _TARGET + op = OP_GATED_MATMUL + + def run(self, region: Any, A: Any, Wg: Any, Wu: Any, + *a: Any, **k: Any) -> tuple[Any, str]: + return _SHARED_RUNNER.run_gated_matmul_region(region, A, Wg, Wu) + + +class NvidiaPointwiseCandidate(Candidate): + """Tier-1 (C5): the synthesized same-shape pointwise-DAG CUDA lane (one thread + per element). Serves any ``PointwiseGraphRegion``.""" + + name = "nvidia_pointwise" + tier = Tier.SYNTHESIZED + target = _TARGET + op = OP_POINTWISE + + def run(self, region: Any, arrays: Any, *a: Any, **k: Any) -> tuple[Any, str]: + return _SHARED_RUNNER.run_pointwise_graph(region, arrays) + + # ── D1 matmul candidates (B1) — bare GEMM, Tier-2 emitted vs Tier-3 shipped ──── # # The arbiter enumerates these per (target="nvidia", op=matmul) and F4-gates each. @@ -493,6 +681,8 @@ def run(self, region: Any, A: Any, B: Any, *a: Any, **k: Any) -> tuple[Any, str] register_candidate(NvidiaGenericCudaCandidate()) register_candidate(NvidiaFlashAttnCandidate()) # C4: synthesized attention +register_candidate(NvidiaGatedCandidate()) # C5: SwiGLU gate +register_candidate(NvidiaPointwiseCandidate()) # C5: pointwise DAG # Bare-GEMM lanes: hand-tuned shipped (Tier 3) + compiler-emitted (Tier 2). register_candidate(NvidiaMmaGemmShippedCandidate()) register_candidate(NvidiaMmaGemmEmittedCandidate()) diff --git a/tests/unit/test_nvidia_e3_escape_hatch.py b/tests/unit/test_nvidia_e3_escape_hatch.py new file mode 100644 index 000000000..fb051fc8f --- /dev/null +++ b/tests/unit/test_nvidia_e3_escape_hatch.py @@ -0,0 +1,125 @@ +"""Workstream E3 — the arbiter escape hatch: a hand-tuned kernel is never orphaned. + +The three-tier / measured-arbiter model (Decision #28) is lead-safe only if the +crown-jewel hand-tuned lane (Tier 3) **wins by default** and can always be +**forced**. This asserts both — host-free with fake candidates, and live on the +NVIDIA shipped mma.sync GEMM (the hand-tuned Tier-3 matmul lane). +""" +from __future__ import annotations + +import os +import shutil + +import numpy as np +import pytest + +from tessera.compiler.emit import candidate as C +from tessera.compiler.emit.candidate import ( + OP_MATMUL, + ArbiterError, + Candidate, + Tier, + arbitrate, + register_candidate, + run_arbitrated, +) + +_TGT = "e3_faketarget" + + +class _FakeReg: + dtype = "bfloat16" + + def reference(self, A, B): + return np.asarray(A, np.float32) @ np.asarray(B, np.float32) + + +class _FakeCand(Candidate): + op = OP_MATMUL + target = _TGT + + def __init__(self, name, tier, tag): + self.name = name + self.tier = tier + self._tag = tag + + def run(self, region, A, B, *a, **k): + return region.reference(A, B), self._tag + + +# ── host-free: the force/tier contract with fake candidates ────────────────── + +def test_e3_hand_tuned_wins_by_default_and_is_forceable(): + hand = _FakeCand("e3_handtuned", Tier.HAND_TUNED, "handtuned_real") + synth = _FakeCand("e3_synth", Tier.SYNTHESIZED, "synth_real") + register_candidate(synth) # register the lower tier FIRST, so a default + register_candidate(hand) # win by tier (not registration order) is real + region = _FakeReg() + A = np.zeros((2, 2), np.float32) + B = np.zeros((2, 2), np.float32) + + # 1) the hand-tuned lane wins by DEFAULT (tier priority) — never orphaned. + win = arbitrate(region, OP_MATMUL, _TGT, verify=False) + assert win is not None and win.name == "e3_handtuned" + + # 2) it CAN be forced, and runs. + _, tag = run_arbitrated(region, OP_MATMUL, _TGT, A, B, + verify=False, force="e3_handtuned") + assert tag == "handtuned_real" + + # 3) a LOWER tier is also forceable — tier priority is overridable (the D2/E3 + # seam a measured cost model uses to displace the default). + _, tag2 = run_arbitrated(region, OP_MATMUL, _TGT, A, B, + verify=False, force="e3_synth") + assert tag2 == "synth_real" + + # 4) forcing an unknown candidate raises honestly (never silently picks another). + with pytest.raises(ArbiterError, match="not available"): + arbitrate(region, OP_MATMUL, _TGT, verify=False, force="e3_nonexistent") + + +# ── live: the NVIDIA shipped mma.sync GEMM is the hand-tuned Tier-3 lane ────── + +def _nvidia_matmul_live() -> bool: + if not (shutil.which("nvcc") or os.path.exists("/usr/local/cuda/bin/nvcc")): + return False + try: + from tessera import runtime as rt + return (rt._nvidia_mma_runtime_available() + and rt._load_nvidia_ptx_launch() is not None) + except Exception: + return False + + +@pytest.mark.slow +@pytest.mark.skipif(not _nvidia_matmul_live(), + reason="live NVIDIA GPU + shipped GEMM + PTX launch bridge required") +def test_e3_live_nvidia_hand_tuned_forced_and_wins(): + import tessera.compiler.emit.nvidia_cuda # noqa: F401 — registers the candidates + from tessera.compiler.fusion_core import MatmulRegion + region = MatmulRegion(dtype="bfloat16") + rng = np.random.default_rng(0) + A = (rng.standard_normal((32, 32)) * 0.4).astype(np.float32) + B = (rng.standard_normal((32, 16)) * 0.4).astype(np.float32) + ref = region.reference(A, B) + + # the shipped hand-tuned GEMM (Tier 3) wins by default over the emitted Tier-2. + win = arbitrate(region, OP_MATMUL, "nvidia") + assert win is not None and win.name == "nvidia_mma_gemm_shipped" + assert int(win.tier) == int(Tier.HAND_TUNED) + + # forceable + runs on-GPU. + out, tag = run_arbitrated(region, OP_MATMUL, "nvidia", A, B, + force="nvidia_mma_gemm_shipped") + assert tag == "nvidia_mma_shipped" + np.testing.assert_allclose(out, ref, atol=5e-3, rtol=0) + + # the emitted Tier-2 lane is forceable too (never orphaned either way). + out2, tag2 = run_arbitrated(region, OP_MATMUL, "nvidia", A, B, + force="nvidia_mma_gemm_emitted") + assert tag2 == "nvidia_ptx_gemm" + np.testing.assert_allclose(out2, ref, atol=5e-3, rtol=0) + + # forcing an unavailable candidate raises honestly. + with pytest.raises(ArbiterError): + arbitrate(region, OP_MATMUL, "nvidia", force="nvidia_does_not_exist") diff --git a/tests/unit/test_nvidia_plugin.py b/tests/unit/test_nvidia_plugin.py index f0fca1422..4741907a2 100644 --- a/tests/unit/test_nvidia_plugin.py +++ b/tests/unit/test_nvidia_plugin.py @@ -2,17 +2,15 @@ Two layers: -1. **Registration + emit + decline paths (host-free)** — a full three-seam plugin - for target "nvidia": the emitter turns a FusedRegion into CUDA source; - attention / gated / pointwise regions (no fused CUDA kernel yet) decline to the - numpy reference; unsupported regions/policies/dtypes raise EmitError. -2. **Live gate (needs a live NVIDIA GPU + nvcc)** — the generically-synthesized - CUDA FusedRegion kernel compiles with nvcc, runs on-device ("nvidia_cuda"), - matches numpy (f32), and passes the same universal F4 oracle as ROCm/x86. - -The generic CUDA lane is the Tier-1 candidate; NVIDIA has no fused hand-tuned -FusedRegion kernel to register as Tier-3 yet (the shipped mma.sync GEMM is a pure -matmul, served by the jit nvidia_mma executor). +1. **Registration + emit paths (host-free)** — a full three-seam plugin for + target "nvidia": the emitter turns fused / attention / gated / pointwise regions + into CUDA source; unsupported regions/policies/dtypes raise EmitError. +2. **Live gates (needs a live NVIDIA GPU + nvcc)** — the generically-synthesized + CUDA kernels compile with nvcc, run on-device ("nvidia_cuda"), match numpy, and + pass the same universal F4 oracle as ROCm/x86. The generic lane covers all four + fusion_core region kinds (FusedRegion, AttentionRegion — C4, GatedMatmulRegion + + PointwiseGraphRegion — C5); bare GEMM is served by the B1 matmul candidates + (shipped + emitted mma.sync). """ from __future__ import annotations @@ -70,22 +68,28 @@ def test_nvidia_emit_is_deterministic(): def test_nvidia_emitter_rejects_unsupported(): e = get_emitter("nvidia") with pytest.raises(EmitError, match="cannot emit"): - e.emit(F.GatedMatmulRegion()) # FusedRegion/AttentionRegion only + e.emit(F.MatmulRegion()) # bare GEMM → the GEMM candidates with pytest.raises(EmitError, match="DYNAMIC"): e.emit(F.FusedRegion(epilogue=("relu",)), spec=SpecPolicy.DYNAMIC) with pytest.raises(EmitError, match="f32"): e.emit(F.FusedRegion(epilogue=("relu",)), dtype="f16") -def test_nvidia_declines_gated_and_pointwise(): - # No single fused CUDA kernel for gated / pointwise yet — always the reference. - # (Attention IS implemented — the flash lane; see the C4 tests below.) - r = get_runner("nvidia") - A = np.zeros((8, 12), np.float32) - _, ex = r.run_gated_matmul_region(F.GatedMatmulRegion(), - A, np.zeros((12, 16), np.float32), - np.zeros((12, 16), np.float32)) - assert ex == "reference" +def test_nvidia_c5_candidates_registered_and_emit(): + # C5: the generic lane now covers gated (SwiGLU) + pointwise-DAG region kinds + # too — Tier-1 candidates + emitter support (attention landed in C4). + from tessera.compiler.emit.candidate import OP_GATED_MATMUL, OP_POINTWISE + gated = {c.name for c in C.candidates_for("nvidia", OP_GATED_MATMUL)} + pw = {c.name for c in C.candidates_for("nvidia", OP_POINTWISE)} + assert "nvidia_gated" in gated and "nvidia_pointwise" in pw + e = get_emitter("nvidia") + gsrc = e.emit(F.GatedMatmulRegion(gate_act="silu")).source + assert "tessera_nvidia_gated" in gsrc and "<<<" in gsrc + region = F.PointwiseGraphRegion( + ops=(("add", ("a", "b"), "s"), ("relu", ("s",), "o")), + inputs=("a", "b"), output="o") + psrc = e.emit(region).source + assert "tessera_nvidia_pointwise" in psrc and "sign" in psrc # sign shim def test_nvidia_flash_attn_candidate_registered_and_emits(): @@ -359,3 +363,81 @@ def test_live_nvidia_flash_attention(scale, causal, shape): aout, atag = C.run_arbitrated(region, OP_ATTENTION, "nvidia", Q, K, V) assert atag == "nvidia_cuda" np.testing.assert_allclose(aout, region.reference(Q, K, V), atol=1e-4) + + +@pytest.mark.slow +@pytest.mark.skipif(not _nvidia_cuda_live(), + reason="live NVIDIA GPU + nvcc required") +@pytest.mark.parametrize("act", ["silu", "gelu", "relu"]) +def test_live_nvidia_gated_swiglu(act): + # C5: the synthesized SwiGLU-gate lane executes on-GPU and matches. + F.clear_verification_cache() + region = F.GatedMatmulRegion(gate_act=act) + rng = np.random.default_rng(hash(act) % 1000) + A = rng.standard_normal((16, 32)).astype(np.float32) + Wg = rng.standard_normal((32, 24)).astype(np.float32) + Wu = rng.standard_normal((32, 24)).astype(np.float32) + out, tag = get_runner("nvidia").run_gated_matmul_region(region, A, Wg, Wu) + assert tag == "nvidia_cuda" + np.testing.assert_allclose(out, region.reference(A, Wg, Wu), atol=1e-3) + assert F.verify_synthesized_gated(region, runner=get_runner("nvidia"), + force=True) is True + + +@pytest.mark.slow +@pytest.mark.skipif(not _nvidia_cuda_live(), + reason="live NVIDIA GPU + nvcc required") +def test_live_nvidia_pointwise_dag(): + # C5: the synthesized pointwise-DAG lane — d = relu(a+b) * sigmoid(c) — runs + # on-GPU (from the POINTWISE_OPS C-expr table) and matches. + F.clear_verification_cache() + region = F.PointwiseGraphRegion( + ops=(("add", ("a", "b"), "s"), ("relu", ("s",), "ra"), + ("sigmoid", ("c",), "sc"), ("mul", ("ra", "sc"), "o")), + inputs=("a", "b", "c"), output="o") + rng = np.random.default_rng(0) + arrs = [rng.standard_normal((6, 7)).astype(np.float32) for _ in range(3)] + out, tag = get_runner("nvidia").run_pointwise_graph(region, arrs) + assert tag == "nvidia_cuda" + np.testing.assert_allclose(out, region.reference(*arrs), atol=1e-5) + assert F.verify_synthesized_pointwise(region, runner=get_runner("nvidia"), + force=True) is True + + +def test_nvidia_pointwise_emits_gelu_and_nan_safe_shims(): + # PR #297 review: gelu's POINTWISE_OPS template uses clamp() (no CUDA builtin), + # and sign must preserve NaN like np.sign. The emitted source defines both + # shims, NaN-aware. + dag = F.PointwiseGraphRegion(ops=(("add", ("a", "b"), "s"), + ("gelu", ("s",), "o")), + inputs=("a", "b"), output="o") + src = get_emitter("nvidia").emit(dag).source + assert "float clamp(float" in src and "float sign(float" in src + assert "isnan(x) ? x" in src # NaN-preserving + + +@pytest.mark.slow +@pytest.mark.skipif(not _nvidia_cuda_live(), + reason="live NVIDIA GPU + nvcc required") +def test_live_nvidia_pointwise_gelu_and_nan_sign(): + # PR #297 review, on-GPU: a gelu DAG must COMPILE + run (was failing on the + # undefined clamp), and a sign chain must preserve NaN (np.sign semantics). + r = get_runner("nvidia") + gelu = F.PointwiseGraphRegion(ops=(("add", ("a", "b"), "s"), + ("gelu", ("s",), "o")), + inputs=("a", "b"), output="o") + rng = np.random.default_rng(1) + a = rng.standard_normal((5, 6)).astype(np.float32) + b = rng.standard_normal((5, 6)).astype(np.float32) + out, tag = r.run_pointwise_graph(gelu, [a, b]) + assert tag == "nvidia_cuda" # compiled + ran (not declined) + np.testing.assert_allclose(out, gelu.reference(a, b), atol=1e-4) + + sign = F.PointwiseGraphRegion(ops=(("abs", ("a",), "aa"), ("sign", ("aa",), "o")), + inputs=("a",), output="o") + x = np.array([[1.0, -2.0, np.nan, 0.0]], np.float32) + sout, stag = r.run_pointwise_graph(sign, [x]) + assert stag == "nvidia_cuda" + ref = sign.reference(x) + np.testing.assert_array_equal(np.isnan(sout), np.isnan(ref)) # NaN preserved + np.testing.assert_allclose(sout[~np.isnan(sout)], ref[~np.isnan(ref)], atol=1e-6)