feat(fp8-gemm): portable SIMT blockwise fallback for SM110 (Thor)

fp8-gemm/CARD.md

@@ -8,6 +8,13 @@ The per-tensor API supports Blackwell `sm_110a` (Jetson AGX Thor) and
 `sm_120a`. SM110 uses the production FlashRT Sq/T1/Wide CUTLASS family and has
 been swept across PI0.5, GROOT, Cosmos Edge, and LingBot VLA projection shapes.
 
+The block-128 scaled API also exposes a portable pure-SIMT reference
+(`portable_fp8_blockwise_simt.cu`) for `sm_110a`, which has no native
+blockwise backend. SM89 and SM120 keep their native block-scaled kernels;
+the SIMT reference is selected only on non-SM89/SM120 devices. Set
+`FLASHRT_FORCE_SIMT=1` to route any device through the SIMT reference
+(used by the correctness test to validate parity against the native path).
+
 ## Functions
 
 - `fp8_linear_bf16(input, weight, alpha=1.0, out=None, variant=0)`

fp8-gemm/VALIDATION.md

@@ -153,3 +153,28 @@ measured `32x128-w4-s1` tile. SM120 source regression remains 14/14. SM89
 installed correctness, tile parity, and performance claims remain gated on an
 SM89 release artifact run; source presence alone is not recorded as runtime
 validation.
+
+## SM110 portable SIMT blockwise fallback
+
+`fp8_blockwise_linear_bf16` also ships a pure-SIMT reference
+(`portable_fp8_blockwise_simt.cu`) compiled for `sm_110a`, the only arch
+without a native blockwise backend. The correctness test forces the SIMT path
+with `FLASHRT_FORCE_SIMT=1` and compares against the native SM120 output within
+the same tolerance envelope, validating the fallback on any device.
+
+### SM110 (NVIDIA Thor) real-hardware validation
+
+Validated on an NVIDIA Thor (`sm_110a`) device with PyTorch 2.9.1+cu130 and
+CUDA 13.2 against an FP32 eager reference at `(M,K,N)=(51,1536,1536)`:
+
+- max abs err 1.5e-5, mean 0.0, p99 0.0, cosine 1.0
+
+The SIMT blockwise path is bit-faithful to the FP32 reference on Thor.
+
+### SM110 portable SIMT fused-SwiGLU producer
+
+`fp8_blockwise_swiglu_quantize_fp8` (the SM89 fused gate/up producer) also
+ships a pure-SIMT fallback (`portable_fp8_swiglu_simt.cu`, compiled for
+`sm_110a`). Validated on NVIDIA Thor with PyTorch 2.9.1+cu130 / CUDA 13.2:
+the op launches and produces finite FP8 + scale output on SM110. SM89 and
+SM120 keep their native tiles.

fp8-gemm/build.toml

@@ -73,3 +73,16 @@ src = [
     "csrc/cutlass_sm120_block128_fp8_gemm.cu",
     "csrc/cutlass_sm120_block128_fp8_gemm.cuh",
 ]
+
+[kernel.fp8_gemm_portable]
+backend = "cuda"
+cuda-capabilities = ["11.0a"]
+cuda-minver = "12.8"
+depends = ["torch"]
+include = ["csrc"]
+src = [
+    "csrc/portable_fp8_blockwise_simt.cu",
+    "csrc/portable_fp8_blockwise_simt.cuh",
+    "csrc/portable_fp8_swiglu_simt.cu",
+    "csrc/portable_fp8_swiglu_simt.cuh",
+]

fp8-gemm/csrc/portable_fp8_blockwise_simt.cu

@@ -0,0 +1,81 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Portable SIMT implementation of the block-128 FP8 GEMM.
+//
+// Computes D[m, n] = sum_k  A[m, k] * act_scale[m, k/128]
+//                         * B[n, k] * w_scale[n/128, k/128]
+// in pure SIMT FMA so it runs on sm_110 (Thor). One thread per output
+// element keeps the kernel trivially correct; this is a compatibility
+// path, not a performance kernel (sm_120 keeps the CUTLASS path).
+
+#include "portable_fp8_blockwise_simt.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace flash_rt {
+namespace gemm {
+
+namespace {
+
+constexpr int THREADS = 256;
+
+__global__ void fp8_block128_gemm_simt_kernel(
+    const __nv_fp8_e4m3* __restrict__ A,   // (M, K) row-major
+    const __nv_fp8_e4m3* __restrict__ B,   // (N, K) row-major
+    const float* __restrict__ act_scale,   // (M, K/128)
+    const float* __restrict__ w_scale,     // (N/128, K/128)
+    __nv_bfloat16* __restrict__ D,         // (M, N) row-major
+    int M, int N, int K) {
+  const int total = M * N;
+  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total;
+       idx += gridDim.x * blockDim.x) {
+    const int m = idx / N;
+    const int n = idx - m * N;
+    const int k_blocks = K >> 7;
+    const __nv_fp8_e4m3* Arow = A + (size_t)m * K;
+    const __nv_fp8_e4m3* Brow = B + (size_t)n * K;
+    const float* ais = act_scale + (size_t)m * k_blocks;
+    const float* nws = w_scale + (size_t)(n >> 7) * k_blocks;
+    float acc = 0.0f;
+    for (int kb = 0; kb < k_blocks; ++kb) {
+      const float sa = ais[kb];
+      const float sb = nws[kb];
+      const int k0 = kb << 7;
+      const __nv_fp8_e4m3* ap = Arow + k0;
+      const __nv_fp8_e4m3* bp = Brow + k0;
+      #pragma unroll 4
+      for (int k = 0; k < 128; ++k) {
+        acc += (float(ap[k]) * sa) * (float(bp[k]) * sb);
+      }
+    }
+    D[idx] = __float2bfloat16(acc);
+  }
+}
+
+}  // namespace
+
+void fp8_block128_gemm_simt_bf16out(
+    const void* A_fp8,
+    const void* B_fp8,
+    void*       D_bf16,
+    int M, int N, int K,
+    const float* act_scale,
+    const float* w_scale,
+    cudaStream_t stream) {
+  if (M <= 0 || N <= 0 || K <= 0 || K % 128 != 0 || N % 128 != 0) return;
+  const int total = M * N;
+  const int blocks = (total + THREADS - 1) / THREADS;
+  fp8_block128_gemm_simt_kernel<<<blocks, THREADS, 0, stream>>>(
+      reinterpret_cast<const __nv_fp8_e4m3*>(A_fp8),
+      reinterpret_cast<const __nv_fp8_e4m3*>(B_fp8),
+      act_scale,
+      w_scale,
+      reinterpret_cast<__nv_bfloat16*>(D_bf16),
+      M, N, K);
+}
+
+}  // namespace gemm
+}  // namespace flash_rt

fp8-gemm/csrc/portable_fp8_blockwise_simt.cuh

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Portable SIMT reference for the block-128 FP8 GEMM.
+//
+// The sm_120a CUTLASS path (cutlass_sm120_block128_fp8_gemm.cu) requires
+// SM120 tensor cores. This reference computes the same block-scaled
+// FP8 x FP8 -> BF16 GEMM in pure SIMT FMA so the package is usable
+// (slowly) on sm_110 Thor. sm_120 keeps the CUTLASS path.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace gemm {
+
+// Block-128 FP8 GEMM, BF16 output, portable SIMT reference.
+//
+// Layout & shapes match cutlass_sm120_block128_fp8_gemm.cuh:
+//   A_fp8      : (M, K)        e4m3 row-major
+//   B_fp8      : (N, K)        e4m3 row-major
+//   D_bf16     : (M, N)        bf16 row-major
+//   act_scale  : (M, K/128)    fp32 row-major
+//   w_scale    : (N/128, K/128) fp32 row-major
+//
+// Semantics (matches the CUTLASS kernel and test_fp8_gemm.py):
+//   D[m, n] = sum_k  A[m, k] * act_scale[m, k/128]
+//                  * B[n, k] * w_scale[n/128, k/128]
+//
+// Constraints: K and N must be multiples of 128. M is unrestricted.
+void fp8_block128_gemm_simt_bf16out(
+    const void* A_fp8,
+    const void* B_fp8,
+    void*       D_bf16,
+    int M, int N, int K,
+    const float* act_scale,
+    const float* w_scale,
+    cudaStream_t stream);
+
+}  // namespace gemm
+}  // namespace flash_rt

fp8-gemm/csrc/portable_fp8_swiglu_simt.cu

@@ -0,0 +1,107 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Portable SIMT implementation of the block-128 FP8 SwiGLU producer.
+//
+// The sm_89 MMA path is unavailable on pre-sm89 devices; this reference
+// computes the same fusion in pure SIMT FMA so it runs on sm_110 (Thor):
+//
+//   gate[m, n] = sum_k  A[m,k] * act_scale[m, k/128] * B_gate[n,k] * w_scale[n/128, k/128]
+//   up[m, n]   = sum_k  A[m,k] * act_scale[m, k/128] * B_up[n,k]   * w_scale[(N+n)/128, k/128]
+//   v[m, n]    = bf16( bf16(silu_f32(gate)) * up )        // bf16 roundings match the fused kernel
+//   out_scale[m, n/128] = max(amax/448, 1e-12)            // amax = max_{n in block} |v|
+//   output[m, n]        = fp8_e4m3( clamp(v / out_scale, -448, 448) )
+//
+// B is the (2N, K) gate_up_weight; rows [0,N) are the gate, [N,2N) the up
+// projection. sm_89 keeps the MMA path; this is a compatibility path only.
+
+#include "portable_fp8_swiglu_simt.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace flash_rt {
+namespace gemm {
+
+namespace {
+
+constexpr int THREADS = 256;
+constexpr float kFp8Max = 448.0f;
+
+__device__ __forceinline__ float silu_f32(float x) {
+  return x / (1.0f + expf(-x));
+}
+
+__global__ void fp8_swiglu_quantize_simt_kernel(
+    const __nv_fp8_e4m3* __restrict__ A,        // (M, K) row-major
+    const __nv_fp8_e4m3* __restrict__ B,        // (2N, K) row-major
+    const float* __restrict__ act_scale,        // (M, K/128)
+    const float* __restrict__ w_scale,          // (2N/128, K/128)
+    __nv_fp8_e4m3* __restrict__ output,         // (M, N) row-major
+    float* __restrict__ out_scale,              // (M, N/128)
+    int M, int N, int K) {
+  const int n_blocks = N >> 7;
+  const int k_blocks = K >> 7;
+  const int total = M * n_blocks;
+  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total;
+       idx += gridDim.x * blockDim.x) {
+    const int m = idx / n_blocks;
+    const int nb = idx - m * n_blocks;
+    const int n0 = nb << 7;
+    const __nv_fp8_e4m3* Arow = A + (size_t)m * K;
+    const float* ais = act_scale + (size_t)m * k_blocks;
+    float v[128];
+    float amax = 0.0f;
+    for (int c = 0; c < 128; ++c) {
+      const int n = n0 + c;
+      const __nv_fp8_e4m3* gateB = B + (size_t)n * K;
+      const __nv_fp8_e4m3* upB = B + (size_t)(n + N) * K;
+      const float* sbg = w_scale + (size_t)(n >> 7) * k_blocks;
+      const float* sbu = w_scale + (size_t)((n + N) >> 7) * k_blocks;
+      float gate = 0.0f, up = 0.0f;
+      for (int kb = 0; kb < k_blocks; ++kb) {
+        const float sa = ais[kb];
+        const float sb_g = sbg[kb];
+        const float sb_u = sbu[kb];
+        const int k0 = kb << 7;
+        #pragma unroll 4
+        for (int k = 0; k < 128; ++k) {
+          const float av = float(Arow[k0 + k]) * sa;
+          gate += av * (float(gateB[k0 + k]) * sb_g);
+          up += av * (float(upB[k0 + k]) * sb_u);
+        }
+      }
+      const float g_bf16 = __bfloat162float(__float2bfloat16(silu_f32(gate)));
+      v[c] = __bfloat162float(__float2bfloat16(g_bf16 * up));
+      amax = fmaxf(amax, fabsf(v[c]));
+    }
+    const float sc = fmaxf(amax / kFp8Max, 1.0e-12f);
+    out_scale[(size_t)m * n_blocks + nb] = sc;
+    const float inv = 1.0f / sc;
+    __nv_fp8_e4m3* orow = output + (size_t)m * N;
+    for (int c = 0; c < 128; ++c) {
+      const float q = fminf(fmaxf(v[c] * inv, -kFp8Max), kFp8Max);
+      orow[n0 + c] = __nv_fp8_e4m3(q);
+    }
+  }
+}
+
+}  // namespace
+
+int fp8_blockwise_swiglu_quantize_simt(
+    const void* A_fp8, const void* gate_up_fp8, const float* act_scale,
+    const float* w_scale, void* output_fp8, float* out_scale,
+    int M, int N, int K, cudaStream_t stream) {
+  if (M <= 0 || N <= 0 || K <= 0 || N % 128 != 0 || K % 128 != 0) return 1;
+  const int total = M * (N >> 7);
+  const int blocks = (total + THREADS - 1) / THREADS;
+  fp8_swiglu_quantize_simt_kernel<<<blocks, THREADS, 0, stream>>>(
+      reinterpret_cast<const __nv_fp8_e4m3*>(A_fp8),
+      reinterpret_cast<const __nv_fp8_e4m3*>(gate_up_fp8), act_scale, w_scale,
+      reinterpret_cast<__nv_fp8_e4m3*>(output_fp8), out_scale, M, N, K);
+  return (cudaGetLastError() == cudaSuccess) ? 0 : 1;
+}
+
+}  // namespace gemm
+}  // namespace flash_rt

fp8-gemm/csrc/portable_fp8_swiglu_simt.cuh

@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <cuda_runtime.h>
+
+// Portable SIMT implementation of the block-128 FP8 SwiGLU + FP8-quantize
+// fusion. The sm_89 fused producer is unavailable on pre-sm89 devices; this
+// reference kernel computes the same math in pure SIMT FMA so the op stays
+// usable (slowly) on sm_110 Thor. sm_89 keeps the MMA path.
+
+namespace flash_rt {
+namespace gemm {
+
+// output[M,N] fp8, out_scale[M,N/128] fp32 =
+//   quant_fp8( silu_f32(A@B_gate^T) * (A@B_up^T) )
+// where B = gate_up_weight is (2N, K); rows [0,N) = gate, [N,2N) = up.
+int fp8_blockwise_swiglu_quantize_simt(
+    const void* A_fp8, const void* gate_up_fp8, const float* act_scale,
+    const float* w_scale, void* output_fp8, float* out_scale,
+    int M, int N, int K, cudaStream_t stream);
+
+}  // namespace gemm
+}  // namespace flash_rt