Optimize gemma4 continue

vllm/custom-esimd-kernels-vllm/csrc/moe_batch/moe.sycl

@@ -4,6 +4,7 @@
 #include <sycl/ext/intel/experimental/esimd/memory.hpp>
 
 #include "moe_topk.h"
+#include "moe_decode_gemv.h"
 #include "../xpu/esimd_kernels/moe_ops.h"
 
 using fp16 = sycl::half;
@@ -1611,6 +1612,152 @@ torch::Tensor moe_forward_full_gelu_tanh(
 }
 
 
+
+// ========== gemma4 decode-only (M==1) expert GEMV (load-width fix) ==========
+// Uses 1D block_load<uint8_t,VL> along K instead of the DPAS 16-wide 2D load.
+// Same vllm weight layout. M==1 only — caller must guarantee n_tokens==1.
+torch::Tensor moe_forward_full_gelu_tanh_routed_decode(
+    torch::Tensor x,
+    torch::Tensor topk_weights, torch::Tensor topk_indices,
+    torch::Tensor gate_up_weight, torch::Tensor gate_up_scale,
+    torch::Tensor down_weight, torch::Tensor down_scale,
+    int64_t top_k, int64_t n_routed_experts) {
+    (void)n_routed_experts;
+    TORCH_CHECK(x.scalar_type() == torch::kHalf);
+    TORCH_CHECK(x.size(0) == 1, "decode kernel requires n_tokens==1");
+    TORCH_CHECK(topk_indices.scalar_type() == torch::kInt32);
+    TORCH_CHECK(topk_weights.scalar_type() == torch::kHalf);
+    int hidden_size = x.size(1);
+    int intermediate_size = gate_up_weight.size(1) / 2;
+    int rows_per_token = top_k;
+    ensure_gemma4_moe_buffers(1, top_k, hidden_size, intermediate_size, x.device());
+
+    const int* idx_ptr = topk_indices.data_ptr<int>();
+    const fp16* w_ptr  = (const fp16*)topk_weights.data_ptr();
+    sycl::queue& q = c10::xpu::getCurrentXPUStream(x.device().index()).queue();
+
+    // ---- Up + gelu_tanh: K = hidden_size ----
+    // gemma: hidden=2816 → VL=256 divides cleanly (tail 0).
+    {
+        const int VL = 256;
+        const int n_routes = top_k;
+        auto k = MoeUpDecodeGeluTanh<VL, 0>{
+            (const fp16*)x.data_ptr(),
+            (const uint8_t*)gate_up_weight.data_ptr(),
+            gate_up_scale.data_ptr<float>(),
+            idx_ptr,
+            (fp16*)sg_intermediates.data_ptr(),
+            hidden_size, intermediate_size, (int)top_k, /*fp8_mode=*/0};
+        q.submit([&](sycl::handler& h){
+            h.parallel_for(sycl::nd_range<2>(
+                sycl::range<2>(n_routes, intermediate_size),
+                sycl::range<2>(1, 1)), k);
+        });
+    }
+    // ---- Down: K = intermediate_size ----
+    // gemma per-card inter=1056 = 256*4 + 32 → VL=256 big chunks + 32 tail.
+    {
+        const int VL = 256;
+        const int VL_TAIL = 32;
+        const int n_routes = top_k;
+        auto k = MoeDownDecode<VL, VL_TAIL>{
+            (const fp16*)sg_intermediates.data_ptr(),
+            (const uint8_t*)down_weight.data_ptr(),
+            down_scale.data_ptr<float>(),
+            w_ptr, idx_ptr,
+            (fp16*)sg_routed_output.data_ptr(),
+            hidden_size, intermediate_size, (int)top_k, /*fp8_mode=*/0};
+        q.submit([&](sycl::handler& h){
+            h.parallel_for(sycl::nd_range<2>(
+                sycl::range<2>(n_routes, hidden_size),
+                sycl::range<2>(1, 1)), k);
+        });
+    }
+    // ---- Accumulate over top_k ----
+    auto& final_out = sg_take_final_output();
+    moe_accumulate_kernel(
+        (const fp16*)sg_routed_output.data_ptr(),
+        (fp16*)final_out.data_ptr(),
+        1, hidden_size, rows_per_token, x.device());
+    return final_out.narrow(0, 0, 1);
+}
+
+// ========== gemma4 fully-fused decode (logits in, topk+scale+expert) =========
+// Takes router logits directly: internal topk (fp32 production kernel, gsm8k-
+// aligned) -> fold per_expert_scale into topk weights (on-device) -> 1D-load
+// gelu_tanh up + down GEMV -> accumulate. Removes the Python-side moe_topk call
+// + torch scale-fold + separate routed_decode dispatch. M==1 only.
+torch::Tensor moe_forward_full_gelu_tanh_decode(
+    torch::Tensor x, torch::Tensor logits,
+    torch::Tensor gate_up_weight, torch::Tensor gate_up_scale,
+    torch::Tensor down_weight, torch::Tensor down_scale,
+    torch::Tensor per_expert_scale,
+    int64_t top_k, int64_t n_routed_experts) {
+    TORCH_CHECK(x.scalar_type() == torch::kHalf);
+    TORCH_CHECK(x.size(0) == 1, "decode kernel requires n_tokens==1");
+    int hidden_size = x.size(1);
+    int intermediate_size = gate_up_weight.size(1) / 2;
+    int rows_per_token = top_k;
+    ensure_gemma4_moe_buffers(1, top_k, hidden_size, intermediate_size, x.device());
+    sycl::queue& q = c10::xpu::getCurrentXPUStream(x.device().index()).queue();
+
+    // ---- TopK (fp32-internal production kernel, norm=true) ----
+    dispatch_moe_topk_forward<class MoeTopKGemma4Decode>(
+        (const fp16*)logits.data_ptr(),
+        sg_topk_idx.data_ptr<int>(),
+        (fp16*)sg_topk_weight.data_ptr(),
+        1, (int)n_routed_experts, (int)top_k, /*norm=*/true,
+        logits.device());
+
+    // ---- Fold per_expert_scale into topk weights (on-device) ----
+    q.submit([&](sycl::handler& h){
+        h.parallel_for(sycl::range<1>(top_k), MoeFoldExpertScale{
+            (fp16*)sg_topk_weight.data_ptr(),
+            sg_topk_idx.data_ptr<int>(),
+            per_expert_scale.data_ptr<float>(),
+            (int)top_k});
+    });
+
+    const int* idx_ptr = sg_topk_idx.data_ptr<int>();
+    const fp16* w_ptr  = (const fp16*)sg_topk_weight.data_ptr();
+
+    // ---- Up + gelu_tanh (1D-load GEMV) ----
+    {
+        const int VL = 256;
+        auto k = MoeUpDecodeGeluTanh<VL, 0>{
+            (const fp16*)x.data_ptr(),
+            (const uint8_t*)gate_up_weight.data_ptr(), gate_up_scale.data_ptr<float>(),
+            idx_ptr, (fp16*)sg_intermediates.data_ptr(),
+            hidden_size, intermediate_size, (int)top_k, /*fp8_mode=*/0};
+        q.submit([&](sycl::handler& h){
+            h.parallel_for(sycl::nd_range<2>(
+                sycl::range<2>(top_k, intermediate_size),
+                sycl::range<2>(1, 1)), k);
+        });
+    }
+    // ---- Down (1D-load GEMV, VL=256 + tail 32) ----
+    {
+        const int VL = 256; const int VL_TAIL = 32;
+        auto k = MoeDownDecode<VL, VL_TAIL>{
+            (const fp16*)sg_intermediates.data_ptr(),
+            (const uint8_t*)down_weight.data_ptr(), down_scale.data_ptr<float>(),
+            w_ptr, idx_ptr, (fp16*)sg_routed_output.data_ptr(),
+            hidden_size, intermediate_size, (int)top_k, /*fp8_mode=*/0};
+        q.submit([&](sycl::handler& h){
+            h.parallel_for(sycl::nd_range<2>(
+                sycl::range<2>(top_k, hidden_size),
+                sycl::range<2>(1, 1)), k);
+        });
+    }
+    // ---- Accumulate over top_k ----
+    auto& final_out = sg_take_final_output();
+    moe_accumulate_kernel(
+        (const fp16*)sg_routed_output.data_ptr(),
+        (fp16*)final_out.data_ptr(),
+        1, hidden_size, rows_per_token, x.device());
+    return final_out.narrow(0, 0, 1);
+}
+
 // ========== gemma4: variant with externally supplied routing ==========
 // Skip the built-in softmax/topk. Caller passes per-token topk indices and
 // weights that already encode model-specific routing logic (e.g. gemma4's
@@ -1789,6 +1936,8 @@ TORCH_LIBRARY_FRAGMENT(moe_ops, m) {
     m.def("moe_forward_full_gelu_tanh(Tensor x, Tensor logits, Tensor gate_up_weight, Tensor gate_up_scale, Tensor down_weight, Tensor down_scale, int top_k, int n_routed_experts) -> Tensor");
     m.def("moe_forward_full_gelu_tanh_routed_no_accum(Tensor x, Tensor topk_weights, Tensor topk_indices, Tensor gate_up_weight, Tensor gate_up_scale, Tensor down_weight, Tensor down_scale, int top_k, int n_routed_experts) -> Tensor");
     m.def("moe_forward_full_gelu_tanh_routed(Tensor x, Tensor topk_weights, Tensor topk_indices, Tensor gate_up_weight, Tensor gate_up_scale, Tensor down_weight, Tensor down_scale, int top_k, int n_routed_experts) -> Tensor");
+    m.def("moe_forward_full_gelu_tanh_routed_decode(Tensor x, Tensor topk_weights, Tensor topk_indices, Tensor gate_up_weight, Tensor gate_up_scale, Tensor down_weight, Tensor down_scale, int top_k, int n_routed_experts) -> Tensor");
+    m.def("moe_forward_full_gelu_tanh_decode(Tensor x, Tensor logits, Tensor gate_up_weight, Tensor gate_up_scale, Tensor down_weight, Tensor down_scale, Tensor per_expert_scale, int top_k, int n_routed_experts) -> Tensor");
     m.def("moe_forward_full(Tensor x, Tensor logits, Tensor gate_up_weight, Tensor gate_up_scale, Tensor shared_gate_up_weight, Tensor shared_gate_up_scale, Tensor down_weight, Tensor down_scale, Tensor shared_down_weight, Tensor shared_down_scale, Tensor shared_expert_gate_weight, int top_k, int num_shared_experts, int n_routed_experts) -> Tensor");
     m.def("moe_forward_full_v2(Tensor x, Tensor logits, Tensor gate_up_weight, Tensor gate_up_scale, Tensor shared_gate_up_weight, Tensor shared_gate_up_scale, Tensor down_weight, Tensor down_scale, Tensor shared_down_weight, Tensor shared_down_scale, Tensor shared_expert_gate_weight, int top_k, int num_shared_experts, int n_routed_experts) -> Tensor");
 }
@@ -1801,6 +1950,8 @@ TORCH_LIBRARY_IMPL(moe_ops, XPU, m) {
     m.impl("moe_topk", &moe_topk);
     m.impl("moe_forward_full_gelu_tanh", &moe_forward_full_gelu_tanh);
     m.impl("moe_forward_full_gelu_tanh_routed", &moe_forward_full_gelu_tanh_routed);
+    m.impl("moe_forward_full_gelu_tanh_routed_decode", &moe_forward_full_gelu_tanh_routed_decode);
+    m.impl("moe_forward_full_gelu_tanh_decode", &moe_forward_full_gelu_tanh_decode);
     m.impl("moe_forward_full_gelu_tanh_routed_no_accum", &moe_forward_full_gelu_tanh_routed_no_accum);
     m.impl("moe_forward_fused", &moe_forward_fused);
     m.impl("moe_forward_full", &moe_forward_full);

vllm/custom-esimd-kernels-vllm/csrc/moe_batch/moe_decode_gemv.h

@@ -0,0 +1,155 @@
+// ============================================================================
+// MoE decode-only expert GEMV (M==1 fast path)
+//
+// Replaces the DPAS GEMM kernels for decode. Root cause fixed: DPAS path uses
+// lsc_load_2d<uint8_t,16,16,1> (16B-wide 2D tiles, 1/4 of BMG's 64B cacheline),
+// capping MoE decode at ~316 GB/s vs dense GEMV's ~574. Expert weight is plain
+// row-major inside each expert (gate_up [E,2*inter,hidden], down [E,hidden,inter]),
+// so 1D block_load<uint8_t,VL> along K reads contiguously like fp8_GEMV_bmg.
+//
+// One work-item computes one output element via VL-strided 1D loads + tail.
+// KS=1 (no K-split) — first version focuses on the load-width fix.
+// ============================================================================
+#pragma once
+#include <sycl/ext/intel/esimd.hpp>
+
+// Forward decl: defined in moe.sycl above the include point's use site.
+template<int N>
+SYCL_ESIMD_FUNCTION simd<sycl::half, N> fp8e4m3_to_half(simd<uint8_t, N> raw);
+
+// Fast E4M3→half: uint16 bit-twiddle for normals + correct subnormal handling.
+// e4m3 bias=7, fp16 bias=15 → exp_fp16 = exp_e4m3 + 8. mant 3b → fp16 mant top.
+// Subnormal (e==0): value = mant * 2^-9 (representable as normal fp16).
+template<int N>
+SYCL_ESIMD_FUNCTION simd<float, N> fp8e4m3_dequant_fast(simd<uint8_t, N> raw) {
+    using namespace sycl::ext::intel::esimd;
+    simd<uint16_t, N> u = convert<uint16_t>(raw);
+    simd<uint16_t, N> sign = (u >> 7) & 1;
+    simd<uint16_t, N> e = (u >> 3) & 0xF;
+    simd<uint16_t, N> m = u & 0x7;
+    // Normal path
+    simd<uint16_t, N> norm_bits = (sign << 15) | ((e + 8) << 10) | (m << 7);
+    simd<fp16, N> hn = norm_bits.template bit_cast_view<fp16>();
+    // Subnormal path: m * 2^-9, with sign
+    simd<fp16, N> hs = convert<fp16>(m) * fp16(1.0f / 512.0f);
+    simd<fp16, N> hs_signed = hs;
+    hs_signed.merge(-hs, sign == 1);
+    // Select subnormal where e==0
+    simd<fp16, N> out = hn;
+    out.merge(hs_signed, e == 0);
+    return simd<float, N>(out);
+}
+
+
+// VL_BIG full chunks + one VL_TAIL chunk (VL_TAIL may be 0). KS=1.
+// Up + gelu_tanh. gate_up_weight [E, 2*inter, hidden], K = hidden.
+template<int VL, int VL_TAIL>
+struct MoeUpDecodeGeluTanh {
+    const fp16*    x;
+    const uint8_t* gate_up_weight;
+    const float*   gate_up_scale;
+    const int*     selected_experts;
+    fp16*          intermediates;   // [top_k, inter]
+    int hidden, inter, top_k, fp8_mode;
+
+    void operator()(sycl::nd_item<2> item) const SYCL_ESIMD_KERNEL {
+        using namespace sycl::ext::intel::esimd;
+        const int route = (int)item.get_global_id(0);
+        const int n     = (int)item.get_global_id(1);
+        if (n >= inter) return;
+
+        const int two_inter = 2 * inter;
+        const int eid = selected_experts[route];
+        const uint8_t* wbase = gate_up_weight + (size_t)eid * two_inter * hidden;
+        const uint8_t* w_gate = wbase + (size_t)n * hidden;
+        const uint8_t* w_up   = wbase + (size_t)(inter + n) * hidden;
+
+        const int kp_full = (hidden / VL) * VL;
+        simd<float, VL> g_acc(0.f), u_acc(0.f);
+        for (int k = 0; k < kp_full; k += VL) {
+            simd<fp16, VL> xv = block_load<fp16, VL>(x + k);
+            simd<float, VL> xf = xv;
+            g_acc += xf * fp8e4m3_dequant_fast<VL>((block_load<uint8_t, VL>(w_gate + k)));
+            u_acc += xf * fp8e4m3_dequant_fast<VL>((block_load<uint8_t, VL>(w_up + k)));
+        }
+        float g_sum = reduce<float>(g_acc, std::plus<>());
+        float u_sum = reduce<float>(u_acc, std::plus<>());
+        if constexpr (VL_TAIL > 0) {
+            int kt = kp_full;
+            simd<fp16, VL_TAIL> xv = block_load<fp16, VL_TAIL>(x + kt);
+            simd<float, VL_TAIL> xf = xv;
+            g_sum += reduce<float>(xf * fp8e4m3_dequant_fast<VL_TAIL>((block_load<uint8_t, VL_TAIL>(w_gate + kt))), std::plus<>());
+            u_sum += reduce<float>(xf * fp8e4m3_dequant_fast<VL_TAIL>((block_load<uint8_t, VL_TAIL>(w_up + kt))), std::plus<>());
+        }
+
+        float scale = gate_up_scale[eid];
+        float gs = g_sum * scale, us = u_sum * scale;
+        constexpr float sqrt_2_over_pi = 0.7978845608f, coeff = 0.044715f;
+        float gs3 = gs*gs*gs;
+        float inner = sqrt_2_over_pi * (gs + coeff*gs3);
+        float e2 = sycl::exp(2.0f*inner);
+        float tanh_v = (e2 - 1.0f)/(e2 + 1.0f);
+        float gelu = 0.5f*gs*(1.0f + tanh_v);
+        intermediates[(size_t)route*inter + n] = fp16(gelu * us);
+    }
+};
+
+// Down. down_weight [E, hidden, inter], K = inter.
+template<int VL, int VL_TAIL>
+struct MoeDownDecode {
+    const fp16*    intermediates;   // [top_k, inter]
+    const uint8_t* down_weight;
+    const float*   down_scale;
+    const fp16*    routing_weights;
+    const int*     selected_experts;
+    fp16*          output;          // [top_k, hidden]
+    int hidden, inter, top_k, fp8_mode;
+
+    void operator()(sycl::nd_item<2> item) const SYCL_ESIMD_KERNEL {
+        using namespace sycl::ext::intel::esimd;
+        const int route = (int)item.get_global_id(0);
+        const int h     = (int)item.get_global_id(1);
+        if (h >= hidden) return;
+
+        const int eid = selected_experts[route];
+        const uint8_t* wrow = down_weight + (size_t)eid * hidden * inter + (size_t)h * inter;
+        const fp16* hi = intermediates + (size_t)route * inter;
+
+        const int kp_full = (inter / VL) * VL;
+        simd<float, VL> acc(0.f);
+        for (int k = 0; k < kp_full; k += VL) {
+            simd<fp16, VL> hv = block_load<fp16, VL>(hi + k);
+            simd<float, VL> hf = hv;
+            acc += hf * fp8e4m3_dequant_fast<VL>((block_load<uint8_t, VL>(wrow + k)));
+        }
+        float s = reduce<float>(acc, std::plus<>());
+        if constexpr (VL_TAIL > 0) {
+            int kt = kp_full;
+            simd<fp16, VL_TAIL> hv = block_load<fp16, VL_TAIL>(hi + kt);
+            simd<float, VL_TAIL> hf = hv;
+            s += reduce<float>(hf * fp8e4m3_dequant_fast<VL_TAIL>((block_load<uint8_t, VL_TAIL>(wrow + kt))), std::plus<>());
+        }
+
+        float w = (float)routing_weights[route];
+        float ds = down_scale[eid];
+        output[(size_t)route*hidden + h] = fp16(s * w * ds);
+    }
+};
+
+
+// ── per_expert_scale fold: topk_weight[r] *= scale[idx[r]] (top_k items) ─────
+// gemma folds a learnable per-expert scale into the routing weights. One
+// work-item per route. Tiny (top_k=8) — pure launch, but stays on-device.
+struct MoeFoldExpertScale {
+    fp16*        topk_weight;   // [top_k] in/out
+    const int*   topk_idx;      // [top_k]
+    const float* per_expert_scale;  // [n_experts]
+    int top_k;
+    void operator()(sycl::id<1> it) const SYCL_ESIMD_KERNEL {
+        using namespace sycl::ext::intel::esimd;
+        const int r = (int)it[0];
+        if (r >= top_k) return;
+        float s = per_expert_scale[topk_idx[r]];
+        topk_weight[r] = fp16((float)topk_weight[r] * s);
+    }
+};

vllm/custom-esimd-kernels-vllm/python/custom_esimd_kernels_vllm/ops.py

@@ -1559,6 +1559,53 @@ def moe_forward_full_gelu_tanh_routed(
         top_k, n_routed_experts)
 
 
+def moe_forward_full_gelu_tanh_routed_decode(
+    x: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_indices: torch.Tensor,
+    gate_up_weight: torch.Tensor,
+    gate_up_scale: torch.Tensor,
+    down_weight: torch.Tensor,
+    down_scale: torch.Tensor,
+    top_k: int,
+    n_routed_experts: int,
+) -> torch.Tensor:
+    """Decode-only (M==1) variant of moe_forward_full_gelu_tanh_routed.
+
+    Uses 1D block_load expert GEMV instead of the 16-wide 2D DPAS load,
+    restoring full HBM bandwidth (~528 vs ~315 GB/s) for the single-token
+    decode case. Requires x.size(0) == 1. Bit-identical to the DPAS path.
+    """
+    return _moe_batch.moe_forward_full_gelu_tanh_routed_decode(
+        x, topk_weights, topk_indices,
+        gate_up_weight, gate_up_scale,
+        down_weight, down_scale,
+        top_k, n_routed_experts)
+
+
+def moe_forward_full_gelu_tanh_decode(
+    x: torch.Tensor,
+    logits: torch.Tensor,
+    gate_up_weight: torch.Tensor,
+    gate_up_scale: torch.Tensor,
+    down_weight: torch.Tensor,
+    down_scale: torch.Tensor,
+    per_expert_scale: torch.Tensor,
+    top_k: int,
+    n_routed_experts: int,
+) -> torch.Tensor:
+    """Fully-fused gemma4 MoE decode (M==1): router logits in, output out.
+
+    Internal topk (fp32 production kernel) + per_expert_scale fold + 1D-load
+    gelu_tanh up/down GEMV + accumulate, all in one op. Removes the Python-side
+    moe_topk call, torch scale-fold, and separate expert dispatch.
+    """
+    return _moe_batch.moe_forward_full_gelu_tanh_decode(
+        x, logits, gate_up_weight, gate_up_scale,
+        down_weight, down_scale, per_expert_scale,
+        top_k, n_routed_experts)
+
+
 def esimd_norm_gemv_norm_fp16(
     residual: torch.Tensor,
     scale_with_root: torch.Tensor,