diff --git a/fp4-gemm/VALIDATION.md b/fp4-gemm/VALIDATION.md index e4dd3fa..0df31f6 100644 --- a/fp4-gemm/VALIDATION.md +++ b/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. diff --git a/fp4-gemm/benchmarks/RESULTS.md b/fp4-gemm/benchmarks/RESULTS.md index ee84b1e..7f0a329 100644 --- a/fp4-gemm/benchmarks/RESULTS.md +++ b/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. diff --git a/fp4-gemm/build.toml b/fp4-gemm/build.toml index 750faad..d7e768f 100644 --- a/fp4-gemm/build.toml +++ b/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", +] diff --git a/fp4-gemm/csrc/gemm/fp4/portable_fp4_gemm_simt.cu b/fp4-gemm/csrc/gemm/fp4/portable_fp4_gemm_simt.cu new file mode 100644 index 0000000..482d10a --- /dev/null +++ b/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 +#include +#include +#include + +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(&scale_q) = sf[off]; + return static_cast(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 +__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(bs); + const int out_n_col_super = (n / 16 + 3) / 4; + out_sfa[sf_off(mi, bi * 16, out_n_col_super)] = + *reinterpret_cast(&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(bias); + const auto* res_b = static_cast(residual); + auto* out_b = static_cast<__nv_bfloat16*>(out); + if (mode == 0) { + gemm_epilogue_simt_kernel<<>>( + a, b, sfa, sfb, bias_b, res_b, out_b, m, n, k, alpha); + } else if (mode == 1) { + gemm_epilogue_simt_kernel<<>>( + a, b, sfa, sfb, bias_b, res_b, out_b, m, n, k, alpha); + } else if (mode == 2) { + gemm_epilogue_simt_kernel<<>>( + a, b, sfa, sfb, bias_b, res_b, out_b, m, n, k, alpha); + } else { + gemm_epilogue_simt_kernel<<>>( + 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(a_packed), static_cast(b_packed), + static_cast(sfa), static_cast(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(a_packed), static_cast(b_packed), + static_cast(sfa), static_cast(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(a_packed), static_cast(b_packed), + static_cast(sfa), static_cast(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(a_packed), static_cast(b_packed), + static_cast(sfa), static_cast(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(a_packed), static_cast(b_packed), + static_cast(sfa), static_cast(sfb), + static_cast(bias_bf16), + static_cast(out_packed), static_cast(out_sfa), + m, n, k, alpha); + return (cudaGetLastError() == cudaSuccess) ? 0 : 1; +} + +} // namespace gemm +} // namespace flash_rt diff --git a/fp4-gemm/csrc/gemm/fp4/portable_fp4_gemm_simt.cuh b/fp4-gemm/csrc/gemm/fp4/portable_fp4_gemm_simt.cuh new file mode 100644 index 0000000..0b6bdd1 --- /dev/null +++ b/fp4-gemm/csrc/gemm/fp4/portable_fp4_gemm_simt.cuh @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include + +// 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 diff --git a/fp4-gemm/torch-ext/torch_binding.cpp b/fp4-gemm/torch-ext/torch_binding.cpp index 4eadccc..9509416 100644 --- a/fp4-gemm/torch-ext/torch_binding.cpp +++ b/fp4-gemm/torch-ext/torch_binding.cpp @@ -19,6 +19,7 @@ #include "gemm/fp4/fp4_w4a4_mma_warpsplit_sm120.cuh" #endif #include "gemm/fp4/sm110_dispatch.cuh" +#include "gemm/fp4/portable_fp4_gemm_simt.cuh" #include "quantize/quantize_fp4_sfa.cuh" #include "registration.h" #include "torch_binding.h" @@ -155,20 +156,29 @@ void fp4_w4a4_gemv_warpsplit_bf16( #if defined(CUDA_KERNEL) at::cuda::CUDAGuard device_guard(a_packed.device()); auto const* props = current_device_properties(a_packed); - TORCH_CHECK(props->major == 12 && props->minor == 0, - "the warp-split GEMV is an SM120 kernel; got SM", - props->major, props->minor); + auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); + if (props->major == 11 && props->minor == 0) { + const int rc = flash_rt::gemm::nvfp4_gemm_linear_simt( + a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), + out.data_ptr(), checked_int(shape.m, "M"), checked_int(shape.n, "N"), + checked_int(shape.k, "K"), static_cast(alpha), stream); + TORCH_CHECK(rc == 0, + "fp4_w4a4_gemv_warpsplit_bf16 SIMT fallback failed rc=", rc); + } else { + TORCH_CHECK(props->major == 12 && props->minor == 0, + "the warp-split GEMV is an SM120 kernel; got SM", + props->major, props->minor); #if defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) - TORCH_CHECK(false, "SM120 FP4 GEMM source is not present in this build"); + TORCH_CHECK(false, "SM120 FP4 GEMM source is not present in this build"); #else - auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); - const int rc = flash_rt::gemm::fp4_w4a4_mma_sm120_warpsplit_bf16out( - a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(), - checked_int(shape.n, "N"), checked_int(shape.k, "K"), - sfa.data_ptr(), sfb.data_ptr(), static_cast(alpha), - static_cast(warps), static_cast(stages), stream); - TORCH_CHECK(rc == 0, "fp4_w4a4_gemv_warpsplit_bf16 failed with rc=", rc); + const int rc = flash_rt::gemm::fp4_w4a4_mma_sm120_warpsplit_bf16out( + a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(), + checked_int(shape.n, "N"), checked_int(shape.k, "K"), + sfa.data_ptr(), sfb.data_ptr(), static_cast(alpha), + static_cast(warps), static_cast(stages), stream); + TORCH_CHECK(rc == 0, "fp4_w4a4_gemv_warpsplit_bf16 failed with rc=", rc); #endif + } #else TORCH_CHECK(false, "fp4-gemm was not built with CUDA support"); #endif @@ -249,15 +259,25 @@ void nvfp4_gemm_residual_bf16( check_same_device(a_packed, out, "a_packed", "out"); #if defined(CUDA_KERNEL) at::cuda::CUDAGuard device_guard(a_packed.device()); - require_sm120(a_packed, "nvfp4_gemm_residual_bf16"); + auto const* props = current_device_properties(a_packed); auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); + if (props->major == 11 && props->minor == 0) { + const int rc = flash_rt::gemm::nvfp4_gemm_residual_simt( + a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), + residual.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"), + checked_int(shape.n, "N"), checked_int(shape.k, "K"), + static_cast(alpha), stream); + TORCH_CHECK(rc == 0, "nvfp4_gemm_residual_bf16 SIMT fallback failed rc=", rc); + } else { + require_sm120(a_packed, "nvfp4_gemm_residual_bf16"); #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) - flash_rt::gemm::fp4_w4a16_gemm_residual_sm120_bf16out( - a_packed.data_ptr(), b_packed.data_ptr(), residual.data_ptr(), - out.data_ptr(), checked_int(shape.m, "M"), checked_int(shape.n, "N"), - checked_int(shape.k, "K"), sfa.data_ptr(), sfb.data_ptr(), - static_cast(alpha), stream); + flash_rt::gemm::fp4_w4a16_gemm_residual_sm120_bf16out( + a_packed.data_ptr(), b_packed.data_ptr(), residual.data_ptr(), + out.data_ptr(), checked_int(shape.m, "M"), checked_int(shape.n, "N"), + checked_int(shape.k, "K"), sfa.data_ptr(), sfb.data_ptr(), + static_cast(alpha), stream); #endif + } #endif } @@ -280,15 +300,25 @@ void nvfp4_gemm_bias_gelu_bf16( check_same_device(a_packed, out, "a_packed", "out"); #if defined(CUDA_KERNEL) at::cuda::CUDAGuard device_guard(a_packed.device()); - require_sm120(a_packed, "nvfp4_gemm_bias_gelu_bf16"); + auto const* props = current_device_properties(a_packed); auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); + if (props->major == 11 && props->minor == 0) { + const int rc = flash_rt::gemm::nvfp4_gemm_bias_gelu_simt( + a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), + bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"), + checked_int(shape.n, "N"), checked_int(shape.k, "K"), + static_cast(alpha), stream); + TORCH_CHECK(rc == 0, "nvfp4_gemm_bias_gelu_bf16 SIMT fallback failed rc=", rc); + } else { + require_sm120(a_packed, "nvfp4_gemm_bias_gelu_bf16"); #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) - flash_rt::gemm::fp4_w4a16_gemm_bias_gelu_bf16out_sm120( - a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), - bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"), - checked_int(shape.n, "N"), checked_int(shape.k, "K"), - static_cast(alpha), stream); + flash_rt::gemm::fp4_w4a16_gemm_bias_gelu_bf16out_sm120( + a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), + bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"), + checked_int(shape.n, "N"), checked_int(shape.k, "K"), + static_cast(alpha), stream); #endif + } #endif } @@ -318,15 +348,25 @@ void nvfp4_gemm_bias_gelu_nvfp4( check_same_device(a_packed, out_sfa, "a_packed", "out_sfa"); #if defined(CUDA_KERNEL) at::cuda::CUDAGuard device_guard(a_packed.device()); - require_sm120(a_packed, "nvfp4_gemm_bias_gelu_nvfp4"); + auto const* props = current_device_properties(a_packed); auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); + if (props->major == 11 && props->minor == 0) { + const int rc = flash_rt::gemm::nvfp4_gemm_bias_gelu_fp4_simt( + a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), + bias.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(), + checked_int(shape.m, "M"), checked_int(shape.n, "N"), + checked_int(shape.k, "K"), static_cast(alpha), stream); + TORCH_CHECK(rc == 0, "nvfp4_gemm_bias_gelu_nvfp4 SIMT fallback failed rc=", rc); + } else { + require_sm120(a_packed, "nvfp4_gemm_bias_gelu_nvfp4"); #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) - flash_rt::gemm::fp4_w4a16_gemm_bias_gelu_fp4out_sm120( - a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), - bias.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(), - checked_int(shape.m, "M"), checked_int(shape.n, "N"), - checked_int(shape.k, "K"), static_cast(alpha), stream); + flash_rt::gemm::fp4_w4a16_gemm_bias_gelu_fp4out_sm120( + a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), + bias.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(), + checked_int(shape.m, "M"), checked_int(shape.n, "N"), + checked_int(shape.k, "K"), static_cast(alpha), stream); #endif + } #endif } @@ -344,14 +384,23 @@ void nvfp4_gemm_streamk_bf16( check_same_device(a_packed, out, "a_packed", "out"); #if defined(CUDA_KERNEL) at::cuda::CUDAGuard device_guard(a_packed.device()); - require_sm120(a_packed, "nvfp4_gemm_streamk_bf16"); + auto const* props = current_device_properties(a_packed); auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); + if (props->major == 11 && props->minor == 0) { + const int rc = flash_rt::gemm::nvfp4_gemm_linear_simt( + a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), + out.data_ptr(), checked_int(shape.m, "M"), checked_int(shape.n, "N"), + checked_int(shape.k, "K"), static_cast(alpha), stream); + TORCH_CHECK(rc == 0, "nvfp4_gemm_streamk_bf16 SIMT fallback failed rc=", rc); + } else { + require_sm120(a_packed, "nvfp4_gemm_streamk_bf16"); #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) - flash_rt::gemm::fp4_w4a16_gemm_dn_streamk_bf16out_sm120( - a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), - out.data_ptr(), checked_int(shape.m, "M"), checked_int(shape.n, "N"), - checked_int(shape.k, "K"), static_cast(alpha), stream); + flash_rt::gemm::fp4_w4a16_gemm_dn_streamk_bf16out_sm120( + a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), + out.data_ptr(), checked_int(shape.m, "M"), checked_int(shape.n, "N"), + checked_int(shape.k, "K"), static_cast(alpha), stream); #endif + } #endif } @@ -374,15 +423,25 @@ void nvfp4_gemm_streamk_bias_bf16( check_same_device(a_packed, out, "a_packed", "out"); #if defined(CUDA_KERNEL) at::cuda::CUDAGuard device_guard(a_packed.device()); - require_sm120(a_packed, "nvfp4_gemm_streamk_bias_bf16"); + auto const* props = current_device_properties(a_packed); auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream(); + if (props->major == 11 && props->minor == 0) { + const int rc = flash_rt::gemm::nvfp4_gemm_bias_simt( + a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), + bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"), + checked_int(shape.n, "N"), checked_int(shape.k, "K"), + static_cast(alpha), stream); + TORCH_CHECK(rc == 0, "nvfp4_gemm_streamk_bias_bf16 SIMT fallback failed rc=", rc); + } else { + require_sm120(a_packed, "nvfp4_gemm_streamk_bias_bf16"); #if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY) - flash_rt::gemm::fp4_w4a16_gemm_dn_streamk_bias_bf16out_sm120( - a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), - bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"), - checked_int(shape.n, "N"), checked_int(shape.k, "K"), - static_cast(alpha), stream); + flash_rt::gemm::fp4_w4a16_gemm_dn_streamk_bias_bf16out_sm120( + a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(), + bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"), + checked_int(shape.n, "N"), checked_int(shape.k, "K"), + static_cast(alpha), stream); #endif + } #endif }