New fusion kernels: QK-norm+RoPE, K0g, SwiGLU requant, FP8 dequant

include/xe-fuse/builder/epilogue_builder.hpp

@@ -137,6 +137,29 @@ template <typename ElementResidual = cutlass::bfloat16_t,
           typename ElementCompute = float>
 using AddResidual = Add<Acc, AuxLoad<ElementResidual>, ElementCompute, ElementCompute>;
 
+// GateAcc: gate[m] * acc — per-token gate (ColBroadcast) applied to accumulator.
+// gate[m] is a scalar per output row; used before residual add in DiT blocks.
+template <typename TileShape,
+          typename ElementGate    = float,
+          typename ElementCompute = float>
+using GateAcc = Mul<ColBroadcast<0, TileShape, ElementGate, ElementCompute>,
+                    Acc,
+                    ElementCompute, ElementCompute>;
+
+// GateResidualGamma: gamma[n] * (gate[m] * acc + residual)  — K0g pattern.
+// Extends K0a with a per-token gate that modulates the GEMM output before
+// the residual add.  Used in FLUX.2-style DiT single-block transformer blocks.
+template <typename TileShape,
+          typename ElementResidual = cutlass::bfloat16_t,
+          typename ElementGate     = float,
+          typename ElementGamma    = float,
+          typename ElementCompute  = float>
+using GateResidualGamma = ScaleCols<
+    Add<GateAcc<TileShape, ElementGate, ElementCompute>,
+        AuxLoad<ElementResidual>,
+        ElementCompute, ElementCompute>,
+    TileShape, ElementGamma, ElementCompute>;
+
 // ============================================================
 // Pairwise Operations — lane-shuffle-based pair computations
 // ============================================================
@@ -279,6 +302,31 @@ template <typename TileShape,
           typename ElementScale = float, typename ElementCompute = float>
 using DequantGeGLU = GeGLU<DequantW8A8<TileShape, ElementScale, ElementCompute>>;
 
+// DequantFP8: float_acc * scale_a[m] * scale_b[n] → bf16
+// FP8×FP8 GEMM dequantization. Identical EVT structure to DequantW8A8 but
+// ElementAcc is float (FP8 upcasts to FP16 before XMX, accumulates in float).
+// scale_a[m] = per-token input scale, scale_b[n] = per-channel weight scale.
+template <typename TileShape,
+          typename ElementScale = float, typename ElementCompute = float>
+using DequantFP8 = Mul<
+    Mul<Acc,
+        ColBroadcast<0, TileShape, ElementScale, ElementCompute>,
+        ElementCompute, ElementCompute>,
+    RowBroadcast<0, TileShape, ElementScale, ElementCompute>,
+    ElementCompute, ElementCompute>;
+
+// DequantFP8SwiGLU: SwiGLU( float_acc * scale_a[m] * scale_b[n] ) → bf16
+// FP8 FFN kernel (K2_FP8): gate+up projection with dequant and SwiGLU fused.
+template <typename TileShape,
+          typename ElementScale = float, typename ElementCompute = float>
+using DequantFP8SwiGLU = SwiGLU<DequantFP8<TileShape, ElementScale, ElementCompute>>;
+
+// DequantFP8GeGLU: GeGLU( float_acc * scale_a[m] * scale_b[n] ) → bf16
+// FP8 FFN kernel for Gemma-style models.
+template <typename TileShape,
+          typename ElementScale = float, typename ElementCompute = float>
+using DequantFP8GeGLU = GeGLU<DequantFP8<TileShape, ElementScale, ElementCompute>>;
+
 // HadamardOutput<InnerEVT, GroupSize>: apply WHT to the output of InnerEVT
 //
 // Used as the final epilogue step in K0_W8A8 (O-projection) for QuaRot:

include/xe-fuse/kernels/compute_rstd.hpp

@@ -3,6 +3,8 @@
 #include <cstdint>
 #include <sycl/sycl.hpp>
 
+#include "cutlass/bfloat16.h"
+
 namespace xe_fuse {
 
 // Standalone rstd reduction kernel.
@@ -124,4 +126,72 @@ void launch_compute_rstd_and_quantize(
   });
 }
 
+// Dual-output RMSNorm + INT8 quantization kernel.
+//
+// Same three-pass algorithm as launch_compute_rstd_and_quantize but also
+// writes a BF16 normed output for the residual path.  Use when both a normed
+// BF16 value (residual path) and an INT8 quantized value (next GEMM input)
+// are needed from the same input.
+template <typename ElementInput, typename ElementNormed = cutlass::bfloat16_t>
+void launch_norm_quantize_dual(
+    sycl::queue&        q,
+    ElementInput const* input_ptr,
+    int8_t*             quant_out_ptr,
+    ElementNormed*      normed_out_ptr,
+    float*              scale_token_ptr,
+    int M, int N, int L,
+    float eps = 1e-6f)
+{
+  constexpr int SG_SIZE = 16;
+  int work_groups = M * L;
+
+  q.submit([&](sycl::handler& cgh) {
+    cgh.parallel_for(
+      sycl::nd_range<1>(static_cast<size_t>(work_groups) * SG_SIZE, SG_SIZE),
+      [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(SG_SIZE)]] {
+        int row  = item.get_group(0);
+        int lane = item.get_local_id(0);
+
+        // ── Pass 1: sum_sq → rstd ──────────────────────────────────────────
+        float sum_sq = 0.f;
+        for (int col = lane; col < N; col += SG_SIZE) {
+          float v = static_cast<float>(input_ptr[row * N + col]);
+          sum_sq += v * v;
+        }
+        auto sg = item.get_sub_group();
+        for (int off = SG_SIZE / 2; off > 0; off /= 2)
+          sum_sq += sycl::shift_group_left(sg, sum_sq, off);
+
+        float rstd = sycl::rsqrt(sum_sq / static_cast<float>(N) + eps);
+        rstd = sycl::group_broadcast(sg, rstd, 0);
+
+        // ── Pass 2: max_abs of normalized values ──────────────────────────
+        float max_abs = 0.f;
+        for (int col = lane; col < N; col += SG_SIZE) {
+          float normed = static_cast<float>(input_ptr[row * N + col]) * rstd;
+          max_abs = sycl::fmax(max_abs, sycl::fabs(normed));
+        }
+        for (int off = SG_SIZE / 2; off > 0; off /= 2)
+          max_abs = sycl::fmax(max_abs, sycl::shift_group_left(sg, max_abs, off));
+
+        float scale_tok = max_abs / 127.f + 1e-8f;
+        scale_tok = sycl::group_broadcast(sg, scale_tok, 0);
+
+        // ── Pass 3: write INT8 quantized + BF16 normed ────────────────────
+        for (int col = lane; col < N; col += SG_SIZE) {
+          float normed = static_cast<float>(input_ptr[row * N + col]) * rstd;
+          normed_out_ptr[row * N + col] = static_cast<ElementNormed>(normed);
+
+          float qval = sycl::round(normed / scale_tok);
+          qval = sycl::fmin(sycl::fmax(qval, -128.f), 127.f);
+          quant_out_ptr[row * N + col] = static_cast<int8_t>(qval);
+        }
+
+        if (lane == 0)
+          scale_token_ptr[row] = scale_tok;
+      }
+    );
+  });
+}
+
 }  // namespace xe_fuse