DGX Spark serving config: source, kernels, receipts, validated env (part 3)

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma128n128k4.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: kitchen 128x128x256 single-N tile == cublasLt.
+
+Profiler name: cutlass3x_sm120_..._128x128x256_1x1x1. Pairing both
+arms doubles B smem; this tile holds one N operand. Same m16n8k64.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_nvfp4_scaled_tma128n128k4,
+    fc1_paired_nvfp4_scaled_tma256k2,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081250)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    full = fc1_nvfp4_scaled_tma128n128k4(qx_live, qxs, qw, qws, alpha)
+    payload = {
+        "identity": (
+            "Kitchen 128x128x256 single-N tile == cublasLt on the "
+            "m16n8k64 atom. One B operand (not paired arms). 2-stage "
+            "K=256. Profiler: 128x128x256_1x1x1 s16864."
+        ),
+        "kitchen_kernel": (
+            "cutlass3x_sm120_bstensorop_s16864gemm_block_scaled_"
+            "ue4m3xe2m1_ue4m3xe2m1_f32_bf16_bf16_128x128x256_1x1x1"
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "kit_vs_kitchen_mismatches": byte_diff(full, kitchen),
+        "kit_vs_kitchen_max_abs": round(max_abs_diff(full, kitchen), 6),
+        "full_finite": bool(torch.isfinite(full.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "kit_full_min_ms": round(
+            time_ms(
+                lambda: fc1_nvfp4_scaled_tma128n128k4(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "k2_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["kit_vs_kitchen_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["kit_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma128n128k4_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma128n128k4_20423.json

@@ -0,0 +1,16 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "Kitchen 128x128x256 single-N tile == cublasLt on the m16n8k64 atom. One B operand (not paired arms). 2-stage K=256. Profiler: 128x128x256_1x1x1 s16864.",
+  "k": 5376,
+  "k2_prod_min_ms": 37.0229,
+  "kit_full_min_ms": 67.7126,
+  "kit_vs_kitchen_max_abs": 0.0,
+  "kit_vs_kitchen_mismatches": 0,
+  "kitchen_gemm_min_ms": 20.1805,
+  "kitchen_kernel": "cutlass3x_sm120_bstensorop_s16864gemm_block_scaled_ue4m3xe2m1_ue4m3xe2m1_f32_bf16_bf16_128x128x256_1x1x1",
+  "n": 28672,
+  "pass": true,
+  "rows": 20423,
+  "vs_kitchen_ms": -47.5321
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma128n128k4_ldm.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: ldmatrix.x4 A fragment == kitchen group/tidg 4xu32.
+
+Probe: ldmatrix from smem[lane&15][(lane>>4)*16] matches the scalar
+kitchen A map (0/32 lanes). Same 128x128x256 single-N tile.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_nvfp4_scaled_tma128n128k4,
+    fc1_nvfp4_scaled_tma128n128k4_ldm,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081261)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    linear = fc1_nvfp4_scaled_tma128n128k4(qx_live, qxs, qw, qws, alpha)
+    ldm = fc1_nvfp4_scaled_tma128n128k4_ldm(qx_live, qxs, qw, qws, alpha)
+    payload = {
+        "identity": (
+            "ldmatrix.x4 from smem[lane&15][(lane>>4)*16] == kitchen "
+            "group/tidg A fragment on 128x128x256. Probe 0/32. Not "
+            "the unit-scale consecutive-K uint4 map (that skips "
+            "ldmatrix redistribute)."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "ldm_vs_kitchen_mismatches": byte_diff(ldm, kitchen),
+        "ldm_vs_linear_mismatches": byte_diff(ldm, linear),
+        "full_finite": bool(torch.isfinite(ldm.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "ldm_full_min_ms": round(
+            time_ms(
+                lambda: fc1_nvfp4_scaled_tma128n128k4_ldm(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "linear_full_min_ms": round(
+            time_ms(
+                lambda: fc1_nvfp4_scaled_tma128n128k4(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["ldm_vs_kitchen_mismatches"] == 0
+        and payload["ldm_vs_linear_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["ldm_full_min_ms"], 4
+        )
+        payload["vs_linear_ms"] = round(
+            payload["linear_full_min_ms"] - payload["ldm_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma128n128k4_ldm_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma128n128k4_ldm_20423.json

@@ -0,0 +1,16 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "ldmatrix.x4 from smem[lane&15][(lane>>4)*16] == kitchen group/tidg A fragment on 128x128x256. Probe 0/32. Not the unit-scale consecutive-K uint4 map (that skips ldmatrix redistribute).",
+  "k": 5376,
+  "kitchen_gemm_min_ms": 20.3012,
+  "ldm_full_min_ms": 67.9157,
+  "ldm_vs_kitchen_mismatches": 0,
+  "ldm_vs_linear_mismatches": 0,
+  "linear_full_min_ms": 67.9085,
+  "n": 28672,
+  "pass": true,
+  "rows": 20423,
+  "vs_kitchen_ms": -47.6145,
+  "vs_linear_ms": -0.0072
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma128n128k4_pipe.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: software-pipelined B fragments == kitchen tile.
+
+Next N-subtile B/SFB loads overlap the current m16n8k64. Same
+128x128x256 single-N atom as the linear clone.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_nvfp4_scaled_tma128n128k4,
+    fc1_nvfp4_scaled_tma128n128k4_pipe,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081259)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    linear = fc1_nvfp4_scaled_tma128n128k4(qx_live, qxs, qw, qws, alpha)
+    piped = fc1_nvfp4_scaled_tma128n128k4_pipe(qx_live, qxs, qw, qws, alpha)
+    payload = {
+        "identity": (
+            "Software-pipelined B fragments on 128x128x256 == linear "
+            "clone == kitchen. Next N-subtile B/SFB overlaps m16n8k64. "
+            "Same atom and tile; ILP only."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "pipe_vs_kitchen_mismatches": byte_diff(piped, kitchen),
+        "pipe_vs_linear_mismatches": byte_diff(piped, linear),
+        "full_finite": bool(torch.isfinite(piped.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "pipe_full_min_ms": round(
+            time_ms(
+                lambda: fc1_nvfp4_scaled_tma128n128k4_pipe(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "linear_full_min_ms": round(
+            time_ms(
+                lambda: fc1_nvfp4_scaled_tma128n128k4(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["pipe_vs_kitchen_mismatches"] == 0
+        and payload["pipe_vs_linear_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["pipe_full_min_ms"], 4
+        )
+        payload["vs_linear_ms"] = round(
+            payload["linear_full_min_ms"] - payload["pipe_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma128n128k4_pipe_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma128n128k4_pipe_20423.json

@@ -0,0 +1,16 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "Software-pipelined B fragments on 128x128x256 == linear clone == kitchen. Next N-subtile B/SFB overlaps m16n8k64. Same atom and tile; ILP only.",
+  "k": 5376,
+  "kitchen_gemm_min_ms": 20.2951,
+  "linear_full_min_ms": 67.6802,
+  "n": 28672,
+  "pass": true,
+  "pipe_full_min_ms": 67.5507,
+  "pipe_vs_kitchen_mismatches": 0,
+  "pipe_vs_linear_mismatches": 0,
+  "rows": 20423,
+  "vs_kitchen_ms": -47.2556,
+  "vs_linear_ms": 0.1295
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma128n128k4_sw.py

@@ -0,0 +1,134 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: SWIZZLE_128B remapped on the 128-byte K-row.
+
+phys_col = col XOR ((row & 7) << 4). Measured on sm_121a. Same
+128x128x256 single-N m16n8k64 tile as the linear kitchen clone.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_nvfp4_scaled_tma128n128k4,
+    fc1_nvfp4_scaled_tma128n128k4_sw,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081256)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    linear = fc1_nvfp4_scaled_tma128n128k4(qx_live, qxs, qw, qws, alpha)
+    swiz = fc1_nvfp4_scaled_tma128n128k4_sw(qx_live, qxs, qw, qws, alpha)
+    payload = {
+        "identity": (
+            "SWIZZLE_128B TMA == linear fragment under "
+            "phys_col = col XOR ((row&7)<<4) on the 128-byte K-row "
+            "of the 128x128x256 single-N m16n8k64 tile. Map measured "
+            "on sm_121a (probe_swizzle128, 0/16384)."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "xor_map": "col XOR ((row & 7) << 4)",
+        "sw_vs_kitchen_mismatches": byte_diff(swiz, kitchen),
+        "sw_vs_kitchen_max_abs": round(max_abs_diff(swiz, kitchen), 6),
+        "sw_vs_linear_mismatches": byte_diff(swiz, linear),
+        "sw_vs_linear_max_abs": round(max_abs_diff(swiz, linear), 6),
+        "full_finite": bool(torch.isfinite(swiz.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "sw_full_min_ms": round(
+            time_ms(
+                lambda: fc1_nvfp4_scaled_tma128n128k4_sw(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "linear_full_min_ms": round(
+            time_ms(
+                lambda: fc1_nvfp4_scaled_tma128n128k4(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["sw_vs_kitchen_mismatches"] == 0
+        and payload["sw_vs_linear_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["sw_full_min_ms"], 4
+        )
+        payload["vs_linear_ms"] = round(
+            payload["linear_full_min_ms"] - payload["sw_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma128n128k4_sw_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma128n128k4_sw_20423.json

@@ -0,0 +1,19 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "SWIZZLE_128B TMA == linear fragment under phys_col = col XOR ((row&7)<<4) on the 128-byte K-row of the 128x128x256 single-N m16n8k64 tile. Map measured on sm_121a (probe_swizzle128, 0/16384).",
+  "k": 5376,
+  "kitchen_gemm_min_ms": 20.1582,
+  "linear_full_min_ms": 67.5818,
+  "n": 28672,
+  "pass": true,
+  "rows": 20423,
+  "sw_full_min_ms": 66.6565,
+  "sw_vs_kitchen_max_abs": 0.0,
+  "sw_vs_kitchen_mismatches": 0,
+  "sw_vs_linear_max_abs": 0.0,
+  "sw_vs_linear_mismatches": 0,
+  "vs_kitchen_ms": -46.4983,
+  "vs_linear_ms": 0.9253,
+  "xor_map": "col XOR ((row & 7) << 4)"
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma128n128k4ws.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: 12-warp 128x128x256 == kitchen / 8-warp clone.
+
+Kitchen launches 384 threads, 88064 B dynamic smem. Four producer
+warps issue A/B/SFA/SFB; eight MMA warps keep the m16n8k64 atom.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_nvfp4_scaled_tma128n128k4,
+    fc1_nvfp4_scaled_tma128n128k4ws,
+    fc1_nvfp4_scaled_tma128n128k4ws_attrs,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081255)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    clone8 = fc1_nvfp4_scaled_tma128n128k4(qx_live, qxs, qw, qws, alpha)
+    ws = fc1_nvfp4_scaled_tma128n128k4ws(qx_live, qxs, qw, qws, alpha)
+    attrs = fc1_nvfp4_scaled_tma128n128k4ws_attrs()
+    payload = {
+        "identity": (
+            "Kitchen 128x128x256 12-warp launch: 4 producer warps "
+            "issue A/B/SFA/SFB TMA; 8 MMA warps keep the m16n8k64 "
+            "atom and 16x4 acc. 88064 B dynamic smem. Same tile as "
+            "the 8-warp clone (already byte-exact)."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "launch": attrs,
+        "ws_vs_kitchen_mismatches": byte_diff(ws, kitchen),
+        "ws_vs_kitchen_max_abs": round(max_abs_diff(ws, kitchen), 6),
+        "ws_vs_clone8_mismatches": byte_diff(ws, clone8),
+        "ws_vs_clone8_max_abs": round(max_abs_diff(ws, clone8), 6),
+        "full_finite": bool(torch.isfinite(ws.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "ws_full_min_ms": round(
+            time_ms(
+                lambda: fc1_nvfp4_scaled_tma128n128k4ws(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "clone8_full_min_ms": round(
+            time_ms(
+                lambda: fc1_nvfp4_scaled_tma128n128k4(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["ws_vs_kitchen_mismatches"] == 0
+        and payload["ws_vs_clone8_mismatches"] == 0
+        and payload["full_finite"]
+        and int(attrs["threads"]) == 384
+        and int(attrs["dynamic_smem"]) == 88064
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["ws_full_min_ms"], 4
+        )
+        payload["vs_clone8_ms"] = round(
+            payload["clone8_full_min_ms"] - payload["ws_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma128n128k4ws_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma128n128k4ws_20423.json

@@ -0,0 +1,30 @@
+{
+  "clone8_full_min_ms": 68.3289,
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "Kitchen 128x128x256 12-warp launch: 4 producer warps issue A/B/SFA/SFB TMA; 8 MMA warps keep the m16n8k64 atom and 16x4 acc. 88064 B dynamic smem. Same tile as the 8-warp clone (already byte-exact).",
+  "k": 5376,
+  "kitchen_gemm_min_ms": 20.134,
+  "launch": {
+    "dynamic_smem": 88064,
+    "local_size_bytes": 8,
+    "max_dynamic_shared_size_bytes": 88064,
+    "mma_warps": 8,
+    "num_regs": 150,
+    "producer_warps": 4,
+    "shared_size_bytes": 0,
+    "smem_struct_bytes": 73856,
+    "threads": 384,
+    "warps": 12
+  },
+  "n": 28672,
+  "pass": true,
+  "rows": 20423,
+  "vs_clone8_ms": 0.9203,
+  "vs_kitchen_ms": -47.2746,
+  "ws_full_min_ms": 67.4086,
+  "ws_vs_clone8_max_abs": 0.0,
+  "ws_vs_clone8_mismatches": 0,
+  "ws_vs_kitchen_max_abs": 0.0,
+  "ws_vs_kitchen_mismatches": 0
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: TMA 256x64 3-stage NVFP4 paired-N vs kitchen.
+
+Same PTX fragment + cuBLAS 128x4 scale slabs. Launch structure: larger
+M tile (two scale slabs) and a 3-stage TMA K pipeline.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    bf16_nvfp4_dynamic,
+    fc1_paired_nvfp4_scaled_tma256,
+    fc1_paired_nvfp4_scaled_tma_sf,
+    load_extension,
+    swiglu_nvfp4_dynamic,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def eager_act(raw: torch.Tensor) -> torch.Tensor:
+    gate, up = raw.chunk(2, dim=-1)
+    return F.silu(gate).mul_(up)
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081240)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    full = fc1_paired_nvfp4_scaled_tma256(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    prod = fc1_paired_nvfp4_scaled_tma256(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    sf = fc1_paired_nvfp4_scaled_tma_sf(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    eager_kit = eager_act(kitchen.contiguous())
+    eager_full = eager_act(full)
+    two_q, two_s, two_g = swiglu_nvfp4_dynamic(kitchen.contiguous())
+    prod_q, prod_s, prod_g = bf16_nvfp4_dynamic(prod)
+    payload = {
+        "identity": (
+            "256x64 3-stage TMA on the kitchen-legal m16n8k64 atom: "
+            "two 128x4 scale slabs cover 256 M; K pipeline depth 3; "
+            "same PTX fragment, UE4M3 map, and eager product"
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "t256_full_vs_kitchen_mismatches": byte_diff(full, kitchen),
+        "t256_full_vs_kitchen_max_abs": round(max_abs_diff(full, kitchen), 6),
+        "t256_prod_vs_sf_prod": byte_diff(prod, sf),
+        "product_vs_eager_full_mismatches": byte_diff(prod, eager_full),
+        "product_vs_eager_kitchen_mismatches": byte_diff(prod, eager_kit),
+        "from_product_vs_kitchen_q": byte_diff(prod_q, two_q),
+        "from_product_vs_kitchen_s": byte_diff(prod_s, two_s),
+        "from_product_vs_kitchen_scale_exact": bool(torch.equal(prod_g, two_g)),
+        "full_finite": bool(torch.isfinite(full.float()).all().item()),
+        "prod_finite": bool(torch.isfinite(prod.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "t256_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+        "t256_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "sf_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma_sf(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "from_product_min_ms": round(
+            time_ms(lambda: bf16_nvfp4_dynamic(prod)), 4
+        ),
+    }
+    payload["pass"] = (
+        payload["t256_full_vs_kitchen_mismatches"] == 0
+        and payload["t256_prod_vs_sf_prod"] == 0
+        and payload["product_vs_eager_full_mismatches"] == 0
+        and payload["product_vs_eager_kitchen_mismatches"] == 0
+        and payload["from_product_vs_kitchen_q"] == 0
+        and payload["from_product_vs_kitchen_s"] == 0
+        and payload["from_product_vs_kitchen_scale_exact"]
+        and payload["full_finite"]
+        and payload["prod_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_sf_ms"] = round(
+            payload["sf_prod_min_ms"] - payload["t256_prod_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256_20423.json

@@ -0,0 +1,24 @@
+{
+  "device": "NVIDIA GB10",
+  "from_product_min_ms": 5.7456,
+  "from_product_vs_kitchen_q": 0,
+  "from_product_vs_kitchen_s": 0,
+  "from_product_vs_kitchen_scale_exact": true,
+  "full_finite": true,
+  "identity": "256x64 3-stage TMA on the kitchen-legal m16n8k64 atom: two 128x4 scale slabs cover 256 M; K pipeline depth 3; same PTX fragment, UE4M3 map, and eager product",
+  "k": 5376,
+  "kitchen_gemm_min_ms": 20.2712,
+  "n": 14336,
+  "pass": true,
+  "prod_finite": true,
+  "product_vs_eager_full_mismatches": 0,
+  "product_vs_eager_kitchen_mismatches": 0,
+  "rows": 20423,
+  "sf_prod_min_ms": 105.2268,
+  "t256_full_min_ms": 57.507,
+  "t256_full_vs_kitchen_max_abs": 0.0,
+  "t256_full_vs_kitchen_mismatches": 0,
+  "t256_prod_min_ms": 54.7271,
+  "t256_prod_vs_sf_prod": 0,
+  "vs_sf_ms": 50.4997
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256_sw.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: SWIZZLE_32B TMA remapped to the PTX fragment.
+
+Same 256x64 3-stage atom. TMA writes 16B chunks swizzled inside each
+32B row; loads XOR ((row&1)<<4) to recover the linear fragment.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    bf16_nvfp4_dynamic,
+    fc1_paired_nvfp4_scaled_tma256,
+    fc1_paired_nvfp4_scaled_tma256_sw,
+    load_extension,
+    swiglu_nvfp4_dynamic,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def eager_act(raw: torch.Tensor) -> torch.Tensor:
+    gate, up = raw.chunk(2, dim=-1)
+    return F.silu(gate).mul_(up)
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081241)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    full = fc1_paired_nvfp4_scaled_tma256_sw(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    prod = fc1_paired_nvfp4_scaled_tma256_sw(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    base = fc1_paired_nvfp4_scaled_tma256(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    eager_kit = eager_act(kitchen.contiguous())
+    eager_full = eager_act(full)
+    two_q, two_s, two_g = swiglu_nvfp4_dynamic(kitchen.contiguous())
+    prod_q, prod_s, prod_g = bf16_nvfp4_dynamic(prod)
+    payload = {
+        "identity": (
+            "SWIZZLE_32B TMA == linear fragment under "
+            "phys_col = col XOR ((row&4)<<2) on the kitchen-legal "
+            "256x64 3-stage m16n8k64 paired-N atom"
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "sw_full_vs_kitchen_mismatches": byte_diff(full, kitchen),
+        "sw_full_vs_kitchen_max_abs": round(max_abs_diff(full, kitchen), 6),
+        "sw_prod_vs_t256_prod": byte_diff(prod, base),
+        "product_vs_eager_full_mismatches": byte_diff(prod, eager_full),
+        "product_vs_eager_kitchen_mismatches": byte_diff(prod, eager_kit),
+        "from_product_vs_kitchen_q": byte_diff(prod_q, two_q),
+        "from_product_vs_kitchen_s": byte_diff(prod_s, two_s),
+        "from_product_vs_kitchen_scale_exact": bool(torch.equal(prod_g, two_g)),
+        "full_finite": bool(torch.isfinite(full.float()).all().item()),
+        "prod_finite": bool(torch.isfinite(prod.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "sw_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256_sw(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+        "sw_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256_sw(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "t256_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "from_product_min_ms": round(
+            time_ms(lambda: bf16_nvfp4_dynamic(prod)), 4
+        ),
+    }
+    payload["pass"] = (
+        payload["sw_full_vs_kitchen_mismatches"] == 0
+        and payload["sw_prod_vs_t256_prod"] == 0
+        and payload["product_vs_eager_full_mismatches"] == 0
+        and payload["product_vs_eager_kitchen_mismatches"] == 0
+        and payload["from_product_vs_kitchen_q"] == 0
+        and payload["from_product_vs_kitchen_s"] == 0
+        and payload["from_product_vs_kitchen_scale_exact"]
+        and payload["full_finite"]
+        and payload["prod_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_t256_ms"] = round(
+            payload["t256_prod_min_ms"] - payload["sw_prod_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256_sw_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256_sw_20423.json

@@ -0,0 +1,24 @@
+{
+  "device": "NVIDIA GB10",
+  "from_product_min_ms": 5.7545,
+  "from_product_vs_kitchen_q": 0,
+  "from_product_vs_kitchen_s": 0,
+  "from_product_vs_kitchen_scale_exact": true,
+  "full_finite": true,
+  "identity": "SWIZZLE_32B TMA == linear fragment under phys_col = col XOR ((row&4)<<2) on the kitchen-legal 256x64 3-stage m16n8k64 paired-N atom",
+  "k": 5376,
+  "kitchen_gemm_min_ms": 20.2633,
+  "n": 14336,
+  "pass": true,
+  "prod_finite": true,
+  "product_vs_eager_full_mismatches": 0,
+  "product_vs_eager_kitchen_mismatches": 0,
+  "rows": 20423,
+  "sw_full_min_ms": 57.6358,
+  "sw_full_vs_kitchen_max_abs": 0.0,
+  "sw_full_vs_kitchen_mismatches": 0,
+  "sw_prod_min_ms": 55.2272,
+  "sw_prod_vs_t256_prod": 0,
+  "t256_prod_min_ms": 54.41,
+  "vs_t256_ms": -0.8172
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: K=128 TMA box == two K=64 MMA steps.
+
+One TMA of 64 packed bytes feeds two m16n8k64 atoms. Same 256x64
+3-stage tile, PTX fragment, and cuBLAS 128x4 scale slabs (two groups).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    bf16_nvfp4_dynamic,
+    fc1_paired_nvfp4_scaled_tma256,
+    fc1_paired_nvfp4_scaled_tma256k2,
+    load_extension,
+    swiglu_nvfp4_dynamic,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def eager_act(raw: torch.Tensor) -> torch.Tensor:
+    gate, up = raw.chunk(2, dim=-1)
+    return F.silu(gate).mul_(up)
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081242)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    full = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    prod = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    base = fc1_paired_nvfp4_scaled_tma256(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    eager_kit = eager_act(kitchen.contiguous())
+    eager_full = eager_act(full)
+    two_q, two_s, two_g = swiglu_nvfp4_dynamic(kitchen.contiguous())
+    prod_q, prod_s, prod_g = bf16_nvfp4_dynamic(prod)
+    payload = {
+        "identity": (
+            "K=128 TMA box == two sequential K=64 MMA steps: one 64-byte "
+            "packed row feeds two m16n8k64 atoms; two 128x4 scale slabs "
+            "cover the 8 K-scale columns. Same 256x64 3-stage tile."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "k2_full_vs_kitchen_mismatches": byte_diff(full, kitchen),
+        "k2_full_vs_kitchen_max_abs": round(max_abs_diff(full, kitchen), 6),
+        "k2_prod_vs_t256_prod": byte_diff(prod, base),
+        "product_vs_eager_full_mismatches": byte_diff(prod, eager_full),
+        "product_vs_eager_kitchen_mismatches": byte_diff(prod, eager_kit),
+        "from_product_vs_kitchen_q": byte_diff(prod_q, two_q),
+        "from_product_vs_kitchen_s": byte_diff(prod_s, two_s),
+        "from_product_vs_kitchen_scale_exact": bool(torch.equal(prod_g, two_g)),
+        "full_finite": bool(torch.isfinite(full.float()).all().item()),
+        "prod_finite": bool(torch.isfinite(prod.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "k2_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+        "k2_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "t256_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "from_product_min_ms": round(
+            time_ms(lambda: bf16_nvfp4_dynamic(prod)), 4
+        ),
+    }
+    payload["pass"] = (
+        payload["k2_full_vs_kitchen_mismatches"] == 0
+        and payload["k2_prod_vs_t256_prod"] == 0
+        and payload["product_vs_eager_full_mismatches"] == 0
+        and payload["product_vs_eager_kitchen_mismatches"] == 0
+        and payload["from_product_vs_kitchen_q"] == 0
+        and payload["from_product_vs_kitchen_s"] == 0
+        and payload["from_product_vs_kitchen_scale_exact"]
+        and payload["full_finite"]
+        and payload["prod_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_t256_ms"] = round(
+            payload["t256_prod_min_ms"] - payload["k2_prod_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_20423.json

@@ -0,0 +1,24 @@
+{
+  "device": "NVIDIA GB10",
+  "from_product_min_ms": 5.7447,
+  "from_product_vs_kitchen_q": 0,
+  "from_product_vs_kitchen_s": 0,
+  "from_product_vs_kitchen_scale_exact": true,
+  "full_finite": true,
+  "identity": "K=128 TMA box == two sequential K=64 MMA steps: one 64-byte packed row feeds two m16n8k64 atoms; two 128x4 scale slabs cover the 8 K-scale columns. Same 256x64 3-stage tile.",
+  "k": 5376,
+  "k2_full_min_ms": 40.1835,
+  "k2_full_vs_kitchen_max_abs": 0.0,
+  "k2_full_vs_kitchen_mismatches": 0,
+  "k2_prod_min_ms": 36.8901,
+  "k2_prod_vs_t256_prod": 0,
+  "kitchen_gemm_min_ms": 20.1649,
+  "n": 14336,
+  "pass": true,
+  "prod_finite": true,
+  "product_vs_eager_full_mismatches": 0,
+  "product_vs_eager_kitchen_mismatches": 0,
+  "rows": 20423,
+  "t256_prod_min_ms": 56.1152,
+  "vs_t256_ms": 19.2251
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_ldm.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: ldmatrix.x4 A on k2 == kitchen / scalar k2.
+
+Best lab mainloop (256x64 K=128 3-stage) with the kitchen A fragment
+via ldmatrix. Same paired-N store.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k2_ldm,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081263)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    ldm = fc1_paired_nvfp4_scaled_tma256k2_ldm(qx_live, qxs, qw, qws, alpha)
+    payload = {
+        "identity": (
+            "ldmatrix.x4 A on k2 (256x64 K=128) == kitchen group/tidg "
+            "fragment == scalar k2. Best lab mainloop + kitchen A load."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "ldm_vs_kitchen_mismatches": byte_diff(ldm, kitchen),
+        "ldm_vs_k2_mismatches": byte_diff(ldm, k2),
+        "full_finite": bool(torch.isfinite(ldm.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "ldm_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2_ldm(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "k2_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["ldm_vs_kitchen_mismatches"] == 0
+        and payload["ldm_vs_k2_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["ldm_full_min_ms"], 4
+        )
+        payload["vs_k2_ms"] = round(
+            payload["k2_full_min_ms"] - payload["ldm_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2_ldm_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_ldm_20423.json

@@ -0,0 +1,16 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "ldmatrix.x4 A on k2 (256x64 K=128) == kitchen group/tidg fragment == scalar k2. Best lab mainloop + kitchen A load.",
+  "k": 5376,
+  "k2_full_min_ms": 38.9594,
+  "kitchen_gemm_min_ms": 20.2181,
+  "ldm_full_min_ms": 38.9408,
+  "ldm_vs_k2_mismatches": 0,
+  "ldm_vs_kitchen_mismatches": 0,
+  "n": 28672,
+  "pass": true,
+  "rows": 20423,
+  "vs_k2_ms": 0.0186,
+  "vs_kitchen_ms": -18.7227
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_ldmb.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: ldmatrix.x4 B on k2 == kitchen / scalar k2.
+
+One x4 covers two 8-N subtiles via the proven A-fragment map
+smem[lane&15][(lane>>4)*16]. Same 256x64 K=128 paired-N store.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k2_ldmb,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081264)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    ldmb = fc1_paired_nvfp4_scaled_tma256k2_ldmb(qx_live, qxs, qw, qws, alpha)
+    payload = {
+        "identity": (
+            "ldmatrix.x4 B on k2 (256x64 K=128) pairs two 8-N subtiles "
+            "via the A-fragment map == kitchen == scalar k2."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "ldmb_vs_kitchen_mismatches": byte_diff(ldmb, kitchen),
+        "ldmb_vs_k2_mismatches": byte_diff(ldmb, k2),
+        "full_finite": bool(torch.isfinite(ldmb.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "ldmb_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2_ldmb(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "k2_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["ldmb_vs_kitchen_mismatches"] == 0
+        and payload["ldmb_vs_k2_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["ldmb_full_min_ms"], 4
+        )
+        payload["vs_k2_ms"] = round(
+            payload["k2_full_min_ms"] - payload["ldmb_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2_ldmb_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_ldmb_20423.json

@@ -0,0 +1,16 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "ldmatrix.x4 B on k2 (256x64 K=128) pairs two 8-N subtiles via the A-fragment map == kitchen == scalar k2.",
+  "k": 5376,
+  "k2_full_min_ms": 40.1911,
+  "kitchen_gemm_min_ms": 20.2815,
+  "ldmb_full_min_ms": 39.4436,
+  "ldmb_vs_k2_mismatches": 0,
+  "ldmb_vs_kitchen_mismatches": 0,
+  "n": 28672,
+  "pass": true,
+  "rows": 20423,
+  "vs_k2_ms": 0.7475,
+  "vs_kitchen_ms": -19.1621
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_leads.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: leader-only SFA/SFB on k2 == kitchen / scalar k2.
+
+scale_vec::4X with selectors {0,0}: SFA from 16 lanes
+((lane>>1)&1)==0; SFB from 8 lanes (lane&3)==0. Other lanes
+pass 0. Same 256x64 K=128 paired-N store.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k2_leads,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081267)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    leads = fc1_paired_nvfp4_scaled_tma256k2_leads(
+        qx_live, qxs, qw, qws, alpha
+    )
+    payload = {
+        "identity": (
+            "leader-only SFA/SFB on k2 (256x64 K=128): scale_vec::4X "
+            "selectors {0,0} read 16 SFA + 8 SFB lanes; others 0 "
+            "== kitchen == scalar k2."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "leads_vs_kitchen_mismatches": byte_diff(leads, kitchen),
+        "leads_vs_k2_mismatches": byte_diff(leads, k2),
+        "full_finite": bool(torch.isfinite(leads.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "leads_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2_leads(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "k2_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["leads_vs_kitchen_mismatches"] == 0
+        and payload["leads_vs_k2_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["leads_full_min_ms"], 4
+        )
+        payload["vs_k2_ms"] = round(
+            payload["k2_full_min_ms"] - payload["leads_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2_leads_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_leads_20423.json

@@ -0,0 +1,16 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "leader-only SFA/SFB on k2 (256x64 K=128): scale_vec::4X selectors {0,0} read 16 SFA + 8 SFB lanes; others 0 == kitchen == scalar k2.",
+  "k": 5376,
+  "k2_full_min_ms": 39.226,
+  "kitchen_gemm_min_ms": 19.4048,
+  "leads_full_min_ms": 39.8251,
+  "leads_vs_k2_mismatches": 0,
+  "leads_vs_kitchen_mismatches": 0,
+  "n": 28672,
+  "pass": true,
+  "rows": 20423,
+  "vs_k2_ms": -0.5991,
+  "vs_kitchen_ms": -20.4203
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_pipe.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: software-pipelined B on k2 == kitchen / scalar k2.
+
+Next N-subtile B/SFB (both arms) overlaps m16n8k64. Same 256x64
+K=128 paired-N store. Not the 128x128x256 pipe.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k2_pipe,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081266)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    pipe = fc1_paired_nvfp4_scaled_tma256k2_pipe(qx_live, qxs, qw, qws, alpha)
+    payload = {
+        "identity": (
+            "software-pipelined B on k2 (256x64 K=128): next N-subtile "
+            "B/SFB both arms overlap m16n8k64 == kitchen == scalar k2."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "pipe_vs_kitchen_mismatches": byte_diff(pipe, kitchen),
+        "pipe_vs_k2_mismatches": byte_diff(pipe, k2),
+        "full_finite": bool(torch.isfinite(pipe.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "pipe_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2_pipe(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "k2_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["pipe_vs_kitchen_mismatches"] == 0
+        and payload["pipe_vs_k2_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["pipe_full_min_ms"], 4
+        )
+        payload["vs_k2_ms"] = round(
+            payload["k2_full_min_ms"] - payload["pipe_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2_pipe_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_pipe_20423.json

@@ -0,0 +1,16 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "software-pipelined B on k2 (256x64 K=128): next N-subtile B/SFB both arms overlap m16n8k64 == kitchen == scalar k2.",
+  "k": 5376,
+  "k2_full_min_ms": 39.9199,
+  "kitchen_gemm_min_ms": 20.1408,
+  "n": 28672,
+  "pass": true,
+  "pipe_full_min_ms": 40.6522,
+  "pipe_vs_k2_mismatches": 0,
+  "pipe_vs_kitchen_mismatches": 0,
+  "rows": 20423,
+  "vs_k2_ms": -0.7323,
+  "vs_kitchen_ms": -20.5114
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_pipea.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: software-pipelined A on k2 == kitchen / scalar k2.
+
+Next K=64 A/SFA overlaps m16n8k64. Same 256x64 K=128 paired-N store.
+Not the B-pipe identity.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k2_pipea,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081269)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    pipea = fc1_paired_nvfp4_scaled_tma256k2_pipea(
+        qx_live, qxs, qw, qws, alpha
+    )
+    payload = {
+        "identity": (
+            "software-pipelined A on k2 (256x64 K=128): next K=64 "
+            "A/SFA overlaps m16n8k64 == kitchen == scalar k2."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "pipea_vs_kitchen_mismatches": byte_diff(pipea, kitchen),
+        "pipea_vs_k2_mismatches": byte_diff(pipea, k2),
+        "full_finite": bool(torch.isfinite(pipea.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "pipea_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2_pipea(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "k2_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["pipea_vs_kitchen_mismatches"] == 0
+        and payload["pipea_vs_k2_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["pipea_full_min_ms"], 4
+        )
+        payload["vs_k2_ms"] = round(
+            payload["k2_full_min_ms"] - payload["pipea_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2_pipea_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_pipea_20423.json

@@ -0,0 +1,16 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "software-pipelined A on k2 (256x64 K=128): next K=64 A/SFA overlaps m16n8k64 == kitchen == scalar k2.",
+  "k": 5376,
+  "k2_full_min_ms": 39.3463,
+  "kitchen_gemm_min_ms": 20.2426,
+  "n": 28672,
+  "pass": true,
+  "pipea_full_min_ms": 39.3357,
+  "pipea_vs_k2_mismatches": 0,
+  "pipea_vs_kitchen_mismatches": 0,
+  "rows": 20423,
+  "vs_k2_ms": 0.0106,
+  "vs_kitchen_ms": -19.0931
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_sw.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: SWIZZLE_64B remapped on the K=128 TMA box.
+
+phys_col = col XOR (((row >> 1) & 3) << 4). Same two m16n8k64 atoms,
+256x64 3-stage tile, and cuBLAS 128x4 scale slabs.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    bf16_nvfp4_dynamic,
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k2_sw,
+    load_extension,
+    swiglu_nvfp4_dynamic,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def eager_act(raw: torch.Tensor) -> torch.Tensor:
+    gate, up = raw.chunk(2, dim=-1)
+    return F.silu(gate).mul_(up)
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081243)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    full = fc1_paired_nvfp4_scaled_tma256k2_sw(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    prod = fc1_paired_nvfp4_scaled_tma256k2_sw(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    base = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    eager_kit = eager_act(kitchen.contiguous())
+    eager_full = eager_act(full)
+    two_q, two_s, two_g = swiglu_nvfp4_dynamic(kitchen.contiguous())
+    prod_q, prod_s, prod_g = bf16_nvfp4_dynamic(prod)
+    payload = {
+        "identity": (
+            "SWIZZLE_64B TMA == linear fragment under "
+            "phys_col = col XOR (((row>>1)&3)<<4) on the kitchen-legal "
+            "K=128 / 256x64 3-stage m16n8k64 paired-N atom"
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "sw_full_vs_kitchen_mismatches": byte_diff(full, kitchen),
+        "sw_full_vs_kitchen_max_abs": round(max_abs_diff(full, kitchen), 6),
+        "sw_prod_vs_k2_prod": byte_diff(prod, base),
+        "product_vs_eager_full_mismatches": byte_diff(prod, eager_full),
+        "product_vs_eager_kitchen_mismatches": byte_diff(prod, eager_kit),
+        "from_product_vs_kitchen_q": byte_diff(prod_q, two_q),
+        "from_product_vs_kitchen_s": byte_diff(prod_s, two_s),
+        "from_product_vs_kitchen_scale_exact": bool(torch.equal(prod_g, two_g)),
+        "full_finite": bool(torch.isfinite(full.float()).all().item()),
+        "prod_finite": bool(torch.isfinite(prod.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "sw_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2_sw(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+        "sw_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2_sw(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "k2_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "from_product_min_ms": round(
+            time_ms(lambda: bf16_nvfp4_dynamic(prod)), 4
+        ),
+    }
+    payload["pass"] = (
+        payload["sw_full_vs_kitchen_mismatches"] == 0
+        and payload["sw_prod_vs_k2_prod"] == 0
+        and payload["product_vs_eager_full_mismatches"] == 0
+        and payload["product_vs_eager_kitchen_mismatches"] == 0
+        and payload["from_product_vs_kitchen_q"] == 0
+        and payload["from_product_vs_kitchen_s"] == 0
+        and payload["from_product_vs_kitchen_scale_exact"]
+        and payload["full_finite"]
+        and payload["prod_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_k2_ms"] = round(
+            payload["k2_prod_min_ms"] - payload["sw_prod_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2_sw_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2_sw_20423.json

@@ -0,0 +1,24 @@
+{
+  "device": "NVIDIA GB10",
+  "from_product_min_ms": 6.1193,
+  "from_product_vs_kitchen_q": 0,
+  "from_product_vs_kitchen_s": 0,
+  "from_product_vs_kitchen_scale_exact": true,
+  "full_finite": true,
+  "identity": "SWIZZLE_64B TMA == linear fragment under phys_col = col XOR (((row>>1)&3)<<4) on the kitchen-legal K=128 / 256x64 3-stage m16n8k64 paired-N atom",
+  "k": 5376,
+  "k2_prod_min_ms": 36.5662,
+  "kitchen_gemm_min_ms": 20.2433,
+  "n": 14336,
+  "pass": true,
+  "prod_finite": true,
+  "product_vs_eager_full_mismatches": 0,
+  "product_vs_eager_kitchen_mismatches": 0,
+  "rows": 20423,
+  "sw_full_min_ms": 39.6867,
+  "sw_full_vs_kitchen_max_abs": 0.0,
+  "sw_full_vs_kitchen_mismatches": 0,
+  "sw_prod_min_ms": 37.5803,
+  "sw_prod_vs_k2_prod": 0,
+  "vs_k2_ms": -1.0141
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2n2.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: sequential N-halves stream A once on 256x64 acc.
+
+One A K=128 TMA feeds two 64-wide N-halves. Acc0 stays in the k2
+2x8x4 register budget; acc1 swaps through opt-in smem. Same PTX
+fragment and cuBLAS 128x4 scale slabs.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_paired_nvfp4_scaled_tma256k2n2,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    props = torch.cuda.get_device_properties(device)
+    optin = int(getattr(props, "shared_memory_per_block_optin", 0))
+    acc_bytes = 256 * 128 * 4
+    tile_smem = 209024
+    torch.manual_seed(26081247)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    launch_error = ""
+    launched = False
+    try:
+        _ = fc1_paired_nvfp4_scaled_tma256k2n2(
+            qx_live, qxs, qw, qws, alpha, product=False
+        )
+        launched = True
+    except RuntimeError as exc:
+        launch_error = str(exc)
+    payload = {
+        "identity": (
+            "Sequential N-halves that stream A once need a second "
+            "256x64 FP32 acc (128 KiB) plus TMA. GB10 block opt-in "
+            "is 99 KiB, so acc1 cannot live in smem. Fail-closed."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "smem_per_block": int(props.shared_memory_per_block),
+        "smem_per_sm": int(props.shared_memory_per_multiprocessor),
+        "smem_optin": optin,
+        "acc1_bytes": acc_bytes,
+        "tile_smem_bytes": tile_smem,
+        "acc1_exceeds_optin": acc_bytes > optin,
+        "tile_exceeds_optin": tile_smem > optin,
+        "launched": launched,
+        "launch_error": launch_error,
+    }
+    payload["pass"] = (
+        not launched
+        and acc_bytes > optin
+        and tile_smem > optin
+        and optin <= 101376
+        and "opt-in" in launch_error
+    )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2n2_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2n2_20423.json

@@ -0,0 +1,17 @@
+{
+  "acc1_bytes": 131072,
+  "acc1_exceeds_optin": true,
+  "device": "NVIDIA GB10",
+  "identity": "Sequential N-halves that stream A once need a second 256x64 FP32 acc (128 KiB) plus TMA. GB10 block opt-in is 99 KiB, so acc1 cannot live in smem. Fail-closed.",
+  "k": 5376,
+  "launch_error": "sequential N-halves need 209024 B smem (131072 B FP32 acc1 + TMA); device opt-in is 101376 B",
+  "launched": false,
+  "n": 14336,
+  "pass": true,
+  "rows": 20423,
+  "smem_optin": 101376,
+  "smem_per_block": 49152,
+  "smem_per_sm": 102400,
+  "tile_exceeds_optin": true,
+  "tile_smem_bytes": 209024
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2p.py

@@ -0,0 +1,171 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: persistent k2 N-walk == grid k2.
+
+One 8-warp CTA owns 256 M and walks every 64-wide N tile. Same
+m16n8k64 atom, 2x8x4 acc, and 3-stage K=128 TMA. A-stationary
+so L2 can reuse the 256-row panel. Smem matches k2 (under 99 KiB).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    bf16_nvfp4_dynamic,
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k2p,
+    load_extension,
+    swiglu_nvfp4_dynamic,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def eager_act(raw: torch.Tensor) -> torch.Tensor:
+    gate, up = raw.chunk(2, dim=-1)
+    return F.silu(gate).mul_(up)
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081249)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    full = fc1_paired_nvfp4_scaled_tma256k2p(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    prod = fc1_paired_nvfp4_scaled_tma256k2p(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    eager_kit = eager_act(kitchen.contiguous())
+    eager_full = eager_act(full)
+    two_q, two_s, two_g = swiglu_nvfp4_dynamic(kitchen.contiguous())
+    prod_q, prod_s, prod_g = bf16_nvfp4_dynamic(prod)
+    payload = {
+        "identity": (
+            "Persistent 8-warp k2 N-walk == grid k2: one CTA owns "
+            "256 M and walks every 64-wide N tile. Same m16n8k64 "
+            "atom, 2x8x4 acc, 3-stage K=128 TMA. A-stationary."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "p_full_vs_kitchen_mismatches": byte_diff(full, kitchen),
+        "p_full_vs_kitchen_max_abs": round(max_abs_diff(full, kitchen), 6),
+        "p_prod_vs_k2_prod": byte_diff(prod, k2),
+        "product_vs_eager_full_mismatches": byte_diff(prod, eager_full),
+        "product_vs_eager_kitchen_mismatches": byte_diff(prod, eager_kit),
+        "from_product_vs_kitchen_q": byte_diff(prod_q, two_q),
+        "from_product_vs_kitchen_s": byte_diff(prod_s, two_s),
+        "from_product_vs_kitchen_scale_exact": bool(torch.equal(prod_g, two_g)),
+        "full_finite": bool(torch.isfinite(full.float()).all().item()),
+        "prod_finite": bool(torch.isfinite(prod.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "p_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2p(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+        "p_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2p(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "k2_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "from_product_min_ms": round(
+            time_ms(lambda: bf16_nvfp4_dynamic(prod)), 4
+        ),
+    }
+    payload["pass"] = (
+        payload["p_full_vs_kitchen_mismatches"] == 0
+        and payload["p_prod_vs_k2_prod"] == 0
+        and payload["product_vs_eager_full_mismatches"] == 0
+        and payload["product_vs_eager_kitchen_mismatches"] == 0
+        and payload["from_product_vs_kitchen_q"] == 0
+        and payload["from_product_vs_kitchen_s"] == 0
+        and payload["from_product_vs_kitchen_scale_exact"]
+        and payload["full_finite"]
+        and payload["prod_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_k2_ms"] = round(
+            payload["k2_prod_min_ms"] - payload["p_prod_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2p_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2p_20423.json

@@ -0,0 +1,24 @@
+{
+  "device": "NVIDIA GB10",
+  "from_product_min_ms": 6.057,
+  "from_product_vs_kitchen_q": 0,
+  "from_product_vs_kitchen_s": 0,
+  "from_product_vs_kitchen_scale_exact": true,
+  "full_finite": true,
+  "identity": "Persistent 8-warp k2 N-walk == grid k2: one CTA owns 256 M and walks every 64-wide N tile. Same m16n8k64 atom, 2x8x4 acc, 3-stage K=128 TMA. A-stationary.",
+  "k": 5376,
+  "k2_prod_min_ms": 38.1653,
+  "kitchen_gemm_min_ms": 20.2363,
+  "n": 14336,
+  "p_full_min_ms": 72.0819,
+  "p_full_vs_kitchen_max_abs": 0.0,
+  "p_full_vs_kitchen_mismatches": 0,
+  "p_prod_min_ms": 64.0328,
+  "p_prod_vs_k2_prod": 0,
+  "pass": true,
+  "prod_finite": true,
+  "product_vs_eager_full_mismatches": 0,
+  "product_vs_eager_kitchen_mismatches": 0,
+  "rows": 20423,
+  "vs_k2_ms": -25.8675
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2s1.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: 1-stage k2 == kitchen / 3-stage k2.
+
+Same 256x64 K=128 paired-N atom. ~29 KiB smem so 2 CTAs/SM
+is legal if the register file allows. Occupancy is the launch
+variable; 3-stage k2 is ~86 KiB and 1-way.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k2s1,
+    fc1_paired_nvfp4_scaled_tma256k2s1_attrs,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081271)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    s1 = fc1_paired_nvfp4_scaled_tma256k2s1(qx_live, qxs, qw, qws, alpha)
+    attrs = fc1_paired_nvfp4_scaled_tma256k2s1_attrs()
+    payload = {
+        "identity": (
+            "1-stage K=128 256x64 paired-N == kitchen == 3-stage k2. "
+            "Smem drop is the occupancy variable."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "s1_vs_kitchen_mismatches": byte_diff(s1, kitchen),
+        "s1_vs_k2_mismatches": byte_diff(s1, k2),
+        "full_finite": bool(torch.isfinite(s1.float()).all().item()),
+        "s1_regs": int(attrs["s1_regs"]),
+        "s1_smem": int(attrs["s1_smem"]),
+        "s1_occupancy": int(attrs["s1_occupancy"]),
+        "k2_regs": int(attrs["k2_regs"]),
+        "k2_smem": int(attrs["k2_smem"]),
+        "k2_occupancy": int(attrs["k2_occupancy"]),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "s1_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2s1(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "k2_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["s1_vs_kitchen_mismatches"] == 0
+        and payload["s1_vs_k2_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["s1_full_min_ms"], 4
+        )
+        payload["vs_k2_ms"] = round(
+            payload["k2_full_min_ms"] - payload["s1_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2s1_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2s1_20423.json

@@ -0,0 +1,22 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "1-stage K=128 256x64 paired-N == kitchen == 3-stage k2. Smem drop is the occupancy variable.",
+  "k": 5376,
+  "k2_full_min_ms": 40.4624,
+  "k2_occupancy": 1,
+  "k2_regs": 167,
+  "k2_smem": 86144,
+  "kitchen_gemm_min_ms": 20.0947,
+  "n": 28672,
+  "pass": true,
+  "rows": 20423,
+  "s1_full_min_ms": 47.8348,
+  "s1_occupancy": 1,
+  "s1_regs": 166,
+  "s1_smem": 28800,
+  "s1_vs_k2_mismatches": 0,
+  "s1_vs_kitchen_mismatches": 0,
+  "vs_k2_ms": -7.3724,
+  "vs_kitchen_ms": -27.7401
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2ws.py

@@ -0,0 +1,171 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: warp-specialized k2 == unified k2.
+
+Producer warp issues TMA; eight MMA warps keep the 256x64
+m16n8k64 atom and 2x8x4 acc. Same smem as k2 (under 99 KiB).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    bf16_nvfp4_dynamic,
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k2ws,
+    load_extension,
+    swiglu_nvfp4_dynamic,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def eager_act(raw: torch.Tensor) -> torch.Tensor:
+    gate, up = raw.chunk(2, dim=-1)
+    return F.silu(gate).mul_(up)
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081248)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    full = fc1_paired_nvfp4_scaled_tma256k2ws(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    prod = fc1_paired_nvfp4_scaled_tma256k2ws(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    eager_kit = eager_act(kitchen.contiguous())
+    eager_full = eager_act(full)
+    two_q, two_s, two_g = swiglu_nvfp4_dynamic(kitchen.contiguous())
+    prod_q, prod_s, prod_g = bf16_nvfp4_dynamic(prod)
+    payload = {
+        "identity": (
+            "Warp-specialized k2 == unified k2: producer warp issues "
+            "TMA; eight MMA warps keep the 256x64 m16n8k64 atom and "
+            "2x8x4 acc. Full/empty mbarriers, no mainloop syncthreads. "
+            "Smem matches k2 (under 99 KiB)."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "ws_full_vs_kitchen_mismatches": byte_diff(full, kitchen),
+        "ws_full_vs_kitchen_max_abs": round(max_abs_diff(full, kitchen), 6),
+        "ws_prod_vs_k2_prod": byte_diff(prod, k2),
+        "product_vs_eager_full_mismatches": byte_diff(prod, eager_full),
+        "product_vs_eager_kitchen_mismatches": byte_diff(prod, eager_kit),
+        "from_product_vs_kitchen_q": byte_diff(prod_q, two_q),
+        "from_product_vs_kitchen_s": byte_diff(prod_s, two_s),
+        "from_product_vs_kitchen_scale_exact": bool(torch.equal(prod_g, two_g)),
+        "full_finite": bool(torch.isfinite(full.float()).all().item()),
+        "prod_finite": bool(torch.isfinite(prod.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "ws_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2ws(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+        "ws_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2ws(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "k2_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "from_product_min_ms": round(
+            time_ms(lambda: bf16_nvfp4_dynamic(prod)), 4
+        ),
+    }
+    payload["pass"] = (
+        payload["ws_full_vs_kitchen_mismatches"] == 0
+        and payload["ws_prod_vs_k2_prod"] == 0
+        and payload["product_vs_eager_full_mismatches"] == 0
+        and payload["product_vs_eager_kitchen_mismatches"] == 0
+        and payload["from_product_vs_kitchen_q"] == 0
+        and payload["from_product_vs_kitchen_s"] == 0
+        and payload["from_product_vs_kitchen_scale_exact"]
+        and payload["full_finite"]
+        and payload["prod_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_k2_ms"] = round(
+            payload["k2_prod_min_ms"] - payload["ws_prod_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2ws_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2ws4.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: 12-warp 4-producer k2 == kitchen == k2.
+
+Kitchen launch structure (4 TMA warps + 8 MMA) on the best lab
+tile (256x64 K=128 paired-N). 9-warp k2ws serializes every TMA
+on one lane; this issues A / Bg / Bu / scales in parallel.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k2ws4,
+    fc1_paired_nvfp4_scaled_tma256k2ws4_attrs,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081274)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    ws4 = fc1_paired_nvfp4_scaled_tma256k2ws4(qx_live, qxs, qw, qws, alpha)
+    attrs = fc1_paired_nvfp4_scaled_tma256k2ws4_attrs()
+    payload = {
+        "identity": (
+            "12-warp 4-producer + 8-MMA on k2 256x64 K=128 == kitchen "
+            "== unified k2. Kitchen launch split: A / Bg / Bu / scales "
+            "issue in parallel. Same atom and smem as k2."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "ws4_vs_kitchen_mismatches": byte_diff(ws4, kitchen),
+        "ws4_vs_kitchen_max_abs": round(max_abs_diff(ws4, kitchen), 6),
+        "ws4_vs_k2_mismatches": byte_diff(ws4, k2),
+        "full_finite": bool(torch.isfinite(ws4.float()).all().item()),
+        "ws4_regs": int(attrs["regs"]),
+        "ws4_smem": int(attrs["smem"]),
+        "ws4_occupancy": int(attrs["occupancy"]),
+        "ws4_threads": int(attrs["threads"]),
+        "ws4_warps": int(attrs["warps"]),
+        "ws4_prod_warps": int(attrs["prod_warps"]),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "ws4_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2ws4(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "k2_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["ws4_vs_kitchen_mismatches"] == 0
+        and payload["ws4_vs_k2_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["ws4_full_min_ms"], 4
+        )
+        payload["vs_k2_ms"] = round(
+            payload["k2_full_min_ms"] - payload["ws4_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k2ws4_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2ws4_20423.json

@@ -0,0 +1,23 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "12-warp 4-producer + 8-MMA on k2 256x64 K=128 == kitchen == unified k2. Kitchen launch split: A / Bg / Bu / scales issue in parallel. Same atom and smem as k2.",
+  "k": 5376,
+  "k2_full_min_ms": 39.2012,
+  "kitchen_gemm_min_ms": 20.1714,
+  "n": 28672,
+  "pass": true,
+  "rows": 20423,
+  "vs_k2_ms": -1.5219,
+  "vs_kitchen_ms": -20.5517,
+  "ws4_full_min_ms": 40.7231,
+  "ws4_occupancy": 1,
+  "ws4_prod_warps": 4,
+  "ws4_regs": 168,
+  "ws4_smem": 86144,
+  "ws4_threads": 384,
+  "ws4_vs_k2_mismatches": 0,
+  "ws4_vs_kitchen_max_abs": 0.0,
+  "ws4_vs_kitchen_mismatches": 0,
+  "ws4_warps": 12
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k2ws_20423.json

@@ -0,0 +1,24 @@
+{
+  "device": "NVIDIA GB10",
+  "from_product_min_ms": 5.7406,
+  "from_product_vs_kitchen_q": 0,
+  "from_product_vs_kitchen_s": 0,
+  "from_product_vs_kitchen_scale_exact": true,
+  "full_finite": true,
+  "identity": "Warp-specialized k2 == unified k2: producer warp issues TMA; eight MMA warps keep the 256x64 m16n8k64 atom and 2x8x4 acc. Full/empty mbarriers, no mainloop syncthreads. Smem matches k2 (under 99 KiB).",
+  "k": 5376,
+  "k2_prod_min_ms": 35.79,
+  "kitchen_gemm_min_ms": 20.1856,
+  "n": 14336,
+  "pass": true,
+  "prod_finite": true,
+  "product_vs_eager_full_mismatches": 0,
+  "product_vs_eager_kitchen_mismatches": 0,
+  "rows": 20423,
+  "vs_k2_ms": -1.5863,
+  "ws_full_min_ms": 48.2706,
+  "ws_full_vs_kitchen_max_abs": 0.0,
+  "ws_full_vs_kitchen_mismatches": 0,
+  "ws_prod_min_ms": 37.3763,
+  "ws_prod_vs_k2_prod": 0
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k4.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: K=256 TMA box on the 256x64 tile.
+
+One TMA of 128 packed bytes feeds four m16n8k64 atoms. Keeps the
+256x64 k2 register budget (2x8x4). 1-stage: 2-stage at this box
+is 112 KiB and exceeds the 99 KiB static smem cap.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    bf16_nvfp4_dynamic,
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256k4,
+    load_extension,
+    swiglu_nvfp4_dynamic,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def eager_act(raw: torch.Tensor) -> torch.Tensor:
+    gate, up = raw.chunk(2, dim=-1)
+    return F.silu(gate).mul_(up)
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081246)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    full = fc1_paired_nvfp4_scaled_tma256k4(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    prod = fc1_paired_nvfp4_scaled_tma256k4(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    eager_kit = eager_act(kitchen.contiguous())
+    eager_full = eager_act(full)
+    two_q, two_s, two_g = swiglu_nvfp4_dynamic(kitchen.contiguous())
+    prod_q, prod_s, prod_g = bf16_nvfp4_dynamic(prod)
+    payload = {
+        "identity": (
+            "K=256 TMA box == four sequential K=64 MMA steps on the "
+            "256x64 tile: one 128-byte packed row feeds four m16n8k64 "
+            "atoms; four 128x4 scale slabs cover the 16 K-scale columns. "
+            "Same 2x8x4 register acc as k2. 1-stage: 2-stage box "
+            "exceeds the 99 KiB static smem cap."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "k4_full_vs_kitchen_mismatches": byte_diff(full, kitchen),
+        "k4_full_vs_kitchen_max_abs": round(max_abs_diff(full, kitchen), 6),
+        "k4_prod_vs_k2_prod": byte_diff(prod, k2),
+        "product_vs_eager_full_mismatches": byte_diff(prod, eager_full),
+        "product_vs_eager_kitchen_mismatches": byte_diff(prod, eager_kit),
+        "from_product_vs_kitchen_q": byte_diff(prod_q, two_q),
+        "from_product_vs_kitchen_s": byte_diff(prod_s, two_s),
+        "from_product_vs_kitchen_scale_exact": bool(torch.equal(prod_g, two_g)),
+        "full_finite": bool(torch.isfinite(full.float()).all().item()),
+        "prod_finite": bool(torch.isfinite(prod.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "k4_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k4(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+        "k4_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k4(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "k2_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "from_product_min_ms": round(
+            time_ms(lambda: bf16_nvfp4_dynamic(prod)), 4
+        ),
+    }
+    payload["pass"] = (
+        payload["k4_full_vs_kitchen_mismatches"] == 0
+        and payload["k4_prod_vs_k2_prod"] == 0
+        and payload["product_vs_eager_full_mismatches"] == 0
+        and payload["product_vs_eager_kitchen_mismatches"] == 0
+        and payload["from_product_vs_kitchen_q"] == 0
+        and payload["from_product_vs_kitchen_s"] == 0
+        and payload["from_product_vs_kitchen_scale_exact"]
+        and payload["full_finite"]
+        and payload["prod_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_k2_ms"] = round(
+            payload["k2_prod_min_ms"] - payload["k4_prod_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k4_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k4_20423.json

@@ -0,0 +1,24 @@
+{
+  "device": "NVIDIA GB10",
+  "from_product_min_ms": 5.7432,
+  "from_product_vs_kitchen_q": 0,
+  "from_product_vs_kitchen_s": 0,
+  "from_product_vs_kitchen_scale_exact": true,
+  "full_finite": true,
+  "identity": "K=256 TMA box == four sequential K=64 MMA steps on the 256x64 tile: one 128-byte packed row feeds four m16n8k64 atoms; four 128x4 scale slabs cover the 16 K-scale columns. Same 2x8x4 register acc as k2. 1-stage: 2-stage box exceeds the 99 KiB static smem cap.",
+  "k": 5376,
+  "k2_prod_min_ms": 36.4231,
+  "k4_full_min_ms": 71.0255,
+  "k4_full_vs_kitchen_max_abs": 0.0,
+  "k4_full_vs_kitchen_mismatches": 0,
+  "k4_prod_min_ms": 65.1102,
+  "k4_prod_vs_k2_prod": 0,
+  "kitchen_gemm_min_ms": 19.9771,
+  "n": 14336,
+  "pass": true,
+  "prod_finite": true,
+  "product_vs_eager_full_mismatches": 0,
+  "product_vs_eager_kitchen_mismatches": 0,
+  "rows": 20423,
+  "vs_k2_ms": -28.6871
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k4n1.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: K=256 2-stage on 256x64 single-N == kitchen / k2.
+
+Paired 256x64 K=256 2-stage is 112 KiB (over the 99 KiB cap).
+Dropping the second B operand (kitchen single-N) is 94 KiB and
+gives a legal pipeline on the k2 M tile + kitchen K box.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_nvfp4_scaled_tma256k4n1,
+    fc1_paired_nvfp4_scaled_tma256k2,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081265)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    n1 = fc1_nvfp4_scaled_tma256k4n1(qx_live, qxs, qw, qws, alpha)
+    payload = {
+        "identity": (
+            "K=256 2-stage TMA on 256x64 single-N == four K=64 MMA "
+            "== kitchen == paired k2. Paired 2-stage is 112 KiB; "
+            "single-N is 94 KiB under the 99 KiB cap."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "smem_bytes": 94336,
+        "device": torch.cuda.get_device_name(device),
+        "n1_vs_kitchen_mismatches": byte_diff(n1, kitchen),
+        "n1_vs_k2_mismatches": byte_diff(n1, k2),
+        "full_finite": bool(torch.isfinite(n1.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "n1_full_min_ms": round(
+            time_ms(
+                lambda: fc1_nvfp4_scaled_tma256k4n1(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "k2_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["n1_vs_kitchen_mismatches"] == 0
+        and payload["n1_vs_k2_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["n1_full_min_ms"], 4
+        )
+        payload["vs_k2_ms"] = round(
+            payload["k2_full_min_ms"] - payload["n1_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256k4n1_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256k4n1_20423.json

@@ -0,0 +1,17 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "K=256 2-stage TMA on 256x64 single-N == four K=64 MMA == kitchen == paired k2. Paired 2-stage is 112 KiB; single-N is 94 KiB under the 99 KiB cap.",
+  "k": 5376,
+  "k2_full_min_ms": 40.1932,
+  "kitchen_gemm_min_ms": 20.1874,
+  "n": 28672,
+  "n1_full_min_ms": 48.913,
+  "n1_vs_k2_mismatches": 0,
+  "n1_vs_kitchen_mismatches": 0,
+  "pass": true,
+  "rows": 20423,
+  "smem_bytes": 94336,
+  "vs_k2_ms": -8.7198,
+  "vs_kitchen_ms": -28.7256
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256n32.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: 256x32 K=128 1-stage == kitchen / k2.
+
+Half-N of k2: 4 n-subtiles, 64-float acc. TMA B box is 64
+(N=32 tile faults). 1-stage so two CTAs/SM can fit.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    fc1_paired_nvfp4_scaled_tma256k2,
+    fc1_paired_nvfp4_scaled_tma256n32,
+    fc1_paired_nvfp4_scaled_tma256n32_attrs,
+    load_extension,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081272)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    k2 = fc1_paired_nvfp4_scaled_tma256k2(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    n32 = fc1_paired_nvfp4_scaled_tma256n32(qx_live, qxs, qw, qws, alpha)
+    attrs = fc1_paired_nvfp4_scaled_tma256n32_attrs()
+    payload = {
+        "identity": (
+            "256x32 K=128 1-stage paired-N == kitchen == k2. "
+            "TMA B is 64; compute is 32. 64-float acc is the occupancy cut."
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": 2 * N,
+        "device": torch.cuda.get_device_name(device),
+        "n32_vs_kitchen_mismatches": byte_diff(n32, kitchen),
+        "n32_vs_k2_mismatches": byte_diff(n32, k2),
+        "full_finite": bool(torch.isfinite(n32.float()).all().item()),
+        "n32_regs": int(attrs["regs"]),
+        "n32_smem": int(attrs["smem"]),
+        "n32_occupancy": int(attrs["occupancy"]),
+        "n32_smem_struct": int(attrs["smem_struct"]),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "n32_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256n32(
+                    qx_live, qxs, qw, qws, alpha
+                )
+            ),
+            4,
+        ),
+        "k2_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma256k2(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+    }
+    payload["pass"] = (
+        payload["n32_vs_kitchen_mismatches"] == 0
+        and payload["n32_vs_k2_mismatches"] == 0
+        and payload["full_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_kitchen_ms"] = round(
+            payload["kitchen_gemm_min_ms"] - payload["n32_full_min_ms"], 4
+        )
+        payload["vs_k2_ms"] = round(
+            payload["k2_full_min_ms"] - payload["n32_full_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma256n32_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma256n32_20423.json

@@ -0,0 +1,18 @@
+{
+  "device": "NVIDIA GB10",
+  "full_finite": true,
+  "identity": "256x32 K=128 1-stage paired-N == kitchen == k2. TMA B is 64; compute is 32. 64-float acc is the occupancy cut.",
+  "k": 5376,
+  "k2_full_min_ms": 39.0139,
+  "kitchen_gemm_min_ms": 19.3708,
+  "n": 28672,
+  "n32_full_min_ms": 44.013,
+  "n32_occupancy": 2,
+  "n32_regs": 125,
+  "n32_smem": 28800,
+  "n32_smem_struct": 28800,
+  "n32_vs_k2_mismatches": 2362273,
+  "n32_vs_kitchen_mismatches": 2362273,
+  "pass": false,
+  "rows": 20423
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma_20423.json

@@ -0,0 +1,24 @@
+{
+  "device": "NVIDIA GB10",
+  "from_product_min_ms": 5.7439,
+  "from_product_vs_kitchen_q": 0,
+  "from_product_vs_kitchen_s": 0,
+  "from_product_vs_kitchen_scale_exact": true,
+  "full_finite": true,
+  "identity": "TMA 2-stage K pipeline on the kitchen-legal m16n8k64 atom: 128x128 2D tile, tensor-map + mbarrier, A reused across 16 n-subtiles and both arms; same PTX fragment, UE4M3 map, and eager product",
+  "k": 5376,
+  "kitchen_gemm_min_ms": 20.1758,
+  "n": 14336,
+  "pass": true,
+  "piped_prod_min_ms": 218.6166,
+  "prod_finite": true,
+  "product_vs_eager_full_mismatches": 0,
+  "product_vs_eager_kitchen_mismatches": 0,
+  "rows": 20423,
+  "tma_full_min_ms": 115.2573,
+  "tma_full_vs_kitchen_max_abs": 0.0,
+  "tma_full_vs_kitchen_mismatches": 0,
+  "tma_prod_min_ms": 104.7121,
+  "tma_prod_vs_piped_prod": 0,
+  "vs_piped_ms": 113.9045
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma_sf.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: TMA scale-slab NVFP4 paired-N vs kitchen.
+
+Same PTX fragment + 128x128 TMA data path. Scales come from one
+cuBLAS 128x4 / 512-byte slab per K-tile, remapped to pack_four_scales.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    bf16_nvfp4_dynamic,
+    fc1_paired_nvfp4_scaled_tma,
+    fc1_paired_nvfp4_scaled_tma_sf,
+    load_extension,
+    swiglu_nvfp4_dynamic,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def kitchen_gemm(qx, sx, qw, sw, scale_x, scale_w) -> torch.Tensor:
+    alpha = (scale_x * scale_w).reshape(1)
+    y = ck.scaled_mm_nvfp4(
+        qx, qw, scale_x, scale_w, sx, sw, out_dtype=torch.bfloat16, alpha=alpha
+    )
+    return y[:ROWS]
+
+
+def eager_act(raw: torch.Tensor) -> torch.Tensor:
+    gate, up = raw.chunk(2, dim=-1)
+    return F.silu(gate).mul_(up)
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def max_abs_diff(left: torch.Tensor, right: torch.Tensor) -> float:
+    return float((left.float() - right.float()).abs().max().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081239)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    sx = nvfp4_scale(x)
+    sw = nvfp4_scale(weight)
+    qx, qxs = ck.quantize_nvfp4(x, sx, pad_16x=True, hi_first=True)
+    qw, qws = ck.quantize_nvfp4(weight, sw, pad_16x=True, hi_first=True)
+    qx_live = qx[:ROWS].contiguous()
+    alpha = (sx * sw).reshape(1).contiguous()
+    kitchen = kitchen_gemm(qx, qxs, qw, qws, sx, sw)
+    full = fc1_paired_nvfp4_scaled_tma_sf(
+        qx_live, qxs, qw, qws, alpha, product=False
+    )
+    prod = fc1_paired_nvfp4_scaled_tma_sf(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    tma = fc1_paired_nvfp4_scaled_tma(
+        qx_live, qxs, qw, qws, alpha, product=True
+    )
+    eager_kit = eager_act(kitchen.contiguous())
+    eager_full = eager_act(full)
+    two_q, two_s, two_g = swiglu_nvfp4_dynamic(kitchen.contiguous())
+    prod_q, prod_s, prod_g = bf16_nvfp4_dynamic(prod)
+    payload = {
+        "identity": (
+            "cuBLAS 128x4 E4M3 scale slab == pack_four_scales: TMA 1D bulk "
+            "of the 512-byte slab, indexed as (row%32)*16+(row/32)*4, on "
+            "the kitchen-legal 128x128 m16n8k64 paired-N atom"
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "sf_full_vs_kitchen_mismatches": byte_diff(full, kitchen),
+        "sf_full_vs_kitchen_max_abs": round(max_abs_diff(full, kitchen), 6),
+        "sf_prod_vs_tma_prod": byte_diff(prod, tma),
+        "product_vs_eager_full_mismatches": byte_diff(prod, eager_full),
+        "product_vs_eager_kitchen_mismatches": byte_diff(prod, eager_kit),
+        "from_product_vs_kitchen_q": byte_diff(prod_q, two_q),
+        "from_product_vs_kitchen_s": byte_diff(prod_s, two_s),
+        "from_product_vs_kitchen_scale_exact": bool(torch.equal(prod_g, two_g)),
+        "full_finite": bool(torch.isfinite(full.float()).all().item()),
+        "prod_finite": bool(torch.isfinite(prod.float()).all().item()),
+        "kitchen_gemm_min_ms": round(
+            time_ms(lambda: kitchen_gemm(qx, qxs, qw, qws, sx, sw)), 4
+        ),
+        "sf_full_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma_sf(
+                    qx_live, qxs, qw, qws, alpha, product=False
+                )
+            ),
+            4,
+        ),
+        "sf_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma_sf(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "tma_prod_min_ms": round(
+            time_ms(
+                lambda: fc1_paired_nvfp4_scaled_tma(
+                    qx_live, qxs, qw, qws, alpha, product=True
+                )
+            ),
+            4,
+        ),
+        "from_product_min_ms": round(
+            time_ms(lambda: bf16_nvfp4_dynamic(prod)), 4
+        ),
+    }
+    payload["pass"] = (
+        payload["sf_full_vs_kitchen_mismatches"] == 0
+        and payload["sf_prod_vs_tma_prod"] == 0
+        and payload["product_vs_eager_full_mismatches"] == 0
+        and payload["product_vs_eager_kitchen_mismatches"] == 0
+        and payload["from_product_vs_kitchen_q"] == 0
+        and payload["from_product_vs_kitchen_s"] == 0
+        and payload["from_product_vs_kitchen_scale_exact"]
+        and payload["full_finite"]
+        and payload["prod_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_tma_ms"] = round(
+            payload["tma_prod_min_ms"] - payload["sf_prod_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name(
+        "gate_fc1_nvfp4_scaled_tma_sf_20423.json"
+    ).write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_scaled_tma_sf_20423.json

@@ -0,0 +1,24 @@
+{
+  "device": "NVIDIA GB10",
+  "from_product_min_ms": 5.7291,
+  "from_product_vs_kitchen_q": 0,
+  "from_product_vs_kitchen_s": 0,
+  "from_product_vs_kitchen_scale_exact": true,
+  "full_finite": true,
+  "identity": "cuBLAS 128x4 E4M3 scale slab == pack_four_scales: TMA 1D bulk of the 512-byte slab, indexed as (row%32)*16+(row/32)*4, on the kitchen-legal 128x128 m16n8k64 paired-N atom",
+  "k": 5376,
+  "kitchen_gemm_min_ms": 20.2346,
+  "n": 14336,
+  "pass": true,
+  "prod_finite": true,
+  "product_vs_eager_full_mismatches": 0,
+  "product_vs_eager_kitchen_mismatches": 0,
+  "rows": 20423,
+  "sf_full_min_ms": 107.2147,
+  "sf_full_vs_kitchen_max_abs": 0.0,
+  "sf_full_vs_kitchen_mismatches": 0,
+  "sf_prod_min_ms": 104.2306,
+  "sf_prod_vs_tma_prod": 0,
+  "tma_prod_min_ms": 104.6651,
+  "vs_tma_ms": 0.4345
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_tiled.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: persistent tiled NVFP4 paired-N vs 1-warp and eager.
+
+Same atom. CTA owns 64 N, streams M, reuses A across the panel.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    bf16_nvfp4_dynamic,
+    fc1_paired_nvfp4,
+    fc1_paired_nvfp4_tiled,
+    load_extension,
+    swiglu_nvfp4_dynamic,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+
+
+def nvfp4_scale(x: torch.Tensor) -> torch.Tensor:
+    return (x.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+
+
+def eager_act(raw: torch.Tensor) -> torch.Tensor:
+    gate, up = raw.chunk(2, dim=-1)
+    return F.silu(gate).mul_(up)
+
+
+def eager_pack(activated: torch.Tensor):
+    scale = (activated.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+    packed, block_scales = ck.quantize_nvfp4(
+        activated, scale, pad_16x=True, hi_first=True
+    )
+    return packed, block_scales, scale
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    fn()
+    torch.cuda.synchronize()
+    starter.record()
+    fn()
+    ender.record()
+    torch.cuda.synchronize()
+    return starter.elapsed_time(ender)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081235)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    qx, _ = ck.quantize_nvfp4(x, nvfp4_scale(x), pad_16x=True, hi_first=True)
+    qw, _ = ck.quantize_nvfp4(
+        weight, nvfp4_scale(weight), pad_16x=True, hi_first=True
+    )
+    qx = qx[:ROWS].contiguous()
+    full = fc1_paired_nvfp4_tiled(qx, qw, product=False)
+    prod = fc1_paired_nvfp4_tiled(qx, qw, product=True)
+    onew = fc1_paired_nvfp4(qx, qw, product=True)
+    eager = eager_act(full)
+    two_q, two_s, two_g = swiglu_nvfp4_dynamic(full)
+    prod_q, prod_s, prod_g = bf16_nvfp4_dynamic(prod)
+    ref_q, ref_s, ref_g = eager_pack(eager)
+    payload = {
+        "identity": (
+            "persistent 64x64 NVFP4 paired-N: A K-slab reused across 8 "
+            "n-subtiles and both arms; eager product epilogue"
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "product_vs_eager_full_mismatches": byte_diff(prod, eager),
+        "tiled_vs_onewarp_prod": byte_diff(prod, onew),
+        "from_product_vs_twopass_q": byte_diff(prod_q, two_q),
+        "from_product_vs_twopass_s": byte_diff(prod_s, two_s),
+        "from_product_vs_twopass_scale_exact": bool(torch.equal(prod_g, two_g)),
+        "from_product_vs_eager_q": byte_diff(prod_q, ref_q),
+        "from_product_vs_eager_s": byte_diff(prod_s, ref_s),
+        "from_product_vs_eager_scale_exact": bool(torch.equal(prod_g, ref_g)),
+        "full_finite": bool(torch.isfinite(full.float()).all().item()),
+        "prod_finite": bool(torch.isfinite(prod.float()).all().item()),
+        "tiled_full_min_ms": round(
+            time_ms(lambda: fc1_paired_nvfp4_tiled(qx, qw, product=False)), 4
+        ),
+        "tiled_prod_min_ms": round(
+            time_ms(lambda: fc1_paired_nvfp4_tiled(qx, qw, product=True)), 4
+        ),
+        "onewarp_prod_min_ms": round(
+            time_ms(lambda: fc1_paired_nvfp4(qx, qw, product=True)), 4
+        ),
+        "twopass_min_ms": round(time_ms(lambda: swiglu_nvfp4_dynamic(full)), 4),
+        "from_product_min_ms": round(time_ms(lambda: bf16_nvfp4_dynamic(prod)), 4),
+    }
+    payload["pass"] = (
+        payload["product_vs_eager_full_mismatches"] == 0
+        and payload["tiled_vs_onewarp_prod"] == 0
+        and payload["from_product_vs_twopass_q"] == 0
+        and payload["from_product_vs_twopass_s"] == 0
+        and payload["from_product_vs_twopass_scale_exact"]
+        and payload["from_product_vs_eager_q"] == 0
+        and payload["from_product_vs_eager_s"] == 0
+        and payload["from_product_vs_eager_scale_exact"]
+        and payload["full_finite"]
+        and payload["prod_finite"]
+    )
+    if payload["pass"]:
+        payload["vs_onewarp_ms"] = round(
+            payload["onewarp_prod_min_ms"] - payload["tiled_prod_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name("gate_fc1_nvfp4_tiled_20423.json").write_text(
+        text + "\n"
+    )
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

labs/swiglu_nvfp4/native_cuda/gate_fc1_nvfp4_tiled_20423.json

@@ -0,0 +1,24 @@
+{
+  "device": "NVIDIA GB10",
+  "from_product_min_ms": 5.8331,
+  "from_product_vs_eager_q": 0,
+  "from_product_vs_eager_s": 0,
+  "from_product_vs_eager_scale_exact": true,
+  "from_product_vs_twopass_q": 0,
+  "from_product_vs_twopass_s": 0,
+  "from_product_vs_twopass_scale_exact": true,
+  "full_finite": true,
+  "identity": "persistent 64x64 NVFP4 paired-N: A K-slab reused across 8 n-subtiles and both arms; eager product epilogue",
+  "k": 5376,
+  "n": 14336,
+  "onewarp_prod_min_ms": 1225.824,
+  "pass": true,
+  "prod_finite": true,
+  "product_vs_eager_full_mismatches": 0,
+  "rows": 20423,
+  "tiled_full_min_ms": 221.7096,
+  "tiled_prod_min_ms": 208.4988,
+  "tiled_vs_onewarp_prod": 0,
+  "twopass_min_ms": 10.8047,
+  "vs_onewarp_ms": 1017.3252
+}

labs/swiglu_nvfp4/native_cuda/gate_fc1_paired.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+"""Byte-exact gate: paired-N FC1 MMA store of eager product vs [gate|up].
+
+Same mainloop, two epilogues. Kitchen GEMM is not modified.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+import comfy_kitchen as ck
+
+from swiglu_nvfp4 import (
+    bf16_nvfp4_dynamic,
+    fc1_paired_store,
+    load_extension,
+    swiglu_nvfp4_dynamic,
+)
+
+
+ROWS = 20423
+K = 5376
+N = 14336
+REPEATS = 1
+
+
+def eager_act(raw: torch.Tensor) -> torch.Tensor:
+    gate, up = raw.chunk(2, dim=-1)
+    return F.silu(gate).mul_(up)
+
+
+def eager_pack(activated: torch.Tensor):
+    scale = (activated.abs().amax() / (6.0 * 448.0)).to(torch.float32).reshape(1)
+    packed, block_scales = ck.quantize_nvfp4(
+        activated, scale, pad_16x=True, hi_first=True
+    )
+    return packed, block_scales, scale
+
+
+def byte_diff(left: torch.Tensor, right: torch.Tensor) -> int:
+    return int((left.view(torch.uint8) != right.view(torch.uint8)).sum().item())
+
+
+def time_ms(fn) -> float:
+    starter = torch.cuda.Event(enable_timing=True)
+    ender = torch.cuda.Event(enable_timing=True)
+    for _ in range(2):
+        fn()
+    torch.cuda.synchronize()
+    samples = []
+    for _ in range(REPEATS):
+        starter.record()
+        fn()
+        ender.record()
+        torch.cuda.synchronize()
+        samples.append(starter.elapsed_time(ender))
+    return min(samples)
+
+
+def main() -> int:
+    load_extension(verbose=False)
+    device = torch.device("cuda")
+    torch.manual_seed(26081228)
+    x = torch.randn(ROWS, K, device=device, dtype=torch.bfloat16)
+    weight = torch.randn(2 * N, K, device=device, dtype=torch.bfloat16) * 0.02
+    full = fc1_paired_store(x, weight, product=False)
+    prod = fc1_paired_store(x, weight, product=True)
+    eager = eager_act(full)
+    two_q, two_s, two_g = swiglu_nvfp4_dynamic(full)
+    prod_q, prod_s, prod_g = bf16_nvfp4_dynamic(prod)
+    ref_q, ref_s, ref_g = eager_pack(eager)
+    payload = {
+        "identity": (
+            "paired-N MMA accumulates gate and up together; "
+            "product epilogue stores eager round(silu(round(g))*round(u)); "
+            "same mainloop as the [gate|up] store"
+        ),
+        "rows": ROWS,
+        "k": K,
+        "n": N,
+        "device": torch.cuda.get_device_name(device),
+        "product_vs_eager_full_mismatches": byte_diff(prod, eager),
+        "from_product_vs_twopass_q": byte_diff(prod_q, two_q),
+        "from_product_vs_twopass_s": byte_diff(prod_s, two_s),
+        "from_product_vs_twopass_scale_exact": bool(torch.equal(prod_g, two_g)),
+        "from_product_vs_eager_q": byte_diff(prod_q, ref_q),
+        "from_product_vs_eager_s": byte_diff(prod_s, ref_s),
+        "from_product_vs_eager_scale_exact": bool(torch.equal(prod_g, ref_g)),
+        "full_store_min_ms": round(
+            time_ms(lambda: fc1_paired_store(x, weight, product=False)), 4
+        ),
+        "product_store_min_ms": round(
+            time_ms(lambda: fc1_paired_store(x, weight, product=True)), 4
+        ),
+        "twopass_min_ms": round(time_ms(lambda: swiglu_nvfp4_dynamic(full)), 4),
+        "from_product_min_ms": round(time_ms(lambda: bf16_nvfp4_dynamic(prod)), 4),
+    }
+    payload["pass"] = (
+        payload["product_vs_eager_full_mismatches"] == 0
+        and payload["from_product_vs_twopass_q"] == 0
+        and payload["from_product_vs_twopass_s"] == 0
+        and payload["from_product_vs_twopass_scale_exact"]
+        and payload["from_product_vs_eager_q"] == 0
+        and payload["from_product_vs_eager_s"] == 0
+        and payload["from_product_vs_eager_scale_exact"]
+    )
+    if payload["pass"]:
+        payload["store_delta_ms"] = round(
+            payload["full_store_min_ms"] - payload["product_store_min_ms"], 4
+        )
+        payload["pack_side_saved_ms"] = round(
+            payload["twopass_min_ms"] - payload["from_product_min_ms"], 4
+        )
+    text = json.dumps(payload, indent=2, sort_keys=True)
+    Path(__file__).with_name("gate_fc1_paired_20423.json").write_text(text + "\n")
+    print(text)
+    return 0 if payload["pass"] else 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())