feat(groot): GR00T N1.6 Thor parity fixes and opt-in 28.5 ms profile

CMakeLists.txt

@@ -1588,7 +1588,8 @@ if(ENABLE_SM100_CUTLASS)
     csrc/kernels/attention_mha_masked.cu
     csrc/kernels/vec_fp16_backbone.cu
     csrc/kernels/attention_seqused_fused.cu
-    csrc/kernels/rope_vec.cu)
+    csrc/kernels/rope_vec.cu
+    csrc/kernels/qk_norm_rope_rotate_half_bf16.cu)
   target_compile_definitions(flash_rt_kernels PRIVATE FLASHRT_HAVE_THOR_VLA_KERNELS=1)
   message(STATUS "Thor VLA helper kernels: ENABLED (sm_${GPU_ARCH})")
 else()
@@ -2056,6 +2057,7 @@ if(ENABLE_SM100_CUTLASS)
     csrc/gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cu
     csrc/quantize/quantize_fp4_sfa_bf16.cu
     csrc/fused_fp4/dit_norm_fp4_sfa.cu
+    csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cu
     csrc/gemm/fp4/cutlass_fp4_gemm_geglu_il_sm100.cu
     csrc/quantize/quantize_e0m3_sfa.cu
     csrc/gemm/fp4/cutlass_fp4_gemm_e0m3w_sm100.cu

csrc/attention/flash_attn_4_src/flashrt_fa4/cute/interface_fwd_sm100.py

@@ -56,9 +56,7 @@ def _raise(*_a, **_k):
 FlashAttentionForwardSm90 = _fa4_trimmed("FlashAttentionForwardSm90")
 FlashAttentionForwardSm120 = _fa4_trimmed("FlashAttentionForwardSm120")
 FlashAttentionMLAForwardSm100 = _fa4_trimmed("FlashAttentionMLAForwardSm100")
-from flashrt_fa4.cute.sm100_hd256_2cta_fmha_forward import (
-    BlackwellFusedMultiHeadAttentionForward,
-)
+BlackwellFusedMultiHeadAttentionForward = _fa4_trimmed("BlackwellFusedMultiHeadAttentionForward")
 
 from flashrt_fa4.cute.block_sparsity import (
     BlockSparseTensorsTorch,
@@ -580,12 +578,8 @@ def _flash_attn_fwd(
         and (tile_m % qhead_per_kvhead == 0 or not pack_gqa)
     )
 
-    # hd=256 2CTA forward uses the dedicated kernel on both SM100 and SM110.
-    use_dedicated_hd256_kernel = (
-        arch // 10 in [10, 11]
-        and head_dim == 256
-        and head_dim_v == 256
-    )
+    # hd=256 2CTA forward uses dedicated kernel (SM100 only; SM110 NVVM ICE)
+    use_dedicated_hd256_kernel = arch // 10 == 10 and head_dim == 256 and head_dim_v == 256
     use_2cta_instrs = use_2cta_instrs or use_dedicated_hd256_kernel
 
     if softcap is not None:

csrc/bindings.cpp

@@ -312,6 +312,9 @@ extern "C" void flash_rt_awq_quant_fp8_static_fp16(
 #include "quantize/bf16_quant_fp8_ncdhw_to_ndhwc.cuh"
 #endif
 #include "quantize/qkv_split_norm_rope_bf16.cuh"
+#ifdef FLASHRT_HAVE_THOR_VLA_KERNELS
+#include "kernels/qk_norm_rope_rotate_half_bf16.cuh"
+#endif
 #include "attention/fmha_dispatch.h"
 #ifdef ENABLE_MOTUS_SAGE2_RAW
 #include "attention/sage2/sage2_attn_raw.cuh"
@@ -1800,6 +1803,18 @@ PYBIND11_MODULE(flash_rt_kernels, m) {
        py::arg("eps") = 1e-5f, py::arg("stream") = 0);
 #endif  // FLASHRT_ENABLE_CHAMELEON
 
+#ifdef FLASHRT_HAVE_THOR_VLA_KERNELS
+    m.def("qk_norm_rope_rotate_half_bf16",
+          [](uintptr_t x, uintptr_t w, uintptr_t cos_t, uintptr_t sin_t,
+             int S, int NH, int HD, float eps, uintptr_t stream) -> int {
+        return flash_rt::kernels::qk_norm_rope_rotate_half_bf16(
+            to_ptr(x), to_ptr(w), to_ptr(cos_t), to_ptr(sin_t),
+            S, NH, HD, eps, to_stream(stream));
+    }, py::arg("x"), py::arg("w"), py::arg("cos_table"), py::arg("sin_table"),
+       py::arg("S"), py::arg("NH"), py::arg("HD"), py::arg("eps") = 1e-6f,
+       py::arg("stream") = 0);
+#endif  // FLASHRT_HAVE_THOR_VLA_KERNELS
+
     m.def("gate_mul_residual_fp16",
           [](uintptr_t residual, uintptr_t x, uintptr_t gate,
              int n, uintptr_t stream) {

csrc/fp4_bindings.cpp

@@ -41,6 +41,7 @@
 #include "gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cuh"
 #include "quantize/quantize_fp4_sfa_bf16.cuh"
 #include "fused_fp4/dit_norm_fp4_sfa.cuh"
+#include "fused_fp4/silu_mul_fp4_sfa_bf16.cuh"
 #include "fused_fp4/layer_norm_fp4_sfa.cuh"
 #include "gemm/fp4/cutlass_fp4_gemm_siglip_ffn_sm100.cuh"
 
@@ -1228,6 +1229,37 @@ same contract as cutlass_fp4_gemm_geglu_il_hw.
         py::arg("stream") = 0,
         "Fused no-affine LayerNorm (bf16) -> NVFP4 packed + SFA.");
 
+  m.def("rms_norm_weight_fp4_sfa_bf16",
+        [](uintptr_t x, uintptr_t weight, uintptr_t packed, uintptr_t sfa,
+           int seq_len, int dim, float eps, uintptr_t stream) -> int {
+          return flash_rt::fused_fp4::rms_norm_weight_fp4_sfa_bf16(
+              reinterpret_cast<void const*>(x),
+              reinterpret_cast<void const*>(weight),
+              reinterpret_cast<void*>(packed),
+              reinterpret_cast<void*>(sfa),
+              seq_len, dim, eps,
+              reinterpret_cast<cudaStream_t>(stream));
+        },
+        py::arg("x"), py::arg("weight"), py::arg("packed"), py::arg("sfa"),
+        py::arg("seq_len"), py::arg("dim"), py::arg("eps") = 1e-5f,
+        py::arg("stream") = 0,
+        "Fused weighted RMSNorm (bf16) -> NVFP4 packed + SFA.");
+
+  m.def("silu_mul_fp4_sfa_bf16",
+        [](uintptr_t gate, uintptr_t up, uintptr_t packed, uintptr_t sfa,
+           int N, int D, bool is_sfb, uintptr_t stream) -> int {
+          return flash_rt::fused_fp4::silu_mul_fp4_sfa_bf16(
+              reinterpret_cast<void const*>(gate),
+              reinterpret_cast<void const*>(up),
+              reinterpret_cast<void*>(packed),
+              reinterpret_cast<void*>(sfa),
+              N, D, is_sfb,
+              reinterpret_cast<cudaStream_t>(stream));
+        },
+        py::arg("gate"), py::arg("up"), py::arg("packed"), py::arg("sfa"),
+        py::arg("N"), py::arg("D"), py::arg("is_sfb"), py::arg("stream") = 0,
+        "Fused SiLU(gate)*up (bf16) -> NVFP4 packed + SFA.");
+
   m.attr("__version__") = "0.1.0-dev";
   m.attr("layout_note") = "scales are linear [N, D/16]; Phase 4 adds tile-interleave conversion";
 }

csrc/fused_fp4/dit_norm_fp4_sfa.cu

@@ -206,3 +206,104 @@ int layer_norm_no_affine_fp4_sfa_bf16(
 
 }  // namespace fused_fp4
 }  // namespace flash_rt
+
+// ----------------------------------------------------------------------------
+//  Weighted RMSNorm (bf16) -> NVFP4 quantize + SFA (GR00T N1.6 Qwen3 tier).
+//  y = x * rsqrt(mean(x^2) + eps) * weight, rounded through bf16 before
+//  quantization (matches the torch RMSNorm(bf16) + quantize chain).
+// ----------------------------------------------------------------------------
+namespace flash_rt {
+namespace fused_fp4 {
+
+#if FV_HAVE_CUTLASS
+
+namespace {
+
+template <class LayoutSF>
+__global__ void rms_norm_w_fp4_sfa_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const __nv_bfloat16* __restrict__ weight,
+    uint2* __restrict__ packed,
+    uint8_t* __restrict__ dst_sfa,
+    LayoutSF layout,
+    int D, float eps) {
+    const int r = blockIdx.x;
+    const __nv_bfloat162* row2 =
+        reinterpret_cast<const __nv_bfloat162*>(x + static_cast<long>(r) * D);
+    const int D2 = D >> 1;
+    __shared__ float sh[32];
+
+    float ssq = 0.f;
+    for (int i = threadIdx.x; i < D2; i += blockDim.x) {
+        const __nv_bfloat162 v = row2[i];
+        const float a = __bfloat162float(v.x);
+        const float b = __bfloat162float(v.y);
+        ssq += a * a + b * b;
+    }
+    const float rstd = rsqrtf(block_sum_dn(ssq, sh) / D + eps);
+
+    const int n_blocks = D >> 4;
+    const int4* x4 = reinterpret_cast<const int4*>(x + static_cast<long>(r) * D);
+    const int4* w4 = reinterpret_cast<const int4*>(weight);
+    for (int blk = threadIdx.x; blk < n_blocks; blk += blockDim.x) {
+        const int4 xr[2] = {x4[2 * blk], x4[2 * blk + 1]};
+        const int4 wr[2] = {w4[2 * blk], w4[2 * blk + 1]};
+        const __nv_bfloat16* xh = reinterpret_cast<const __nv_bfloat16*>(xr);
+        const __nv_bfloat16* wh = reinterpret_cast<const __nv_bfloat16*>(wr);
+        float vals[16];
+        float amax = 0.f;
+        #pragma unroll
+        for (int i = 0; i < 16; ++i) {
+            const float normed =
+                __bfloat162float(xh[i]) * rstd * __bfloat162float(wh[i]);
+            vals[i] = __bfloat162float(__float2bfloat16(normed));
+            const float a = fabsf(vals[i]);
+            if (a > amax) amax = a;
+        }
+        float desired = amax / 6.f;
+        if (desired < 1e-12f) desired = 1e-12f;
+        __nv_fp8_e4m3 bs_q = __nv_fp8_e4m3(fmaxf(desired, 0.f));
+        const float bs_dq = static_cast<float>(bs_q);
+        dst_sfa[layout(r, blk * 16, 0)] = *reinterpret_cast<uint8_t*>(&bs_q);
+        const float inv_bs = 1.f / bs_dq;
+        uint2 out;
+        uint8_t* ob = reinterpret_cast<uint8_t*>(&out);
+        #pragma unroll
+        for (int p = 0; p < 8; ++p) {
+            const uint8_t lo = fp32_to_e2m1_dn(vals[2 * p] * inv_bs);
+            const uint8_t hi = fp32_to_e2m1_dn(vals[2 * p + 1] * inv_bs);
+            ob[p] = static_cast<uint8_t>(lo | (hi << 4));
+        }
+        packed[static_cast<long>(r) * n_blocks + blk] = out;
+    }
+}
+
+}  // namespace
+
+#endif  // FV_HAVE_CUTLASS
+
+int rms_norm_weight_fp4_sfa_bf16(
+    const void* x, const void* weight, void* packed, void* sfa,
+    int seq_len, int dim, float eps, cudaStream_t stream) {
+#if FV_HAVE_CUTLASS
+    if (check_dn_args(x, packed, dim) != 0) return -1;
+    if (reinterpret_cast<uintptr_t>(weight) & 15) return -1;
+    auto shape = cute::make_shape(seq_len, 1, dim, 1);
+    auto layout = CfgDN::tile_atom_to_shape_SFA(shape);
+    rms_norm_w_fp4_sfa_kernel<<<seq_len, 128, 0, stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(x),
+        reinterpret_cast<const __nv_bfloat16*>(weight),
+        reinterpret_cast<uint2*>(packed),
+        reinterpret_cast<uint8_t*>(sfa),
+        layout, dim, eps);
+    const cudaError_t e = cudaGetLastError();
+    return (e == cudaSuccess) ? 0 : -static_cast<int>(e);
+#else
+    (void)x; (void)weight; (void)packed; (void)sfa;
+    (void)seq_len; (void)dim; (void)eps; (void)stream;
+    return -2;
+#endif
+}
+
+}  // namespace fused_fp4
+}  // namespace flash_rt

csrc/fused_fp4/dit_norm_fp4_sfa.cuh

@@ -29,5 +29,11 @@ int layer_norm_no_affine_fp4_sfa_bf16(
     const void* x, void* packed, void* sfa,
     int seq_len, int dim, float eps, cudaStream_t stream);
 
+// packed/sfa = quantize(RMSNorm(x[row]) * weight)  — no mean removal.
+// x: bf16 [S, D]; weight: bf16 [D]. Qwen3 pre-attn / pre-FF norms.
+int rms_norm_weight_fp4_sfa_bf16(
+    const void* x, const void* weight, void* packed, void* sfa,
+    int seq_len, int dim, float eps, cudaStream_t stream);
+
 }  // namespace fused_fp4
 }  // namespace flash_rt

csrc/fused_fp4/silu_mul_fp4_sfa_bf16.cu

@@ -0,0 +1,152 @@
+// ============================================================================
+//  Fused SiLU(gate) * up (bf16) + NVFP4 quantize + SFA write.
+//  One thread per 16-element block: two int4 loads per operand, silu in
+//  fp32 rounded to bf16, bf16 multiply, then the standard per-block scale
+//  selection + e2m1 rounding + tile-interleaved SFA byte.
+// ============================================================================
+#include "silu_mul_fp4_sfa_bf16.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(__CUDA_ARCH__)
+#  include "cutlass/cutlass.h"
+#  include "cutlass/detail/sm100_blockscaled_layout.hpp"
+#  include "cute/tensor.hpp"
+#  define FV_HAVE_CUTLASS 1
+#else
+#  define FV_HAVE_CUTLASS 0
+#endif
+
+namespace flash_rt {
+namespace fused_fp4 {
+
+#if FV_HAVE_CUTLASS
+
+namespace {
+
+using CfgSM = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+__device__ __forceinline__ uint8_t fp32_to_e2m1_sm(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;
+}
+
+template <class LayoutSF>
+__global__ void kernel_silu_mul_fp4_sfa_bf16(
+    const int4* __restrict__ gate,
+    const int4* __restrict__ up,
+    uint2* __restrict__ dst_packed,
+    uint8_t* __restrict__ dst_sfa,
+    LayoutSF layout,
+    int N, int D8) {
+  const int block_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  const int row = blockIdx.y;
+  const int n_blocks = D8 >> 1;
+  if (row >= N || block_idx >= n_blocks) return;
+
+  const int4 g0 = gate[row * D8 + 2 * block_idx];
+  const int4 g1 = gate[row * D8 + 2 * block_idx + 1];
+  const int4 u0 = up[row * D8 + 2 * block_idx];
+  const int4 u1 = up[row * D8 + 2 * block_idx + 1];
+  const __nv_bfloat16* gh0 = reinterpret_cast<const __nv_bfloat16*>(&g0);
+  const __nv_bfloat16* gh1 = reinterpret_cast<const __nv_bfloat16*>(&g1);
+  const __nv_bfloat16* uh0 = reinterpret_cast<const __nv_bfloat16*>(&u0);
+  const __nv_bfloat16* uh1 = reinterpret_cast<const __nv_bfloat16*>(&u1);
+
+  float vals[16];
+  float amax = 0.f;
+  #pragma unroll
+  for (int i = 0; i < 8; ++i) {
+    const float g[2] = {__bfloat162float(gh0[i]), __bfloat162float(gh1[i])};
+    const __nv_bfloat16 u[2] = {uh0[i], uh1[i]};
+    #pragma unroll
+    for (int h = 0; h < 2; ++h) {
+      const float s = g[h] / (1.f + expf(-g[h]));         // silu fp32
+      const __nv_bfloat16 sb = __float2bfloat16(s);       // round like torch
+      const __nv_bfloat16 prod = __hmul(sb, u[h]);        // bf16 multiply
+      vals[h * 8 + i] = __bfloat162float(prod);
+      const float a = fabsf(vals[h * 8 + i]);
+      if (a > amax) amax = a;
+    }
+  }
+
+  float desired = amax / 6.f;
+  if (desired < 1e-12f) desired = 1e-12f;
+  __nv_fp8_e4m3 bs_q = __nv_fp8_e4m3(fmaxf(desired, 0.f));
+  const float bs_dq = static_cast<float>(bs_q);
+
+  dst_sfa[layout(row, block_idx * 16, 0)] =
+      *reinterpret_cast<uint8_t*>(&bs_q);
+
+  const float inv_bs = 1.f / bs_dq;
+  uint2 out;
+  uint8_t* ob = reinterpret_cast<uint8_t*>(&out);
+  #pragma unroll
+  for (int p = 0; p < 8; ++p) {
+    const uint8_t lo = fp32_to_e2m1_sm(vals[2 * p] * inv_bs);
+    const uint8_t hi = fp32_to_e2m1_sm(vals[2 * p + 1] * inv_bs);
+    ob[p] = static_cast<uint8_t>(lo | (hi << 4));
+  }
+  dst_packed[row * n_blocks + block_idx] = out;
+}
+
+}  // namespace
+
+#endif  // FV_HAVE_CUTLASS
+
+int silu_mul_fp4_sfa_bf16(
+    const void* gate, const void* up, void* packed, void* sfa,
+    int N, int D, bool is_sfb, cudaStream_t stream) {
+#if FV_HAVE_CUTLASS
+  if (D % 16 != 0) return -1;
+  if ((reinterpret_cast<uintptr_t>(gate) & 15) ||
+      (reinterpret_cast<uintptr_t>(up) & 15) ||
+      (reinterpret_cast<uintptr_t>(packed) & 7)) return -1;
+  const int n_blocks = D / 16;
+  const int threads = 128;
+  dim3 grid((n_blocks + threads - 1) / threads, N);
+
+  auto shape = cute::make_shape(
+      is_sfb ? 1 : N,
+      is_sfb ? N : 1,
+      D, 1);
+
+  if (is_sfb) {
+    auto layout = CfgSM::tile_atom_to_shape_SFB(shape);
+    kernel_silu_mul_fp4_sfa_bf16<<<grid, threads, 0, stream>>>(
+        reinterpret_cast<const int4*>(gate),
+        reinterpret_cast<const int4*>(up),
+        reinterpret_cast<uint2*>(packed),
+        reinterpret_cast<uint8_t*>(sfa),
+        layout, N, D >> 3);
+  } else {
+    auto layout = CfgSM::tile_atom_to_shape_SFA(shape);
+    kernel_silu_mul_fp4_sfa_bf16<<<grid, threads, 0, stream>>>(
+        reinterpret_cast<const int4*>(gate),
+        reinterpret_cast<const int4*>(up),
+        reinterpret_cast<uint2*>(packed),
+        reinterpret_cast<uint8_t*>(sfa),
+        layout, N, D >> 3);
+  }
+  const cudaError_t e = cudaGetLastError();
+  return (e == cudaSuccess) ? 0 : -static_cast<int>(e);
+#else
+  (void)gate; (void)up; (void)packed; (void)sfa;
+  (void)N; (void)D; (void)is_sfb; (void)stream;
+  return -2;
+#endif
+}
+
+}  // namespace fused_fp4
+}  // namespace flash_rt