Correct M attention precision and validate every temporal transformer block

tools/windows/cuda_mma_probe.cpp

@@ -0,0 +1,165 @@
+#include "hip_api.h"
+#include "cuda_api.h"
+#include <array>
+#include <cmath>
+#include <cstring>
+
+// Public synthetic matrix fragments, independent of NVIDIA DLLs and weights.
+static constexpr char source[] = R"PTX(
+.version 7.0
+.target sm_75
+.address_size 64
+.visible .entry d4r_f16_mma_reference(
+ .param .u64 A, .param .u64 B, .param .u64 C, .param .u64 D)
+{
+ .reg .b32 %r<24>;
+ .reg .b64 %rd<12>;
+ .reg .b16 %h<4>;
+ ld.param.u64 %rd0, [A];
+ ld.param.u64 %rd1, [B];
+ ld.param.u64 %rd2, [C];
+ ld.param.u64 %rd3, [D];
+ mov.u32 %r0, %ctaid.x;
+ mov.u32 %r1, %tid.x;
+ shr.u32 %r2, %r1, 2;
+ and.b32 %r3, %r1, 3;
+ shl.b32 %r3, %r3, 1;
+ mul.wide.u32 %rd4, %r0, 512;
+ add.u64 %rd0, %rd0, %rd4;
+ mul.wide.u32 %rd4, %r0, 256;
+ add.u64 %rd1, %rd1, %rd4;
+ add.u64 %rd2, %rd2, %rd4;
+ add.u64 %rd3, %rd3, %rd4;
+ mad.lo.u32 %r4, %r2, 16, %r3;
+ mul.wide.u32 %rd4, %r4, 2;
+ add.u64 %rd5, %rd0, %rd4;
+ ld.global.u32 %r8, [%rd5];
+ ld.global.u32 %r9, [%rd5+256];
+ ld.global.u32 %r10, [%rd5+16];
+ ld.global.u32 %r11, [%rd5+272];
+ mad.lo.u32 %r4, %r3, 8, %r2;
+ mul.wide.u32 %rd4, %r4, 2;
+ add.u64 %rd5, %rd1, %rd4;
+ ld.global.b16 %h0, [%rd5];
+ ld.global.b16 %h1, [%rd5+16];
+ mov.b32 %r12, {%h0,%h1};
+ ld.global.b16 %h2, [%rd5+128];
+ ld.global.b16 %h3, [%rd5+144];
+ mov.b32 %r13, {%h2,%h3};
+ mad.lo.u32 %r4, %r2, 8, %r3;
+ mul.wide.u32 %rd4, %r4, 2;
+ add.u64 %rd5, %rd2, %rd4;
+ ld.global.u32 %r14, [%rd5];
+ ld.global.u32 %r15, [%rd5+128];
+ mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16
+  {%r16,%r17}, {%r8,%r9,%r10,%r11}, {%r12,%r13}, {%r14,%r15};
+ add.u64 %rd5, %rd3, %rd4;
+ st.global.u32 [%rd5], %r16;
+ st.global.u32 [%rd5+128], %r17;
+ ret;
+}
+)PTX";
+
+static uint16_t half_bits(double value) {
+    _Float16 half = static_cast<_Float16>(value);
+    uint16_t bits;
+    std::memcpy(&bits, &half, sizeof(bits));
+    return bits;
+}
+static double half_value(uint16_t bits) {
+    _Float16 half;
+    std::memcpy(&half, &bits, sizeof(bits));
+    return static_cast<double>(half);
+}
+
+int main(int argc, char** argv) {
+    using namespace d4r::diag;
+    start();
+    try {
+        Args args(argc, argv);
+        HipApi hip(args.hip_root);
+        hipDeviceProp_t props{};
+        hip.select_gfx1201(args.device, props);
+        CudaApi cuda(args.cuda_dll);
+        cuda.check(cuda.cuInit(0), "cuInit");
+        CUdevice device = -1;
+        int count = 0;
+        cuda.check(cuda.cuDeviceGetCount(&count), "cuDeviceGetCount");
+        for (int i = 0; i < count; ++i) {
+            CUdevice candidate;
+            int bus, pci, domain;
+            cuda.check(cuda.cuDeviceGet(&candidate, i), "cuDeviceGet");
+            cuda.check(cuda.cuDeviceGetAttribute(&bus, 33, candidate), "PCI bus");
+            cuda.check(cuda.cuDeviceGetAttribute(&pci, 34, candidate), "PCI device");
+            cuda.check(cuda.cuDeviceGetAttribute(&domain, 50, candidate), "PCI domain");
+            if (bus == props.pciBusID && pci == props.pciDeviceID && domain == props.pciDomainID) device = candidate;
+        }
+        if (device < 0) throw std::runtime_error("CUDA/HIP PCI identity mismatch");
+        CUcontext context = nullptr;
+        cuda.check(cuda.cuDevicePrimaryCtxRetain(&context, device), "primary context");
+        struct Context { CudaApi& api; CUdevice device; ~Context() { api.cuCtxSetCurrent(nullptr); api.cuDevicePrimaryCtxRelease_v2(device); } } context_owner{cuda, device};
+        cuda.check(cuda.cuCtxSetCurrent(context), "cuCtxSetCurrent");
+        CUmodule module = nullptr;
+        cuda.check(cuda.cuModuleLoadData(&module, source), "MMA PTX JIT");
+        struct Module { CudaApi& api; CUmodule module; ~Module() { api.cuModuleUnload(module); } } module_owner{cuda, module};
+        CUfunction kernel = nullptr;
+        cuda.check(cuda.cuModuleGetFunction(&kernel, module, "d4r_f16_mma_reference"), "MMA kernel");
+        constexpr size_t cases = 512;
+        std::vector<uint16_t> a(cases * 256), b(cases * 128), c(cases * 128), d(cases * 128);
+        uint32_t random = 0x192a321u;
+        auto next = [&]() { random = random * 1664525u + 1013904223u; return random; };
+        for (size_t n = 0; n < cases; ++n) {
+            for (size_t i = 0; i < 256; ++i) a[n * 256 + i] = uint16_t(0x2800u + (next() >> 20));
+            for (size_t i = 0; i < 128; ++i) {
+                b[n * 128 + i] = uint16_t((next() & 0x8000u) | (0x3000u + (next() >> 20)));
+                c[n * 128 + i] = half_bits(n % 3 ? 0. : 0.015625);
+            }
+            // Sweep large cancellation terms and a tiny surviving product
+            // across K positions; this exposes loss hidden by integer fixtures.
+            if (n < 256) {
+                for (size_t i = 0; i < 256; ++i) a[n * 256 + i] = half_bits(1.);
+                std::fill(b.begin() + n * 128, b.begin() + (n + 1) * 128, half_bits(0.));
+                for (size_t col = 0; col < 8; ++col) {
+                    b[n * 128 + (n % 16) * 8 + col] = half_bits(32768.);
+                    b[n * 128 + ((n % 16 + 1) % 16) * 8 + col] = half_bits(-32768.);
+                    b[n * 128 + ((n % 16 + 2 + n / 16 % 14) % 16) * 8 + col] = half_bits(0.0009765625);
+                }
+            }
+        }
+        struct Allocation {
+            CudaApi& api; CUdeviceptr pointer = 0;
+            Allocation(CudaApi& api, size_t bytes) : api(api) { api.check(api.cuMemAlloc_v2(&pointer, bytes), "MMA allocate"); }
+            ~Allocation() { if (pointer) api.cuMemFree_v2(pointer); }
+        } da(cuda, a.size() * 2), db(cuda, b.size() * 2), dc(cuda, c.size() * 2), dd(cuda, d.size() * 2);
+        cuda.check(cuda.cuMemcpyHtoD_v2(da.pointer, a.data(), a.size() * 2), "MMA A upload");
+        cuda.check(cuda.cuMemcpyHtoD_v2(db.pointer, b.data(), b.size() * 2), "MMA B upload");
+        cuda.check(cuda.cuMemcpyHtoD_v2(dc.pointer, c.data(), c.size() * 2), "MMA C upload");
+        std::fill(d.begin(), d.end(), uint16_t(0x7e00));
+        cuda.check(cuda.cuMemcpyHtoD_v2(dd.pointer, d.data(), d.size() * 2), "MMA NaN sentinel");
+        void* parameters[] = {&da.pointer, &db.pointer, &dc.pointer, &dd.pointer};
+        cuda.check(cuda.cuLaunchKernel(kernel, cases, 1, 1, 32, 1, 1, 0, nullptr, parameters, nullptr), "MMA launch");
+        cuda.check(cuda.cuCtxSynchronize(), "MMA synchronize");
+        cuda.check(cuda.cuMemcpyDtoH_v2(d.data(), dd.pointer, d.size() * 2), "MMA download");
+        size_t mismatch = 0;
+        for (size_t n = 0; n < cases; ++n)
+            for (size_t row = 0; row < 16; ++row)
+                for (size_t col = 0; col < 8; ++col) {
+                    double sum = half_value(c[n * 128 + row * 8 + col]);
+                    for (size_t k = 0; k < 16; ++k)
+                        sum += half_value(a[n * 256 + row * 16 + k]) * half_value(b[n * 128 + k * 8 + col]);
+                    const uint16_t expected = half_bits(sum);
+                    const uint16_t actual = d[n * 128 + row * 8 + col];
+                    if (expected != actual) {
+                        if (mismatch++ < 8) std::fprintf(stderr, "MMA_DIFFERENT case=%zu row=%zu column=%zu expected=%04x actual=%04x\n", n, row, col, expected, actual);
+                    }
+                }
+        std::printf("MMA_REFERENCE cases=%zu elements=%zu differing=%zu\n", cases, d.size(), mismatch);
+        if (mismatch) return 5;
+        std::printf("PASS CUDA_F16_MMA architecture=gfx1201 cases=%zu exact_f16=1\n", cases);
+        return 0;
+    } catch (const std::exception& error) {
+        std::fprintf(stderr, "FAIL CUDA_F16_MMA %s\n", error.what());
+        loaded_modules();
+        return 4;
+    }
+}

tools/windows/m_capture_compare.py

@@ -0,0 +1,54 @@
+"""Compare logical M inputs/outputs between private temporal capture sequences."""
+import argparse
+import pathlib
+
+import numpy as np
+
+from m_replay_validate import me, LAYERS
+
+
+def logical(parameters, pointer, count, output=False, directory=None):
+    base, allocation = parameters.alloc_of(pointer)
+    if output:
+        identifier = next(i for i, (b, _) in parameters.allocs.items() if b == base)
+        allocation = np.fromfile(directory / 'gpu-output' / f'alloc-{identifier}.bin', np.uint8)
+    return allocation[pointer - base:pointer - base + count]
+
+
+def compare(reference, actual):
+    paths = []
+    order = ['enc1', 'enc2', 'enc3_tube', 'dec2', 'dec1']
+    for path in reference.glob('replay-*-rrlite_*_4x4'):
+        layer = path.name.split('rrlite_', 1)[1][:-4]
+        ordinal = int(path.name.split('-')[1])
+        frame = (ordinal - 1) // (6 if layer == 'enc3_tube' else 1)
+        paths.append((frame, order.index(layer), ordinal, path, layer))
+    for frame, _, ordinal, path, layer in sorted(paths):
+        other = actual / path.name
+        a, b = me.Params(str(path)), me.Params(str(other))
+        c, _, merged, cin = LAYERS[layer]
+        if (a.W, a.H, a.sx, a.sy, a.grid) != (b.W, b.H, b.sx, b.sy, b.grid):
+            raise RuntimeError(f'Launch geometry differs: {path.name}')
+        fields = [('input', 'inp', (a.W // 2) * (a.H // 2) * cin if cin else a.W * a.H * c, False)]
+        if cin:
+            fields.append(('skip', 'p32', a.W * a.H * c, False))
+        fields.append(('output', 'out', a.W * a.H * c, True))
+        if merged:
+            fields.append(('merge', 'p48', (a.W // 2) * (a.H // 2) * merged, True))
+        for label, field, count, output in fields:
+            x = logical(a, getattr(a, field), count, output, path)
+            y = logical(b, getattr(b, field), count, output, other)
+            if x.size != count or y.size != count:
+                raise RuntimeError(f'Truncated {path.name} {label}')
+            differing = np.flatnonzero(x != y)
+            print(f'M_CAPTURE_COMPARE frame={frame} layer={layer} ordinal={ordinal} field={label} '
+                  f'codes={count} differing={differing.size} '
+                  f'first={int(differing[0]) if differing.size else -1}', flush=True)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--reference', type=pathlib.Path, required=True)
+    parser.add_argument('--actual', type=pathlib.Path, required=True)
+    args = parser.parse_args()
+    compare(args.reference, args.actual)

tools/windows/m_ptx_validate.py

@@ -18,7 +18,10 @@ def validate(args):
         raise RuntimeError('Capture kernel does not match requested M layer')
     native = {i: np.fromfile(args.native_dir / f'alloc-{i}.bin', np.uint8) for i in parameters.allocs}
     candidates = []
-    for bx, by in blocks(parameters.grid, 12):
+    coordinates = args.block or blocks(parameters.grid, args.candidate_blocks)
+    for bx, by in coordinates:
+        if not (0 <= bx < parameters.grid[0] and 0 <= by < parameters.grid[1]):
+            raise RuntimeError(f'Block {bx},{by} is outside the capture grid')
         model = model_block(parameters, args.layer, bx, by)
         expected, actual = stored(parameters, args.layer, bx, by, native, model)
         difference = float(np.abs(expected - actual).max())
@@ -52,8 +55,12 @@ if __name__ == '__main__':
     parser.add_argument('--dlss-dll', type=pathlib.Path, required=True)
     parser.add_argument('--output-dir', type=pathlib.Path, required=True)
     parser.add_argument('--blocks', type=int, default=2)
+    parser.add_argument('--candidate-blocks', type=int, default=12,
+                        help='NumPy scan size used to choose the worst blocks for independent PTX')
+    parser.add_argument('--block', type=lambda value: tuple(map(int, value.split(','))),
+                        action='append', help='Explicit x,y regression block (repeatable)')
     args = parser.parse_args()
-    if args.blocks < 1:
+    if args.blocks < 1 or args.candidate_blocks < 1:
         parser.error('--blocks must be positive')
     try:
         validate(args)

tools/windows/m_replay_validate.py

@@ -142,7 +142,7 @@ def metrics(expected, actual, label):
     return maximum, peak, psnr, steps
 
 
-def validate(layer, directory, output, block_count, real_capture):
+def validate(layer, directory, output, block_count, real_capture, exact=False):
     parameters = me.Params(str(directory))
     if parameters.kernel != f'rrlite_{layer}_4x4':
         raise RuntimeError('Replay kernel does not match requested M layer')
@@ -151,7 +151,7 @@ def validate(layer, directory, output, block_count, real_capture):
              for bx, by in blocks(parameters.grid, block_count)]
     expected, actual = (np.concatenate(values) for values in zip(*pairs))
     maximum, peak, psnr, steps = metrics(expected, actual, f'M_REPLAY_REFERENCE kernel={layer}')
-    if (not real_capture and maximum != 0) or (real_capture and (steps.max() > 1 or psnr < 50)):
+    if ((exact or not real_capture) and maximum != 0) or (real_capture and (steps.max() > 1 or psnr < 50)):
         raise RuntimeError(f'M mismatch: exact synthetic / <=1 FP8 step and >=50 dB real gate; peak={peak}')
     print(f'PASS M_REPLAY_REFERENCE kernel={layer} blocks={len(pairs)} architecture=gfx1201 nan_inf=0')
 
@@ -165,6 +165,7 @@ if __name__ == '__main__':
     parser.add_argument('--size', type=int, choices=[8, 16, 24, 32], default=16)
     parser.add_argument('--blocks', type=int, default=12)
     parser.add_argument('--real-capture', action='store_true')
+    parser.add_argument('--exact', action='store_true', help='Require every stored real FP8 value to match')
     args = parser.parse_args()
     try:
         if args.mode == 'generate':
@@ -172,7 +173,7 @@ if __name__ == '__main__':
         else:
             if not args.output_dir or args.blocks < 1:
                 parser.error('validate requires --output-dir and positive --blocks')
-            validate(args.layer, args.fixture_dir, args.output_dir, args.blocks, args.real_capture)
+            validate(args.layer, args.fixture_dir, args.output_dir, args.blocks, args.real_capture, args.exact)
     except Exception as error:
         print(f'FAIL M_REPLAY {error}', file=sys.stderr)
         raise SystemExit(5)

tools/windows/m_temporal_validate.py

@@ -0,0 +1,53 @@
+"""Validate every captured M launch and spatial block; all allocations stay private."""
+import argparse
+import contextlib
+import json
+import pathlib
+import sys
+import traceback
+
+from m_replay_validate import LAYERS, validate
+
+
+def main(args):
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    summary = {'passed': False, 'captured_outputs': True, 'exact': args.exact, 'launches': []}
+    status = 1
+    try:
+        captures = sorted(args.capture_dir.glob('replay-*-rrlite_*_4x4'))
+        if not captures:
+            raise RuntimeError('No audited M captures were found')
+        for capture in captures:
+            layer = capture.name.split('rrlite_', 1)[1][:-4]
+            if layer not in LAYERS:
+                raise RuntimeError(f'Unknown captured M layer: {layer}')
+            logfile = args.output_dir / (capture.name + '.log')
+            with logfile.open('w', encoding='utf-8') as output, contextlib.redirect_stdout(output):
+                validate(layer, capture, capture / 'gpu-output', args.blocks, True, args.exact)
+            print(capture.name + ' ' + logfile.read_text().splitlines()[0], flush=True)
+            summary['launches'].append({'capture': capture.name, 'layer': layer, 'passed': True})
+        present = {item['layer'] for item in summary['launches']}
+        if present != set(LAYERS):
+            raise RuntimeError('Capture set does not contain all five M layers')
+        summary['passed'] = True
+        print(f'PASS M_TEMPORAL_REFERENCE launches={len(captures)} blocks={args.blocks} exact={int(args.exact)}', flush=True)
+        status = 0
+    except Exception as error:
+        summary['error'] = str(error)
+        summary['failed_capture'] = capture.name if 'capture' in locals() else None
+        traceback.print_exc()
+    finally:
+        (args.output_dir / 'summary.json').write_text(json.dumps(summary, indent=2) + '\n', encoding='utf-8')
+    return status
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--capture-dir', type=pathlib.Path, required=True)
+    parser.add_argument('--output-dir', type=pathlib.Path, required=True)
+    parser.add_argument('--blocks', type=int, default=10000)
+    parser.add_argument('--exact', action='store_true')
+    args = parser.parse_args()
+    if args.blocks < 1:
+        parser.error('--blocks must be positive')
+    sys.exit(main(args))

tools/windows/native_replay_probe.cpp

@@ -82,6 +82,14 @@ int main(int argc, char** argv) {
         hipModule_t module = nullptr;
         hip.check(hip.hipModuleLoad(&module, args.module.c_str()), "hipModuleLoad(replay gfx1201)");
         struct Cleanup { HipApi& hip; hipModule_t module; ~Cleanup() { (void)hip.hipModuleUnload(module); } } cleanup{hip, module};
+        if (const char* debugBlock = std::getenv("D4R_SWIN_DEBUG_BLOCK")) {
+            int coordinates[2]{};
+            if (std::sscanf(debugBlock, "%d,%d", &coordinates[0], &coordinates[1]) != 2) throw std::runtime_error("D4R_SWIN_DEBUG_BLOCK must be x,y");
+            hipDeviceptr_t address = nullptr; size_t size = 0;
+            hip.check(hip.hipModuleGetGlobal(&address, &size, module, "d4r_swin_debug_block"), "Find Swin diagnostic block");
+            if (size != sizeof(coordinates)) throw std::runtime_error("Swin diagnostic coordinates ABI");
+            hip.check(hip.hipMemcpy(address, coordinates, size, hipMemcpyHostToDevice), "Set Swin diagnostic block");
+        }
         auto scalar = [&](const char* name, unsigned fallback) {
             hipDeviceptr_t address = nullptr; size_t size = 0;
             if (hip.hipModuleGetGlobal(&address, &size, module, name) != hipSuccess) return fallback;
@@ -108,6 +116,15 @@ int main(int argc, char** argv) {
                 nullptr, launchArguments, nullptr), "hipModuleLaunchKernel(replay)");
         hip.check(hip.hipDeviceSynchronize(), "hipDeviceSynchronize(replay)");
         std::filesystem::create_directories(output);
+        if (std::getenv("D4R_SWIN_DEBUG_BLOCK")) {
+            hipDeviceptr_t address = nullptr; size_t size = 0;
+            hip.check(hip.hipModuleGetGlobal(&address, &size, module, "d4r_swin_debug_values"), "Find Swin diagnostic stages");
+            std::vector<uint8_t> bytes(size);
+            hip.check(hip.hipMemcpy(bytes.data(), address, size, hipMemcpyDeviceToHost), "Download Swin diagnostic stages");
+            std::ofstream file(output / "swin-stages.f16", std::ios::binary);
+            file.write(reinterpret_cast<const char*>(bytes.data()), std::streamsize(bytes.size()));
+            if (!file) throw std::runtime_error("Cannot save Swin diagnostic stages");
+        }
         for (const auto& allocation : allocations) {
             std::vector<uint8_t> bytes(allocation->size);
             hip.check(hip.hipMemcpy(bytes.data(), allocation->device, bytes.size(), hipMemcpyDeviceToHost), "hipMemcpy(replay result)");

tools/windows/swin_debug_compare.py

@@ -0,0 +1,26 @@
+"""Compare compile-time Swin diagnostic snapshots with the NumPy intermediate values."""
+import argparse
+import numpy as np
+from m_replay_validate import me, sm, LAYERS, metrics
+
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument('--capture-dir', required=True)
+parser.add_argument('--stages', required=True)
+parser.add_argument('--block', default='0,4')
+parser.add_argument('--layer', choices=list(LAYERS), default='enc1')
+args = parser.parse_args()
+p = me.Params(args.capture_dir)
+c, heads, _, cin = LAYERS[args.layer]
+bx, by = map(int, args.block.split(','))
+pre = 4 * c * cin + 8 * c if cin else 0
+x0 = sm.decoder_input(p, bx, by, c, cin, heads, p.p32) if cin else sm.load_tokens(p, bx, by, c // 32)
+stages = {}
+sm.swin_block(p, x0, c, heads, pre, stages=stages)
+actual = np.fromfile(args.stages, np.float16).reshape(4, 64, 128).astype(np.float64)
+for index, name in enumerate(('h1', 'x1', 'h2', 'O')):
+    expected = stages[name]
+    values = actual[index, :, :expected.shape[1]]
+    metrics(expected, values, f'SWIN_STAGE {name}')
+    different = np.argwhere(expected != values)
+    for row, channel in different[:8]:
+        print(f'DIFFERENT stage={name} token={row} channel={channel} reference={expected[row, channel]} native={values[row, channel]}')