feat(fp4-gemm): portable SIMT NVFP4 GEMM fallback for SM110 (Thor)

fp4-gemm/VALIDATION.md

@@ -81,3 +81,19 @@ The Tensor wrapper was compared against the same native FlashRT launchers on
 Thor with 20 warmup and 100 measured iterations. For production auto-dispatch
 across the six model shapes, wrapper/native latency ratio had median `1.019`
 and maximum `1.086`. Correctness was exact (`max_abs=mean_abs=p99_abs=0`).
+
+## SM110 portable SIMT fallback (Thor real hardware)
+
+On NVIDIA Thor (`sm_110a`) the SM120-only fused-epilogue GEMMs and the
+warp-split GEMV route to `portable_fp4_gemm_simt` (compiled for `sm_110a`):
+
+- `fp4_w4a4_gemv_warpsplit_bf16`
+- `nvfp4_gemm_residual_bf16`
+- `nvfp4_gemm_bias_gelu_bf16`
+- `nvfp4_gemm_bias_gelu_nvfp4`
+- `nvfp4_gemm_streamk_bf16`
+- `nvfp4_gemm_streamk_bias_bf16`
+
+Validated on Thor with PyTorch 2.9.1+cu130 / CUDA 13.2: all six ops launch,
+produce finite BF16/FP4 output, and match the FP32 reference within the BF16
+contract. SM120 keeps the CUTLASS fused-epilogue path.

fp4-gemm/benchmarks/RESULTS.md

@@ -55,3 +55,25 @@ The direct entry is byte-exact against the package's established
 BF16-to-FP16 plus FP16-producer contract. The native timing is reported as a
 performance reference only because that producer uses a distinct quantization
 strategy.
+
+## NVIDIA Thor SM110 Results (installed artifact, 2026-08-04)
+
+Measured against `/data/test_thor/fp4-gemm` (installed artifact) on Thor
+`sm_110a`, torch `2.9.1+cu130`. `nvfp4_gemm_bf16` routes SM110 through the
+dedicated CUTLASS `sm110_gemm_dispatch` path (not the SIMT fallback), so the
+production GEMM is fast on Thor. Reference = PyTorch fp32 GEMM over the same
+dequantized FP4/SFA+FP4/SFB inputs.
+
+| Workload `(M,N,K)` | FlashRT us | Eager us | Compile us | vs eager | Max abs | Cosine |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| pi0.5 action gate/up `(51,16384,2048)` | 38.710 | 2400.608 | 2117.904 | 62.0x | 0.0 | 1.000000 |
+| pi0.5 action down `(51,2048,8192)` | 18.074 | 1330.918 | 1268.806 | 73.6x | 0.0 | 1.000000 |
+| GROOT DiT QKV `(51,4608,1536)` | 12.752 | 545.200 | 500.822 | 42.8x | 0.0 | 1.000000 |
+| GROOT backbone gate/up `(277,16384,2048)` | 79.242 | 7006.390 | 6748.576 | 88.4x | 0.0 | 1.000000 |
+| Cosmos Edge action `(64,9216,2048)` | 21.398 | 1446.109 | 1312.918 | 67.6x | 0.0 | 1.000000 |
+| LingBot action gate/up `(105,16384,2048)` | 31.331 | 2947.763 | 2856.307 | 94.1x | 0.0 | 1.000000 |
+
+The SM110 CUTLASS GEMM is 42-94x faster than the dequantized PyTorch reference
+with bitwise-equal output. The portable SIMT fallbacks (`nvfp4_gemm_linear_simt`
+et al.) remain the compatibility path for fused-epilogue ops without an SM110
+CUTLASS implementation and are exercised by the on-device correctness suite.

fp4-gemm/build.toml

@@ -63,3 +63,15 @@ src = [
   "csrc/gemm/fp4/cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cu",
   "csrc/gemm/fp4/cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cuh",
 ]
+
+[kernel.fp4_gemm_portable]
+backend = "cuda"
+depends = ["torch"]
+include = ["csrc"]
+cuda-minver = "12.8"
+cuda-capabilities = ["11.0a"]
+cuda-flags = ["--expt-relaxed-constexpr", "-O3"]
+src = [
+  "csrc/gemm/fp4/portable_fp4_gemm_simt.cu",
+  "csrc/gemm/fp4/portable_fp4_gemm_simt.cuh",
+]

fp4-gemm/csrc/gemm/fp4/portable_fp4_gemm_simt.cu

@@ -0,0 +1,238 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Portable SIMT reference for the FP4 block-scaled GEMM fused epilogues.
+//
+// The sm_120a CUTLASS kernels use block-scaled MMA paths that are not
+// available on pre-sm120 devices. These kernels compute the same
+//   D[m,n] = epilogue( alpha * sum_k fp4(A[m,k])*ue4m3(SFA[m,k])
+//                                     * fp4(B[n,k])*ue4m3(SFB[n,k]) )
+// in pure SIMT FMA so the fused epilogue ops remain usable (slowly) on
+// sm_110 Thor. sm_120 keeps the CUTLASS path.
+//
+// Packed layout: byte k/2 holds element (k&1) in the high nibble, element
+// (k&1)==0 in the low nibble, e2m1 (NVFP4) code points.
+// Scale layout: NVFP4 128-row super-block swizzle used by the SFA/SFB
+// producers (Sm1xxBlockScaledConfig<16>); identical for SFA and SFB.
+
+#include "portable_fp4_gemm_simt.cuh"
+
+#include <cuda_fp8.h>
+#include <cuda_bf16.h>
+#include <cstdint>
+#include <cmath>
+
+namespace flash_rt {
+namespace gemm {
+
+namespace {
+
+constexpr int THREADS = 256;
+
+__device__ __forceinline__ float e2m1_to_float(uint8_t v) {
+  static constexpr float mags[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f};
+  float mag = mags[v & 0x7];
+  return (v & 0x8) ? -mag : mag;
+}
+
+__device__ __forceinline__ float sf_read(const uint8_t* sf, int off) {
+  __nv_fp8_e4m3 scale_q;
+  *reinterpret_cast<uint8_t*>(&scale_q) = sf[off];
+  return static_cast<float>(scale_q);
+}
+
+// NVFP4 super-block swizzle byte offset for (row, k) within a flat SFA/SFB buf.
+__device__ __forceinline__ int sf_off(int row, int k, int n_col_super) {
+  int rb = row >> 7;
+  int ri = row & 127;
+  int kt = k >> 6;
+  int cb = (k >> 4) & 3;
+  return (rb * n_col_super + kt) * 512 + (ri & 31) * 16 + ((ri >> 5) & 3) * 4 + cb;
+}
+
+__device__ __forceinline__ float read_val(
+    const uint8_t* __restrict__ packed, int row, int k, int K,
+    const uint8_t* __restrict__ sf, int n_col_super) {
+  uint8_t byte = packed[(size_t)row * (K / 2) + (k >> 1)];
+  uint8_t nib = (k & 1) ? (byte >> 4) : (byte & 0xF);
+  return e2m1_to_float(nib) * sf_read(sf, sf_off(row, k, n_col_super));
+}
+
+__device__ __forceinline__ float gelu_taylor(float x) {
+  const float k0 = 0.7978845608028654f;
+  return 0.5f * x * (1.f + tanhf(k0 * x * (1.f + 0.044715f * x * x)));
+}
+
+__device__ __forceinline__ uint8_t fp32_to_e2m1(float x) {
+  uint8_t sign = (x < 0.f) ? 0x8u : 0x0u;
+  float ax = fabsf(x);
+  uint8_t mant;
+  if      (ax <= 0.25f) mant = 0u;
+  else if (ax <= 0.75f) mant = 1u;
+  else if (ax <= 1.25f) mant = 2u;
+  else if (ax <= 1.75f) mant = 3u;
+  else if (ax <= 2.5f)  mant = 4u;
+  else if (ax <= 3.5f)  mant = 5u;
+  else if (ax <= 5.0f)  mant = 6u;
+  else                  mant = 7u;
+  return sign | mant;
+}
+
+enum class Ep : int { Linear, Residual, Bias, BiasGelu };
+
+template <Ep EPI>
+__global__ void gemm_epilogue_simt_kernel(
+    const uint8_t* __restrict__ a, const uint8_t* __restrict__ b,
+    const uint8_t* __restrict__ sfa, const uint8_t* __restrict__ sfb,
+    const __nv_bfloat16* __restrict__ bias,
+    const __nv_bfloat16* __restrict__ residual,
+    __nv_bfloat16* __restrict__ out,
+    int m, int n, int k, float alpha) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= m * n) return;
+  int mi = idx / n, ni = idx - mi * n;
+  const int n_col_super = (k / 16 + 3) / 4;
+  float acc = 0.f;
+  for (int kk = 0; kk < k; ++kk) {
+    float av = read_val(a, mi, kk, k, sfa, n_col_super);
+    float bv = read_val(b, ni, kk, k, sfb, n_col_super);
+    acc += av * bv;
+  }
+  float r = acc * alpha;
+  if (EPI == Ep::Residual) r += __bfloat162float(residual[idx]);
+  if (EPI == Ep::Bias || EPI == Ep::BiasGelu) r += __bfloat162float(bias[ni]);
+  if (EPI == Ep::BiasGelu) r = gelu_taylor(r);
+  out[idx] = __float2bfloat16(r);
+}
+
+// One thread per (row, 16-element column block): computes the 16 gelu values
+// then quantizes them to NVFP4 (e2m1 packed + per-16-block scale), matching
+// the fp4 producer's quantization.
+__global__ void gemm_bias_gelu_fp4_simt_kernel(
+    const uint8_t* __restrict__ a, const uint8_t* __restrict__ b,
+    const uint8_t* __restrict__ sfa, const uint8_t* __restrict__ sfb,
+    const __nv_bfloat16* __restrict__ bias,
+    uint8_t* __restrict__ out_packed, uint8_t* __restrict__ out_sfa,
+    int m, int n, int k, float alpha) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  const int n_blocks = n / 16;
+  if (idx >= m * n_blocks) return;
+  const int mi = idx / n_blocks, bi = idx - mi * n_blocks;
+  const int n_col_super = (k / 16 + 3) / 4;
+  float vals[16];
+#pragma unroll
+  for (int cc = 0; cc < 16; ++cc) {
+    const int ni = bi * 16 + cc;
+    float acc = 0.f;
+    for (int kk = 0; kk < k; ++kk) {
+      float av = read_val(a, mi, kk, k, sfa, n_col_super);
+      float bv = read_val(b, ni, kk, k, sfb, n_col_super);
+      acc += av * bv;
+    }
+    vals[cc] = gelu_taylor(acc * alpha + __bfloat162float(bias[ni]));
+  }
+  float amax = 0.f;
+#pragma unroll
+  for (int cc = 0; cc < 16; ++cc) amax = fmaxf(amax, fabsf(vals[cc]));
+  float desired = amax / 6.f;
+  if (desired < 1e-12f) desired = 1e-12f;
+  __nv_fp8_e4m3 bs = __nv_fp8_e4m3(desired);
+  const float bs_dq = static_cast<float>(bs);
+  const int out_n_col_super = (n / 16 + 3) / 4;
+  out_sfa[sf_off(mi, bi * 16, out_n_col_super)] =
+      *reinterpret_cast<uint8_t*>(&bs);
+  const float inv = 1.f / bs_dq;
+  uint8_t* op = out_packed + (size_t)mi * (n / 2) + bi * 8;
+#pragma unroll
+  for (int p = 0; p < 8; ++p) {
+    uint8_t lo = fp32_to_e2m1(vals[2 * p] * inv);
+    uint8_t hi = fp32_to_e2m1(vals[2 * p + 1] * inv);
+    op[p] = lo | (hi << 4);
+  }
+}
+
+int launch_epilogue(const uint8_t* a, const uint8_t* b, const uint8_t* sfa,
+                    const uint8_t* sfb, const void* bias, const void* residual,
+                    void* out, int m, int n, int k, float alpha, int mode,
+                    cudaStream_t stream) {
+  const int total = m * n;
+  const int blocks = (total + THREADS - 1) / THREADS;
+  const auto* bias_b = static_cast<const __nv_bfloat16*>(bias);
+  const auto* res_b = static_cast<const __nv_bfloat16*>(residual);
+  auto* out_b = static_cast<__nv_bfloat16*>(out);
+  if (mode == 0) {
+    gemm_epilogue_simt_kernel<Ep::Linear><<<blocks, THREADS, 0, stream>>>(
+        a, b, sfa, sfb, bias_b, res_b, out_b, m, n, k, alpha);
+  } else if (mode == 1) {
+    gemm_epilogue_simt_kernel<Ep::Residual><<<blocks, THREADS, 0, stream>>>(
+        a, b, sfa, sfb, bias_b, res_b, out_b, m, n, k, alpha);
+  } else if (mode == 2) {
+    gemm_epilogue_simt_kernel<Ep::Bias><<<blocks, THREADS, 0, stream>>>(
+        a, b, sfa, sfb, bias_b, res_b, out_b, m, n, k, alpha);
+  } else {
+    gemm_epilogue_simt_kernel<Ep::BiasGelu><<<blocks, THREADS, 0, stream>>>(
+        a, b, sfa, sfb, bias_b, res_b, out_b, m, n, k, alpha);
+  }
+  return (cudaGetLastError() == cudaSuccess) ? 0 : 1;
+}
+
+}  // namespace
+
+int nvfp4_gemm_linear_simt(
+    const void* a_packed, const void* b_packed, const void* sfa,
+    const void* sfb, void* out_bf16, int m, int n, int k, float alpha,
+    cudaStream_t stream) {
+  return launch_epilogue(
+      static_cast<const uint8_t*>(a_packed), static_cast<const uint8_t*>(b_packed),
+      static_cast<const uint8_t*>(sfa), static_cast<const uint8_t*>(sfb),
+      nullptr, nullptr, out_bf16, m, n, k, alpha, 0, stream);
+}
+
+int nvfp4_gemm_residual_simt(
+    const void* a_packed, const void* b_packed, const void* sfa,
+    const void* sfb, const void* residual_bf16, void* out_bf16,
+    int m, int n, int k, float alpha, cudaStream_t stream) {
+  return launch_epilogue(
+      static_cast<const uint8_t*>(a_packed), static_cast<const uint8_t*>(b_packed),
+      static_cast<const uint8_t*>(sfa), static_cast<const uint8_t*>(sfb),
+      nullptr, residual_bf16, out_bf16, m, n, k, alpha, 1, stream);
+}
+
+int nvfp4_gemm_bias_simt(
+    const void* a_packed, const void* b_packed, const void* sfa,
+    const void* sfb, const void* bias_bf16, void* out_bf16,
+    int m, int n, int k, float alpha, cudaStream_t stream) {
+  return launch_epilogue(
+      static_cast<const uint8_t*>(a_packed), static_cast<const uint8_t*>(b_packed),
+      static_cast<const uint8_t*>(sfa), static_cast<const uint8_t*>(sfb),
+      bias_bf16, nullptr, out_bf16, m, n, k, alpha, 2, stream);
+}
+
+int nvfp4_gemm_bias_gelu_simt(
+    const void* a_packed, const void* b_packed, const void* sfa,
+    const void* sfb, const void* bias_bf16, void* out_bf16,
+    int m, int n, int k, float alpha, cudaStream_t stream) {
+  return launch_epilogue(
+      static_cast<const uint8_t*>(a_packed), static_cast<const uint8_t*>(b_packed),
+      static_cast<const uint8_t*>(sfa), static_cast<const uint8_t*>(sfb),
+      bias_bf16, nullptr, out_bf16, m, n, k, alpha, 3, stream);
+}
+
+int nvfp4_gemm_bias_gelu_fp4_simt(
+    const void* a_packed, const void* b_packed, const void* sfa,
+    const void* sfb, const void* bias_bf16, void* out_packed, void* out_sfa,
+    int m, int n, int k, float alpha, cudaStream_t stream) {
+  if (m <= 0 || n <= 0 || k <= 0 || n % 16 != 0) return 1;
+  const int n_blocks = n / 16;
+  const int total = m * n_blocks;
+  gemm_bias_gelu_fp4_simt_kernel<<<(total + THREADS - 1) / THREADS, THREADS, 0,
+                                   stream>>>(
+      static_cast<const uint8_t*>(a_packed), static_cast<const uint8_t*>(b_packed),
+      static_cast<const uint8_t*>(sfa), static_cast<const uint8_t*>(sfb),
+      static_cast<const __nv_bfloat16*>(bias_bf16),
+      static_cast<uint8_t*>(out_packed), static_cast<uint8_t*>(out_sfa),
+      m, n, k, alpha);
+  return (cudaGetLastError() == cudaSuccess) ? 0 : 1;
+}
+
+}  // namespace gemm
+}  // namespace flash_rt

fp4-gemm/csrc/gemm/fp4/portable_fp4_gemm_simt.cuh

@@ -0,0 +1,45 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <cuda_runtime.h>
+
+// Portable SIMT fallbacks for the FP4 block-scaled GEMM fused epilogues.
+// The sm_120a CUTLASS kernels are unavailable on pre-sm120 devices; these
+// reference kernels compute the same GEMM + epilogue in pure SIMT FMA so the
+// ops stay usable (slowly) on sm_110 Thor. sm_120 keeps the CUTLASS path.
+
+namespace flash_rt {
+namespace gemm {
+
+// out[M,N] = alpha * (A@B^T)
+int nvfp4_gemm_linear_simt(
+    const void* a_packed, const void* b_packed, const void* sfa,
+    const void* sfb, void* out_bf16, int m, int n, int k, float alpha,
+    cudaStream_t stream);
+
+// out[M,N] = alpha * (A@B^T) + residual[M,N]
+int nvfp4_gemm_residual_simt(
+    const void* a_packed, const void* b_packed, const void* sfa,
+    const void* sfb, const void* residual_bf16, void* out_bf16,
+    int m, int n, int k, float alpha, cudaStream_t stream);
+
+// out[M,N] = alpha * (A@B^T) + bias[N]
+int nvfp4_gemm_bias_simt(
+    const void* a_packed, const void* b_packed, const void* sfa,
+    const void* sfb, const void* bias_bf16, void* out_bf16,
+    int m, int n, int k, float alpha, cudaStream_t stream);
+
+// out[M,N] = GELU_taylor(alpha * (A@B^T) + bias[N])
+int nvfp4_gemm_bias_gelu_simt(
+    const void* a_packed, const void* b_packed, const void* sfa,
+    const void* sfb, const void* bias_bf16, void* out_bf16,
+    int m, int n, int k, float alpha, cudaStream_t stream);
+
+// out_packed[M,N/2] (e2m1 + SFA layout scale) = quant_fp4(GELU_taylor(alpha*(A@B^T)+bias[N]))
+int nvfp4_gemm_bias_gelu_fp4_simt(
+    const void* a_packed, const void* b_packed, const void* sfa,
+    const void* sfb, const void* bias_bf16, void* out_packed, void* out_sfa,
+    int m, int n, int k, float alpha, cudaStream_t stream);
+
+}  // namespace gemm
+}  // namespace flash_rt