diff --git a/sglang/README.md b/sglang/README.md index ffe41746..138ec646 100644 --- a/sglang/README.md +++ b/sglang/README.md @@ -1,7 +1,9 @@ # SGLang on Intel BMG -End-to-end recipe for running Qwen3.6-35B-A3B online fp8 inference on Intel -Battlemage (BMG) GPUs with the optimized ESIMD kernel fast-paths. +End-to-end recipe for running Qwen3.6-35B-A3B online fp8 (e5m2) inference on +Intel Battlemage (BMG) GPUs with the optimized ESIMD kernel fast-paths. Decode +runs eager (XPU graph disabled for accuracy); the e5m2 fused decode kernels +recover the per-step host-dispatch cost. ## What's in here @@ -11,7 +13,7 @@ sglang/ │ └── Dockerfile # builds the full image ├── scripts/ │ ├── build_image.sh # wrapper around `docker buildx build` -│ ├── run_qwen3_6.sh # launches the TP=2 fp8 server +│ ├── start_qwen3_6_service.sh # launches the TP=2 e5m2 fp8 server │ └── run_gsm8k.py # standalone GSM8K accuracy harness ├── patches/ # sglang / sgl-kernel-xpu source patches └── custom-esimd-kernels/ # merged ESIMD kernel package: @@ -30,7 +32,7 @@ The script resolves `docker/Dockerfile` relative to itself, forwards `http_proxy` / `https_proxy` from the environment, and bumps `SGLANG_CACHEBUST` each run. Override the tag with `IMAGE_TAG=...`. -Time: ~25 min on a workstation (cold), dominated by the ESIMD AOT compile +Cold builds take a while, dominated by the ESIMD AOT compile and the sgl-kernel-xpu cmake build. ## Run @@ -42,26 +44,39 @@ docker run --rm -it \ -v /home/intel/LLM/models/Qwen3.6-35B-A3B:/models/Qwen3.6-35B-A3B:ro \ -p 30000:30000 \ llm-scaler-sgl:bmg \ - /workspace/scripts/run_qwen3_6.sh + /llm-scaler/sglang/scripts/start_qwen3_6_service.sh ``` ## Fast-paths enabled -Each is gated by an env var (set by `run_qwen3_6.sh`): +Each is gated by an env var (set by `start_qwen3_6_service.sh`): | Env var | Path | |------------------------------------|----------------------------------------| | `SGL_XPU_ESIMD_DECODE` | Decode SDPA (split-K, flat NHD KV) | | `SGL_XPU_ESIMD_MOE` | FP8 MoE silu routed kernel | +| `SGL_XPU_ESIMD_MOE_FULL` | Full decode MoE fusion (router+routed+shared+gate, e5m2, native N-major w13) | | `SGL_XPU_ESIMD_MOE_PREFILL` | FP8 MoE prefill (M-tiled DPAS) | | `SGL_XPU_FA_ESIMD_QKV` | Full-attention fused QKV+RMSNorm+RoPE | +| `SGL_XPU_FA_RESADD_NORM` | Fuse FA input_layernorm (resadd+rmsnorm) into qkv_proj (decode) | | `SGL_XPU_GDN_ESIMD` | GDN conv fused_seq decode | | `SGL_XPU_GDN_EXTEND_ESIMD` | GDN chunk_gated_delta_rule prefill | +| `SGL_XPU_GDN_NORM_GEMV` | GDN gated-RMSNorm as ESIMD GEMV (decode) | +| `SGL_XPU_GDN_RESADD_NORM` | Fuse GDN input_layernorm + in_proj (qkvz+ba) into one GEMV | +| `SGL_XPU_MOE_ROUTER_FP8` | MoE router as fp8 ESIMD GEMV (vs fp16 aten::mm) | | `SGL_XPU_PREFILL_DPAS` | Prefill SDPA via DPAS/XMX | -| `SGL_XPU_ENABLE_GRAPH` | XPU device-graph capture/replay | +| `SGL_XPU_ENABLE_GRAPH` | XPU device-graph capture/replay (kept **0** here) | > **Note:** all ESIMD/XPU fast-path gates use the `SGL_XPU_*` prefix. +The full decode MoE fusion (`SGL_XPU_ESIMD_MOE_FULL`) and the MoE router fp8 +path require online fp8 to be quantized as **e5m2** — set `SGLANG_FP8_DTYPE=e5m2` +(the script does). The e5m2 fused MoE kernel reads the native N-major `w13` +weight directly (no transposed weight copy), so it needs no extra device +memory for a transposed copy. `SGL_XPU_MOE_ROUTER_FP8=1` perturbs top-8 routing +on a fraction of tokens — A/B against GSM8K before trusting it (set to 0 for the +accurate fp16 gate). + In addition `SGLANG_MAMBA_{CONV,SSM}_DTYPE=float16` is required when running the model with `--dtype float16` so the mamba state pool matches activation dtype (the triton causal_conv1d_update kernel rejects mismatches). diff --git a/sglang/custom-esimd-kernels/.gitignore b/sglang/custom-esimd-kernels/.gitignore index 067da1e1..27d7f2a4 100644 --- a/sglang/custom-esimd-kernels/.gitignore +++ b/sglang/custom-esimd-kernels/.gitignore @@ -1,5 +1,5 @@ *.so -build/ +build*/ __pycache__/ *.egg-info/ *.egg-link diff --git a/sglang/custom-esimd-kernels/csrc/eagle/eagle.sycl b/sglang/custom-esimd-kernels/csrc/eagle/eagle.sycl index 577db5de..4d9ce798 100644 --- a/sglang/custom-esimd-kernels/csrc/eagle/eagle.sycl +++ b/sglang/custom-esimd-kernels/csrc/eagle/eagle.sycl @@ -503,7 +503,7 @@ void page_attn_decode( // (out, last_state) — last_state aliases initial_state (mutated in place) // and is returned as a convenience so the caller can // scatter ssm_states[cache_indices] = last_state. -std::tuple chunk_gated_delta_rule_extend( +std::tuple chunk_gated_delta_rule_extend( torch::Tensor& q, torch::Tensor& k, torch::Tensor& v, @@ -511,7 +511,8 @@ std::tuple chunk_gated_delta_rule_extend( torch::Tensor& beta, torch::Tensor& initial_state, torch::Tensor& cu_seqlens, - double scale) + double scale, + int64_t h_chunk_size) { TORCH_CHECK(q.scalar_type() == torch::kHalf || q.scalar_type() == torch::kBFloat16); TORCH_CHECK(k.scalar_type() == q.scalar_type()); @@ -550,6 +551,19 @@ std::tuple chunk_gated_delta_rule_extend( auto out = torch::empty({1, T_total, H_v, V}, q.options()); + // Optional per-chunk intermediate states. Sequence s owns ceil(L_s/chunk) + // entries; sum_s ceil(L_s/chunk) <= T_total/chunk + n_seqs, so allocate that + // upper bound (trailing rows are simply never indexed by the caller). + torch::Tensor h; + uint8_t* hPtr = nullptr; + if (h_chunk_size > 0) { + const int64_t max_chunks = T_total / h_chunk_size + n_seqs; + h = torch::empty({1, max_chunks, H_v, V, K}, initial_state.options()); + hPtr = (uint8_t*)h.data_ptr(); + } else { + h = torch::empty({0}, initial_state.options()); + } + sycl::range<2> globalRange(H_v * 16, n_seqs); sycl::range<2> localRange(16, 1); sycl::nd_range<2> range(globalRange, localRange); @@ -559,6 +573,7 @@ std::tuple chunk_gated_delta_rule_extend( const uint32_t headV = (uint32_t)H_v; const uint32_t headDim = (uint32_t)K; const float qScale = (float)scale; + const uint32_t hChunk = (uint32_t)(h_chunk_size > 0 ? h_chunk_size : 0); uint8_t* qPtr = (uint8_t*)q.data_ptr(); uint8_t* kPtr = (uint8_t*)k.data_ptr(); @@ -580,8 +595,8 @@ std::tuple chunk_gated_delta_rule_extend( cgh.parallel_for( range, [=](sycl::nd_item<2> ndi) SYCL_ESIMD_KERNEL { chunkGatedDeltaRuleExtendFp16( - qPtr, kPtr, vPtr, gPtr, bPtr, sPtr, oPtr, cuPtr, - headQk, headV, headDim, qScale, ndi); + qPtr, kPtr, vPtr, gPtr, bPtr, sPtr, oPtr, hPtr, cuPtr, + headQk, headV, headDim, hChunk, qScale, ndi); }); }; } else { @@ -589,8 +604,8 @@ std::tuple chunk_gated_delta_rule_extend( cgh.parallel_for( range, [=](sycl::nd_item<2> ndi) SYCL_ESIMD_KERNEL { chunkGatedDeltaRuleExtendFp16( - qPtr, kPtr, vPtr, gPtr, bPtr, sPtr, oPtr, cuPtr, - headQk, headV, headDim, qScale, ndi); + qPtr, kPtr, vPtr, gPtr, bPtr, sPtr, oPtr, hPtr, cuPtr, + headQk, headV, headDim, hChunk, qScale, ndi); }); }; } @@ -602,8 +617,8 @@ std::tuple chunk_gated_delta_rule_extend( cgh.parallel_for( range, [=](sycl::nd_item<2> ndi) SYCL_ESIMD_KERNEL { chunkGatedDeltaRuleExtendBf16( - qPtr, kPtr, vPtr, gPtr, bPtr, sPtr, oPtr, cuPtr, - headQk, headV, headDim, qScale, ndi); + qPtr, kPtr, vPtr, gPtr, bPtr, sPtr, oPtr, hPtr, cuPtr, + headQk, headV, headDim, hChunk, qScale, ndi); }); }; } else { @@ -611,8 +626,8 @@ std::tuple chunk_gated_delta_rule_extend( cgh.parallel_for( range, [=](sycl::nd_item<2> ndi) SYCL_ESIMD_KERNEL { chunkGatedDeltaRuleExtendBf16( - qPtr, kPtr, vPtr, gPtr, bPtr, sPtr, oPtr, cuPtr, - headQk, headV, headDim, qScale, ndi); + qPtr, kPtr, vPtr, gPtr, bPtr, sPtr, oPtr, hPtr, cuPtr, + headQk, headV, headDim, hChunk, qScale, ndi); }); }; } @@ -622,7 +637,7 @@ std::tuple chunk_gated_delta_rule_extend( TORCH_CHECK(false, "unreachable dtype check"); } - return std::make_tuple(out, initial_state); + return std::make_tuple(out, initial_state, h); } // =========================================================================== @@ -1844,7 +1859,7 @@ TORCH_LIBRARY_FRAGMENT(eagle_ops, m) { m.def("verify_tree_greedy(Tensor! predicts, Tensor! accept_index, Tensor! accept_token_num, Tensor candidates, Tensor retrive_index, Tensor retrive_next_token, Tensor retrive_next_sibling, Tensor target_predict) -> ()"); m.def("gdn_eagle(Tensor! qkvz, Tensor! z_out, Tensor conv_w, Tensor? conv_b, Tensor! conv_state, Tensor accepted_tokens, Tensor ba, Tensor a_log, Tensor dt_bias, Tensor! state_in, Tensor ssm_state_idx, Tensor norm_w, int max_query_len) -> Tensor"); m.def("page_attn_decode(Tensor query, Tensor key_cache, Tensor value_cache, Tensor block_table, Tensor seq_lens, Tensor! out, int max_query_len, int max_seq_len, Tensor(a!)? temp_p=None) -> ()"); - m.def("chunk_gated_delta_rule_extend(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, Tensor! initial_state, Tensor cu_seqlens, float scale) -> (Tensor, Tensor)"); + m.def("chunk_gated_delta_rule_extend(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, Tensor! initial_state, Tensor cu_seqlens, float scale, int h_chunk_size=0) -> (Tensor, Tensor, Tensor)"); m.def("mamba_state_scatter(Tensor! dst, Tensor src, Tensor dst_indices, Tensor step_indices) -> ()"); // Contract: q/k/v/initial_state are fp16 or bf16 (all same); g and beta are fp32 // (both emitted by fused_gdn_gating). cu_seqlens is int32. diff --git a/sglang/custom-esimd-kernels/csrc/eagle/extend.kernels.bf16.h b/sglang/custom-esimd-kernels/csrc/eagle/extend.kernels.bf16.h index 0215d2af..7e56189d 100644 --- a/sglang/custom-esimd-kernels/csrc/eagle/extend.kernels.bf16.h +++ b/sglang/custom-esimd-kernels/csrc/eagle/extend.kernels.bf16.h @@ -16,10 +16,12 @@ ESIMD_INLINE void chunkGatedDeltaRuleExtendBf16( uint8_t* betaState, // fp32 uint8_t* stateBuf, // StateT in/out uint8_t* oState, // bf16 out + uint8_t* hState, // StateT out, may be null: [total_chunks, H_v, V, K] uint32_t* cuSeqlens, uint32_t headQk, // H_k: headV must be a multiple of this (GQA on GDN) uint32_t headV, uint32_t headDim, + uint32_t hChunkSize, // token stride between h snapshots (0 = disabled) float qScale, sycl::nd_item<2>& ndi) { @@ -69,10 +71,38 @@ ESIMD_INLINE void chunkGatedDeltaRuleExtendBf16( fp32InS_persistent.select<128, 1>(r * 128) = raw; } + // Per-chunk snapshot base (see extend.kernels.fp16.h for the layout). + const bool emitH = (hState != nullptr) && (hChunkSize > 0); + StateT* hMyRows = nullptr; + if (emitH) { + uint32_t hSeqOffset = 0; + for (uint32_t s = 0; s < seqIdx; s++) { + const uint32_t len = cuSeqlens[s + 1] - cuSeqlens[s]; + hSeqOffset += (len + hChunkSize - 1) / hChunkSize; + } + hMyRows = (StateT*)hState + + (size_t)hSeqOffset * stateSeqElems + + headIdx * stateHeadElems + + (hh * 8) * headDim; + } + uint32_t hNextSnapTok = 0; + uint32_t hSnapIdx = 0; + namespace xens = sycl::ext::intel::experimental::esimd; for (uint32_t tRel = 0; tRel < nTokSeq; tRel++) { const uint32_t t = tStart + tRel; + if (emitH && tRel == hNextSnapTok) { + StateT* hDst = hMyRows + (size_t)hSnapIdx * stateSeqElems; + #pragma unroll + for (int r = 0; r < 8; r++) { + simd hv = fp32InS_persistent.select<128, 1>(r * 128); + block_store(hDst + r * headDim, hv); + } + hSnapIdx += 1; + hNextSnapTok += hChunkSize; + } + simd q_bf16 = block_load(qPtr + t * qkTokStride + kHeadIdx * headDim); simd k_bf16 = block_load(kPtr + t * qkTokStride + kHeadIdx * headDim); simd v_bf16 = block_load(vPtr + t * vTokStride + headIdx * headDim); diff --git a/sglang/custom-esimd-kernels/csrc/eagle/extend.kernels.fp16.h b/sglang/custom-esimd-kernels/csrc/eagle/extend.kernels.fp16.h index e77fc32c..a7eccc17 100644 --- a/sglang/custom-esimd-kernels/csrc/eagle/extend.kernels.fp16.h +++ b/sglang/custom-esimd-kernels/csrc/eagle/extend.kernels.fp16.h @@ -44,10 +44,13 @@ ESIMD_INLINE void chunkGatedDeltaRuleExtendFp16( uint8_t* betaState, // fp32 post-sigmoid uint8_t* stateBuf, // StateT in/out: initial_state on entry, last_state on exit uint8_t* oState, + uint8_t* hState, // StateT out, may be null: per-chunk intermediate states + // [total_chunks, H_v, V, K]; h[i] = state after i chunks uint32_t* cuSeqlens, uint32_t headQk, // H_k: headV must be a multiple of this (GQA on GDN) uint32_t headV, uint32_t headDim, // = 128 + uint32_t hChunkSize, // token stride between successive h snapshots (0 = disabled) float qScale, sycl::nd_item<2>& ndi) { @@ -106,11 +109,44 @@ ESIMD_INLINE void chunkGatedDeltaRuleExtendFp16( fp32InS_persistent.select<128, 1>(r * 128) = raw; } + // Base pointer to this thread's 8-row slab inside h, when the caller asked + // for per-chunk snapshots. h is packed per sequence: sequence s owns + // ceil(L_s / hChunkSize) entries, so our base is the running sum over the + // preceding sequences (n_seqs is small, so the serial scan is cheap). + const bool emitH = (hState != nullptr) && (hChunkSize > 0); + StateT* hMyRows = nullptr; + if (emitH) { + uint32_t hSeqOffset = 0; + for (uint32_t s = 0; s < seqIdx; s++) { + const uint32_t len = cuSeqlens[s + 1] - cuSeqlens[s]; + hSeqOffset += (len + hChunkSize - 1) / hChunkSize; + } + hMyRows = (StateT*)hState + + (size_t)hSeqOffset * stateSeqElems + + headIdx * stateHeadElems + + (hh * 8) * headDim; + } + // Snapshot cursor: avoids a division per token. + uint32_t hNextSnapTok = 0; + uint32_t hSnapIdx = 0; + // ---- Per-token loop -------------------------------------------------- namespace xens = sycl::ext::intel::experimental::esimd; for (uint32_t tRel = 0; tRel < nTokSeq; tRel++) { const uint32_t t = tStart + tRel; + // --- h snapshot: state *before* this token, i.e. after hSnapIdx chunks -- + if (emitH && tRel == hNextSnapTok) { + StateT* hDst = hMyRows + (size_t)hSnapIdx * stateSeqElems; + #pragma unroll + for (int r = 0; r < 8; r++) { + simd hv = fp32InS_persistent.select<128, 1>(r * 128); + block_store(hDst + r * headDim, hv); + } + hSnapIdx += 1; + hNextSnapTok += hChunkSize; + } + // --- Load q, k, v for the token's head ----------------------------- simd q_fp16 = block_load(qPtr + t * qkTokStride + kHeadIdx * headDim); simd k_fp16 = block_load(kPtr + t * qkTokStride + kHeadIdx * headDim); diff --git a/sglang/custom-esimd-kernels/csrc/moe_batch/moe.sycl b/sglang/custom-esimd-kernels/csrc/moe_batch/moe.sycl index 66b7872e..1a10c816 100644 --- a/sglang/custom-esimd-kernels/csrc/moe_batch/moe.sycl +++ b/sglang/custom-esimd-kernels/csrc/moe_batch/moe.sycl @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -240,6 +241,323 @@ void moe_router_forward_e5m2_kernel( submit_kernel(cgf, device, "moe router forward e5m2"); } +// ═══════════════════════════════════════════════════════════════════════════ +// FUSED: Router GEMV (e4m3 gate) + softmax TopK in ONE dispatch. +// One work-group per token, SG work-items per group. Each work-item computes a +// strided subset of the NE expert logits (e4m3 gate GEMV, matching +// moe_router_forward_e4m3_kernel), stores them into SLM. After a barrier, +// local id 0 performs the exact softmax + top-K + renormalize of +// MoE_TopK_V2_Kernel (norm=true) reading the logits from SLM. Removes the +// standalone router_forward dispatch (topk was already a dispatch inside +// moe_forward_full) → −1 dispatch/layer. +// +// Logits are rounded through fp16 before the topk to bit-match the split path +// (router writes fp16 logits, topk reads them back). +// ═══════════════════════════════════════════════════════════════════════════ +template class MoeRouterTopKFusedE4M3; +template +void moe_router_topk_fused_e4m3_kernel( + const fp16* x, + const uint8_t* rweight, const float* rscale, + fp16* top_values, int32_t* top_indices, + const int n_tokens, const int hidden_size, + const torch::Device& device) { + static_assert(NE % 64 == 0, "NE must be a multiple of 64"); + static_assert(TOPK <= 16, "TOPK must be <= 16"); + constexpr int C = 64; + constexpr int N_CHUNKS = NE / C; + + auto cgf = [&](sycl::handler& cgh) { + cgh.parallel_for>( + sycl::nd_range<1>(sycl::range<1>((size_t)n_tokens * SG), + sycl::range<1>(SG)), + [=](sycl::nd_item<1> item) SYCL_ESIMD_KERNEL { + slm_init(); + const int token = (int)item.get_group(0); + const int lid = (int)item.get_local_id(0); + if (token >= n_tokens) return; + + const fp16* x_row = x + (size_t)token * hidden_size; + const fp16 s = fp16(rscale[0]); + + // Each work-item computes logits for experts e = lid, lid+SG, ... + for (int e = lid; e < NE; e += SG) { + const uint8_t* w = rweight + (size_t)e * hidden_size; + simd acc(0.f); + for (int k = 0; k < hidden_size; k += C) { + simd wv = + fp8e4m3_to_half(block_load(w + k)) * s; + simd xv = block_load(x_row + k); + acc += convert(xv * wv); + } + // Round through fp16 to match the split router (fp16 logits). + float logit = (float)fp16( + sycl::ext::intel::esimd::detail::sum(acc)); + slm_block_store((uint32_t)e * 4u, simd(logit)); + } + + barrier(); + if (lid != 0) return; + + // ── softmax + top-K + renormalize (== MoE_TopK_V2_Kernel) ── + simd probs[N_CHUNKS]; + #pragma unroll + for (int c = 0; c < N_CHUNKS; c++) + probs[c] = slm_block_load((uint32_t)(c * C) * 4u); + + float row_max = h_max(probs[0]); + #pragma unroll + for (int c = 1; c < N_CHUNKS; c++) { + float m = h_max(probs[c]); + if (m > row_max) row_max = m; + } + float total = 0.0f; + #pragma unroll + for (int c = 0; c < N_CHUNKS; c++) { + probs[c] -= row_max; + probs[c] = esimd_math::exp(probs[c]); + total += h_sum(probs[c]); + } + float inv = 1.0f / total; + #pragma unroll + for (int c = 0; c < N_CHUNKS; c++) probs[c] *= inv; + + float tv[16]; + int32_t ti[16]; + #pragma unroll + for (int k = 0; k < TOPK; k++) { + float bv = -1.0f; int bc = 0; + #pragma unroll + for (int c = 0; c < N_CHUNKS; c++) { + float m = h_max(probs[c]); + if (m > bv) { bv = m; bc = c; } + } + float fv; int32_t fi; + chunk_argmax_and_zero(probs[bc], bc * C, fv, fi); + tv[k] = fv; ti[k] = fi; + } + float top_sum = 0.0f; + #pragma unroll + for (int k = 0; k < TOPK; k++) top_sum += tv[k]; + float inv_top = 1.0f / top_sum; + #pragma unroll + for (int k = 0; k < TOPK; k++) tv[k] *= inv_top; + + fp16* vp = top_values + (size_t)token * TOPK; + int32_t* ip = top_indices + (size_t)token * TOPK; + #pragma unroll + for (int i = 0; i < TOPK; i++) { + block_store(vp + i, simd((fp16)tv[i])); + block_store(ip + i, simd(ti[i])); + } + }); + }; + submit_kernel(cgf, device, "moe router+topk fused e4m3"); +} + +// Host dispatcher: pick the (NE, TOPK) specialization for the fused router+topk. +static inline bool dispatch_moe_router_topk_fused_e4m3( + const fp16* x, const uint8_t* rweight, const float* rscale, + fp16* top_values, int32_t* top_indices, + const int n_tokens, const int hidden_size, + const int num_experts, const int top_k, + const torch::Device& device) { + constexpr int SG = 32; + if (num_experts == 256 && top_k == 8) { + moe_router_topk_fused_e4m3_kernel<256, 8, SG>( + x, rweight, rscale, top_values, top_indices, n_tokens, hidden_size, device); + return true; + } else if (num_experts == 512 && top_k == 10) { + moe_router_topk_fused_e4m3_kernel<512, 10, SG>( + x, rweight, rscale, top_values, top_indices, n_tokens, hidden_size, device); + return true; + } else if (num_experts == 512 && top_k == 8) { + moe_router_topk_fused_e4m3_kernel<512, 8, SG>( + x, rweight, rscale, top_values, top_indices, n_tokens, hidden_size, device); + return true; + } else if (num_experts == 128 && top_k == 8) { + moe_router_topk_fused_e4m3_kernel<128, 8, SG>( + x, rweight, rscale, top_values, top_indices, n_tokens, hidden_size, device); + return true; + } + return false; +} + +// ═══════════════════════════════════════════════════════════════════════════ +// PHASE 3 FUSED: GemmaRMSNorm(post_attention_layernorm) resadd + rmsnorm folded +// into the head of the router+topk fused kernel. One work-group per token, +// SG work-items. Head phase: +// 1. new_residual[j] = hidden[j] + residual[j] (fp16, → global new_residual) +// 2. variance = mean(new_residual.float()^2) (cross-WI reduction via SLM) +// 3. normed[j] = new_residual.float() * rsqrt(var+eps) * norm_weight[j] (fp16, +// → global normed_x). norm_weight is the pre-folded Gemma (1+weight). +// Then the SAME router GEMV + softmax topk as moe_router_topk_fused_e4m3_kernel, +// but reading the freshly-written normed_x row instead of x. This removes the +// standalone gemma_fused_add_rmsnorm dispatch (−1/layer) AND writes new_residual +// so the caller needs no separate python add. +// +// H (hidden_size) must be a multiple of HB=64; the caller guarantees this. +// ═══════════════════════════════════════════════════════════════════════════ +template class MoeRouterTopKNormFusedE4M3; +template +void moe_router_topk_norm_fused_e4m3_kernel( + const fp16* hidden, const fp16* residual, const fp16* norm_weight, + fp16* normed_x, fp16* new_residual, + const uint8_t* rweight, const float* rscale, + fp16* top_values, int32_t* top_indices, + const int n_tokens, const int hidden_size, const float eps, + const torch::Device& device) { + static_assert(NE % 64 == 0, "NE must be a multiple of 64"); + static_assert(TOPK <= 16, "TOPK must be <= 16"); + constexpr int C = 64; + constexpr int N_CHUNKS = NE / C; + constexpr int HB = 64; + + auto cgf = [&](sycl::handler& cgh) { + cgh.parallel_for>( + sycl::nd_range<1>(sycl::range<1>((size_t)n_tokens * SG), + sycl::range<1>(SG)), + [=](sycl::nd_item<1> item) SYCL_ESIMD_KERNEL { + slm_init(); + const int token = (int)item.get_group(0); + const int lid = (int)item.get_local_id(0); + if (token >= n_tokens) return; + + const fp16* h_row = hidden + (size_t)token * hidden_size; + const fp16* r_row = residual + (size_t)token * hidden_size; + fp16* nr_row = new_residual + (size_t)token * hidden_size; + fp16* nx_row = normed_x + (size_t)token * hidden_size; + + // ── Phase A: resadd + partial sum-of-squares ── + simd pss(0.f); + for (int off = lid * HB; off < hidden_size; off += SG * HB) { + simd hv = block_load(h_row + off); + simd rv = block_load(r_row + off); + simd sv = hv + rv; + block_store(nr_row + off, sv); + simd sf = convert(sv); + pss += sf * sf; + } + float my_ss = + sycl::ext::intel::esimd::detail::sum(pss); + slm_block_store((uint32_t)lid * 4u, simd(my_ss)); + barrier(); + + simd parts = slm_block_load(0); + float total_ss = + sycl::ext::intel::esimd::detail::sum(parts); + float variance = total_ss / (float)hidden_size; + float rms = 1.0f / sycl::sqrt(variance + eps); + barrier(); // all WIs done reading partials before SLM reused for logits + + // ── Phase B: normed = new_residual * rms * (1+weight) ── + for (int off = lid * HB; off < hidden_size; off += SG * HB) { + simd sv = block_load(nr_row + off); + simd wv = block_load(norm_weight + off); + simd nf = + convert(sv) * rms * convert(wv); + simd no = convert(nf); + block_store(nx_row + off, no); + } + barrier(); // full normed row visible before router GEMV reads it + + // ── Phase C: router GEMV + softmax topk on normed_x ── + const fp16* x_row = nx_row; + const fp16 s = fp16(rscale[0]); + for (int e = lid; e < NE; e += SG) { + const uint8_t* w = rweight + (size_t)e * hidden_size; + simd acc(0.f); + for (int k = 0; k < hidden_size; k += C) { + simd wv = + fp8e4m3_to_half(block_load(w + k)) * s; + simd xv = block_load(x_row + k); + acc += convert(xv * wv); + } + float logit = (float)fp16( + sycl::ext::intel::esimd::detail::sum(acc)); + slm_block_store((uint32_t)e * 4u, simd(logit)); + } + barrier(); + if (lid != 0) return; + + simd probs[N_CHUNKS]; + #pragma unroll + for (int c = 0; c < N_CHUNKS; c++) + probs[c] = slm_block_load((uint32_t)(c * C) * 4u); + + float row_max = h_max(probs[0]); + #pragma unroll + for (int c = 1; c < N_CHUNKS; c++) { + float m = h_max(probs[c]); + if (m > row_max) row_max = m; + } + float total = 0.0f; + #pragma unroll + for (int c = 0; c < N_CHUNKS; c++) { + probs[c] -= row_max; + probs[c] = esimd_math::exp(probs[c]); + total += h_sum(probs[c]); + } + float inv = 1.0f / total; + #pragma unroll + for (int c = 0; c < N_CHUNKS; c++) probs[c] *= inv; + + float tv[16]; + int32_t ti[16]; + #pragma unroll + for (int k = 0; k < TOPK; k++) { + float bv = -1.0f; int bc = 0; + #pragma unroll + for (int c = 0; c < N_CHUNKS; c++) { + float m = h_max(probs[c]); + if (m > bv) { bv = m; bc = c; } + } + float fv; int32_t fi; + chunk_argmax_and_zero(probs[bc], bc * C, fv, fi); + tv[k] = fv; ti[k] = fi; + } + float top_sum = 0.0f; + #pragma unroll + for (int k = 0; k < TOPK; k++) top_sum += tv[k]; + float inv_top = 1.0f / top_sum; + #pragma unroll + for (int k = 0; k < TOPK; k++) tv[k] *= inv_top; + + fp16* vp = top_values + (size_t)token * TOPK; + int32_t* ip = top_indices + (size_t)token * TOPK; + #pragma unroll + for (int i = 0; i < TOPK; i++) { + block_store(vp + i, simd((fp16)tv[i])); + block_store(ip + i, simd(ti[i])); + } + }); + }; + submit_kernel(cgf, device, "moe router+topk+norm fused e4m3"); +} + +static inline bool dispatch_moe_router_topk_norm_fused_e4m3( + const fp16* hidden, const fp16* residual, const fp16* norm_weight, + fp16* normed_x, fp16* new_residual, + const uint8_t* rweight, const float* rscale, + fp16* top_values, int32_t* top_indices, + const int n_tokens, const int hidden_size, const float eps, + const int num_experts, const int top_k, + const torch::Device& device) { + constexpr int SG = 32; + if (hidden_size % 64 != 0) return false; + #define RTN_CASE(NE_, TK_) \ + moe_router_topk_norm_fused_e4m3_kernel( \ + hidden, residual, norm_weight, normed_x, new_residual, \ + rweight, rscale, top_values, top_indices, \ + n_tokens, hidden_size, eps, device); return true; + if (num_experts == 256 && top_k == 8) { RTN_CASE(256, 8); } + else if (num_experts == 512 && top_k == 10) { RTN_CASE(512, 10); } + else if (num_experts == 512 && top_k == 8) { RTN_CASE(512, 8); } + else if (num_experts == 128 && top_k == 8) { RTN_CASE(128, 8); } + #undef RTN_CASE + return false; +} + // -- Option C: 1 WI per (token, expert), n_tokens > 4 -- void moe_router_forward_e4m3_wide_kernel( const fp16* x, @@ -745,6 +1063,103 @@ void moe_up_routed_e5m2_kernel( submit_kernel(cgf, device, "moe up routed e5m2"); } +// ──────────────────────────────────────────────────────────────────────────── +// E5M2 up kernel for sglang FusedMoE w13 layout: [E, 2*intermediate, hidden] +// (N-major). Mirror of moe_up_routed_e4m3_nmajor_kernel but for e5m2 weights. +// Reads each 16×16 weight block as [N=16 rows × K=16 cols] via lsc_load_2d and +// uses fp8e5m2_block_to_vnni_nk to permute into the VNNI b_tile DPAS expects. +// This lets callers pass sglang's native w13 directly — NO transposed copy. +// Numerically identical to moe_up_routed_e5m2_kernel fed w13.transpose(1,2). +// ──────────────────────────────────────────────────────────────────────────── +void moe_up_routed_e5m2_nmajor_kernel( + const fp16* x, + const uint8_t* gate_up_weight, const float* gate_up_scale, + const int* selected_experts, + fp16* intermediates, + const int n_tokens, const int hidden_size, const int intermediate_size, + const int top_k, const int rows_per_token, + const torch::Device& device) { + + constexpr int GS = 16; + const int two_inter = 2 * intermediate_size; + const int n_tiles = intermediate_size / 16; + + auto cgf = [&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<2>( + sycl::range<2>(n_tokens * top_k, n_tiles * GS), + sycl::range<2>(1, GS)), + [=](sycl::nd_item<2> item) SYCL_ESIMD_KERNEL { + slm_init(); + + const int route_idx = (int)item.get_global_id(0); + const int token = route_idx / top_k; + const int k_idx = route_idx % top_k; + const int n_tile = (int)item.get_group(1); + const int tid = (int)item.get_local_id(1); + + const fp16* x_row = x + (size_t)token * hidden_size; + const int eid = selected_experts[route_idx]; + // weight layout: [E, 2*inter, hidden] → base + eid * two_inter * hidden + const uint8_t* base = gate_up_weight + (size_t)eid * two_inter * hidden_size; + const int ng = n_tile * 16; // N row offset (gate half) + const int nu = intermediate_size + ng; // N row offset (up half) + simd g_acc(0.f), u_acc(0.f); + + for (int k = 16 * tid; k < hidden_size; k += 16 * GS) { + simd a_tile = block_load(x_row + k); + + // [N=16, K=16] block at row=ng, col=k + // surface: width=hidden, height=two_inter, pitch=hidden + xesimd::config_2d_mem_access pgW( + base, (uint32_t)hidden_size - 1u, (uint32_t)two_inter - 1u, + (uint32_t)hidden_size - 1u, (uint32_t)k, (uint32_t)ng); + auto rg = xesimd::lsc_load_2d(pgW); + simd bg = fp8e5m2_block_to_vnni_nk(rg); + g_acc = dpas<8, 1, float, float, fp16, fp16>(g_acc, bg, a_tile); + + xesimd::config_2d_mem_access puW( + base, (uint32_t)hidden_size - 1u, (uint32_t)two_inter - 1u, + (uint32_t)hidden_size - 1u, (uint32_t)k, (uint32_t)nu); + auto ru = xesimd::lsc_load_2d(puW); + simd bu = fp8e5m2_block_to_vnni_nk(ru); + u_acc = dpas<8, 1, float, float, fp16, fp16>(u_acc, bu, a_tile); + } + + uint32_t slm_off = (uint32_t)(tid * 32) * 4u; + slm_block_store(slm_off, g_acc); + slm_block_store(slm_off + 64u, u_acc); + + barrier(); + + if (tid == 0) { + simd gs(0.f), us(0.f); + #pragma unroll + for (int i = 0; i < GS; i++) { + gs += slm_block_load((uint32_t)(i * 128)); + us += slm_block_load((uint32_t)(i * 128 + 64)); + } + + float scale = gate_up_scale[eid]; + gs *= scale; + us *= scale; + + simd silu = gs / (1.f + exp(-gs)); + simd result = silu * us; + + const int routed_row = token * rows_per_token + k_idx; + block_store( + intermediates + (size_t)routed_row * intermediate_size + ng, + convert(result)); + } + }); + }; + + submit_kernel(cgf, device, "moe up routed e5m2 nmajor (sglang layout)"); +} + // -- Shared expert up kernels -- void moe_up_shared_e4m3_kernel( const fp16* x, @@ -840,6 +1255,218 @@ void moe_up_shared_e5m2_kernel( submit_kernel(cgf, device, "moe up shared e5m2"); } +// -- Merged up kernel (routed + shared in ONE dispatch) -- +// Both routed and shared gate_up weights must be in DPAS/transposed layout +// [E, hidden, 2*inter] / [NS, hidden, 2*inter]. Each row (slot) of a token is +// either a routed expert (slot < top_k) or a shared expert (slot >= top_k); +// the only per-slot difference is the weight base pointer, the per-tensor +// scale, and the expert-id lookup. GPU work is identical to running the two +// kernels separately — this just collapses two host dispatches into one. +void moe_up_merged_e5m2_kernel( + const fp16* x, + const uint8_t* gate_up_weight, const float* gate_up_scale, + const uint8_t* shared_gate_up_weight, const float* shared_gate_up_scale, + const int* selected_experts, + fp16* intermediates, + const int n_tokens, const int hidden_size, const int intermediate_size, + const int top_k, const int num_shared_experts, const int rows_per_token, + const torch::Device& device) { + + constexpr int GS = 16; + const int two_inter = 2 * intermediate_size; + const int n_tiles = intermediate_size / 16; + + auto cgf = [&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<2>( + sycl::range<2>(n_tokens * rows_per_token, n_tiles * GS), + sycl::range<2>(1, GS)), + [=](sycl::nd_item<2> item) SYCL_ESIMD_KERNEL { + slm_init(); + + const int route_idx = (int)item.get_global_id(0); + const int token = route_idx / rows_per_token; + const int slot = route_idx % rows_per_token; + const int n_tile = (int)item.get_group(1); + const int tid = (int)item.get_local_id(1); + + const fp16* x_row = x + (size_t)token * hidden_size; + + // Select routed vs shared weight base + scale. + const uint8_t* base; + float scale; + if (slot < top_k) { + const int eid = selected_experts[token * top_k + slot]; + base = gate_up_weight + (size_t)eid * hidden_size * two_inter; + scale = gate_up_scale[eid]; + } else { + const int sid = slot - top_k; + base = shared_gate_up_weight + (size_t)sid * hidden_size * two_inter; + scale = shared_gate_up_scale[sid]; + } + + const int ng = n_tile * 16; + const int nu = intermediate_size + ng; + simd g_acc(0.f), u_acc(0.f); + + for (int k = 16 * tid; k < hidden_size; k += 16 * GS) { + simd a_tile = block_load(x_row + k); + + xesimd::config_2d_mem_access pgW( + base, (uint32_t)two_inter - 1u, (uint32_t)hidden_size - 1u, + (uint32_t)two_inter - 1u, (uint32_t)ng, (uint32_t)k); + auto rg = xesimd::lsc_load_2d(pgW); + simd bg = fp8e5m2_block_to_vnni(rg); + g_acc = dpas<8, 1, float, float, fp16, fp16>(g_acc, bg, a_tile); + + xesimd::config_2d_mem_access puW( + base, (uint32_t)two_inter - 1u, (uint32_t)hidden_size - 1u, + (uint32_t)two_inter - 1u, (uint32_t)nu, (uint32_t)k); + auto ru = xesimd::lsc_load_2d(puW); + simd bu = fp8e5m2_block_to_vnni(ru); + u_acc = dpas<8, 1, float, float, fp16, fp16>(u_acc, bu, a_tile); + } + + uint32_t slm_off = (uint32_t)(tid * 32) * 4u; + slm_block_store(slm_off, g_acc); + slm_block_store(slm_off + 64u, u_acc); + + barrier(); + + if (tid == 0) { + simd gs(0.f), us(0.f); + #pragma unroll + for (int i = 0; i < GS; i++) { + gs += slm_block_load((uint32_t)(i * 128)); + us += slm_block_load((uint32_t)(i * 128 + 64)); + } + + gs *= scale; + us *= scale; + + simd silu = gs / (1.f + exp(-gs)); + simd result = silu * us; + + const int out_row = token * rows_per_token + slot; + block_store( + intermediates + (size_t)out_row * intermediate_size + ng, + convert(result)); + } + }); + }; + + submit_kernel(cgf, device, "moe up merged e5m2"); +} + +// -- Merged up kernel, N-major variant (routed + shared in ONE dispatch) -- +// Both routed and shared gate_up weights are in sglang's native N-major layout +// [E, 2*inter, hidden] / [NS, 2*inter, hidden]. Same fusion as +// moe_up_merged_e5m2_kernel but reads weights via lsc_load_2d + nk permute so +// callers pass w13 (and the shared N-major gate_up) directly — NO transposed +// copy. Numerically identical to the K-major merged kernel fed transposed w13. +void moe_up_merged_e5m2_nmajor_kernel( + const fp16* x, + const uint8_t* gate_up_weight, const float* gate_up_scale, + const uint8_t* shared_gate_up_weight, const float* shared_gate_up_scale, + const int* selected_experts, + fp16* intermediates, + const int n_tokens, const int hidden_size, const int intermediate_size, + const int top_k, const int num_shared_experts, const int rows_per_token, + const torch::Device& device) { + + constexpr int GS = 16; + const int two_inter = 2 * intermediate_size; + const int n_tiles = intermediate_size / 16; + + auto cgf = [&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<2>( + sycl::range<2>(n_tokens * rows_per_token, n_tiles * GS), + sycl::range<2>(1, GS)), + [=](sycl::nd_item<2> item) SYCL_ESIMD_KERNEL { + slm_init(); + + const int route_idx = (int)item.get_global_id(0); + const int token = route_idx / rows_per_token; + const int slot = route_idx % rows_per_token; + const int n_tile = (int)item.get_group(1); + const int tid = (int)item.get_local_id(1); + + const fp16* x_row = x + (size_t)token * hidden_size; + + // Select routed vs shared weight base + scale. Both are N-major + // [*, 2*inter, hidden], so the per-expert stride is the same. + const uint8_t* base; + float scale; + if (slot < top_k) { + const int eid = selected_experts[token * top_k + slot]; + base = gate_up_weight + (size_t)eid * two_inter * hidden_size; + scale = gate_up_scale[eid]; + } else { + const int sid = slot - top_k; + base = shared_gate_up_weight + (size_t)sid * two_inter * hidden_size; + scale = shared_gate_up_scale[sid]; + } + + const int ng = n_tile * 16; + const int nu = intermediate_size + ng; + simd g_acc(0.f), u_acc(0.f); + + for (int k = 16 * tid; k < hidden_size; k += 16 * GS) { + simd a_tile = block_load(x_row + k); + + // [N=16, K=16] block at row=ng, col=k + // surface: width=hidden, height=two_inter, pitch=hidden + xesimd::config_2d_mem_access pgW( + base, (uint32_t)hidden_size - 1u, (uint32_t)two_inter - 1u, + (uint32_t)hidden_size - 1u, (uint32_t)k, (uint32_t)ng); + auto rg = xesimd::lsc_load_2d(pgW); + simd bg = fp8e5m2_block_to_vnni_nk(rg); + g_acc = dpas<8, 1, float, float, fp16, fp16>(g_acc, bg, a_tile); + + xesimd::config_2d_mem_access puW( + base, (uint32_t)hidden_size - 1u, (uint32_t)two_inter - 1u, + (uint32_t)hidden_size - 1u, (uint32_t)k, (uint32_t)nu); + auto ru = xesimd::lsc_load_2d(puW); + simd bu = fp8e5m2_block_to_vnni_nk(ru); + u_acc = dpas<8, 1, float, float, fp16, fp16>(u_acc, bu, a_tile); + } + + uint32_t slm_off = (uint32_t)(tid * 32) * 4u; + slm_block_store(slm_off, g_acc); + slm_block_store(slm_off + 64u, u_acc); + + barrier(); + + if (tid == 0) { + simd gs(0.f), us(0.f); + #pragma unroll + for (int i = 0; i < GS; i++) { + gs += slm_block_load((uint32_t)(i * 128)); + us += slm_block_load((uint32_t)(i * 128 + 64)); + } + + gs *= scale; + us *= scale; + + simd silu = gs / (1.f + exp(-gs)); + simd result = silu * us; + + const int out_row = token * rows_per_token + slot; + block_store( + intermediates + (size_t)out_row * intermediate_size + ng, + convert(result)); + } + }); + }; + + submit_kernel(cgf, device, "moe up merged e5m2 nmajor"); +} + torch::Tensor moe_up_forward( torch::Tensor x, torch::Tensor gate_up_weight, torch::Tensor gate_up_scale, @@ -1539,6 +2166,68 @@ void moe_down_finalize_e5m2_kernel( submit_kernel(cgf, device, "moe down finalize e5m2"); } +// Fused down finalize (4c): computes the routed-expert down projection INLINE +// per (token, j) instead of reading a precomputed routed_output buffer, folding +// the separate moe_down_routed dispatch into this kernel. Routed down uses the +// vllm layout down_weight [E, hidden, inter] (for fixed (eid,j) the inter values +// are contiguous, so a plain block_load GEMV replaces the DPAS pass). At BSZ=1 +// decode the extra per-work-item GPU work is free under the host-dispatch stall +// envelope, and we save one host dispatch + the s_routed_output buffer. +void moe_down_finalize_fused_e5m2_kernel( + const fp16* x, const fp16* intermediates, + const uint8_t* down_weight, const float* down_scale, + const fp16* routing_weights, const int* selected_experts, + const uint8_t* shared_down_weight, const float* shared_down_scale, + const fp16* shared_expert_gate_weight, fp16* final_output, + const int n_tokens, const int hidden_size, const int intermediate_size, + const int top_k, const int num_shared_experts, const int rows_per_token, + const torch::Device& device) { + auto cgf = [&](sycl::handler& cgh) { + cgh.parallel_for(sycl::range<2>(n_tokens, hidden_size), + [=](sycl::id<2> idx) SYCL_ESIMD_KERNEL { + const int token = (int)idx[0], j = (int)idx[1]; + float sum = 0.f; + // Routed experts: inline down GEMV (vllm layout [E, hidden, inter]). + for (int k = 0; k < top_k; k++) { + const int route_idx = token * top_k + k; + const int eid = selected_experts[route_idx]; + const int routed_row = token * rows_per_token + k; + const fp16* hi = intermediates + (size_t)routed_row * intermediate_size; + const uint8_t* dw = down_weight + + (size_t)eid * hidden_size * intermediate_size + + (size_t)j * intermediate_size; + fp16 ds = fp16(down_scale[eid]); + simd acc(0.f); + for (int kk = 0; kk < intermediate_size; kk += 64) { + auto d_dq = fp8e5m2_to_half<64>(block_load(dw+kk)) * ds; + acc += convert(d_dq * block_load(hi+kk)); + } + float w = (float)routing_weights[route_idx]; + sum += sycl::ext::intel::esimd::detail::sum(acc) * w; + } + const fp16* x_row = x + (size_t)token * hidden_size; + for (int sid = 0; sid < num_shared_experts; sid++) { + simd gate_acc(0.f); + for (int gk = 0; gk < hidden_size; gk += 64) + gate_acc += convert(block_load(x_row+gk)) * convert(block_load(shared_expert_gate_weight+sid*hidden_size+gk)); + float gate_w = 1.f/(1.f+sycl::exp(-sycl::ext::intel::esimd::detail::sum(gate_acc))); + const int shared_row = token * rows_per_token + top_k + sid; + const fp16* hi = intermediates + (size_t)shared_row * intermediate_size; + const uint8_t* dw = shared_down_weight + (size_t)sid*hidden_size*intermediate_size + (size_t)j*intermediate_size; + fp16 ds = fp16(shared_down_scale[sid]); + simd acc(0.f); + for (int k = 0; k < intermediate_size; k += 64) { + auto d_dq = fp8e5m2_to_half<64>(block_load(dw+k)) * ds; + acc += convert(d_dq * block_load(hi+k)); + } + sum += sycl::ext::intel::esimd::detail::sum(acc) * gate_w; + } + final_output[(size_t)token * hidden_size + j] = fp16(sum); + }); + }; + submit_kernel(cgf, device, "moe down finalize fused e5m2"); +} + torch::Tensor moe_forward_fused( torch::Tensor x, torch::Tensor gate_up_weight, torch::Tensor gate_up_scale, torch::Tensor shared_gate_up_weight, torch::Tensor shared_gate_up_scale, @@ -1579,6 +2268,7 @@ torch::Tensor moe_forward_fused( static thread_local torch::Tensor s_topk_idx, s_topk_weight; static thread_local torch::Tensor s_routed_output, s_final_output; static thread_local torch::Tensor s_intermediates; +static thread_local torch::Tensor s_normed_x, s_new_residual; static thread_local int s_cached_ntokens = -1; static void ensure_moe_buffers(int n_tokens, int top_k, int num_shared_experts, @@ -1596,6 +2286,10 @@ static void ensure_moe_buffers(int n_tokens, int top_k, int num_shared_experts, torch::device(dev).dtype(torch::kHalf)); s_intermediates = torch::empty({n_tokens * rows_per_token, intermediate_size}, torch::device(dev).dtype(torch::kHalf)); + s_normed_x = torch::empty({n_tokens, hidden_size}, + torch::device(dev).dtype(torch::kHalf)); + s_new_residual = torch::empty({n_tokens, hidden_size}, + torch::device(dev).dtype(torch::kHalf)); s_cached_ntokens = n_tokens; } } @@ -1610,7 +2304,7 @@ torch::Tensor moe_forward_full( int64_t top_k, int64_t num_shared_experts, int64_t n_routed_experts) { TORCH_CHECK(x.scalar_type() == torch::kHalf); int n_tokens = x.size(0), hidden_size = x.size(1); - int intermediate_size = gate_up_weight.size(2) / 2; + int intermediate_size = gate_up_weight.size(1) / 2; int rows_per_token = top_k + num_shared_experts; // Ensure pre-allocated buffers exist @@ -1629,7 +2323,7 @@ torch::Tensor moe_forward_full( // Up forward (inlined, using pre-allocated s_intermediates) const fp16* x_ptr = (const fp16*)x.data_ptr(); fp16* inter_ptr = (fp16*)s_intermediates.data_ptr(); - moe_up_routed_e5m2_kernel( + moe_up_routed_e5m2_nmajor_kernel( x_ptr, (const uint8_t*)gate_up_weight.data_ptr(), gate_up_scale.data_ptr(), s_topk_idx.data_ptr(), @@ -1643,8 +2337,13 @@ torch::Tensor moe_forward_full( n_tokens, hidden_size, intermediate_size, top_k, num_shared_experts, rows_per_token, x.device()); - // Down routed - moe_down_routed_e5m2_kernel((const fp16*)s_intermediates.data_ptr(), + // Down routed — use the vllm/sglang-natural down layout [E, hidden, inter] + // so we consume the routed down weight WITHOUT a transposed copy (the + // transposed w2_t replica would OOM at load; see moe_forward_full_v2). + // The output row layout (token*rows_per_token + k_idx, routing weight + // applied in-kernel) is identical to what moe_down_finalize_e5m2_kernel + // below sums over, so this is a drop-in for the transposed variant. + moe_down_routed_e5m2_vllm_layout_kernel((const fp16*)s_intermediates.data_ptr(), (const uint8_t*)down_weight.data_ptr(), down_scale.data_ptr(), (const fp16*)s_topk_weight.data_ptr(), s_topk_idx.data_ptr(), (fp16*)s_routed_output.data_ptr(), @@ -1660,6 +2359,168 @@ torch::Tensor moe_forward_full( return s_final_output; } +// Router-fused variant of moe_forward_full: takes the e4m3 gate weight + scale +// instead of precomputed logits, and fuses the router GEMV + softmax topk into a +// SINGLE dispatch (moe_router_topk_fused_e4m3_kernel), removing the standalone +// router_forward dispatch. Everything after topk is identical to +// moe_forward_full. BSZ=1 decode path. Falls back (returns undefined tensor via +// TORCH_CHECK-less bool) — caller must ensure the (NE, top_k) combo is supported. +torch::Tensor moe_forward_full_rtfused( + torch::Tensor x, + torch::Tensor router_weight, torch::Tensor router_scale, + torch::Tensor gate_up_weight, torch::Tensor gate_up_scale, + torch::Tensor shared_gate_up_weight, torch::Tensor shared_gate_up_scale, + torch::Tensor down_weight, torch::Tensor down_scale, + torch::Tensor shared_down_weight, torch::Tensor shared_down_scale, + torch::Tensor shared_expert_gate_weight, + int64_t top_k, int64_t num_shared_experts, int64_t n_routed_experts) { + TORCH_CHECK(x.scalar_type() == torch::kHalf); + int n_tokens = x.size(0), hidden_size = x.size(1); + int intermediate_size = gate_up_weight.size(1) / 2; + int rows_per_token = top_k + num_shared_experts; + + ensure_moe_buffers(n_tokens, top_k, num_shared_experts, hidden_size, + intermediate_size, x.device()); + + // Fused router GEMV + topk (one dispatch). If the (NE, top_k) combo is not + // specialized, fall back to the standalone dense e4m3 router + topk so the + // op is always correct. + bool fused = dispatch_moe_router_topk_fused_e4m3( + (const fp16*)x.data_ptr(), + (const uint8_t*)router_weight.data_ptr(), router_scale.data_ptr(), + (fp16*)s_topk_weight.data_ptr(), s_topk_idx.data_ptr(), + n_tokens, hidden_size, (int)n_routed_experts, (int)top_k, x.device()); + if (!fused) { + // Fallback: dense router GEMV into a temp logits buffer, then topk. + torch::Tensor logits = torch::empty({n_tokens, (int)n_routed_experts}, + torch::device(x.device()).dtype(torch::kHalf)); + moe_router_forward_e4m3_kernel( + (const fp16*)x.data_ptr(), + (const uint8_t*)router_weight.data_ptr(), router_scale.data_ptr(), + (fp16*)logits.data_ptr(), + n_tokens, hidden_size, (int)n_routed_experts, x.device()); + dispatch_moe_topk_forward( + (const fp16*)logits.data_ptr(), + s_topk_idx.data_ptr(), (fp16*)s_topk_weight.data_ptr(), + n_tokens, (int)n_routed_experts, (int)top_k, /*norm=*/true, x.device()); + } + + const fp16* x_ptr = (const fp16*)x.data_ptr(); + fp16* inter_ptr = (fp16*)s_intermediates.data_ptr(); + moe_up_routed_e5m2_nmajor_kernel( + x_ptr, + (const uint8_t*)gate_up_weight.data_ptr(), gate_up_scale.data_ptr(), + s_topk_idx.data_ptr(), + inter_ptr, + n_tokens, hidden_size, intermediate_size, + top_k, rows_per_token, x.device()); + moe_up_shared_e5m2_kernel( + x_ptr, + (const uint8_t*)shared_gate_up_weight.data_ptr(), shared_gate_up_scale.data_ptr(), + inter_ptr, + n_tokens, hidden_size, intermediate_size, + top_k, num_shared_experts, rows_per_token, x.device()); + + moe_down_routed_e5m2_vllm_layout_kernel((const fp16*)s_intermediates.data_ptr(), + (const uint8_t*)down_weight.data_ptr(), down_scale.data_ptr(), + (const fp16*)s_topk_weight.data_ptr(), s_topk_idx.data_ptr(), + (fp16*)s_routed_output.data_ptr(), + n_tokens, hidden_size, intermediate_size, top_k, rows_per_token, x.device()); + + moe_down_finalize_e5m2_kernel((const fp16*)x.data_ptr(), + (const fp16*)s_intermediates.data_ptr(), (const fp16*)s_routed_output.data_ptr(), + (const uint8_t*)shared_down_weight.data_ptr(), shared_down_scale.data_ptr(), + (const fp16*)shared_expert_gate_weight.data_ptr(), (fp16*)s_final_output.data_ptr(), + n_tokens, hidden_size, intermediate_size, top_k, num_shared_experts, rows_per_token, x.device()); + + return s_final_output; +} + +// Phase 3: router+topk+norm fused variant. Folds the post_attention_layernorm +// (GemmaRMSNorm resadd + rmsnorm) into the head of the fused router kernel and +// returns BOTH the MoE output and the new residual (hidden + residual), so the +// caller drops the standalone gemma_fused_add_rmsnorm dispatch AND needs no +// separate python residual add. `hidden`/`residual` are the pre-norm attention +// output and residual stream; `norm_weight` is the pre-folded Gemma (1+weight) +// in fp16. All internal MoE kernels consume the freshly-normed activation. +std::vector moe_forward_full_rtfused_norm( + torch::Tensor hidden, torch::Tensor residual, torch::Tensor norm_weight, + double eps, + torch::Tensor router_weight, torch::Tensor router_scale, + torch::Tensor gate_up_weight, torch::Tensor gate_up_scale, + torch::Tensor shared_gate_up_weight, torch::Tensor shared_gate_up_scale, + torch::Tensor down_weight, torch::Tensor down_scale, + torch::Tensor shared_down_weight, torch::Tensor shared_down_scale, + torch::Tensor shared_expert_gate_weight, + int64_t top_k, int64_t num_shared_experts, int64_t n_routed_experts) { + TORCH_CHECK(hidden.scalar_type() == torch::kHalf); + TORCH_CHECK(residual.scalar_type() == torch::kHalf); + int n_tokens = hidden.size(0), hidden_size = hidden.size(1); + int intermediate_size = gate_up_weight.size(1) / 2; + int rows_per_token = top_k + num_shared_experts; + + ensure_moe_buffers(n_tokens, top_k, num_shared_experts, hidden_size, + intermediate_size, hidden.device()); + + // Fused resadd+rmsnorm head + router GEMV + topk (one dispatch). Writes + // s_normed_x (fp16 normed activation) and s_new_residual (hidden+residual). + bool fused = dispatch_moe_router_topk_norm_fused_e4m3( + (const fp16*)hidden.data_ptr(), (const fp16*)residual.data_ptr(), + (const fp16*)norm_weight.data_ptr(), + (fp16*)s_normed_x.data_ptr(), (fp16*)s_new_residual.data_ptr(), + (const uint8_t*)router_weight.data_ptr(), router_scale.data_ptr(), + (fp16*)s_topk_weight.data_ptr(), s_topk_idx.data_ptr(), + n_tokens, hidden_size, (float)eps, (int)n_routed_experts, (int)top_k, + hidden.device()); + TORCH_CHECK(fused, + "moe_forward_full_rtfused_norm: unsupported (n_experts, top_k) combo"); + + const fp16* x_ptr = (const fp16*)s_normed_x.data_ptr(); + fp16* inter_ptr = (fp16*)s_intermediates.data_ptr(); + // Merged up: routed + shared experts in ONE dispatch (4b). Both routed and + // shared gate_up are in sglang N-major layout [*, 2*inter, hidden]. + moe_up_merged_e5m2_nmajor_kernel( + x_ptr, + (const uint8_t*)gate_up_weight.data_ptr(), gate_up_scale.data_ptr(), + (const uint8_t*)shared_gate_up_weight.data_ptr(), shared_gate_up_scale.data_ptr(), + s_topk_idx.data_ptr(), + inter_ptr, + n_tokens, hidden_size, intermediate_size, + top_k, num_shared_experts, rows_per_token, hidden.device()); + + // 4c (env-gated, SGL_XPU_MOE_DOWN_FUSED=1): fold moe_down_routed into the + // finalize kernel (routed down GEMV computed inline), dropping one dispatch + + // the s_routed_output buffer. Default OFF (two-dispatch DPAS path) so behaviour + // is unchanged unless explicitly enabled for A/B. + static const bool s_down_fused = [](){ + const char* e = std::getenv("SGL_XPU_MOE_DOWN_FUSED"); + return e && e[0] == '1'; + }(); + if (s_down_fused) { + moe_down_finalize_fused_e5m2_kernel(x_ptr, + (const fp16*)s_intermediates.data_ptr(), + (const uint8_t*)down_weight.data_ptr(), down_scale.data_ptr(), + (const fp16*)s_topk_weight.data_ptr(), s_topk_idx.data_ptr(), + (const uint8_t*)shared_down_weight.data_ptr(), shared_down_scale.data_ptr(), + (const fp16*)shared_expert_gate_weight.data_ptr(), (fp16*)s_final_output.data_ptr(), + n_tokens, hidden_size, intermediate_size, top_k, num_shared_experts, rows_per_token, hidden.device()); + } else { + moe_down_routed_e5m2_vllm_layout_kernel((const fp16*)s_intermediates.data_ptr(), + (const uint8_t*)down_weight.data_ptr(), down_scale.data_ptr(), + (const fp16*)s_topk_weight.data_ptr(), s_topk_idx.data_ptr(), + (fp16*)s_routed_output.data_ptr(), + n_tokens, hidden_size, intermediate_size, top_k, rows_per_token, hidden.device()); + + moe_down_finalize_e5m2_kernel(x_ptr, + (const fp16*)s_intermediates.data_ptr(), (const fp16*)s_routed_output.data_ptr(), + (const uint8_t*)shared_down_weight.data_ptr(), shared_down_scale.data_ptr(), + (const fp16*)shared_expert_gate_weight.data_ptr(), (fp16*)s_final_output.data_ptr(), + n_tokens, hidden_size, intermediate_size, top_k, num_shared_experts, rows_per_token, hidden.device()); + } + + return {s_final_output, s_new_residual}; +} + // V2: For BSZ>1 — uses standalone down path (gate_precompute + down_routed + // down_shared + accumulate) instead of down_finalize. down_finalize fuses // shared-expert GEMV per work-item with no weight reuse, which is optimal for @@ -1908,9 +2769,9 @@ torch::Tensor moe_forward_full_silu_routed_e5m2( TORCH_CHECK(down_weight.scalar_type() == at::kFloat8_e5m2, "down_weight must be fp8_e5m2"); int n_tokens = x.size(0), hidden_size = x.size(1); - TORCH_CHECK(gate_up_weight.size(1) == hidden_size, - "gate_up_weight must be [E, hidden, 2*intermediate]"); - int intermediate_size = gate_up_weight.size(2) / 2; + TORCH_CHECK(gate_up_weight.size(2) == hidden_size, + "gate_up_weight must be [E, 2*intermediate, hidden]"); + int intermediate_size = gate_up_weight.size(1) / 2; int rows_per_token = top_k; ensure_gemma4_moe_buffers(n_tokens, top_k, hidden_size, intermediate_size, x.device()); @@ -1919,7 +2780,7 @@ torch::Tensor moe_forward_full_silu_routed_e5m2( const fp16* w_ptr = (const fp16*)topk_weights.data_ptr(); // Up + silu_mul (e5m2 silu kernel — built-in silu, not gelu_tanh) - moe_up_routed_e5m2_kernel( + moe_up_routed_e5m2_nmajor_kernel( (const fp16*)x.data_ptr(), (const uint8_t*)gate_up_weight.data_ptr(), gate_up_scale.data_ptr(), idx_ptr, @@ -2088,7 +2949,7 @@ torch::Tensor moe_forward_full_v2( int64_t top_k, int64_t num_shared_experts, int64_t n_routed_experts) { TORCH_CHECK(x.scalar_type() == torch::kHalf); int n_tokens = x.size(0), hidden_size = x.size(1); - int intermediate_size = gate_up_weight.size(2) / 2; + int intermediate_size = gate_up_weight.size(1) / 2; int rows_per_token = top_k + num_shared_experts; ensure_moe_buffers_v2(n_tokens, top_k, num_shared_experts, hidden_size, @@ -2106,7 +2967,7 @@ torch::Tensor moe_forward_full_v2( // Up: routed + shared const fp16* x_ptr = (const fp16*)x.data_ptr(); fp16* inter_ptr = (fp16*)s2_intermediates.data_ptr(); - moe_up_routed_e5m2_kernel( + moe_up_routed_e5m2_nmajor_kernel( x_ptr, (const uint8_t*)gate_up_weight.data_ptr(), gate_up_scale.data_ptr(), s2_topk_idx.data_ptr(), @@ -2129,7 +2990,10 @@ torch::Tensor moe_forward_full_v2( s2_gate_values.data_ptr(), n_tokens, hidden_size, num_shared_experts, x.device()); - moe_down_routed_e5m2_kernel( + // Use the vllm/sglang-natural down layout [E, hidden, inter] so v2 can + // consume the routed down weight WITHOUT a transposed copy (saves one + // full weight replica per rank; the transposed variant would OOM at load). + moe_down_routed_e5m2_vllm_layout_kernel( (const fp16*)s2_intermediates.data_ptr(), (const uint8_t*)down_weight.data_ptr(), down_scale.data_ptr(), (const fp16*)s2_topk_weight.data_ptr(), s2_topk_idx.data_ptr(), @@ -2167,6 +3031,8 @@ TORCH_LIBRARY_FRAGMENT(moe_ops, m) { m.def("moe_forward_full_silu_routed_sglang(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(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"); + m.def("moe_forward_full_rtfused(Tensor x, Tensor router_weight, Tensor router_scale, 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_rtfused_norm(Tensor hidden, Tensor residual, Tensor norm_weight, float eps, Tensor router_weight, Tensor router_scale, 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[]"); } TORCH_LIBRARY_IMPL(moe_ops, XPU, m) { @@ -2184,6 +3050,8 @@ TORCH_LIBRARY_IMPL(moe_ops, XPU, m) { m.impl("moe_forward_fused", &moe_forward_fused); m.impl("moe_forward_full", &moe_forward_full); m.impl("moe_forward_full_v2", &moe_forward_full_v2); + m.impl("moe_forward_full_rtfused", &moe_forward_full_rtfused); + m.impl("moe_forward_full_rtfused_norm", &moe_forward_full_rtfused_norm); } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { @@ -2195,6 +3063,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("moe_forward_fused", &moe_forward_fused); m.def("moe_forward_full", &moe_forward_full); m.def("moe_forward_full_v2", &moe_forward_full_v2); + m.def("moe_forward_full_rtfused", &moe_forward_full_rtfused); + m.def("moe_forward_full_rtfused_norm", &moe_forward_full_rtfused_norm); m.def("moe_forward_full_silu_routed", &moe_forward_full_silu_routed); m.def("moe_forward_full_silu_routed_e5m2", &moe_forward_full_silu_routed_e5m2); m.def("moe_forward_full_silu_routed_sglang", &moe_forward_full_silu_routed_sglang); diff --git a/sglang/custom-esimd-kernels/csrc/moe_prefill/moe_prefill_fp8.sycl b/sglang/custom-esimd-kernels/csrc/moe_prefill/moe_prefill_fp8.sycl index ca7014e4..26c8e1c9 100644 --- a/sglang/custom-esimd-kernels/csrc/moe_prefill/moe_prefill_fp8.sycl +++ b/sglang/custom-esimd-kernels/csrc/moe_prefill/moe_prefill_fp8.sycl @@ -58,6 +58,66 @@ namespace xmx_ns = sycl::ext::intel::esimd::xmx; #include "../xpu/esimd_kernels/moe_ops.h" +// Build-configurable DPAS tile sizes (bit-exact across choices: per-output +// fp32 K-accumulation order is unchanged; only token/N grouping differs). +// Override via -DMP_UP_MAXM / -DMP_UP_N / -DMP_DOWN_MAXM / -DMP_DOWN_N. +#ifndef MP_UP_MAXM +#define MP_UP_MAXM 32 +#endif +#ifndef MP_UP_N +#define MP_UP_N 16 +#endif +#ifndef MP_DOWN_MAXM +#define MP_DOWN_MAXM 32 +#endif +#ifndef MP_DOWN_N +#define MP_DOWN_N 64 +#endif + +// ── Chunk-based load balancing ─────────────────────────────────────────────── +// The (expert, N-tile) grid serialized each expert's tokens inside a WI via an +// m_base loop, so a hot expert (2000+ routed tokens) ran ~70 serial MAX_M blocks +// while cold experts' WIs finished immediately → the GPU idled on the hot tail. +// Instead we pre-split every expert's token range into uniform MAX_M chunks so +// each WI does exactly one block; the launch grid is (chunk, N-tile). This is +// BIT-EXACT: each output element is still summed over the full K by one WI in +// the same order — only the WI→token-block mapping changed. +struct MpChunk { int eid; int t0; int nt; }; + +// Upper bound on #chunks for a MAX_M split: each expert adds at most +// ceil(nt_e/MAX_M) ≤ nt_e/MAX_M + 1, summed = total_seqlen/MAX_M + num_experts. +static inline int mp_chunk_capacity(int total_seqlen, int num_experts, int max_m) { + return total_seqlen / max_m + num_experts + 1; +} + +// Single-WI fill from expert_offsets (num_experts ~256, trivial). Trailing +// unused entries get nt=0 so their WIs early-return. +template +sycl::event moe_prefill_build_chunks_kernel( + sycl::queue& q, const int* expert_offsets, MpChunk* chunks, + int num_experts, int total_seqlen, int capacity) +{ + return q.submit([=](sycl::handler& cgh) { + cgh.parallel_for(sycl::range<1>(1), [=](sycl::id<1>) { + int c = 0; + for (int eid = 0; eid < num_experts; eid++) { + int lo = expert_offsets[eid]; + int hi = (eid + 1 < num_experts) ? expert_offsets[eid + 1] : total_seqlen; + for (int mb = lo; mb < hi; mb += MAX_M) { + if (c < capacity) { + chunks[c].eid = eid; + chunks[c].t0 = mb; + int nt = hi - mb; if (nt > MAX_M) nt = MAX_M; + chunks[c].nt = nt; + } + c++; + } + } + for (; c < capacity; c++) { chunks[c].eid = 0; chunks[c].t0 = 0; chunks[c].nt = 0; } + }); + }); +} + // ── FP8 → half scalar-vector dequant (identical math to the helpers in ────── // csrc/moe_batch/moe.sycl; kept local so this TU is self-contained) ─────── template @@ -169,6 +229,7 @@ sycl::event moe_prefill_up_forward_kernel_fp8( sycl::queue& q, const IT* input, const uint8_t* gate_up_weight, const float* gate_up_scale, IT* intermediate, const int* expert_offsets, const int* expert_tokens, + const MpChunk* chunks, int num_chunks, int num_experts, int total_seqlen, int hidden_size, int intermediate_size, int top_k) { static_assert(MAX_M % 16 == 0, "MAX_M must be multiple of 16"); @@ -179,16 +240,18 @@ sycl::event moe_prefill_up_forward_kernel_fp8( return q.submit([&](sycl::handler& cgh) { cgh.parallel_for( - sycl::range<2>(num_experts, intermediate_size / N), + sycl::range<2>(num_chunks, intermediate_size / N), [=](sycl::id<2> id) SYCL_ESIMD_KERNEL { const int two_inter = 2 * intermediate_size; - const int eid = (int)id[0]; + // One WI = one uniform MAX_M token-chunk × one N-tile (load balanced). + const MpChunk ck = chunks[(int)id[0]]; + const int eid = ck.eid; + const int m_base = ck.t0; + const int nt = ck.nt; + if (nt <= 0) return; const int tid = (int)id[1]; const int n_start = tid * N; // gate N-start; up = + intermediate_size - const int t0 = expert_offsets[eid]; - const int t1 = (eid + 1 < num_experts) ? expert_offsets[eid + 1] : total_seqlen; - if (t0 == t1) return; // w13 = [E, 2*I, H]: base for this expert, K = hidden_size = last dim const uint8_t* w_base = gate_up_weight + (size_t)eid * two_inter * hidden_size; @@ -198,10 +261,10 @@ sycl::event moe_prefill_up_forward_kernel_fp8( const float scale_gate = gate_up_scale[(size_t)eid * 2 + 0]; const float scale_up = gate_up_scale[(size_t)eid * 2 + 1]; - for (int m_base = t0; m_base < t1; m_base += MAX_M) { + { simd sorted_idxs = min(simd(0u, 1u) + (uint32_t)m_base, - simd((uint32_t)(t1 - 1))); + simd((uint32_t)(m_base + nt - 1))); simd pair_idxs = convert(gather(expert_tokens, sorted_idxs * 4u)); simd in_off = @@ -235,20 +298,20 @@ sycl::event moe_prefill_up_forward_kernel_fp8( // col_base : N-row offset into the 2*I weight rows. auto fp8_dpas = [&](int col_base, simd& acc) SYCL_ESIMD_FUNCTION { + // One 2D block load fetches all N weight rows × 16 K for this + // gate/up tile in row-major [row][k] — replaces N/8 * 8 tiny + // block_load scalar-row loads. Rows are strided by + // hidden_size; byte-x = k, row-y = col_base. Bit-identical bytes. + simd wraw_all = + xesimd::lsc_load_2d( + w_base, (unsigned)(hidden_size - 1), (unsigned)(two_inter - 1), + (unsigned)(hidden_size - 1), k, col_base); #pragma unroll for (int ns = 0; ns < NS; ns++) { - // Load FP8 weight [8 N-rows × 16 K] = 128 bytes, row-major. - // w13[eid][col_base + ns*8 + r][k .. k+15] - simd wraw; - #pragma unroll - for (int r = 0; r < 8; r++) { - const size_t row = (size_t)(col_base + ns * 8 + r); - wraw.template select<16, 1>(r * 16) = - block_load( - w_base + row * (size_t)hidden_size + k); - } // Dequant → plain [N=8, K=16] = 128 a_tile (DPAS operand A). - simd a_tile = mp_fp8_to_half(wraw); + simd a_tile = mp_fp8_to_half( + wraw_all.template select<128, 1>(ns * 128)); #pragma unroll for (int ms = 0; ms < MS; ms++) { @@ -294,7 +357,7 @@ sycl::event moe_prefill_up_forward_kernel_fp8( } } } - } // m_base + } // single chunk }); }); } @@ -307,6 +370,7 @@ sycl::event moe_prefill_down_forward_kernel_fp8( sycl::queue& q, const IT* intermediate, const uint8_t* down_weight, const float* down_scale, IT* expert_output, const int* expert_offsets, const int* expert_tokens, + const MpChunk* chunks, int num_chunks, int num_experts, int total_seqlen, int intermediate_size, int hidden_size) { static_assert(MAX_M % 16 == 0, "MAX_M must be multiple of 16"); @@ -317,24 +381,25 @@ sycl::event moe_prefill_down_forward_kernel_fp8( return q.submit([&](sycl::handler& cgh) { cgh.parallel_for( - sycl::range<2>(num_experts, hidden_size / N), + sycl::range<2>(num_chunks, hidden_size / N), [=](sycl::id<2> id) SYCL_ESIMD_KERNEL { - const int eid = (int)id[0]; + const MpChunk ck = chunks[(int)id[0]]; + const int eid = ck.eid; + const int m_base = ck.t0; + const int nt = ck.nt; + if (nt <= 0) return; const int tid = (int)id[1]; const int n_start = tid * N; - const int t0 = expert_offsets[eid]; - const int t1 = (eid + 1 < num_experts) ? expert_offsets[eid + 1] : total_seqlen; - if (t0 == t1) return; // w2 = [E, H, I]: K = intermediate_size = last dim const uint8_t* w_base = down_weight + (size_t)eid * hidden_size * intermediate_size; const float scale = down_scale[eid]; - for (int m_base = t0; m_base < t1; m_base += MAX_M) { + { simd sorted_idxs = min(simd(0u, 1u) + (uint32_t)m_base, - simd((uint32_t)(t1 - 1))); + simd((uint32_t)(m_base + nt - 1))); simd pair_idxs = convert(gather(expert_tokens, sorted_idxs * 4u)); simd in_off = @@ -358,17 +423,24 @@ sycl::event moe_prefill_down_forward_kernel_fp8( b_tile.template select<256, 1>(ms * 256) = convert(btmp); } + // 2D block load(s): N weight rows × 16 K, row-major, rows + // strided by intermediate_size. lsc_load_2d block height is + // capped at 32 rows, so for N > 32 we issue N/32 loads. + constexpr int WH = (N <= 32) ? N : 32; + constexpr int WNB = N / WH; + simd wraw_all; + #pragma unroll + for (int wb = 0; wb < WNB; wb++) { + wraw_all.template select(wb * WH * 16) = + xesimd::lsc_load_2d( + w_base, (unsigned)(intermediate_size - 1), (unsigned)(hidden_size - 1), + (unsigned)(intermediate_size - 1), k, n_start + wb * WH); + } #pragma unroll for (int ns = 0; ns < NS; ns++) { - simd wraw; - #pragma unroll - for (int r = 0; r < 8; r++) { - const size_t row = (size_t)(n_start + ns * 8 + r); - wraw.template select<16, 1>(r * 16) = - block_load( - w_base + row * (size_t)intermediate_size + k); - } - simd a_tile = mp_fp8_to_half(wraw); + simd a_tile = mp_fp8_to_half( + wraw_all.template select<128, 1>(ns * 128)); #pragma unroll for (int ms = 0; ms < MS; ms++) { @@ -404,7 +476,7 @@ sycl::event moe_prefill_down_forward_kernel_fp8( } } } - } // m_base + } // single chunk }); }); } @@ -477,13 +549,18 @@ std::vector moe_prefill_gather_forward_fp8( } // w13 = [E, 2*I, H] FP8 ; gate_up_scale = float[E, 2] (gate, up per expert) -torch::Tensor moe_prefill_up_forward_fp8( +// Core impl. If chunks_ext is defined, reuse the caller-built load-balancing +// chunks (built with MP_UP_MAXM) instead of building fresh — lets the fused +// driver build once and share between up+down (requires MP_UP_MAXM==MP_DOWN_MAXM). +static torch::Tensor moe_prefill_up_core( torch::Tensor x, torch::Tensor gate_up_weight, torch::Tensor gate_up_scale, torch::Tensor expert_offsets, torch::Tensor expert_tokens, - int64_t top_k) { + int64_t top_k, + torch::Tensor chunks_ext, + int cap_ext) { TORCH_CHECK(x.dim() == 2); TORCH_CHECK(gate_up_weight.dim() == 3, "gate_up_weight must be [E, 2*I, H]"); @@ -515,27 +592,63 @@ torch::Tensor moe_prefill_up_forward_fp8( const float* s_ptr = gate_up_scale.data_ptr(); fp16* o_ptr = (fp16*)intermediate.data_ptr(); + // Build load-balancing chunks (uniform MAX_M token blocks) from offsets, + // unless the caller passed prebuilt chunks (shared up+down). + constexpr int UP_MAXM = MP_UP_MAXM; + torch::Tensor chunks_t; + MpChunk* chunks; + int cap; + if (chunks_ext.defined()) { + chunks_t = chunks_ext; + cap = cap_ext; + chunks = reinterpret_cast(chunks_t.data_ptr()); + } else { + cap = mp_chunk_capacity(total_seqlen, num_experts, UP_MAXM); + chunks_t = torch::empty( + {cap * 3}, torch::device(x.device()).dtype(torch::kInt32)); + chunks = reinterpret_cast(chunks_t.data_ptr()); + moe_prefill_build_chunks_kernel( + queue, expert_offsets.data_ptr(), chunks, num_experts, total_seqlen, cap); + } + if (fp8_mode == 0) - moe_prefill_up_forward_kernel_fp8( + moe_prefill_up_forward_kernel_fp8( queue, x_ptr, w_ptr, s_ptr, o_ptr, expert_offsets.data_ptr(), expert_tokens.data_ptr(), + chunks, cap, num_experts, total_seqlen, hidden_size, intermediate_size, top_k); else - moe_prefill_up_forward_kernel_fp8( + moe_prefill_up_forward_kernel_fp8( queue, x_ptr, w_ptr, s_ptr, o_ptr, expert_offsets.data_ptr(), expert_tokens.data_ptr(), + chunks, cap, num_experts, total_seqlen, hidden_size, intermediate_size, top_k); return intermediate; } +// Public op: builds its own chunks. +torch::Tensor moe_prefill_up_forward_fp8( + torch::Tensor x, + torch::Tensor gate_up_weight, + torch::Tensor gate_up_scale, + torch::Tensor expert_offsets, + torch::Tensor expert_tokens, + int64_t top_k) { + return moe_prefill_up_core(x, gate_up_weight, gate_up_scale, + expert_offsets, expert_tokens, top_k, + torch::Tensor(), -1); +} + // w2 = [E, H, I] FP8 ; down_scale = float[E] (per-expert scalar) -torch::Tensor moe_prefill_down_forward_fp8( +static torch::Tensor moe_prefill_down_core( torch::Tensor intermediate, torch::Tensor down_weight, torch::Tensor down_scale, torch::Tensor expert_offsets, - torch::Tensor expert_tokens) { + torch::Tensor expert_tokens, + torch::Tensor chunks_ext, + int cap_ext) { TORCH_CHECK(intermediate.dim() == 2); TORCH_CHECK(down_weight.dim() == 3, "down_weight must be [E, H, I]"); @@ -567,20 +680,51 @@ torch::Tensor moe_prefill_down_forward_fp8( const float* s_ptr = down_scale.data_ptr(); fp16* o_ptr = (fp16*)expert_output.data_ptr(); + constexpr int DOWN_MAXM = MP_DOWN_MAXM; + torch::Tensor chunks_t; + MpChunk* chunks; + int cap; + if (chunks_ext.defined()) { + chunks_t = chunks_ext; + cap = cap_ext; + chunks = reinterpret_cast(chunks_t.data_ptr()); + } else { + cap = mp_chunk_capacity(total_seqlen, num_experts, DOWN_MAXM); + chunks_t = torch::empty( + {cap * 3}, torch::device(intermediate.device()).dtype(torch::kInt32)); + chunks = reinterpret_cast(chunks_t.data_ptr()); + moe_prefill_build_chunks_kernel( + queue, expert_offsets.data_ptr(), chunks, num_experts, total_seqlen, cap); + } + if (fp8_mode == 0) - moe_prefill_down_forward_kernel_fp8( + moe_prefill_down_forward_kernel_fp8( queue, i_ptr, w_ptr, s_ptr, o_ptr, expert_offsets.data_ptr(), expert_tokens.data_ptr(), + chunks, cap, num_experts, total_seqlen, intermediate_size, hidden_size); else - moe_prefill_down_forward_kernel_fp8( + moe_prefill_down_forward_kernel_fp8( queue, i_ptr, w_ptr, s_ptr, o_ptr, expert_offsets.data_ptr(), expert_tokens.data_ptr(), + chunks, cap, num_experts, total_seqlen, intermediate_size, hidden_size); return expert_output; } +// Public op: builds its own chunks. +torch::Tensor moe_prefill_down_forward_fp8( + torch::Tensor intermediate, + torch::Tensor down_weight, + torch::Tensor down_scale, + torch::Tensor expert_offsets, + torch::Tensor expert_tokens) { + return moe_prefill_down_core(intermediate, down_weight, down_scale, + expert_offsets, expert_tokens, + torch::Tensor(), -1); +} + torch::Tensor moe_prefill_accumulate_forward_fp8( torch::Tensor output, torch::Tensor routing_weights) { @@ -634,15 +778,34 @@ torch::Tensor moe_prefill_full_fp8( auto expert_offsets = gather_result[0]; auto expert_tokens = gather_result[1]; - // 2. Up GEMM: gate_up projection + SiLU. - auto intermediate = moe_prefill_up_forward_fp8( - hidden_states, w13, w13_scale, expert_offsets, expert_tokens, top_k); + // 2. Build load-balancing chunks ONCE and share between up+down. + // Valid only when MP_UP_MAXM == MP_DOWN_MAXM (identical split); otherwise + // each core builds its own. + torch::Tensor shared_chunks; + int shared_cap = -1; + if (MP_UP_MAXM == MP_DOWN_MAXM) { + int total_seqlen = expert_tokens.numel(); + shared_cap = mp_chunk_capacity(total_seqlen, (int)num_experts, MP_UP_MAXM); + shared_chunks = torch::empty( + {shared_cap * 3}, torch::device(hidden_states.device()).dtype(torch::kInt32)); + sycl::queue& q = c10::xpu::getCurrentXPUStream(hidden_states.device().index()).queue(); + moe_prefill_build_chunks_kernel( + q, expert_offsets.data_ptr(), + reinterpret_cast(shared_chunks.data_ptr()), + (int)num_experts, total_seqlen, shared_cap); + } + + // 3. Up GEMM: gate_up projection + SiLU. + auto intermediate = moe_prefill_up_core( + hidden_states, w13, w13_scale, expert_offsets, expert_tokens, top_k, + shared_chunks, shared_cap); - // 3. Down GEMM. - auto expert_output = moe_prefill_down_forward_fp8( - intermediate, w2, w2_scale, expert_offsets, expert_tokens); + // 4. Down GEMM. + auto expert_output = moe_prefill_down_core( + intermediate, w2, w2_scale, expert_offsets, expert_tokens, + shared_chunks, shared_cap); - // 4. Accumulate: weighted gather over top_k. + // 5. Accumulate: weighted gather over top_k. auto final_output = moe_prefill_accumulate_forward_fp8(expert_output, topk_weights); return final_output; diff --git a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernel.sycl b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernel.sycl index f91f8ee6..3f3572d1 100644 --- a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernel.sycl +++ b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernel.sycl @@ -308,6 +308,9 @@ at::Tensor esimd_gemv_q6_k_m( } #include "esimd_kernels/moe_kquant_GEMV.h" +#include "esimd_kernels/moe_shared_q8_GEMV.h" +#include "esimd_kernels/moe_up_fused.h" +#include "esimd_kernels/moe_norm_router.h" // Fused GGUF k-quant MoE up/gate stage (Q4_K gate + Q4_K up -> silu*up). // gate_ql/up_ql [E,inter,hidden/2] u8; gate_sc/mn,up_sc/mn [E,inter,hidden/32] @@ -373,6 +376,229 @@ at::Tensor esimd_moe_down_q6k( return out_partial; } +// Fused GGUF Q8_0 SHARED-expert MLP (decode): one op = one python dispatch. +// g[t] = sigmoid(dot(x[t], wg)) +// inter[t,c] = silu() * (Q8_0 gate/up merged) +// out[t,n] = g[t] * sum_c inter[t,c]*W_down[n,c] (Q8_0 down) +// gu_qs [2*inter_s, hidden] int8 (rows gate then up); gu_sc [2*inter_s, hidden/32] +// fp16; d_qs [hidden, inter_s] int8; d_sc [hidden, inter_s/32] fp16; +// wg [hidden] fp16 (shared_expert_gate weight row 0). x/out [M, hidden] fp16. +at::Tensor esimd_shared_expert_q8( + at::Tensor x, at::Tensor gu_qs, at::Tensor gu_sc, + at::Tensor d_qs, at::Tensor d_sc, at::Tensor wg, int64_t inter_s) { + int64_t M = x.size(0); + int64_t hidden = x.size(1); + auto out = at::empty({M, hidden}, x.options()); + auto inter = at::empty({M, inter_s}, x.options()); + auto g = at::empty({M}, x.options()); + auto p_x = reinterpret_cast(x.data_ptr()); + auto p_gq = reinterpret_cast(gu_qs.data_ptr()); + auto p_gs = reinterpret_cast(gu_sc.data_ptr()); + auto p_dq = reinterpret_cast(d_qs.data_ptr()); + auto p_ds = reinterpret_cast(d_sc.data_ptr()); + auto p_wg = reinterpret_cast(wg.data_ptr()); + auto p_g = reinterpret_cast(g.data_ptr()); + auto p_in = reinterpret_cast(inter.data_ptr()); + auto p_o = reinterpret_cast(out.data_ptr()); + auto& dpcpp_queue = get_device_queue(x); + moe_shared_up_q8_host(p_x, p_gq, p_gs, p_wg, p_in, p_g, + (int)M, (int)hidden, (int)inter_s, dpcpp_queue); + moe_shared_down_q8_host(p_in, p_dq, p_ds, p_g, p_o, + (int)M, (int)hidden, (int)inter_s, dpcpp_queue); + return out; +} + +// Fused GGUF FULL MoE (decode): router topk + routed experts + shared expert +// -> ONE op = ONE python dispatch. Router GEMV stays in python (logits input). +// Chain: topk -> routed up (Q4_K) -> shared gate/up (Q8_0) -> routed down +// (Q5_K or Q6_K) -> finalize (routed combine + shared down + gate*sigmoid). +// Routed: gate_ql/up_ql [E,inter,hidden/2] u8; *_sc/*_mn [E,inter,hidden/32]; +// down_ql [E,hidden,inter/2] u8; down_qh [E,hidden,inter/32] u8; +// down_sc/down_mn [E,hidden,inter/32] fp16 (q6k: down_mn unused). +// Shared: gu_qs [2*inter_s,hidden] i8; gu_sc [2*inter_s,hidden/32] fp16; +// d_qs [hidden,inter_s] i8; d_sc [hidden,inter_s/32] fp16; wg [hidden]. +// Shared body of the GGUF decode MoE op. `p_x` is the (already normalised) +// activation and `p_log` the router logits; the norm-fused entry point below +// produces both on device instead of taking them from python. + +// Pooled decode scratch for the GGUF MoE op (see the comment in the body). +// One rank == one process == one device, and decode is single-threaded, so a +// single function-local instance is enough; a shape/dtype/device change simply +// reallocates. +struct GgufMoeWs { + int64_t M = -1, hidden = -1, n_routed = -1, intermediate = -1, inter_s = -1; + at::ScalarType dtype = at::ScalarType::Undefined; + c10::DeviceIndex dev = -1; + at::Tensor sel, tw, inter_r, inter_sh, g, out_part; + at::Tensor xn, logits; // norm-fused variant only +}; + +static GgufMoeWs& gguf_moe_ws() { + static GgufMoeWs ws; + return ws; +} + +static at::Tensor gguf_moe_full_body( + const at::Tensor& x, const fp16* p_x, const fp16* p_log, + at::Tensor gate_ql, at::Tensor gate_sc, at::Tensor gate_mn, + at::Tensor up_ql, at::Tensor up_sc, at::Tensor up_mn, + at::Tensor down_ql, at::Tensor down_qh, at::Tensor down_sc, at::Tensor down_mn, + at::Tensor gu_qs, at::Tensor gu_sc, at::Tensor d_qs, at::Tensor d_sc, + at::Tensor wg, + int64_t n_experts, int64_t top_k, int64_t intermediate, int64_t inter_s, + bool down_is_q6, bool renorm) { + int64_t M = x.size(0); + int64_t hidden = x.size(1); + int64_t n_routed = M * top_k; + // Scratch is pooled across calls. At decode every one of the 40 MoE layers + // asks for the identical shapes, and 6 at::empty per call x 40 layers was + // ~0.5ms of pure allocator dispatch per step on this host-bound path. + // Safe because all six kernels below go on the SAME in-order queue, so + // layer i+1's first kernel cannot start before layer i's last one retired + // (the existing code already relies on that ordering between its own + // kernels). `out` is NOT pooled: it escapes to python (all-reduce, residual + // add) and must stay live past this call. + auto& ws = gguf_moe_ws(); + if (ws.M != M || ws.hidden != hidden || ws.n_routed != n_routed || + ws.intermediate != intermediate || ws.inter_s != inter_s || + ws.dtype != x.scalar_type() || ws.dev != x.device().index()) { + auto i32opt = x.options().dtype(at::kInt); + ws.sel = at::empty({n_routed}, i32opt); + ws.tw = at::empty({n_routed}, x.options()); + ws.inter_r = at::empty({n_routed, intermediate}, x.options()); + ws.inter_sh = at::empty({M, inter_s}, x.options()); + ws.g = at::empty({M}, x.options()); + ws.out_part = at::empty({n_routed, hidden}, x.options()); + ws.M = M; ws.hidden = hidden; ws.n_routed = n_routed; + ws.intermediate = intermediate; ws.inter_s = inter_s; + ws.dtype = x.scalar_type(); ws.dev = x.device().index(); + } + auto out = at::empty({M, hidden}, x.options()); + + auto p_sel = reinterpret_cast(ws.sel.data_ptr()); + auto p_tw = reinterpret_cast(ws.tw.data_ptr()); + auto p_ir = reinterpret_cast(ws.inter_r.data_ptr()); + auto p_ish = reinterpret_cast(ws.inter_sh.data_ptr()); + auto p_g = reinterpret_cast(ws.g.data_ptr()); + auto p_op = reinterpret_cast(ws.out_part.data_ptr()); + auto p_o = reinterpret_cast(out.data_ptr()); + // routed reps + auto p_gq = reinterpret_cast(gate_ql.data_ptr()); + auto p_gs = reinterpret_cast(gate_sc.data_ptr()); + auto p_gm = reinterpret_cast(gate_mn.data_ptr()); + auto p_uq = reinterpret_cast(up_ql.data_ptr()); + auto p_us = reinterpret_cast(up_sc.data_ptr()); + auto p_um = reinterpret_cast(up_mn.data_ptr()); + auto p_dql = reinterpret_cast(down_ql.data_ptr()); + auto p_dqh = reinterpret_cast(down_qh.data_ptr()); + auto p_dsc = reinterpret_cast(down_sc.data_ptr()); + auto p_dmn = reinterpret_cast(down_mn.data_ptr()); + // shared reps + auto p_guq = reinterpret_cast(gu_qs.data_ptr()); + auto p_gus = reinterpret_cast(gu_sc.data_ptr()); + auto p_ddq = reinterpret_cast(d_qs.data_ptr()); + auto p_dds = reinterpret_cast(d_sc.data_ptr()); + auto p_wg = reinterpret_cast(wg.data_ptr()); + + auto& q = get_device_queue(x); + // The topk kernel keeps all expert scores in one simd, and the + // host below only instantiates PAD = 256 or 512, so a larger expert count + // would index past the end of that vector. + TORCH_CHECK(n_experts > 0 && n_experts <= 512, + "gguf MoE: n_experts ", n_experts, " outside the supported 1..512"); + moe_topk_gguf_host(p_log, p_sel, p_tw, (int)M, (int)n_experts, (int)top_k, + renorm, q); + // One enqueue for both up stages when the shape allows it (see + // moe_up_fused.h); otherwise the original pair of launches. + if (!moe_up_fused_host(p_x, p_gq, p_gs, p_gm, p_uq, p_us, p_um, p_sel, p_ir, + p_guq, p_gus, p_wg, p_ish, p_g, + (int)M, (int)hidden, (int)intermediate, (int)top_k, + (int)inter_s, q)) { + moe_up_q4k_host(p_x, p_gq, p_gs, p_gm, p_uq, p_us, p_um, p_sel, p_ir, + (int)M, (int)hidden, (int)intermediate, (int)top_k, q); + moe_shared_up_q8_host(p_x, p_guq, p_gus, p_wg, p_ish, p_g, + (int)M, (int)hidden, (int)inter_s, q); + } + if (down_is_q6) + moe_down_q6k_host(p_ir, p_dql, p_dqh, p_dsc, p_sel, p_tw, p_op, + (int)M, (int)hidden, (int)intermediate, (int)top_k, q); + else + moe_down_q5k_host(p_ir, p_dql, p_dqh, p_dsc, p_dmn, p_sel, p_tw, p_op, + (int)M, (int)hidden, (int)intermediate, (int)top_k, q); + moe_finalize_gguf_host(p_op, p_ish, p_ddq, p_dds, p_g, p_o, + (int)M, (int)hidden, (int)inter_s, (int)top_k, q); + return out; +} + +#define GGUF_MOE_FULL_ARGS \ + gate_ql, gate_sc, gate_mn, up_ql, up_sc, up_mn, \ + down_ql, down_qh, down_sc, down_mn, gu_qs, gu_sc, d_qs, d_sc, wg, \ + n_experts, top_k, intermediate, inter_s, down_is_q6, renorm + +at::Tensor esimd_moe_forward_full_gguf( + at::Tensor x, at::Tensor logits, + at::Tensor gate_ql, at::Tensor gate_sc, at::Tensor gate_mn, + at::Tensor up_ql, at::Tensor up_sc, at::Tensor up_mn, + at::Tensor down_ql, at::Tensor down_qh, at::Tensor down_sc, at::Tensor down_mn, + at::Tensor gu_qs, at::Tensor gu_sc, at::Tensor d_qs, at::Tensor d_sc, + at::Tensor wg, + int64_t n_experts, int64_t top_k, int64_t intermediate, int64_t inter_s, + bool down_is_q6, bool renorm) { + return gguf_moe_full_body( + x, reinterpret_cast(x.data_ptr()), + reinterpret_cast(logits.data_ptr()), GGUF_MOE_FULL_ARGS); +} + +// Norm-fused variant: absorbs the post-attention GemmaRMSNorm (residual add + +// RMS norm) and the fp16 router GEMV into the MoE op, removing two per-layer +// python/dispatcher round-trips on the host-bound decode path. +// `residual` is updated in place (== h + residual, matching +// gemma_fused_add_rmsnorm) and returned alongside the MoE output. +std::vector esimd_moe_forward_full_gguf_norm( + at::Tensor h, at::Tensor residual, at::Tensor norm_w, double eps, + at::Tensor router_w, + at::Tensor gate_ql, at::Tensor gate_sc, at::Tensor gate_mn, + at::Tensor up_ql, at::Tensor up_sc, at::Tensor up_mn, + at::Tensor down_ql, at::Tensor down_qh, at::Tensor down_sc, at::Tensor down_mn, + at::Tensor gu_qs, at::Tensor gu_sc, at::Tensor d_qs, at::Tensor d_sc, + at::Tensor wg, + int64_t n_experts, int64_t top_k, int64_t intermediate, int64_t inter_s, + bool down_is_q6, bool renorm) { + const int64_t M = h.size(0); + const int64_t hidden = h.size(1); + // xn/logits never leave this op, so they are pooled like the other scratch. + auto& nws = gguf_moe_ws(); + if (nws.xn.defined() == false || nws.xn.size(0) != M || + nws.xn.size(1) != hidden || nws.logits.size(1) != n_experts || + nws.xn.scalar_type() != h.scalar_type() || + nws.xn.device().index() != h.device().index()) { + nws.xn = at::empty({M, hidden}, h.options()); + nws.logits = at::empty({M, n_experts}, h.options()); + } + auto& xn = nws.xn; + auto& logits = nws.logits; + // Distinct output buffer: block 0 stores the post-add residual while the + // other row-blocks are still loading the OLD residual, so writing in place + // would be a data race. + auto nr = at::empty({M, hidden}, h.options()); + auto& q = get_device_queue(h); + const bool ok = moe_norm_router_host( + reinterpret_cast(h.data_ptr()), + reinterpret_cast(residual.data_ptr()), + reinterpret_cast(nr.data_ptr()), + reinterpret_cast(norm_w.data_ptr()), + reinterpret_cast(router_w.data_ptr()), + reinterpret_cast(xn.data_ptr()), + reinterpret_cast(logits.data_ptr()), + (float)eps, (int)M, (int)hidden, (int)n_experts, q); + TORCH_CHECK(ok, "esimd_moe_forward_full_gguf_norm: unsupported hidden ", hidden); + auto out = gguf_moe_full_body( + xn, reinterpret_cast(xn.data_ptr()), + reinterpret_cast(logits.data_ptr()), GGUF_MOE_FULL_ARGS); + return {out, nr}; +} +#undef GGUF_MOE_FULL_ARGS + #include "esimd_kernels/q4_0_GEMM.h" // GGUF q4_0 GEMM (prefill / batched M>=2), non-DPAS SIMD (PTL Xe3 has no XMX). @@ -394,6 +620,85 @@ at::Tensor esimd_gemm_q4_0( return output; } +// Fused (residual-add + GemmaRMSNorm) + GGUF q8_0 GEMV [+ optional fp16 dense +// GEMV] for the decode attention input projection. +// +// nr = fp16(h + residual) (new residual, distinct +// buffer: block 0 stores +// while others still read) +// xn = fp16(nr * rsqrt(mean(nr^2)+eps) * nw) (scratch) +// o0 = xn @ w0^T (q8_0: qs [N0,K] int8, sc [N0,K/32] fp16) +// o1 = xn @ w1^T (fp16 dense [N1,K]; skipped when w1 is empty) +// +// Collapses 2-3 torch op dispatches per attention layer into one. Decode is +// host-bound (~13-25us of dispatch per call vs ~4us of enqueue), so removing +// the CALL is the win; the two kernels are still enqueued back to back on the +// same in-order queue, which serialises them correctly (o0's kernel reads the +// xn that the norm kernel wrote). +void esimd_resadd_norm_gemv_q8_ba( + at::Tensor h, at::Tensor residual, at::Tensor nw, double eps, + at::Tensor xn, at::Tensor nr, + at::Tensor w0, at::Tensor s0, at::Tensor o0, + at::Tensor w1, at::Tensor o1) { + const int M = (int)h.size(0); + const int hidden = (int)h.size(1); + const bool has_ba = w1.numel() > 0; + const int E = has_ba ? (int)w1.size(0) : 0; + + EXTRACT_PTR(p_h, h); EXTRACT_PTR(p_res, residual); EXTRACT_PTR(p_nw, nw); + EXTRACT_PTR(p_xn, xn); EXTRACT_PTR(p_nr, nr); + auto& dpcpp_queue = get_device_queue(h); + + // Single enqueue: the WG grid covers the fp16 `ba` rows and the q8_0 rows, + // every WG recomputing the (register-local) norm. + const bool ok = moe_norm_q8_host( + reinterpret_cast(p_h), + reinterpret_cast(p_res), reinterpret_cast(p_nr), + reinterpret_cast(p_nw), + has_ba ? reinterpret_cast(w1.data_ptr()) : nullptr, + has_ba ? reinterpret_cast(o1.data_ptr()) : nullptr, + reinterpret_cast(w0.data_ptr()), + reinterpret_cast(s0.data_ptr()), + reinterpret_cast(o0.data_ptr()), + nullptr, (float)eps, M, hidden, E, (int)w0.size(0), dpcpp_queue); + TORCH_CHECK(ok, "esimd_resadd_norm_gemv_q8_ba: unsupported hidden ", hidden); + (void)p_xn; +} + +#include "esimd_kernels/gdn_norm_gated.h" + +// Fused RMSNormGated + GGUF q8_0 GEMV for the GDN out_proj (decode). +// +// y[t,h,:] = rmsnorm(x[t,h,:]) * nw * silu(z[t,h,:]) (per token, per head) +// out[t,:] = y_flat[t] @ dequant(w_q8)^T +// +// The GGUF counterpart of esimd_norm_gemv_fp8_pert: collapses the standalone +// gdn_rms_norm_gated launch and the out_proj GEMV into one op dispatch. +// +// x/z are [M*HV, V] and out is [M, N]; M is inferred from the row count so a +// batched decode step is one launch, same as the single-token case. +void esimd_norm_gemv_q8_0( + at::Tensor x, at::Tensor z, at::Tensor nw, at::Tensor y, + at::Tensor w0, at::Tensor s0, at::Tensor o0, + int64_t HV, int64_t V, double eps) { + EXTRACT_PTR(p_x, x); EXTRACT_PTR(p_z, z); EXTRACT_PTR(p_nw, nw); + EXTRACT_PTR(p_y, y); + auto& dpcpp_queue = get_device_queue(x); + + TORCH_CHECK(HV > 0 && x.size(0) % HV == 0, + "esimd_norm_gemv_q8_0: rows ", x.size(0), " not a multiple of HV ", HV); + const int M = (int)(x.size(0) / HV); + + const bool ok = gdn_norm_gated_q8_host( + reinterpret_cast(p_x), reinterpret_cast(p_z), + reinterpret_cast(p_nw), reinterpret_cast(p_y), + reinterpret_cast(w0.data_ptr()), + reinterpret_cast(s0.data_ptr()), + reinterpret_cast(o0.data_ptr()), + (float)eps, (int)HV, (int)V, (int)w0.size(0), M, dpcpp_queue); + TORCH_CHECK(ok, "esimd_norm_gemv_q8_0: unsupported head_v_dim ", V); +} + // Fused 2-matrix INT4 GEMV: two GEMVs sharing the same input vector. // Used for GDN input projection (in_proj_qkvz + in_proj_ba). at::Tensor esimd_gemv_int4_fused2( @@ -556,6 +861,7 @@ at::Tensor esimd_resadd_norm_gemv2_fp8_pert( at::Tensor norm_weight, at::Tensor w0, at::Tensor s0, at::Tensor o0, at::Tensor w1, at::Tensor s1, at::Tensor o1, + at::Tensor new_residual, double eps) { int N0 = (int)w0.size(0); @@ -564,6 +870,9 @@ at::Tensor esimd_resadd_norm_gemv2_fp8_pert( int fp8_mode = get_fp8_mode(w0); auto& dpcpp_queue = get_device_queue(hidden_states); + fp16* new_residual_ptr = new_residual.defined() + ? reinterpret_cast(new_residual.data_ptr()) : nullptr; + resadd_norm_gemv2_fp8_pert_host( reinterpret_cast(hidden_states.data_ptr()), reinterpret_cast(residual.data_ptr()), @@ -574,6 +883,7 @@ at::Tensor esimd_resadd_norm_gemv2_fp8_pert( reinterpret_cast(w1.data_ptr()), s1.data_ptr(), reinterpret_cast(o1.data_ptr()), + new_residual_ptr, N0, N1, K, (float)eps, fp8_mode, dpcpp_queue); diff --git a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernel_lgrf.sycl b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernel_lgrf.sycl index 0b9e310c..8b19efb2 100644 --- a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernel_lgrf.sycl +++ b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernel_lgrf.sycl @@ -105,13 +105,20 @@ at::Tensor esimd_gdn_conv_fused_seq( auto* p_out = reinterpret_cast(output.data_ptr()); auto* p_zout = reinterpret_cast(z_out.data_ptr()); + // Detect conv_state memory layout. The native MambaPool stores conv as + // (cache, conv_dim, W-1) whose last dim (W-1, e.g. 3) is tiny, so taps for + // each channel are interleaved (stride = W-1). The legacy path passes a + // transposed (cache, W-1, conv_dim) copy whose last dim is conv_dim (large, + // contiguous per tap). Distinguish by the size of the last dim. + int conv_native = (conv_state.dim() >= 2 && conv_state.size(-1) <= 8) ? 1 : 0; + gdn_conv_fused_seq_host( p_qkvz, qkvz_stride0, p_cstate, p_cweight, p_cbias, p_csidx, p_alog, p_dtbias, p_ba, ba_stride0, p_sstate, p_ssidx, p_out, p_zout, (int)N, (int)H, (int)HV, (int)K, (int)V, (float)scale, conv_stride0, ssm_stride0, - dpcpp_queue); + conv_native, dpcpp_queue); return output; } diff --git a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/gdn_conv_fused_seq.h b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/gdn_conv_fused_seq.h index 9397566a..e8481410 100644 --- a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/gdn_conv_fused_seq.h +++ b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/gdn_conv_fused_seq.h @@ -122,6 +122,7 @@ ESIMD_INLINE void gdn_conv_fused_seq_kernel( int N, int H, int HV, int gdn_K, int gdn_V, float attn_scale, int64_t conv_stride0, int64_t ssm_stride0, int inline_conv_shift, // 1 = do conv_state shift inline (safe when N*HV<=32) + int conv_native, // 1 = conv_state is native pool layout (cache, conv_dim, W-1=3) nd_item<3>& ndi) { slm_init<2048>(); @@ -208,9 +209,21 @@ ESIMD_INLINE void gdn_conv_fused_seq_kernel( simd x_fp16 = block_load(qkvz_row + qkvz_offset); simd x_f32 = x_fp16; - simd s0 = block_load(cstate_base + 0 * dim + chunk_start); - simd s1 = block_load(cstate_base + 1 * dim + chunk_start); - simd s2 = block_load(cstate_base + 2 * dim + chunk_start); + simd s0, s1, s2; + if (conv_native) { + // Native pool layout (cache, conv_dim, W-1=3): taps for each channel + // are stored interleaved with stride 3. Load 64 channels x 3 taps as a + // contiguous 192-vector starting at chunk_start*3, then deinterleave. + simd raw = block_load(cstate_base + (int64_t)chunk_start * 3); + simd t0 = raw.select<64, 3>(0); + simd t1 = raw.select<64, 3>(1); + simd t2 = raw.select<64, 3>(2); + s0 = t0; s1 = t1; s2 = t2; + } else { + s0 = block_load(cstate_base + 0 * dim + chunk_start); + s1 = block_load(cstate_base + 1 * dim + chunk_start); + s2 = block_load(cstate_base + 2 * dim + chunk_start); + } simd w_raw = block_load(conv_weight_ptr + (int64_t)chunk_start * 4); simd conv_result = @@ -232,9 +245,17 @@ ESIMD_INLINE void gdn_conv_fused_seq_kernel( x_fp16_hi = block_load(qkvz_row + qkvz_offset_hi); simd x_f32_hi = x_fp16_hi; - s0_hi = block_load(cstate_base + 0 * dim + chunk_start_hi); - s1_hi = block_load(cstate_base + 1 * dim + chunk_start_hi); - s2_hi = block_load(cstate_base + 2 * dim + chunk_start_hi); + if (conv_native) { + simd raw_hi = block_load(cstate_base + (int64_t)chunk_start_hi * 3); + simd h0 = raw_hi.select<64, 3>(0); + simd h1 = raw_hi.select<64, 3>(1); + simd h2 = raw_hi.select<64, 3>(2); + s0_hi = h0; s1_hi = h1; s2_hi = h2; + } else { + s0_hi = block_load(cstate_base + 0 * dim + chunk_start_hi); + s1_hi = block_load(cstate_base + 1 * dim + chunk_start_hi); + s2_hi = block_load(cstate_base + 2 * dim + chunk_start_hi); + } simd w_raw_hi = block_load( conv_weight_ptr + (int64_t)chunk_start_hi * 4); @@ -419,16 +440,33 @@ ESIMD_INLINE void gdn_conv_fused_seq_kernel( // Phase 1 reads for the same seq_idx. // Uses register-cached s1, s2, x_fp16 from Phase 1 (not re-read from memory). if (inline_conv_shift && conv_idx >= 0 && hv == 0 && !v_oob) { - // lo chunk (all threads) - block_store(cstate_base + 0 * dim + chunk_start, simd(s1)); - block_store(cstate_base + 1 * dim + chunk_start, simd(s2)); - block_store(cstate_base + 2 * dim + chunk_start, x_fp16); - - // hi chunk (v-threads only, when double_v) - if (double_v && tid >= 4 * H) { - block_store(cstate_base + 0 * dim + chunk_start_hi, simd(s1_hi)); - block_store(cstate_base + 1 * dim + chunk_start_hi, simd(s2_hi)); - block_store(cstate_base + 2 * dim + chunk_start_hi, x_fp16_hi); + if (conv_native) { + // Native layout: shift row0<-s1, row1<-s2, row2<-x, interleaved + // (stride 3) then one 192-wide store per chunk. + simd out; + out.select<64, 3>(0) = simd(s1); + out.select<64, 3>(1) = simd(s2); + out.select<64, 3>(2) = x_fp16; + block_store(cstate_base + (int64_t)chunk_start * 3, out); + if (double_v && tid >= 4 * H) { + simd out_hi; + out_hi.select<64, 3>(0) = simd(s1_hi); + out_hi.select<64, 3>(1) = simd(s2_hi); + out_hi.select<64, 3>(2) = x_fp16_hi; + block_store(cstate_base + (int64_t)chunk_start_hi * 3, out_hi); + } + } else { + // lo chunk (all threads) + block_store(cstate_base + 0 * dim + chunk_start, simd(s1)); + block_store(cstate_base + 1 * dim + chunk_start, simd(s2)); + block_store(cstate_base + 2 * dim + chunk_start, x_fp16); + + // hi chunk (v-threads only, when double_v) + if (double_v && tid >= 4 * H) { + block_store(cstate_base + 0 * dim + chunk_start_hi, simd(s1_hi)); + block_store(cstate_base + 1 * dim + chunk_start_hi, simd(s2_hi)); + block_store(cstate_base + 2 * dim + chunk_start_hi, x_fp16_hi); + } } } @@ -478,6 +516,7 @@ ESIMD_INLINE void conv_state_shift_seq_kernel( const int* __restrict__ conv_state_indices_ptr, int N, int H, int HV, int gdn_K, int gdn_V, int64_t conv_stride0, + int conv_native, nd_item<3>& ndi) { const int seq_idx = ndi.get_group(0); @@ -530,23 +569,43 @@ ESIMD_INLINE void conv_state_shift_seq_kernel( fp16* cstate_base = conv_state_ptr + (int64_t)conv_idx * conv_stride0; // lo chunk (all threads) - simd s1_val = block_load(cstate_base + 1 * dim + chunk_start); - simd s2_val = block_load(cstate_base + 2 * dim + chunk_start); simd x_new = block_load(qkvz_row + qkvz_offset); - - block_store(cstate_base + 0 * dim + chunk_start, simd(s1_val)); - block_store(cstate_base + 1 * dim + chunk_start, simd(s2_val)); - block_store(cstate_base + 2 * dim + chunk_start, x_new); + if (conv_native) { + simd raw = block_load(cstate_base + (int64_t)chunk_start * 3); + simd s1_val = raw.select<64, 3>(1); + simd s2_val = raw.select<64, 3>(2); + simd out; + out.select<64, 3>(0) = s1_val; + out.select<64, 3>(1) = s2_val; + out.select<64, 3>(2) = x_new; + block_store(cstate_base + (int64_t)chunk_start * 3, out); + } else { + simd s1_val = block_load(cstate_base + 1 * dim + chunk_start); + simd s2_val = block_load(cstate_base + 2 * dim + chunk_start); + block_store(cstate_base + 0 * dim + chunk_start, simd(s1_val)); + block_store(cstate_base + 1 * dim + chunk_start, simd(s2_val)); + block_store(cstate_base + 2 * dim + chunk_start, x_new); + } // hi chunk (v-threads only, when double_v) if (double_v && tid >= 4 * H) { - simd s1_hi = block_load(cstate_base + 1 * dim + chunk_start_hi); - simd s2_hi = block_load(cstate_base + 2 * dim + chunk_start_hi); simd x_hi = block_load(qkvz_row + qkvz_offset_hi); - - block_store(cstate_base + 0 * dim + chunk_start_hi, simd(s1_hi)); - block_store(cstate_base + 1 * dim + chunk_start_hi, simd(s2_hi)); - block_store(cstate_base + 2 * dim + chunk_start_hi, x_hi); + if (conv_native) { + simd raw_hi = block_load(cstate_base + (int64_t)chunk_start_hi * 3); + simd s1_hi = raw_hi.select<64, 3>(1); + simd s2_hi = raw_hi.select<64, 3>(2); + simd out_hi; + out_hi.select<64, 3>(0) = s1_hi; + out_hi.select<64, 3>(1) = s2_hi; + out_hi.select<64, 3>(2) = x_hi; + block_store(cstate_base + (int64_t)chunk_start_hi * 3, out_hi); + } else { + simd s1_hi = block_load(cstate_base + 1 * dim + chunk_start_hi); + simd s2_hi = block_load(cstate_base + 2 * dim + chunk_start_hi); + block_store(cstate_base + 0 * dim + chunk_start_hi, simd(s1_hi)); + block_store(cstate_base + 1 * dim + chunk_start_hi, simd(s2_hi)); + block_store(cstate_base + 2 * dim + chunk_start_hi, x_hi); + } } } @@ -564,6 +623,7 @@ inline void gdn_conv_fused_seq_dispatch( fp16* output_ptr, fp16* z_out_ptr, int N, int H, int HV, int K, int V, float scale, int64_t conv_stride0, int64_t ssm_stride0, + int conv_native, sycl::queue& q) { const int total_wgs = N * HV; @@ -582,7 +642,7 @@ inline void gdn_conv_fused_seq_dispatch( ssm_state_ptr, ssm_state_indices_ptr, output_ptr, z_out_ptr, N, H, HV, K, V, scale, conv_stride0, ssm_stride0, - inline_shift, ndi); + inline_shift, conv_native, ndi); }); }); @@ -597,7 +657,7 @@ inline void gdn_conv_fused_seq_dispatch( qkvz_ptr, qkvz_stride0, conv_state_ptr, conv_state_indices_ptr, N, H, HV, K, V, - conv_stride0, ndi); + conv_stride0, conv_native, ndi); }); }); } @@ -625,6 +685,7 @@ inline void gdn_conv_fused_seq_host( float scale, int64_t conv_stride0, int64_t ssm_stride0, + int conv_native, sycl::queue& q) { TORCH_CHECK(HV > 0 && HV % H == 0, @@ -642,7 +703,7 @@ inline void gdn_conv_fused_seq_host( A_log_ptr, dt_bias_ptr, ba_ptr, ba_stride0, ssm_state_ptr, ssm_state_indices_ptr, output_ptr, z_out_ptr, - N, H, HV, K, V, scale, conv_stride0, ssm_stride0, q); + N, H, HV, K, V, scale, conv_stride0, ssm_stride0, conv_native, q); } else { const int v_slots_64 = 64 - 4 * H; TORCH_CHECK(v_slots_64 > 0 && HV <= v_slots_64, @@ -654,6 +715,6 @@ inline void gdn_conv_fused_seq_host( A_log_ptr, dt_bias_ptr, ba_ptr, ba_stride0, ssm_state_ptr, ssm_state_indices_ptr, output_ptr, z_out_ptr, - N, H, HV, K, V, scale, conv_stride0, ssm_stride0, q); + N, H, HV, K, V, scale, conv_stride0, ssm_stride0, conv_native, q); } } diff --git a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/gdn_norm_gated.h b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/gdn_norm_gated.h new file mode 100644 index 00000000..e98fd011 --- /dev/null +++ b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/gdn_norm_gated.h @@ -0,0 +1,189 @@ +/* gdn_norm_gated.h — standalone RMSNormGated head for the GGUF GDN out_proj. + * + * y[h,:] = fp16( x[h,:] * rsqrt(mean(x[h,:]^2) + eps) * nw[:] * silu(z[h,:]) ) + * + * i.e. RMSNormGated(norm_before_gate=True, activation="swish"), per value head. + * + * `norm_gemv_fused.h` already fuses this into an fp8 GEMV, but the GGUF build's + * out_proj is an ESIMD q8_0 GEMV, so that kernel never applies. Rather than + * duplicate the q8_0 dequant inside a second fused kernel, this one only emits + * the normed activation; the caller enqueues the existing, already-validated + * q8_0 GEMV right behind it on the same in-order queue. Two kernels, but still + * ONE torch op dispatch — and at bs=1 decode the dispatch is what costs + * (~13-25us) while a kernel of this size costs ~2us. + * + * The RMS reduction is per head (V dims), so one work-item owns a whole head + * and no cross-lane/SLM reduction is needed at all. + * + * Math copied from NormGEMV_fp8_pert_kernel so both paths stay bit-comparable. + */ +#pragma once + +template +struct Gdn_norm_gated_kernel { + const fp16* x; // [HV, V] + const fp16* z; // [HV, V] + const fp16* nw; // [V] + fp16* y; // [HV, V] out + float eps; + + void operator()(sycl::item<1> item) const SYCL_ESIMD_KERNEL { + const size_t off = (size_t)item.get_id(0) * V; + + simd x_f = block_load(x + off); + float mean_sq = reduce(x_f * x_f, std::plus<>()) * (1.0f / (float)V); + float inv_rms = + sycl::ext::intel::esimd::rsqrt(simd(mean_sq + eps))[0]; + + simd z_f = block_load(z + off); + simd silu_z = z_f / (1.0f + sycl::ext::intel::esimd::exp(-z_f)); + + simd normed = + x_f * inv_rms * simd(block_load(nw)) * silu_z; + block_store(y + off, convert(normed)); + } +}; + +// Returns false for an unsupported head width so the caller can fall back. +inline bool gdn_norm_gated_host( + const fp16* x, const fp16* z, const fp16* nw, fp16* y, + float eps, int HV, int V, sycl::queue& q) { +#define LAUNCH_GDN_NG(VV) \ + q.submit([&](sycl::handler& hd) { \ + hd.parallel_for(sycl::range<1>((size_t)HV), \ + Gdn_norm_gated_kernel{x, z, nw, y, eps}); \ + }); \ + return true; + + switch (V) { + case 64: LAUNCH_GDN_NG(64) + case 128: LAUNCH_GDN_NG(128) + case 256: LAUNCH_GDN_NG(256) + default: return false; + } +#undef LAUNCH_GDN_NG +} + +/* --------------------------------------------------------------------------- + * Gdn_norm_gated_q8_kernel — RMSNormGated fused with the q8_0 out_proj GEMV. + * + * One work-group = HVT lanes; lane `l` owns value head `l` (V elements), so the + * gated norm needs no cross-lane communication and the normalised head stays in + * registers. The GEMV then contracts over the flattened [HVT*V] vector with each + * lane supplying the partial dot of its own head, combined through SLM — the + * same K-split shape q8_0_GEMV.h uses, with K_SPLIT == HVT and KP == V. + * + * Saves one enqueue per GDN layer; at M=1 decode we are launch-bound, so the + * redundant per-WG norm recompute (register/L2 local) is cheaper than the + * second kernel launch. + * + * Batched decode: the WG grid is M*blocks groups, `token = gid / blocks` + * selecting the row of x/z/y/out and `blk = gid % blocks` the band of output + * rows, exactly as Moe_norm_q8_kernel does. The weight offset within K stays + * `lid*V` (shared by all tokens) while the activation offset carries the extra + * `token*K`, so the weights are read once per band and reused across tokens. + * ------------------------------------------------------------------------- */ +template +struct Gdn_norm_gated_q8_kernel { + static constexpr int ROWS = HVT < 8 ? HVT : 8; + const fp16* x; // [M*HV, V] + const fp16* z; // [M*HV, V] + const fp16* nw; // [V] + fp16* y; // [M*HV, V] may be null + const int8_t* qs; // [N, HVT*V] + const fp16* sc; // [N, HVT*V/32] + fp16* out; // [M, N] + float eps; + int N, K, M, blocks; + + void operator()(sycl::nd_item<1> item) const SYCL_ESIMD_KERNEL { + slm_init(); + const int gid = (int)item.get_group(0); + const int lid = (int)item.get_local_id(0); + const int token = gid / blocks; + const int blk = gid % blocks; + // Offset of this lane's head inside the contraction dim K (shared by + // every token), and the same offset inside this token's activation row. + const size_t koff = (size_t)lid * V; + const size_t xoff = (size_t)token * K + koff; + + simd x_f = block_load(x + xoff); + float mean_sq = reduce(x_f * x_f, std::plus<>()) * (1.0f / (float)V); + float inv_rms = + sycl::ext::intel::esimd::rsqrt(simd(mean_sq + eps))[0]; + + simd z_f = block_load(z + xoff); + simd silu_z = z_f / (1.0f + sycl::ext::intel::esimd::exp(-z_f)); + simd yv = convert( + x_f * inv_rms * simd(block_load(nw)) * silu_z); + if (blk == 0 && y) block_store(y + xoff, yv); + + simd xf = simd(yv); + const int rbase = blk * ROWS; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const int row = rbase + r; + float p = 0.0f; + if (row < N) { + const int8_t* wrow = qs + (size_t)row * K + koff; + const fp16* srow = sc + (size_t)row * (K / 32) + (koff / 32); + simd acc = 0.0f; + for (int k = 0; k < V; k += 32) { + simd wf = + convert(block_load(wrow + k)); + float s = static_cast(srow[k >> 5]); + acc += xf.template select<32, 1>(k) * (wf * s); + } + p = reduce(acc, std::plus<>()); + } + slm_block_store((HVT * (1 + r) + lid) * sizeof(float), + simd(p)); + } + barrier(); + if (lid < ROWS) { + const int row = rbase + lid; + if (row < N) { + simd d = + slm_block_load(HVT * (1 + lid) * sizeof(float)); + out[(size_t)token * N + row] = fp16(reduce(d, std::plus<>())); + } + } + } +}; + +inline bool gdn_norm_gated_q8_host( + const fp16* x, const fp16* z, const fp16* nw, fp16* y, + const int8_t* qs, const fp16* sc, fp16* out, + float eps, int HV, int V, int N, int M, sycl::queue& q) { + const int rows = HV < 8 ? HV : 8; + const int K = HV * V; + const int blocks = (N + rows - 1) / rows; + if (M < 1) return false; + +#define LAUNCH_GDN_NQ(VV, HH) \ + q.submit([&](sycl::handler& hd) { \ + hd.parallel_for(sycl::nd_range<1>((size_t)M * blocks * HH, HH), \ + Gdn_norm_gated_q8_kernel{ \ + x, z, nw, y, qs, sc, out, eps, N, K, M, blocks}); \ + }); \ + return true; +#define DISPATCH_HV(VV) \ + switch (HV) { \ + case 2: LAUNCH_GDN_NQ(VV, 2) \ + case 4: LAUNCH_GDN_NQ(VV, 4) \ + case 8: LAUNCH_GDN_NQ(VV, 8) \ + case 16: LAUNCH_GDN_NQ(VV, 16) \ + case 32: LAUNCH_GDN_NQ(VV, 32) \ + default: return false; \ + } + + if (V % 32 != 0) return false; + switch (V) { + case 64: DISPATCH_HV(64) + case 128: DISPATCH_HV(128) + case 256: DISPATCH_HV(256) + default: return false; + } +#undef DISPATCH_HV +#undef LAUNCH_GDN_NQ +} diff --git a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_kquant_GEMV.h b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_kquant_GEMV.h index f88bd332..f42148e1 100644 --- a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_kquant_GEMV.h +++ b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_kquant_GEMV.h @@ -32,6 +32,10 @@ static constexpr int MOE_Q4K_HALF = 16; // ── Up/gate stage: Q4_K gate + Q4_K up -> silu(gate)*up -> intermediate ────── // grid (n_routed = n_tokens*top_k, intermediate_size). Each WI: one n_col. +// VL elements per iteration: the original stepped 32 at a time, issuing 16-byte +// nibble loads (below LSC granularity) and re-reading a scalar scale/min per +// block. Wider tiles cut the load count by VL/32 and keep one flat accumulator. +template struct Moe_up_q4k_kernel { const fp16* x; // [n_tokens, hidden] const uint8_t* gate_ql; // [E, inter, hidden/2] @@ -45,12 +49,18 @@ struct Moe_up_q4k_kernel { int n_tokens, hidden, intermediate, top_k; void operator()(sycl::id<2> idx) const SYCL_ESIMD_KERNEL { - const int route = (int)idx[0]; - const int n_col = (int)idx[1]; + run((int)idx[0], (int)idx[1]); + } + + // Split out so Moe_up_fused_kernel can reuse the exact same body without + // duplicating the q4_k dequant. + void run(int route, int n_col) const { const int token = route / top_k; const int eid = sel_experts[route]; - const int Kh = hidden / 2; // packed bytes per row + constexpr int VH = VL / 2; // packed bytes per tile + constexpr int NSC = VL / MOE_Q4K_GROUP; // scale/min entries per tile + const int Kh = hidden / 2; const int Kg = hidden / MOE_Q4K_GROUP; const fp16* x_row = x + (size_t)token * hidden; const uint8_t* gq = gate_ql + ((size_t)eid * intermediate + n_col) * Kh; @@ -60,28 +70,42 @@ struct Moe_up_q4k_kernel { const fp16* us = up_sc + ((size_t)eid * intermediate + n_col) * Kg; const fp16* um = up_mn + ((size_t)eid * intermediate + n_col) * Kg; - simd ag_e = 0.0f, ag_o = 0.0f, au_e = 0.0f, au_o = 0.0f; - int gi = 0; - for (int k = 0; k < hidden; k += MOE_Q4K_GROUP) { - simd iv = block_load(x_row + k); - simd ie = iv.template select(0); - simd io = iv.template select(1); - - simd graw = block_load(gq + k / 2); - simd g16 = convert(graw); - float gsc = (float)gs[gi], gmn = (float)gm[gi]; - ag_e += ie * (convert(g16 & 0x000F) * gsc - gmn); - ag_o += io * (convert((g16 >> 4) & 0x000F) * gsc - gmn); - - simd uraw = block_load(uq + k / 2); - simd u16 = convert(uraw); - float usc = (float)us[gi], umn = (float)um[gi]; - au_e += ie * (convert(u16 & 0x000F) * usc - umn); - au_o += io * (convert((u16 >> 4) & 0x000F) * usc - umn); - gi++; + simd ag = 0.0f, au = 0.0f; + for (int k = 0; k < hidden; k += VL) { + const int gi = k / MOE_Q4K_GROUP; + simd iv = block_load(x_row + k); + simd xf = simd(iv); + + simd graw = block_load(gq + k / 2); + simd uraw = block_load(uq + k / 2); + simd gsv = block_load(gs + gi); + simd gmv = block_load(gm + gi); + simd usv = block_load(us + gi); + simd umv = block_load(um + gi); + + simd g16 = convert(graw); + simd u16 = convert(uraw); + simd gw, uw; + gw.template select(0) = convert(g16 & 0x000F); + gw.template select(1) = convert((g16 >> 4) & 0x000F); + uw.template select(0) = convert(u16 & 0x000F); + uw.template select(1) = convert((u16 >> 4) & 0x000F); + + #pragma unroll + for (int c = 0; c < NSC; c++) { + fp16 gsc = gsv[c], gmn = gmv[c], usc = usv[c], umn = umv[c]; + gw.template select(c * MOE_Q4K_GROUP) = + gw.template select(c * MOE_Q4K_GROUP) + * (float)gsc - (float)gmn; + uw.template select(c * MOE_Q4K_GROUP) = + uw.template select(c * MOE_Q4K_GROUP) + * (float)usc - (float)umn; + } + ag += xf * gw; + au += xf * uw; } - float g = reduce(ag_e, std::plus<>()) + reduce(ag_o, std::plus<>()); - float u = reduce(au_e, std::plus<>()) + reduce(au_o, std::plus<>()); + float g = reduce(ag, std::plus<>()); + float u = reduce(au, std::plus<>()); float silu = g / (1.0f + sycl::exp(-g)); inter[(size_t)route * intermediate + n_col] = fp16(silu * u); } @@ -92,11 +116,17 @@ inline void moe_up_q4k_host( const uint8_t* uq, const fp16* us, const fp16* um, const int* sel, fp16* inter, int n_tokens, int hidden, int intermediate, int top_k, sycl::queue& q) { - q.submit([&](sycl::handler& h) { - h.parallel_for(sycl::range<2>((size_t)n_tokens * top_k, intermediate), - Moe_up_q4k_kernel{x, gq, gs, gm, uq, us, um, sel, inter, - n_tokens, hidden, intermediate, top_k}); +#define LAUNCH_MOE_UP_Q4K(V) \ + q.submit([&](sycl::handler& h) { \ + h.parallel_for(sycl::range<2>((size_t)n_tokens * top_k, intermediate),\ + Moe_up_q4k_kernel{x, gq, gs, gm, uq, us, um, sel, inter, \ + n_tokens, hidden, intermediate, top_k}); \ }); + + if (hidden % 128 == 0) { LAUNCH_MOE_UP_Q4K(128) } + else if (hidden % 64 == 0) { LAUNCH_MOE_UP_Q4K(64) } + else { LAUNCH_MOE_UP_Q4K(32) } +#undef LAUNCH_MOE_UP_Q4K } // ── Down stage: PACKED (zero extra memory, mirrors q5_k/q6_k_GEMV.h) ───────── @@ -108,6 +138,14 @@ inline void moe_up_q4k_host( // Q5_K down: ql [E,N,K/2] nibble + qh [E,N,K/8] PLAIN 1-bit (byte j bit b = elem // 8j+b) + scale,min [E,N,K/32] fp16. v5 = nibble|(qh<<4); w = v5*scale - min. +// VL elements per iteration, ROWS output cols per work-item. The original +// stepped one 32-elem superblock at a time (16B nibble + 4B qh loads, plus a +// full horizontal sum folded into a serial `dot` chain every block) and only +// reached ~180GB/s. Wider tiles remove the tiny loads and the serial reduction; +// ROWS>1 additionally amortises the `inter` row load, which every work-item +// otherwise re-reads in full (N=2048 work-items x 512B = 1MB of cache traffic +// per route for a 4KB tensor). +template struct Moe_down_q5k_kernel { const fp16* inter; // [n_routed, K] const uint8_t* ql; // [E, N, K/2] @@ -121,38 +159,54 @@ struct Moe_down_q5k_kernel { void operator()(sycl::id<2> idx) const SYCL_ESIMD_KERNEL { const int route = (int)idx[0]; - const int n = (int)idx[1]; // hidden output col + const int n0 = (int)idx[1] * ROWS; // first hidden output col const int eid = sel_experts[route]; const int Kh = K / 2, Kq = K / 8, Kg = K / 32; - const fp16* i_row = inter + (size_t)route * K; - const uint8_t* qlr = ql + ((size_t)eid * N + n) * Kh; - const uint8_t* qhr = qh + ((size_t)eid * N + n) * Kq; - const fp16* scr = sc + ((size_t)eid * N + n) * Kg; - const fp16* mnr = mn + ((size_t)eid * N + n) * Kg; + const fp16* i_row = inter + (size_t)route * K; + const size_t rbase = (size_t)eid * N + n0; - simd acc; // per-superblock scratch - const simd lane(0u, 1u); // 0..31 = bit index within word - float dot = 0.0f; - for (int blk = 0; blk < Kg; blk++) { - const int e0 = blk * 32; - simd iv = block_load(i_row + e0); - simd qd = block_load(qlr + e0 / 2); - simd qhw = - block_load( - reinterpret_cast(qhr + e0 / 8)); - const uint32_t qh32 = qhw[0]; - simd wf; - wf.template select<16, 2>(0) = qd & 0x0F; // low nibble -> 2j - wf.template select<16, 2>(1) = (qd >> 4) & 0x0F; // high nibble -> 2j+1 - // 5th bit: element e high = bit e of the little-endian 4-byte qh word. - simd hb = (simd(qh32) >> lane) & 1u; - wf += simd(hb) * 16.0f; - const float s = (float)scr[blk], m = (float)mnr[blk]; - wf = wf * s - m; - acc = simd(iv) * wf; - dot += esimd_detail2::sum(acc); + constexpr int VH = VL / 2; // nibble bytes per tile + constexpr int VW = VL / 32; // qh dwords / scale entries per tile + const simd lane(0u, 1u); // bit index within a qh dword + + simd dots = 0.0f; + for (int k = 0; k < K; k += VL) { + const int gi = k / 32; + simd iv = block_load(i_row + k); + simd xf = simd(iv); + + #pragma unroll + for (int r = 0; r < ROWS; r++) { + const uint8_t* qlr = ql + (rbase + r) * Kh; + const uint8_t* qhr = qh + (rbase + r) * Kq; + const fp16* scr = sc + (rbase + r) * Kg; + const fp16* mnr = mn + (rbase + r) * Kg; + + simd qd = block_load(qlr + k / 2); + simd qhw = block_load( + reinterpret_cast(qhr + k / 8)); + simd scv = block_load(scr + gi); + simd mnv = block_load(mnr + gi); + + simd wf; + wf.template select(0) = convert(qd & 0x0F); + wf.template select(1) = convert((qd >> 4) & 0x0F); + + #pragma unroll + for (int c = 0; c < VW; c++) { + uint32_t qh32 = qhw[c]; + simd hb = (simd(qh32) >> lane) & 1u; + fp16 s = scv[c], m = mnv[c]; + auto blk = wf.template select<32, 1>(c * 32); + blk = (blk + simd(hb) * 16.0f) * (float)s - (float)m; + } + dots[r] += reduce(xf * wf, std::plus<>()); + } } - out[(size_t)route * N + n] = fp16(dot * (float)topk_w[route]); + const float tw = (float)topk_w[route]; + #pragma unroll + for (int r = 0; r < ROWS; r++) + out[(size_t)route * N + n0 + r] = fp16(dots[r] * tw); } }; @@ -212,11 +266,24 @@ inline void moe_down_q5k_host( const fp16* inter, const uint8_t* ql, const uint8_t* qh, const fp16* sc, const fp16* mn, const int* sel, const fp16* topk_w, fp16* out_partial, int n_tokens, int N, int K, int top_k, sycl::queue& q) { - q.submit([&](sycl::handler& h) { - h.parallel_for(sycl::range<2>((size_t)n_tokens * top_k, N), - Moe_down_q5k_kernel{inter, ql, qh, sc, mn, sel, topk_w, out_partial, - n_tokens, N, K, top_k}); +#define LAUNCH_MOE_DOWN_Q5K(V, R) \ + q.submit([&](sycl::handler& h) { \ + h.parallel_for(sycl::range<2>((size_t)n_tokens * top_k, N / (R)), \ + Moe_down_q5k_kernel{inter, ql, qh, sc, mn, sel, topk_w, \ + out_partial, n_tokens, N, K, top_k}); \ }); + + const int R = (N % 4 == 0) ? 4 : 1; + if (K % 256 == 0) { + if (R == 4) { LAUNCH_MOE_DOWN_Q5K(256, 4) } else { LAUNCH_MOE_DOWN_Q5K(256, 1) } + } else if (K % 128 == 0) { + if (R == 4) { LAUNCH_MOE_DOWN_Q5K(128, 4) } else { LAUNCH_MOE_DOWN_Q5K(128, 1) } + } else if (K % 64 == 0) { + if (R == 4) { LAUNCH_MOE_DOWN_Q5K(64, 4) } else { LAUNCH_MOE_DOWN_Q5K(64, 1) } + } else { + if (R == 4) { LAUNCH_MOE_DOWN_Q5K(32, 4) } else { LAUNCH_MOE_DOWN_Q5K(32, 1) } + } +#undef LAUNCH_MOE_DOWN_Q5K } inline void moe_down_q6k_host( diff --git a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_norm_router.h b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_norm_router.h new file mode 100644 index 00000000..bd247145 --- /dev/null +++ b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_norm_router.h @@ -0,0 +1,249 @@ +/* moe_norm_router.h — fused (residual-add + GemmaRMSNorm + router GEMV) head + * for the GGUF decode MoE op (Intel XPU, ESIMD). + * + * res_out[t,:]= fp16( h[t,:] + res[t,:] ) + * v = float(res_out[t,:]) + * xn[t,:] = fp16( v * rsqrt(mean(v^2) + eps) * nw ) (nw = 1 + gemma w) + * logits[t,e] = sum_k xn[t,k] * rw[e,k] + * + * Motivation: at M=1 decode the model is HOST-bound (~13-25us of torch dispatch + * per op call vs ~4us of actual enqueue), so the win is in collapsing op CALLS, + * not kernel time. Folding `gemma_fused_add_rmsnorm` + the fp16 router `linear` + * into the MoE op removes two python/dispatcher round-trips per layer. + * + * One work-group per (token, block of ROWS expert rows); K_SPLIT threads split + * the hidden reduction and combine through SLM. Every WG recomputes the norm + * from values it has already loaded into registers (no extra DRAM traffic), and + * only block 0 of a token stores xn/res back, so there is no cross-WG race. + * + * Included into esimd_kernel.sycl alongside the other moe_* headers. + */ +#pragma once + +template +struct Moe_norm_router_kernel { + const fp16* h; // [M, hidden] + const fp16* res; // [M, hidden] in + fp16* res_out; // [M, hidden] out (MUST NOT alias `res`: block 0 + // stores while the other blocks are still loading) + const fp16* nw; // [hidden] (1 + gemma weight) + const fp16* rw; // [E, hidden] + fp16* xn; // [M, hidden] out + fp16* logits; // [M, E] out + float eps; + int M, hidden, E, blocks; + + void operator()(sycl::nd_item<1> item) const SYCL_ESIMD_KERNEL { + slm_init(); + const int gid = (int)item.get_group(0); + const int lid = (int)item.get_local_id(0); + const int token = gid / blocks; + const int blk = gid % blocks; + const int kbeg = lid * KP; + + const fp16* hr = h + (size_t)token * hidden; + const fp16* rr = res + (size_t)token * hidden; + + // fp16 add first: the reference rounds the residual to fp16 BEFORE the + // fp32 variance, so accumulating in fp32 here would not match. + simd vh = block_load(hr + kbeg) + + block_load(rr + kbeg); + simd v = simd(vh); + + slm_block_store(lid * sizeof(float), + simd(reduce(v * v, std::plus<>()))); + barrier(); + simd parts = slm_block_load(0); + float rstd = 1.0f / sycl::sqrt(reduce(parts, std::plus<>()) / (float)hidden + eps); + + simd xv = convert( + v * rstd * simd(block_load(nw + kbeg))); + if (blk == 0) { + block_store(res_out + (size_t)token * hidden + kbeg, vh); + block_store(xn + (size_t)token * hidden + kbeg, xv); + } + simd xf = simd(xv); + + const int rbase = blk * ROWS; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const int row = rbase + r; + float p = 0.0f; + if (row < E) + p = reduce( + xf * simd( + block_load(rw + (size_t)row * hidden + kbeg)), + std::plus<>()); + slm_block_store((K_SPLIT * (1 + r) + lid) * sizeof(float), + simd(p)); + } + barrier(); + if (lid < ROWS) { + const int row = rbase + lid; + if (row < E) { + simd d = slm_block_load( + K_SPLIT * (1 + lid) * sizeof(float)); + logits[(size_t)token * E + row] = + fp16(reduce(d, std::plus<>())); + } + } + } +}; + +// Returns false when the shape is unsupported, so the caller can fall back to +// the separate norm + router path. +inline bool moe_norm_router_host( + const fp16* h, const fp16* res, fp16* res_out, const fp16* nw, const fp16* rw, + fp16* xn, fp16* logits, float eps, + int M, int hidden, int E, sycl::queue& q) { + constexpr int KS = 8; + constexpr int ROWS = 4; + // E == 0 means "norm only, no GEMV": still launch one block per token so + // xn / res are written (the `row < E` guards make the GEMV half a no-op). + const int blocks = E > 0 ? (E + ROWS - 1) / ROWS : 1; + +#define LAUNCH_MOE_NR(KP) \ + q.submit([&](sycl::handler& hd) { \ + hd.parallel_for( \ + sycl::nd_range<1>((size_t)M * blocks * KS, KS), \ + Moe_norm_router_kernel{h, res, res_out, nw, rw, xn, \ + logits, eps, M, hidden, E, \ + blocks}); \ + }); \ + return true; + + switch (hidden / KS) { + case 128: LAUNCH_MOE_NR(128) + case 256: LAUNCH_MOE_NR(256) + case 512: LAUNCH_MOE_NR(512) + default: return false; + } +#undef LAUNCH_MOE_NR +} + +/* --------------------------------------------------------------------------- + * Moe_norm_q8_kernel — same head as above, but the WG grid is split into two + * ranges so ONE kernel launch produces both GEMVs: + * + * blk < nb_r : rows of the fp16 matrix `rw` -> logits[M, E] + * blk >= nb_r : rows of the q8_0 matrix qs/sc -> out0[M, N0] + * + * At M=1 decode we are launch-bound (~4-8us of host per enqueue vs ~2-20us of + * GPU), so collapsing the (norm-kernel + gemv-kernel) pair into a single + * enqueue is worth the redundant norm recompute, which is register/L2 local. + * ------------------------------------------------------------------------- */ +template +struct Moe_norm_q8_kernel { + const fp16* h; + const fp16* res; + fp16* res_out; // may be null + const fp16* nw; + const fp16* rw; // [E, hidden] fp16 (may be null when E == 0) + fp16* logits; // [M, E] + const int8_t* qs; // [N0, hidden] int8 (may be null when N0 == 0) + const fp16* sc; // [N0, hidden/32] + fp16* out0; // [M, N0] + fp16* xn; // [M, hidden] may be null + float eps; + int M, hidden, E, N0, nb_r, blocks; + + void operator()(sycl::nd_item<1> item) const SYCL_ESIMD_KERNEL { + slm_init(); + const int gid = (int)item.get_group(0); + const int lid = (int)item.get_local_id(0); + const int token = gid / blocks; + const int blk = gid % blocks; + const int kbeg = lid * KP; + + simd vh = + block_load(h + (size_t)token * hidden + kbeg) + + block_load(res + (size_t)token * hidden + kbeg); + simd v = simd(vh); + + slm_block_store(lid * sizeof(float), + simd(reduce(v * v, std::plus<>()))); + barrier(); + simd parts = slm_block_load(0); + float rstd = 1.0f / sycl::sqrt(reduce(parts, std::plus<>()) / (float)hidden + eps); + + simd xv = convert( + v * rstd * simd(block_load(nw + kbeg))); + if (blk == 0) { + if (res_out) block_store(res_out + (size_t)token * hidden + kbeg, vh); + if (xn) block_store(xn + (size_t)token * hidden + kbeg, xv); + } + simd xf = simd(xv); + + const bool q8_range = (blk >= nb_r); + const int rbase = (q8_range ? (blk - nb_r) : blk) * ROWS; + const int nrow = q8_range ? N0 : E; + +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const int row = rbase + r; + float p = 0.0f; + if (row < nrow) { + if (q8_range) { + const int8_t* wrow = qs + (size_t)row * hidden + kbeg; + const fp16* srow = sc + (size_t)row * (hidden / 32) + (kbeg / 32); + simd acc = 0.0f; + for (int k = 0; k < KP; k += 32) { + simd wf = + convert(block_load(wrow + k)); + float s = static_cast(srow[k >> 5]); + acc += xf.template select<32, 1>(k) * (wf * s); + } + p = reduce(acc, std::plus<>()); + } else { + p = reduce( + xf * simd( + block_load(rw + (size_t)row * hidden + kbeg)), + std::plus<>()); + } + } + slm_block_store((K_SPLIT * (1 + r) + lid) * sizeof(float), + simd(p)); + } + barrier(); + if (lid < ROWS) { + const int row = rbase + lid; + if (row < nrow) { + simd d = slm_block_load( + K_SPLIT * (1 + lid) * sizeof(float)); + fp16 val = fp16(reduce(d, std::plus<>())); + if (q8_range) out0[(size_t)token * N0 + row] = val; + else logits[(size_t)token * E + row] = val; + } + } + } +}; + +inline bool moe_norm_q8_host( + const fp16* h, const fp16* res, fp16* res_out, const fp16* nw, + const fp16* rw, fp16* logits, + const int8_t* qs, const fp16* sc, fp16* out0, + fp16* xn, float eps, int M, int hidden, int E, int N0, sycl::queue& q) { + constexpr int KS = 8; + constexpr int ROWS = 4; + if (hidden % (KS * 32) != 0) return false; + const int nb_r = E > 0 ? (E + ROWS - 1) / ROWS : 0; + const int nb_q = N0 > 0 ? (N0 + ROWS - 1) / ROWS : 0; + const int blocks = (nb_r + nb_q) > 0 ? (nb_r + nb_q) : 1; + +#define LAUNCH_MOE_NQ(KP) \ + q.submit([&](sycl::handler& hd) { \ + hd.parallel_for(sycl::nd_range<1>((size_t)M * blocks * KS, KS), \ + Moe_norm_q8_kernel{ \ + h, res, res_out, nw, rw, logits, qs, sc, out0, xn, \ + eps, M, hidden, E, N0, nb_r, blocks}); \ + }); \ + return true; + + switch (hidden / KS) { + case 128: LAUNCH_MOE_NQ(128) + case 256: LAUNCH_MOE_NQ(256) + case 512: LAUNCH_MOE_NQ(512) + default: return false; + } +#undef LAUNCH_MOE_NQ +} diff --git a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_shared_q8_GEMV.h b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_shared_q8_GEMV.h new file mode 100644 index 00000000..ebc703a2 --- /dev/null +++ b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_shared_q8_GEMV.h @@ -0,0 +1,379 @@ +/* moe_shared_q8_GEMV.h — fused GGUF Q8_0 SHARED-expert MLP for Intel XPU + * (ESIMD), decode. Collapses the ~6 python dispatches of the Qwen2MoE shared + * expert (gate_up GEMV + SiLU_and_mul + down GEMV + shared_expert_gate GEMV + + * sigmoid + mul) into ONE host op = ONE python dispatch: + * + * g[t] = sigmoid( sum_k x[t,k] * wg[k] ) (gate scalar) + * inter[t,c] = silu( ) * (Q8_0 gate/up) + * out[t,n] = g[t] * sum_c inter[t,c] * W_down[n,c] (Q8_0 down) + * + * Q8_0 rep (matches gguf.py _xpu_repack_q8_0): symmetric, group-32, + * qs [N, K] int8 (per-element quant) + * scale[N, K/32] fp16 (per-group); w[n,k] = scale[n,k/32] * qs[n,k]. + * gate_up is the MERGED rep: rows [0, inter_s) = gate, [inter_s, 2*inter_s) = up + * (MergedColumnParallelLinear shard order gate,up). down K = inter_s. + * K (=hidden for gate/up, =inter_s for down) is a multiple of 32 (Q8_0 block). + * + * Included into esimd_kernel.sycl after moe_kquant_GEMV.h (utils.h namespaces). + */ +#pragma once + +static constexpr int MOE_Q8_GROUP = 32; + +// ── Shared gate scalar: g[t] = sigmoid( dot(x[t], wg) ) ────────────────────── +struct Moe_shared_gate_kernel { + const fp16* x; // [M, hidden] + const fp16* wg; // [hidden] (shared_expert_gate weight row 0) + fp16* g; // [M] + int hidden; + + void operator()(sycl::id<1> idx) const SYCL_ESIMD_KERNEL { + const int token = (int)idx[0]; + const fp16* xr = x + (size_t)token * hidden; + simd acc = 0.0f; + for (int k = 0; k < hidden; k += MOE_Q8_GROUP) { + simd xv = block_load(xr + k); + simd wv = block_load(wg + k); + acc += simd(xv) * simd(wv); + } + float d = reduce(acc, std::plus<>()); + g[token] = fp16(1.0f / (1.0f + sycl::exp(-d))); + } +}; + +inline void moe_shared_gate_host( + const fp16* x, const fp16* wg, fp16* g, int M, int hidden, sycl::queue& q) { + q.submit([&](sycl::handler& h) { + h.parallel_for(sycl::range<1>((size_t)M), + Moe_shared_gate_kernel{x, wg, g, hidden}); + }); +} + +// ── Up/gate stage: Q8_0 gate + Q8_0 up -> silu(gate)*up -> intermediate ─────── +// One work-group per (token, col); K_SPLIT threads split the `hidden` reduction +// and combine through SLM (same structure as q8_0_gemv, which sustains ~470GB/s +// while the old flat range<2>(M, inter_s) launch only reached ~120GB/s because +// inter_s=256 work-items cannot fill the machine). +// +// The shared-expert gate scalar rides along as one extra column (col == +// inter_s), which removes a whole kernel launch whose single work-item spent +// ~7.7us on a serial 64-iteration dependent load chain. +template +struct Moe_shared_up_q8_kernel { + const fp16* x; // [M, hidden] + const int8_t* gu_qs; // [2*inter_s, hidden] (rows: gate then up) + const fp16* gu_sc; // [2*inter_s, hidden/32] + const fp16* wg; // [hidden] shared_expert_gate row (may be null) + fp16* inter; // [M, inter_s] + fp16* g; // [M] shared gate scalar (may be null) + int M, hidden, inter_s, n_cols; + + void operator()(sycl::nd_item<1> item) const SYCL_ESIMD_KERNEL { + slm_init<2 * K_SPLIT * sizeof(float)>(); + run((int)item.get_group(0), (int)item.get_local_id(0)); + } + + // Body without slm_init, so Moe_up_fused_kernel can call it after doing its + // own (single, uniform) slm_init. Contains barriers -> must be entered by + // every lane of the work-group, which holds because the caller branches on + // the group id only. + void run(int wgid, int lid) const { + const int token = wgid / n_cols; + const int col = wgid % n_cols; + const int Kg = hidden / MOE_Q8_GROUP; + const int kp = hidden / K_SPLIT; + const int kbeg = lid * kp; + const int kend = kbeg + kp; + const fp16* xr = x + (size_t)token * hidden; + + if (col == inter_s) { // shared-expert gate scalar + simd acc = 0.0f; + for (int k = kbeg; k < kend; k += MOE_Q8_GROUP) { + simd xv = block_load(xr + k); + simd wv = block_load(wg + k); + acc += simd(xv) * simd(wv); + } + float part = reduce(acc, std::plus<>()); + slm_block_store(lid * sizeof(float), simd(part)); + barrier(); + if (lid == 0) { + simd p = slm_block_load(0); + float d = reduce(p, std::plus<>()); + g[token] = fp16(1.0f / (1.0f + sycl::exp(-d))); + } + return; + } + + const int grow = col; + const int urow = inter_s + col; + const int8_t* gq = gu_qs + (size_t)grow * hidden; + const fp16* gs = gu_sc + (size_t)grow * Kg; + const int8_t* uq = gu_qs + (size_t)urow * hidden; + const fp16* us = gu_sc + (size_t)urow * Kg; + + simd ag = 0.0f, au = 0.0f; + int gi = kbeg / MOE_Q8_GROUP; + for (int k = kbeg; k < kend; k += MOE_Q8_GROUP) { + simd xv = block_load(xr + k); + simd xf = simd(xv); + simd gr = block_load(gq + k); + simd ur = block_load(uq + k); + float gsc = (float)gs[gi], usc = (float)us[gi]; + ag += xf * (convert(gr) * gsc); + au += xf * (convert(ur) * usc); + gi++; + } + float pg = reduce(ag, std::plus<>()); + float pu = reduce(au, std::plus<>()); + + if constexpr (K_SPLIT == 1) { + float silu = pg / (1.0f + sycl::exp(-pg)); + inter[(size_t)token * inter_s + col] = fp16(silu * pu); + } else { + slm_block_store(lid * sizeof(float), simd(pg)); + slm_block_store((K_SPLIT + lid) * sizeof(float), + simd(pu)); + barrier(); + if (lid == 0) { + simd vg = slm_block_load(0); + simd vu = + slm_block_load(K_SPLIT * sizeof(float)); + float gate = reduce(vg, std::plus<>()); + float up = reduce(vu, std::plus<>()); + float silu = gate / (1.0f + sycl::exp(-gate)); + inter[(size_t)token * inter_s + col] = fp16(silu * up); + } + } + } +}; + +// wg/g may be null: then no gate column is launched (standalone use). +inline void moe_shared_up_q8_host( + const fp16* x, const int8_t* gu_qs, const fp16* gu_sc, + const fp16* wg, fp16* inter, fp16* g, + int M, int hidden, int inter_s, sycl::queue& q) { + const int n_cols = inter_s + (g != nullptr ? 1 : 0); + int ks = 8; + while (ks > 1 && (hidden / ks) % MOE_Q8_GROUP != 0) ks /= 2; + +#define LAUNCH_MOE_SH_UP(S) \ + q.submit([&](sycl::handler& h) { \ + h.parallel_for( \ + sycl::nd_range<1>((size_t)M * n_cols * (S), (S)), \ + Moe_shared_up_q8_kernel{x, gu_qs, gu_sc, wg, inter, g, \ + M, hidden, inter_s, n_cols}); \ + }); + + if (ks == 8) { LAUNCH_MOE_SH_UP(8) } + else if (ks == 4) { LAUNCH_MOE_SH_UP(4) } + else if (ks == 2) { LAUNCH_MOE_SH_UP(2) } + else { LAUNCH_MOE_SH_UP(1) } +#undef LAUNCH_MOE_SH_UP +} + +// ── Down stage: Q8_0 down -> * gate scalar -> shared output ─────────────────── +// grid (M, hidden). Each WI: one hidden output col. K = inter_s. +struct Moe_shared_down_q8_kernel { + const fp16* inter; // [M, inter_s] + const int8_t* d_qs; // [hidden, inter_s] + const fp16* d_sc; // [hidden, inter_s/32] + const fp16* g; // [M] + fp16* out; // [M, hidden] + int M, hidden, inter_s; + + void operator()(sycl::id<2> idx) const SYCL_ESIMD_KERNEL { + const int token = (int)idx[0]; + const int n = (int)idx[1]; + const int Kg = inter_s / MOE_Q8_GROUP; + const fp16* ir = inter + (size_t)token * inter_s; + const int8_t* dq = d_qs + (size_t)n * inter_s; + const fp16* ds = d_sc + (size_t)n * Kg; + + simd acc = 0.0f; + int gi = 0; + for (int k = 0; k < inter_s; k += MOE_Q8_GROUP) { + simd iv = block_load(ir + k); + simd dr = block_load(dq + k); + float dsc = (float)ds[gi]; + acc += simd(iv) * (convert(dr) * dsc); + gi++; + } + float dot = reduce(acc, std::plus<>()); + out[(size_t)token * hidden + n] = fp16(dot * (float)g[token]); + } +}; + +inline void moe_shared_down_q8_host( + const fp16* inter, const int8_t* d_qs, const fp16* d_sc, const fp16* g, + fp16* out, int M, int hidden, int inter_s, sycl::queue& q) { + q.submit([&](sycl::handler& h) { + h.parallel_for(sycl::range<2>((size_t)M, hidden), + Moe_shared_down_q8_kernel{inter, d_qs, d_sc, g, out, + M, hidden, inter_s}); + }); +} + +// ═══════════════ FULL fusion: topk + routed + shared + finalize ═══════════════ +// Mirrors the fp8 moe_forward_full: ONE host op = ONE python dispatch. Router +// GEMV stays in python (logits input); everything else is chained on the queue. + +// ── Softmax over all experts + top-k heap select + renorm (per token) ───────── +// Ported from moe_batch/moe_topk.h (validated). n_experts padded to 512. +static constexpr int MOE_TOPK_PAD = 512; + +// Fast top-k: the original kernel ran a 32-slot min-heap with one serial +// iteration per expert (n_experts=256 -> 256 dependent hmin/pack_mask/merge +// chains), costing ~20us for a single work-item. Instead we pack each score +// into an order-preserving uint32 key (monotonic fp16 bits << 16 | inverted +// lane id) and extract the top-k with `top_k` hmax reductions, i.e. 8 steps +// instead of 256. +// +// Weights are mathematically identical to "full softmax -> select -> renorm": +// norm=true -> softmax restricted to the selected logits +// norm=false -> full softmax probability of the selected logits +// Both use the global max (== the top-1 logit) as the shift. +template +struct Moe_topk_gguf_kernel { + const fp16* logits; // [M, n_experts] + int* sel; // [M*top_k] expert ids + fp16* tw; // [M*top_k] renormalized weights + int n_experts, top_k; + bool norm; + + void operator()(sycl::id<1> idx) const SYCL_ESIMD_KERNEL { + const int nid = (int)idx[0]; + const fp16* row_ptr = logits + (size_t)nid * n_experts; + + simd scores(fp16(-65504.f)); + // Load whole 32-blocks, then the tail one element at a time. Reading a + // full block past the end of the row would both run off the logits + // allocation on the last token and clobber the -65504 padding, letting a + // garbage lane win the hmax below and yield a selected expert id >= + // n_experts -- which then indexes the expert weights out of range. The + // tail loop is dead code whenever n_experts % 32 == 0, which is the case + // for every shipped config (Qwen3.5-35B-A3B has 256 routed experts), so + // the fast path is unchanged. + int i = 0; + for (; i + 32 <= n_experts; i += 32) + scores.template select<32, 1>(i) = block_load(row_ptr + i); + for (; i < n_experts; ++i) + scores[i] = row_ptr[i]; + + // fp16 bits -> monotonically ordered uint16 key. + simd bits = scores.template bit_cast_view(); + simd key = bits | (uint16_t)0x8000; + key.merge(~bits, (bits >> 15) != 0); + + // Low 16 bits break ties towards the smaller expert id. + simd lane(0, 1); + simd keys = + (simd(key) << 16) | ((uint32_t)(PAD - 1) - lane); + + int top_i[32]; + float top_l[32]; + for (int k = 0; k < top_k; ++k) { + uint32_t mx = hmax(keys); + int i = (int)((uint32_t)(PAD - 1) - (mx & 0xFFFFu)); + top_i[k] = i; + fp16 sv = scores[i]; + top_l[k] = (float)sv; + keys.merge(simd(0), keys == mx); + } + + const float mx_l = top_l[0]; // global max logit + float denom; + if (norm) { + denom = 0.f; + for (int k = 0; k < top_k; ++k) denom += sycl::exp(top_l[k] - mx_l); + } else { + simd e = exp(convert(scores) - mx_l); + denom = reduce(e, std::plus<>()); + } + const float inv = 1.0f / denom; + + int* idx_base = sel + (size_t)nid * top_k; + fp16* w_base = tw + (size_t)nid * top_k; + for (int k = 0; k < top_k; ++k) { + idx_base[k] = top_i[k]; + w_base[k] = fp16(sycl::exp(top_l[k] - mx_l) * inv); + } + } +}; + +inline void moe_topk_gguf_host( + const fp16* logits, int* sel, fp16* tw, + int M, int n_experts, int top_k, bool norm, sycl::queue& q) { + q.submit([&](sycl::handler& h) { + if (n_experts <= 256) + h.parallel_for(sycl::range<1>((size_t)M), + Moe_topk_gguf_kernel<256>{logits, sel, tw, n_experts, top_k, norm}); + else + h.parallel_for(sycl::range<1>((size_t)M), + Moe_topk_gguf_kernel<512>{logits, sel, tw, n_experts, top_k, norm}); + }); +} + +// ── Finalize: routed top_k combine + shared Q8_0 down + gate*sigmoid ────────── +// grid (M, hidden). final[t,n] = sum_{k +struct Moe_finalize_gguf_kernel { + const fp16* out_partial; // [M*top_k, hidden] (routed, topk_w already applied) + const fp16* inter_sh; // [M, inter_s] + const int8_t* d_qs; // [hidden, inter_s] shared down Q8_0 + const fp16* d_sc; // [hidden, inter_s/32] + const fp16* g; // [M] shared gate scalar + fp16* out; // [M, hidden] + int M, hidden, inter_s, top_k; + + void operator()(sycl::id<2> idx) const SYCL_ESIMD_KERNEL { + const int token = (int)idx[0]; + const int n = (int)idx[1]; + float acc = 0.0f; + for (int k = 0; k < top_k; k++) + acc += (float)out_partial[((size_t)token * top_k + k) * hidden + n]; + + constexpr int NSC = VL / MOE_Q8_GROUP; + const int Kg = inter_s / MOE_Q8_GROUP; + const fp16* ir = inter_sh + (size_t)token * inter_s; + const int8_t* dq = d_qs + (size_t)n * inter_s; + const fp16* ds = d_sc + (size_t)n * Kg; + + simd sacc = 0.0f; + for (int kk = 0; kk < inter_s; kk += VL) { + simd iv = block_load(ir + kk); + simd dr = block_load(dq + kk); + simd sh = block_load(ds + kk / MOE_Q8_GROUP); + simd wf = convert(dr); + #pragma unroll + for (int sb = 0; sb < NSC; sb++) { + fp16 sv = sh[sb]; + wf.template select(sb * MOE_Q8_GROUP) = + wf.template select(sb * MOE_Q8_GROUP) * (float)sv; + } + sacc += simd(iv) * wf; + } + float sdot = reduce(sacc, std::plus<>()); + out[(size_t)token * hidden + n] = fp16(acc + (float)g[token] * sdot); + } +}; + +inline void moe_finalize_gguf_host( + const fp16* out_partial, const fp16* inter_sh, const int8_t* d_qs, + const fp16* d_sc, const fp16* g, fp16* out, + int M, int hidden, int inter_s, int top_k, sycl::queue& q) { +#define LAUNCH_MOE_FIN(V) \ + q.submit([&](sycl::handler& h) { \ + h.parallel_for(sycl::range<2>((size_t)M, hidden), \ + Moe_finalize_gguf_kernel{out_partial, inter_sh, d_qs, d_sc, g, \ + out, M, hidden, inter_s, top_k}); \ + }); + + if (inter_s % 256 == 0) { LAUNCH_MOE_FIN(256) } + else if (inter_s % 128 == 0) { LAUNCH_MOE_FIN(128) } + else if (inter_s % 64 == 0) { LAUNCH_MOE_FIN(64) } + else { LAUNCH_MOE_FIN(32) } +#undef LAUNCH_MOE_FIN +} diff --git a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_up_fused.h b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_up_fused.h new file mode 100644 index 00000000..1cd6c00f --- /dev/null +++ b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/moe_up_fused.h @@ -0,0 +1,78 @@ +/* moe_up_fused.h — single-launch routed-up (Q4_K) + shared-up (Q8_0) stage. + * + * The two stages are independent: both only read the layer input `x`, and they + * write disjoint buffers (inter_r / inter_sh+g). At M=1 decode the MoE op is + * launch-bound (~30-48us of host per enqueue against ~5-12us of GPU per + * kernel), so folding them into one enqueue removes 40 launches per step + * without changing the work done. + * + * The grid is a concatenation of the two original grids: + * + * group < split : routed q4_k, work-item (group*K_SPLIT + lid) -> (route,col) + * group >= split : shared q8_0, work-group (group - split) exactly as before + * + * Every lane of a group takes the same side (the branch is on the group id + * only), so the barriers inside the shared body stay uniform. + * + * Both bodies are reused verbatim through Moe_*_kernel::run(), so this file + * carries no copy of the dequant math. + */ +#pragma once +#include + +template +struct Moe_up_fused_kernel { + Moe_up_q4k_kernel routed; + Moe_shared_up_q8_kernel shared; + int split; // first group index belonging to the shared stage + int n_routed_wi; // M * top_k * intermediate + + void operator()(sycl::nd_item<1> item) const SYCL_ESIMD_KERNEL { + slm_init<2 * K_SPLIT * sizeof(float)>(); + const int grp = (int)item.get_group(0); + const int lid = (int)item.get_local_id(0); + if (grp < split) { + const int wi = grp * K_SPLIT + lid; + if (wi < n_routed_wi) + routed.run(wi / routed.intermediate, wi % routed.intermediate); + } else { + shared.run(grp - split, lid); + } + } +}; + +// Returns false when the shape/tuning combination is outside the fused fast +// path; the caller then falls back to the two separate launches. +inline bool moe_up_fused_host( + const fp16* x, + const uint8_t* gq, const fp16* gs, const fp16* gm, + const uint8_t* uq, const fp16* us, const fp16* um, + const int* sel, fp16* inter_r, + const int8_t* gu_qs, const fp16* gu_sc, const fp16* wg, + fp16* inter_sh, fp16* g, + int M, int hidden, int intermediate, int top_k, int inter_s, + sycl::queue& q) { + // Escape hatch for A/B testing the fused path against the original pair. + static const bool disabled = getenv("SGL_ESIMD_NO_MOE_UP_FUSE") != nullptr; + if (disabled) return false; + if (hidden % 128 != 0) return false; // routed wants VL=128 + if ((hidden / 8) % MOE_Q8_GROUP != 0) return false; // shared wants K_SPLIT=8 + constexpr int VL = 128, KS = 8; + + const int n_cols = inter_s + (g != nullptr ? 1 : 0); + const int n_routed_wi = M * top_k * intermediate; + const int split = (n_routed_wi + KS - 1) / KS; + const int groups = split + M * n_cols; + + q.submit([&](sycl::handler& h) { + h.parallel_for( + sycl::nd_range<1>((size_t)groups * KS, KS), + Moe_up_fused_kernel{ + Moe_up_q4k_kernel{x, gq, gs, gm, uq, us, um, sel, inter_r, + M, hidden, intermediate, top_k}, + Moe_shared_up_q8_kernel{x, gu_qs, gu_sc, wg, inter_sh, g, + M, hidden, inter_s, n_cols}, + split, n_routed_wi}); + }); + return true; +} diff --git a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/prefill_dpas.h b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/prefill_dpas.h index d5b02075..40787caf 100644 --- a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/prefill_dpas.h +++ b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/prefill_dpas.h @@ -58,6 +58,22 @@ static constexpr uint32_t PF_MAX_SLM_BASE = 0x18000; // 4 KB static constexpr uint32_t PF_SUM_SLM_BASE = 0x19000; // 4 KB static constexpr uint32_t PF_TOTAL_SLM = 0x1A000; // 104 KB total +// Cross-subgroup exchange of the running max / partial sums / S tiles goes +// through SLM. A named barrier only orders execution, it does not make the +// SLM stores visible to the other subgroups, so the split arrive/wait form +// leaves the exchange racy. barrier() carries the required memory fence. +// Set PF_SPLIT_BARRIER=1 to get the old (racy) split form back. +#ifndef PF_SPLIT_BARRIER +#define PF_SPLIT_BARRIER 0 +#endif +#if PF_SPLIT_BARRIER +#define PF_BARRIER_ARRIVE() __esimd_nbarrier_arrive(0, 0, 32, 32) +#define PF_BARRIER_WAIT() __esimd_nbarrier(0, 0, 32) +#else +#define PF_BARRIER_ARRIVE() ((void)0) +#define PF_BARRIER_WAIT() barrier() +#endif + template ESIMD_INLINE void sdp_paged_prefill_dpas( const unsigned short* __restrict__ query_ptr, @@ -146,8 +162,22 @@ ESIMD_INLINE void sdp_paged_prefill_dpas( // Inputs: kv_row (absolute KV position in this sequence). // Returns: the Y coordinate to feed into payload.set_y(), accounting for // the row's containing phys block. +// +// The final PF_KV_CHUNK iteration covers up to 127 KV rows past `seq_len` +// (their scores are masked to FP32_MIN, so they are numerically inert) but the +// address computation still runs for them. Without a clamp those rows index +// block_table_row[] past this request's own blocks. block_table has exactly +// ceil(max_seq_len_k / page_size) columns, i.e. the longest request in the +// batch fills its row completely, so for that request the read runs off the end +// of the row -- and for the last request off the end of the whole tensor, which +// is an out-of-bounds device read (observed as UR_RESULT_ERROR_DEVICE_LOST). +// Clamping to the last block that actually belongs to this request keeps the +// address in range; the loaded values are discarded by the mask either way. +#define BLK_LOGICAL_CLAMP(idx) \ + ((int32_t)(idx) < max_valid_blk_idx ? (int32_t)(idx) : max_valid_blk_idx) + #define KV_PHYS_Y(kv_row) \ - ((uint32_t)((BLK_TABLE_LOAD((int32_t)((kv_row) >> block_size_shift)) \ + ((uint32_t)((BLK_TABLE_LOAD(BLK_LOGICAL_CLAMP((kv_row) >> block_size_shift)) \ << phys_block_shift) \ + ((kv_row) & block_size_mask))) @@ -256,7 +286,7 @@ ESIMD_INLINE void sdp_paged_prefill_dpas( } } if (kvOuterLoops > 1) { - int32_t pf1_logical = (int32_t)PF_KV_CHUNK >> block_size_shift; + int32_t pf1_logical = BLK_LOGICAL_CLAMP((int32_t)PF_KV_CHUNK >> block_size_shift); int32_t pf1_off = (int32_t)PF_KV_CHUNK & block_size_mask; int32_t pf1_phys = BLK_TABLE_LOAD(pf1_logical); payloadKpf.set_y((uint32_t)((pf1_phys << phys_block_shift) + pf1_off + sg_i * (int32_t)PF_KV_PER_SG)); @@ -440,14 +470,14 @@ ESIMD_INLINE void sdp_paged_prefill_dpas( // ======================================== // BARRIER A: arrive, QK[k+1] overlap, wait // ======================================== - __esimd_nbarrier_arrive(0, 0, 32, 32); + PF_BARRIER_ARRIVE(); if (outerIter < kvOuterLoops - 1) { uint32_t next_kv_start = (outerIter + 1) * PF_KV_CHUNK; // Per-row K lookup for sg_i's slice of the next chunk. uint32_t next_Y_base_K = KV_PHYS_Y(next_kv_start + (uint32_t)(sg_i * PF_KV_PER_SG)); // V prefetch base — best-effort, use first row's phys. - int32_t next_logical = next_kv_start >> block_size_shift; + int32_t next_logical = BLK_LOGICAL_CLAMP(next_kv_start >> block_size_shift); int32_t next_off = next_kv_start & block_size_mask; int32_t next_phys = BLK_TABLE_LOAD(next_logical); uint32_t next_Y_base_V = (uint32_t)((next_phys << phys_block_shift) + next_off); @@ -536,7 +566,7 @@ ESIMD_INLINE void sdp_paged_prefill_dpas( } } - __esimd_nbarrier(0, 0, 32); + PF_BARRIER_WAIT(); // ======================================== // SOFTMAX SECOND HALF @@ -635,7 +665,7 @@ ESIMD_INLINE void sdp_paged_prefill_dpas( // ======================================== // BARRIER B: arrive, V loads + compensation, wait // ======================================== - __esimd_nbarrier_arrive(0, 0, 32, 32); + PF_BARRIER_ARRIVE(); fp32_sum = fp32_sum * delta + local_sum; @@ -666,7 +696,7 @@ ESIMD_INLINE void sdp_paged_prefill_dpas( } } - __esimd_nbarrier(0, 0, 32); + PF_BARRIER_WAIT(); // ======================================== // VS PHASE + K PREFETCH (remaining tiles) @@ -816,6 +846,7 @@ ESIMD_INLINE void sdp_paged_prefill_dpas( } } #undef BLK_TABLE_LOAD +#undef BLK_LOGICAL_CLAMP } // ============================================================ diff --git a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/resadd_norm_gemv2_fused.h b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/resadd_norm_gemv2_fused.h index 1dfc5ac8..18e77a7f 100644 --- a/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/resadd_norm_gemv2_fused.h +++ b/sglang/custom-esimd-kernels/csrc/xpu/esimd_kernels/resadd_norm_gemv2_fused.h @@ -46,6 +46,7 @@ struct ResAddNormGEMV2_fp8_pert_kernel { const uint8_t* w1_ptr; // [N1, K] FP8 const float* s1_ptr; // [1] fp16* o1_ptr; // [1, N1] + fp16* new_residual_ptr; // [1, K] — written by gid==0 (= hidden+residual); may be null int N0, N1, K; float eps; int fp8_mode; @@ -61,6 +62,18 @@ struct ResAddNormGEMV2_fp8_pert_kernel { constexpr int VL = 512; int n_chunks = K / VL; + // gid 0 additionally writes the post-add residual (hidden+residual) so + // the caller doesn't need a separate aten::add dispatch. fp16 add to + // match the reference (h.half + r.half). Only one group writes → no race. + if (gid == 0 && new_residual_ptr != nullptr) { + for (int c = 0; c < n_chunks; c++) { + int offset = c * VL; + simd hh = block_load(hidden_ptr + offset); + simd rr = block_load(residual_ptr + offset); + block_store(new_residual_ptr + offset, hh + rr); + } + } + // Pass 1: compute sum_sq for RMS float sum_sq = 0.0f; for (int c = 0; c < n_chunks; c++) { @@ -121,6 +134,7 @@ inline void resadd_norm_gemv2_fp8_pert_host( const fp16* hidden_ptr, const fp16* residual_ptr, const fp16* norm_w_ptr, const uint8_t* w0, const float* s0, fp16* o0, const uint8_t* w1, const float* s1, fp16* o1, + fp16* new_residual_ptr, int N0, int N1, int K, float eps, int fp8_mode, sycl::queue& q) { @@ -130,7 +144,7 @@ inline void resadd_norm_gemv2_fp8_pert_host( sycl::nd_range<1>(total_N, 1), ResAddNormGEMV2_fp8_pert_kernel{ hidden_ptr, residual_ptr, norm_w_ptr, - w0, s0, o0, w1, s1, o1, + w0, s0, o0, w1, s1, o1, new_residual_ptr, N0, N1, K, eps, fp8_mode}); }); } diff --git a/sglang/custom-esimd-kernels/csrc/xpu/torch_extension.cc b/sglang/custom-esimd-kernels/csrc/xpu/torch_extension.cc index 19b421e6..708b0b97 100644 --- a/sglang/custom-esimd-kernels/csrc/xpu/torch_extension.cc +++ b/sglang/custom-esimd-kernels/csrc/xpu/torch_extension.cc @@ -100,6 +100,45 @@ TORCH_LIBRARY(custom_esimd_kernels_sglang, m) { "int n_tokens, int hidden, int intermediate, int top_k) -> Tensor"); m.impl("esimd_moe_down_q6k", torch::kXPU, &esimd_moe_down_q6k); + m.def("esimd_shared_expert_q8(Tensor x, Tensor gu_qs, Tensor gu_sc, " + "Tensor d_qs, Tensor d_sc, Tensor wg, int inter_s) -> Tensor"); + m.impl("esimd_shared_expert_q8", torch::kXPU, &esimd_shared_expert_q8); + + m.def("esimd_moe_forward_full_gguf(Tensor x, Tensor logits, " + "Tensor gate_ql, Tensor gate_sc, Tensor gate_mn, " + "Tensor up_ql, Tensor up_sc, Tensor up_mn, " + "Tensor down_ql, Tensor down_qh, Tensor down_sc, Tensor down_mn, " + "Tensor gu_qs, Tensor gu_sc, Tensor d_qs, Tensor d_sc, Tensor wg, " + "int n_experts, int top_k, int intermediate, int inter_s, " + "bool down_is_q6, bool renorm) -> Tensor"); + m.impl("esimd_moe_forward_full_gguf", torch::kXPU, &esimd_moe_forward_full_gguf); + + m.def("esimd_moe_forward_full_gguf_norm(Tensor h, Tensor residual, " + "Tensor norm_w, float eps, Tensor router_w, " + "Tensor gate_ql, Tensor gate_sc, Tensor gate_mn, " + "Tensor up_ql, Tensor up_sc, Tensor up_mn, " + "Tensor down_ql, Tensor down_qh, Tensor down_sc, Tensor down_mn, " + "Tensor gu_qs, Tensor gu_sc, Tensor d_qs, Tensor d_sc, Tensor wg, " + "int n_experts, int top_k, int intermediate, int inter_s, " + "bool down_is_q6, bool renorm) -> Tensor[]"); + m.impl("esimd_moe_forward_full_gguf_norm", torch::kXPU, + &esimd_moe_forward_full_gguf_norm); + + // Fused (resadd + GemmaRMSNorm) + q8_0 GEMV [+ optional fp16 GEMV] for the + // decode attention input projection. `w1`/`o1` empty => no second GEMV. + m.def("esimd_resadd_norm_gemv_q8_ba(Tensor h, Tensor residual, " + "Tensor nw, float eps, Tensor(b!) xn, Tensor(e!) nr, " + "Tensor w0, Tensor s0, Tensor(c!) o0, " + "Tensor w1, Tensor(d!) o1) -> ()"); + m.impl("esimd_resadd_norm_gemv_q8_ba", torch::kXPU, + &esimd_resadd_norm_gemv_q8_ba); + + // Fused RMSNormGated + q8_0 GEMV for the GGUF GDN out_proj. + m.def("esimd_norm_gemv_q8_0(Tensor x, Tensor z, Tensor nw, Tensor(a!) y, " + "Tensor w0, Tensor s0, Tensor(b!) o0, " + "int HV, int V, float eps) -> ()"); + m.impl("esimd_norm_gemv_q8_0", torch::kXPU, &esimd_norm_gemv_q8_0); + // GGUF q4_0 GEMM (prefill / M>=2) via DPAS. Same interleaved weight layout. m.def("esimd_gemm_q4_0(Tensor input, Tensor weight, Tensor weight_scale, " "Tensor output) -> Tensor"); @@ -130,6 +169,7 @@ TORCH_LIBRARY(custom_esimd_kernels_sglang, m) { "Tensor norm_weight, " "Tensor w0, Tensor s0, Tensor o0, " "Tensor w1, Tensor s1, Tensor o1, " + "Tensor new_residual, " "float eps) -> Tensor"); m.impl("esimd_resadd_norm_gemv2_fp8_pert", torch::kXPU, &esimd_resadd_norm_gemv2_fp8_pert); diff --git a/sglang/custom-esimd-kernels/include/kernel_ops.h b/sglang/custom-esimd-kernels/include/kernel_ops.h index 45540297..42115bfa 100644 --- a/sglang/custom-esimd-kernels/include/kernel_ops.h +++ b/sglang/custom-esimd-kernels/include/kernel_ops.h @@ -124,6 +124,7 @@ at::Tensor esimd_resadd_norm_gemv2_fp8_pert( at::Tensor hidden_states, at::Tensor residual, at::Tensor norm_weight, at::Tensor w0, at::Tensor s0, at::Tensor o0, at::Tensor w1, at::Tensor s1, at::Tensor o1, + at::Tensor new_residual, double eps); // Fused RMSNormGated + FP8 GEMV for GDN out_proj decode path @@ -231,6 +232,17 @@ at::Tensor esimd_gemv_q8_0( at::Tensor output); // M-tiled q8_0 dense GEMV (small M, MTP verify): input [M,K], output [M,N]. +void esimd_norm_gemv_q8_0( + at::Tensor x, at::Tensor z, at::Tensor nw, at::Tensor y, + at::Tensor w0, at::Tensor s0, at::Tensor o0, + int64_t HV, int64_t V, double eps); + +void esimd_resadd_norm_gemv_q8_ba( + at::Tensor h, at::Tensor residual, at::Tensor nw, double eps, + at::Tensor xn, at::Tensor nr, + at::Tensor w0, at::Tensor s0, at::Tensor o0, + at::Tensor w1, at::Tensor o1); + at::Tensor esimd_gemv_q8_0_m( at::Tensor input, at::Tensor weight, at::Tensor weight_scale, at::Tensor output); @@ -276,6 +288,36 @@ at::Tensor esimd_moe_down_q6k( at::Tensor sel, at::Tensor topk_w, at::Tensor out_partial, int64_t n_tokens, int64_t hidden, int64_t intermediate, int64_t top_k); +// Fused GGUF Q8_0 shared-expert MLP (gate_up + silu + down + gate*sigmoid). +at::Tensor esimd_shared_expert_q8( + at::Tensor x, at::Tensor gu_qs, at::Tensor gu_sc, + at::Tensor d_qs, at::Tensor d_sc, at::Tensor wg, int64_t inter_s); + +// Fused GGUF FULL MoE (decode): topk + routed + shared -> one op. +at::Tensor esimd_moe_forward_full_gguf( + at::Tensor x, at::Tensor logits, + at::Tensor gate_ql, at::Tensor gate_sc, at::Tensor gate_mn, + at::Tensor up_ql, at::Tensor up_sc, at::Tensor up_mn, + at::Tensor down_ql, at::Tensor down_qh, at::Tensor down_sc, at::Tensor down_mn, + at::Tensor gu_qs, at::Tensor gu_sc, at::Tensor d_qs, at::Tensor d_sc, + at::Tensor wg, + int64_t n_experts, int64_t top_k, int64_t intermediate, int64_t inter_s, + bool down_is_q6, bool renorm); + +// Norm-fused variant: absorbs the post-attention GemmaRMSNorm (residual add + +// RMS norm, `norm_w` = 1 + gemma weight) and the fp16 router GEMV into the MoE +// op. `residual` is updated in place; returns {moe_out, residual}. +std::vector esimd_moe_forward_full_gguf_norm( + at::Tensor h, at::Tensor residual, at::Tensor norm_w, double eps, + at::Tensor router_w, + at::Tensor gate_ql, at::Tensor gate_sc, at::Tensor gate_mn, + at::Tensor up_ql, at::Tensor up_sc, at::Tensor up_mn, + at::Tensor down_ql, at::Tensor down_qh, at::Tensor down_sc, at::Tensor down_mn, + at::Tensor gu_qs, at::Tensor gu_sc, at::Tensor d_qs, at::Tensor d_sc, + at::Tensor wg, + int64_t n_experts, int64_t top_k, int64_t intermediate, int64_t inter_s, + bool down_is_q6, bool renorm); + // GGUF q4_0 GEMM (prefill / M>=2) via DPAS. Same interleaved weight layout. at::Tensor esimd_gemm_q4_0( at::Tensor input, at::Tensor weight, at::Tensor weight_scale, diff --git a/sglang/custom-esimd-kernels/python/custom_esimd_kernels_sglang/ops.py b/sglang/custom-esimd-kernels/python/custom_esimd_kernels_sglang/ops.py index 93cddd0f..a2cf1bdf 100644 --- a/sglang/custom-esimd-kernels/python/custom_esimd_kernels_sglang/ops.py +++ b/sglang/custom-esimd-kernels/python/custom_esimd_kernels_sglang/ops.py @@ -531,16 +531,19 @@ def esimd_resadd_norm_gemv2_fp8_pert( norm_weight: torch.Tensor, w0: torch.Tensor, s0: torch.Tensor, o0: torch.Tensor, w1: torch.Tensor, s1: torch.Tensor, o1: torch.Tensor, + new_residual: torch.Tensor, eps: float, ) -> torch.Tensor: """Fused ResidualAdd + RMSNorm + 2-matrix FP8 GEMV. For input_layernorm + GDN in_proj (qkvz + ba projections). - residual updated in-place. o0/o1 are output buffers. + o0/o1 are output buffers. ``new_residual`` is written with + ``hidden + residual`` (fp16) by the kernel, removing the separate + aten::add dispatch on the caller side. """ return _ops.esimd_resadd_norm_gemv2_fp8_pert( hidden_states, residual, norm_weight, - w0, s0, o0, w1, s1, o1, eps) + w0, s0, o0, w1, s1, o1, new_residual, eps) def esimd_norm_gemv_fp8_pert( diff --git a/sglang/custom-esimd-kernels/setup.py b/sglang/custom-esimd-kernels/setup.py index ccd7a05c..8563593f 100644 --- a/sglang/custom-esimd-kernels/setup.py +++ b/sglang/custom-esimd-kernels/setup.py @@ -353,7 +353,7 @@ extra_compile_args={ "cxx": ["-O3", "-std=c++20"], "sycl": ["-fsycl", "-ffast-math", "-fsycl-device-code-split=per_kernel", - "-fsycl-targets=spir64_gen", "-Xs", "-device bmg", + "-fsycl-targets=spir64_gen", "-Xs", "-device bmg -options -doubleGRF", f"-I{torch_include}"], }, extra_link_args=["-Wl,-rpath,$ORIGIN/../../torch/lib"], diff --git a/sglang/docker/Dockerfile b/sglang/docker/Dockerfile index df816419..ac96308a 100644 --- a/sglang/docker/Dockerfile +++ b/sglang/docker/Dockerfile @@ -95,6 +95,8 @@ ARG http_proxy ARG https_proxy ARG no_proxy=localhost,127.0.0.1,::1,intel.com,.intel.com ARG DPCPP_SYCL_TARGET=bmg +ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh" +ARG ONECCL_INSTALLER_SHA256="f7ab81b6ed1b10dd35fadec366a78046d8af214888dfd625047ce8953d5aa4ef" ENV PIP_NO_CACHE_DIR=1 PIP_DISABLE_PIP_VERSION_CHECK=1 ENV http_proxy=${http_proxy} https_proxy=${https_proxy} no_proxy=${no_proxy} ENV TORCH_XPU_ARCH_LIST=${DPCPP_SYCL_TARGET} @@ -141,6 +143,16 @@ RUN --mount=type=bind,from=builder,source=/root/wheels,target=/tmp/wheels \ # --- runtime helpers (eval scripts: gsm8k, etc.) --- && uv pip install --python /opt/venv/bin/python \ pytest expecttest ray huggingface_hub \ + # --- Replace the torch-bundled oneCCL in /opt/venv/lib with links to + # oneCCL 2021.15 (libtorch_xpu.so resolves libccl.so.1 through its + # RPATH into /opt/venv/lib). --- + && uv pip uninstall --python /opt/venv/bin/python oneccl oneccl-devel || true \ + && rm -f /opt/venv/lib/libccl.so /opt/venv/lib/libccl.so.1 /opt/venv/lib/libccl.so.1.0 \ + /opt/venv/lib/libccl_openmp.so /opt/venv/lib/libccl_openmp.so.0 /opt/venv/lib/libccl_openmp.so.0.1 \ + && ln -sf /opt/intel/oneapi/ccl/2021.15/lib/libccl.so.1.0 /opt/venv/lib/libccl.so.1 \ + && ln -sf /opt/venv/lib/libccl.so.1 /opt/venv/lib/libccl.so \ + && ln -sf /opt/intel/oneapi/ccl/2021.15/lib/libccl_openmp.so.0.1 /opt/venv/lib/libccl_openmp.so.0 \ + && ln -sf /opt/venv/lib/libccl_openmp.so.0 /opt/venv/lib/libccl_openmp.so \ && chmod +x /llm-scaler/sglang/scripts/*.sh \ # --- Size cleanup: drop caches + bytecode --- && find /opt/venv /usr/local/lib/python3* -name "*.pyc" -delete 2>/dev/null || true \ @@ -150,4 +162,12 @@ RUN --mount=type=bind,from=builder,source=/root/wheels,target=/tmp/wheels \ && rm -rf /root/.cache /root/.cmake /var/tmp/* /var/lib/apt/lists/* WORKDIR /llm-scaler/sglang + +# NOTE: no global LD_PRELOAD for oneCCL. The 2021.15 libccl.so.1 carries no +# RPATH/RUNPATH, so preloading it into *every* process (including /bin/bash) +# fails with "libsvml.so: cannot open shared object file" — its oneAPI runtime +# deps only live in /opt/venv/lib, which is not on the loader search path. +# The /opt/venv/lib/libccl.so.1 -> 2021.15 symlink above is what makes torch +# pick up this oneCCL: libtorch_xpu.so has DT_RPATH=$ORIGIN/../../../.. +# (= /opt/venv/lib), which is inherited by its transitive dependencies. ENTRYPOINT ["bash", "-c", "exec sglang serve \"$@\"", "--"] diff --git a/sglang/docker/Dockerfile.dev b/sglang/docker/Dockerfile.dev index 001540eb..9e34d0f5 100644 --- a/sglang/docker/Dockerfile.dev +++ b/sglang/docker/Dockerfile.dev @@ -46,6 +46,8 @@ ARG SGL_KERNEL_UPSTREAM_REPO=https://github.com/sgl-project/sgl-kernel-xpu.git # sgl-kernel-xpu has no release tags; pin to the exact upstream/main commit # the analytics-zoo/dev-bmg fork was branched from. ARG SGL_KERNEL_UPSTREAM_COMMIT=ea5c70f0909bcd55ceaf1803302651fc0593b64d +ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh" +ARG ONECCL_INSTALLER_SHA256="f7ab81b6ed1b10dd35fadec366a78046d8af214888dfd625047ce8953d5aa4ef" SHELL ["bash", "-c"] @@ -71,6 +73,48 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ curl -LsSf https://astral.sh/uv/install.sh | sh && \ mkdir -p ${WHEELS} +# --------------------------------------------------------------------------- +# Intel(R) oneAPI Collective Communications Library (oneCCL) 2021.15.9, which +# contains several enhancements for Intel(R) Arc(TM) Pro graphics. +# For details, please refer to: +# https://github.com/uxlfoundation/oneCCL/releases/tag/2021.15.9 +# +# Only the runtime payload is extracted from the offline installer (a full +# `-a --silent` install would drag in the whole oneAPI installer layout), and +# it replaces the oneCCL shipped by the OMIX base image. +# --------------------------------------------------------------------------- +RUN set -eo pipefail && \ + curl -fL --retry 3 -o "${ONECCL_INSTALLER}" "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \ + printf "%s %s\n" "${ONECCL_INSTALLER_SHA256}" "${ONECCL_INSTALLER}" > /tmp/oneccl.sha256 && \ + sha256sum -c /tmp/oneccl.sha256 && \ + rm -f /tmp/oneccl.sha256 && \ + bash "${ONECCL_INSTALLER}" -s -x -f /tmp && \ + ONECCL_CUP="/tmp/${ONECCL_INSTALLER%.sh}/packages/intel.oneapi.lin.ccl.runtime,v=2021.15.9+13/cupPayload.cup" && \ + ONECCL_DST="/tmp/oneccl-runtime" && \ + python3 -m zipfile -e "${ONECCL_CUP}" "${ONECCL_DST}" && \ + ONECCL_LIB="${ONECCL_DST}/_installdir/ccl/2021.15/lib" && \ + rm -f "${ONECCL_LIB}/libccl.so" "${ONECCL_LIB}/libccl.so.1" "${ONECCL_LIB}/libccl_openmp.so" "${ONECCL_LIB}/libccl_openmp.so.0" && \ + ln -sf libccl.so.1 "${ONECCL_LIB}/libccl.so" && \ + ln -sf libccl.so.1.0 "${ONECCL_LIB}/libccl.so.1" && \ + ln -sf libccl_openmp.so.0 "${ONECCL_LIB}/libccl_openmp.so" && \ + ln -sf libccl_openmp.so.0.1 "${ONECCL_LIB}/libccl_openmp.so.0" && \ + ONECCL_CPU_LIB="${ONECCL_DST}/_installdir/ccl/2021.15/lib/ccl/cpu/lib" && \ + rm -f "${ONECCL_CPU_LIB}/libccl.so" "${ONECCL_CPU_LIB}/libccl.so.1" "${ONECCL_CPU_LIB}/libccl_openmp.so" "${ONECCL_CPU_LIB}/libccl_openmp.so.0" && \ + ln -sf libccl.so.1 "${ONECCL_CPU_LIB}/libccl.so" && \ + ln -sf libccl.so.1.0 "${ONECCL_CPU_LIB}/libccl.so.1" && \ + ln -sf libccl_openmp.so.0 "${ONECCL_CPU_LIB}/libccl_openmp.so" && \ + ln -sf libccl_openmp.so.0.1 "${ONECCL_CPU_LIB}/libccl_openmp.so.0" && \ + mkdir -p /opt/intel/oneapi/ccl && \ + rm -rf /opt/intel/oneapi/ccl/2021.15 && \ + cp -a /tmp/oneccl-runtime/_installdir/ccl/2021.15 /opt/intel/oneapi/ccl/ && \ + rm "${ONECCL_INSTALLER}" && \ + echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \ + echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --ccl-bundled-mpi=no" >> /root/.bashrc && \ + rm -f /opt/intel/oneapi/ccl/latest && \ + ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest && \ + find /opt/intel/oneapi/ccl -mindepth 1 -maxdepth 1 -type d -not -name 2021.15 -exec rm -rf {} + && \ + rm -rf "/tmp/${ONECCL_INSTALLER%.sh}" /tmp/oneccl-runtime + # Install the same Torch version used by the production runtime. RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install --system \ @@ -162,6 +206,18 @@ RUN --mount=type=cache,target=/root/.cache/uv \ huggingface_hub && \ chmod +x /llm-scaler/sglang/scripts/*.sh +# Replace the torch-bundled oneCCL with links to oneCCL 2021.15. With +# `uv pip install --system` on this base the wheel data files land in +# /usr/local/lib, which is where libtorch_xpu.so resolves libccl.so.1 from. +RUN set -eo pipefail && \ + uv pip uninstall --system oneccl oneccl-devel || true && \ + rm -f /usr/local/lib/libccl.so /usr/local/lib/libccl.so.1 /usr/local/lib/libccl.so.1.0 \ + /usr/local/lib/libccl_openmp.so /usr/local/lib/libccl_openmp.so.0 /usr/local/lib/libccl_openmp.so.0.1 && \ + ln -sf /opt/intel/oneapi/ccl/2021.15/lib/libccl.so.1.0 /usr/local/lib/libccl.so.1 && \ + ln -sf /usr/local/lib/libccl.so.1 /usr/local/lib/libccl.so && \ + ln -sf /opt/intel/oneapi/ccl/2021.15/lib/libccl_openmp.so.0.1 /usr/local/lib/libccl_openmp.so.0 && \ + ln -sf /usr/local/lib/libccl_openmp.so.0 /usr/local/lib/libccl_openmp.so + # Keep wheels, source trees, build directories, caches, bytecode, and debug # symbols. Print an inventory for build-log validation. RUN echo "Built wheels:" && \ @@ -171,6 +227,12 @@ RUN echo "Built wheels:" && \ grep -iE "^(sglang|sgl-kernel|custom-esimd|torch|triton|transformers|xgrammar)" \ || true +# NOTE: no global LD_PRELOAD for oneCCL — the 2021.15 libccl.so.1 has no +# RPATH/RUNPATH, so preloading it into every process breaks plain binaries +# (e.g. /bin/bash) with "libsvml.so: cannot open shared object file". The +# /usr/local/lib/libccl.so.1 -> 2021.15 symlink above is what makes torch load +# this oneCCL, via the libtorch_xpu.so DT_RPATH. + # Interactive development shell. The OMIX base configures the oneAPI # environment for login-shell workflows. CMD ["/bin/bash", "-l"] \ No newline at end of file diff --git a/sglang/patches/sglang_for_multi_arc.patch b/sglang/patches/sglang_for_multi_arc.patch index d515f6b2..5b728a98 100644 --- a/sglang/patches/sglang_for_multi_arc.patch +++ b/sglang/patches/sglang_for_multi_arc.patch @@ -208,10 +208,10 @@ index 44ee91a1a..38de4d83a 100644 REMOTE = "remote" diff --git a/python/sglang/srt/layers/attention/fla/chunk_torch_xpu.py b/python/sglang/srt/layers/attention/fla/chunk_torch_xpu.py new file mode 100644 -index 000000000..378fea0cc +index 000000000..7c2b07690 --- /dev/null +++ b/python/sglang/srt/layers/attention/fla/chunk_torch_xpu.py -@@ -0,0 +1,144 @@ +@@ -0,0 +1,167 @@ +"""Pure-PyTorch fallback for chunk_gated_delta_rule on Intel XPU. + +Triton-XPU 3.7.0 cannot compile the block-pointer-heavy GDN kernels in @@ -248,7 +248,8 @@ index 000000000..378fea0cc + cu_seqlens: Optional[torch.LongTensor] = None, + head_first: bool = False, + use_qk_l2norm_in_kernel: bool = False, -+) -> Tuple[torch.Tensor, None, Optional[torch.Tensor]]: ++ intermediate_chunk_size: Optional[int] = None, ++) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Drop-in for sglang.srt.layers.attention.fla.chunk.chunk_gated_delta_rule. + + Shapes (head_first=False, the only mode sglang uses today): @@ -301,6 +302,23 @@ index 000000000..378fea0cc + (n_seqs, H, V, K), dtype=torch.float32, device=q.device + ) + ++ # Optional per-chunk intermediate states. h[i] is the recurrent state after ++ # i full chunks (h[0] == the initial state), packed per sequence with ++ # ceil(L_s / chunk) entries each -- the layout _init_track_ssm_indices ++ # assumes. Only materialised when the caller asks for it. ++ chunk = int(intermediate_chunk_size or 0) ++ h_out = None ++ h_offsets = None ++ if chunk > 0: ++ seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() ++ n_h_per_seq = [(L + chunk - 1) // chunk for L in seq_lens] ++ h_offsets = [0] ++ for c in n_h_per_seq[:-1]: ++ h_offsets.append(h_offsets[-1] + c) ++ h_out = torch.empty( ++ (sum(n_h_per_seq), H, V, K), dtype=torch.float32, device=q.device ++ ) ++ + cu = cu_seqlens.tolist() + idx_list = ( + initial_state_indices.tolist() if initial_state_indices is not None else None @@ -321,6 +339,9 @@ index 000000000..378fea0cc + state = torch.zeros((H, V, K), dtype=torch.float32, device=q.device) + + for t in range(t0, t1): ++ if chunk > 0 and (t - t0) % chunk == 0: ++ h_out[h_offsets[s] + (t - t0) // chunk] = state ++ + q_t = q[t].to(torch.float32) # [H_k, K] + k_t = k[t].to(torch.float32) # [H_k, K] + v_t = v[t].to(torch.float32) # [H, V] @@ -355,7 +376,9 @@ index 000000000..378fea0cc + # gdn_backend.forward_extend scatters last_recurrent_state back into + # ssm_states[cache_indices] on non-CUDA backends; returning it as the + # middle element (not as h_aux) is required for that scatter to run. -+ return o, last_state, None ++ if h_out is not None: ++ h_out = h_out.unsqueeze(0) # [1, total_chunks, H, V, K] ++ return o, last_state, h_out diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/srt/layers/attention/fla/layernorm_gated.py index 5a8fda41b..99c56f9ff 100644 --- a/python/sglang/srt/layers/attention/fla/layernorm_gated.py @@ -447,10 +470,59 @@ index 5a8fda41b..99c56f9ff 100644 + activation=self.activation, + ) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py -index d0a8ff7a9..7c0d37076 100644 +index d0a8ff7a9..32d5c8b67 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py -@@ -533,7 +533,14 @@ class MambaAttnBackendBase(AttentionBackend): +@@ -1,4 +1,5 @@ + import logging ++import os + from typing import Optional, Union + + import torch +@@ -25,7 +26,6 @@ from sglang.srt.speculative.spec_info import SpecInput + + logger = logging.getLogger(__name__) + +- + # Kernel to track mamba states if needed based on track mask + @triton.jit + def track_mamba_state_if_needed_kernel( +@@ -399,6 +399,34 @@ class MambaAttnBackendBase(AttentionBackend): + ) + track_ssm_h_dst = dst_masked[not_aligned] + ++ # Bounds guard. Everything here is already on CPU, so this costs no ++ # device sync. `h` has one row per chunk per sequence; an index past ++ # the end would be an out-of-bounds *device* read in ++ # `_track_mamba_state_extend`, which shows up as a GPU page fault ++ # (UR_RESULT_ERROR_DEVICE_LOST) at the next synchronisation point ++ # rather than as a Python IndexError. ++ if track_ssm_h_src.numel() > 0: ++ total_h_rows = int(num_h_states.sum()) ++ bad = (track_ssm_h_src < 0) | (track_ssm_h_src >= total_h_rows) ++ if bool(bad.any()): ++ logger.error( ++ "Mamba extend tracking: %d/%d h index(es) out of range " ++ "[0, %d); clamping. extend_seq_lens=%s prefix_lens=%s " ++ "mamba_track_seqlens=%s mamba_track_mask=%s " ++ "lens_to_track=%s offsets=%s src=%s", ++ int(bad.sum()), ++ track_ssm_h_src.numel(), ++ total_h_rows, ++ extend_seq_lens.tolist(), ++ prefix_lens.tolist(), ++ mamba_track_seqlens.tolist(), ++ mamba_track_mask.tolist(), ++ lens_to_track.tolist(), ++ track_ssm_src_offset.tolist(), ++ track_ssm_h_src.tolist(), ++ ) ++ track_ssm_h_src = track_ssm_h_src.clamp(0, total_h_rows - 1) ++ + # Move back to GPU + return ( + track_ssm_h_src.to(self.device, non_blocking=True), +@@ -533,7 +561,14 @@ class MambaAttnBackendBase(AttentionBackend): # Make sure forward metadata is correctly handled for padding reqs req_pool_indices[bs - num_padding :] = 0 mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) @@ -466,21 +538,104 @@ index d0a8ff7a9..7c0d37076 100644 self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices) if forward_mode.is_decode_or_idle(): if num_padding == 0: +@@ -624,10 +659,23 @@ class MambaAttnBackendBase(AttentionBackend): + forward_batch.batch_size, + ) + ++ def _warn_missing_intermediate_states(self, n_rows: int): ++ """Warn once when the kernel gave us no per-chunk states to snapshot.""" ++ if not getattr(self, "_warned_missing_h", False): ++ self._warned_missing_h = True ++ logger.warning( ++ "Mamba extend tracking: kernel returned no intermediate chunk " ++ "states, so %d unaligned snapshot row(s) per batch cannot be " ++ "written. Prefix-cached prefills that resume from those " ++ "snapshots will restore a stale SSM state. Aligned snapshots " ++ "are unaffected.", ++ n_rows, ++ ) ++ + def _track_mamba_state_extend( + self, + forward_batch: ForwardBatch, +- h: torch.Tensor, ++ h: Optional[torch.Tensor], + ssm_states: torch.Tensor, + forward_metadata: ForwardMetadata, + ): +@@ -640,16 +688,25 @@ class MambaAttnBackendBase(AttentionBackend): + to the chunk size. See `_init_track_ssm_indices` for more details on how + the source and destination indices are computed. + ++ ``h`` (the per-chunk intermediate states) may be None when the kernel ++ backend does not produce them. Only the *unaligned* branch needs it; the ++ aligned branch copies the final state out of ``ssm_states`` and must ++ still run, otherwise the snapshot inserted into the radix cache keeps a ++ stale SSM state while its conv state is fresh. ++ + Note: Conv state tracking for extend is handled separately via gather operations + using indices computed by `_init_track_conv_indices`. + """ + if forward_metadata.has_mamba_track_mask: +- h = h.squeeze(0) +- +- if forward_metadata.track_ssm_h_src.numel() > 0: +- ssm_states[forward_metadata.track_ssm_h_dst] = h[ +- forward_metadata.track_ssm_h_src +- ].to(ssm_states.dtype, copy=False) ++ n_h = forward_metadata.track_ssm_h_src.numel() ++ if n_h > 0: ++ if h is None: ++ self._warn_missing_intermediate_states(n_h) ++ else: ++ h = h.squeeze(0) ++ ssm_states[forward_metadata.track_ssm_h_dst] = h[ ++ forward_metadata.track_ssm_h_src ++ ].to(ssm_states.dtype, copy=False) + if forward_metadata.track_ssm_final_src.numel() > 0: + ssm_states[forward_metadata.track_ssm_final_dst] = ssm_states[ + forward_metadata.track_ssm_final_src diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py -index ade156e98..a430c6695 100644 +index ade156e98..122e952f7 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py -@@ -18,7 +18,7 @@ from sglang.srt.layers.radix_linear_attention import RadixLinearAttention +@@ -1,3 +1,5 @@ ++import logging ++import os + from typing import Optional, Tuple, Union + + import torch +@@ -18,7 +20,8 @@ from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.mem_cache.memory_pool import MambaPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.model_runner import ModelRunner -from sglang.srt.utils import is_cpu, is_cuda, is_npu ++from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import is_cpu, is_cuda, is_npu, is_xpu from sglang.srt.utils.common import rank0_log if not is_cpu(): -@@ -501,7 +501,9 @@ class GDNAttnBackend(MambaAttnBackendBase): +@@ -490,6 +493,13 @@ class GDNAttnBackend(MambaAttnBackendBase): + ) + else: + g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) ++ # Per-chunk intermediate states are only needed when this batch has ++ # to snapshot at a position that is not chunk-aligned; asking for ++ # them unconditionally would materialise a large tensor per layer. ++ need_h = ( ++ forward_metadata.has_mamba_track_mask ++ and forward_metadata.track_ssm_h_src.numel() > 0 ++ ) + core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend( + q=query, + k=key, +@@ -499,17 +509,26 @@ class GDNAttnBackend(MambaAttnBackendBase): + ssm_states=ssm_states, + cache_indices=cache_indices, query_start_loc=query_start_loc, ++ intermediate_chunk_size=( ++ get_global_server_args().mamba_cache_chunk_size if need_h else 0 ++ ), ) - if (is_npu() or is_cpu()) and last_recurrent_state is not None: @@ -490,8 +645,24 @@ index ade156e98..a430c6695 100644 last_recurrent_state = last_recurrent_state.to( ssm_states.dtype, copy=False ) + ssm_states[cache_indices] = last_recurrent_state + +- if h is not None: +- self._track_mamba_state_extend( +- forward_batch, h, ssm_states, forward_metadata +- ) ++ ++ # `h` may be None: some kernel backends (XPU) do not emit the ++ # per-chunk intermediate states. Only the unaligned branch of the ++ # tracking needs them, so the call must not be gated on `h` -- doing ++ # so silently dropped *every* extend-time SSM snapshot on XPU. ++ self._track_mamba_state_extend( ++ forward_batch, h, ssm_states, forward_metadata ++ ) + + return core_attn_out diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py -index 0daf8f4e2..cab43381b 100644 +index 0daf8f4e2..93ce0defc 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py @@ -34,6 +34,17 @@ elif is_xpu(): @@ -512,7 +683,7 @@ index 0daf8f4e2..cab43381b 100644 class TritonGDNKernel(LinearAttnKernelBase): """Triton-based kernel for GDN (Gated Delta Network) linear attention.""" -@@ -141,9 +152,42 @@ class TritonGDNKernel(LinearAttnKernelBase): +@@ -141,9 +152,45 @@ class TritonGDNKernel(LinearAttnKernelBase): ) -> tuple: recurrent_state = ssm_states recurrent_state_indices_args = {"initial_state_indices": cache_indices} @@ -539,23 +710,38 @@ index 0daf8f4e2..cab43381b 100644 + and hasattr(torch.ops.eagle_ops, "chunk_gated_delta_rule_extend") + ): + scale = float(q.size(-1)) ** -0.5 -+ # Kernel contract: (q, k, v, g, beta, initial_state, cu_seqlens, scale) -+ # Returns (out [1, T, H_v, V], last_state [n_seqs, H_v, V, K]). ++ # Kernel contract: (q, k, v, g, beta, initial_state, cu_seqlens, ++ # scale, h_chunk_size) -> (out, last_state, h). ++ # h (per-chunk intermediate states) is only materialised when ++ # h_chunk_size > 0, i.e. when this batch actually has to write a ++ # mamba track snapshot at a non-chunk-aligned position. + # g is fp32 log-space decay; kernel expects exactly that. + # initial_state is IN/OUT: kernel mutates it to last_state. ++ h_chunk = int(kwargs.get("intermediate_chunk_size") or 0) + state_in = recurrent_state.contiguous() -+ out, last_state = torch.ops.eagle_ops.chunk_gated_delta_rule_extend( ++ out, last_state, h = torch.ops.eagle_ops.chunk_gated_delta_rule_extend( + q.contiguous(), k.contiguous(), v.contiguous(), + g.contiguous(), beta.contiguous(), + state_in, query_start_loc.to(torch.int32).contiguous(), -+ scale, ++ scale, h_chunk, + ) -+ # Match chunk_gated_delta_rule_torch return: (o, last_recurrent_state, h_aux) -+ return out, last_state, None ++ return out, last_state, (h if h_chunk > 0 else None) + return chunk_gated_delta_rule( q=q, k=k, +@@ -154,6 +201,11 @@ class TritonGDNKernel(LinearAttnKernelBase): + cu_seqlens=query_start_loc, + head_first=False, + use_qk_l2norm_in_kernel=True, ++ **( ++ {"intermediate_chunk_size": kwargs["intermediate_chunk_size"]} ++ if is_xpu() and kwargs.get("intermediate_chunk_size") ++ else {} ++ ), + **recurrent_state_indices_args, + ) + diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index ae5365e8d..0f54d2a8c 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py @@ -667,7 +853,7 @@ index ae5365e8d..0f54d2a8c 100644 q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), self.token_to_kv_pool.get_key_buffer(layer.layer_id), diff --git a/python/sglang/srt/layers/attention/xpu_backend.py b/python/sglang/srt/layers/attention/xpu_backend.py -index ea02e001f..5c1ef4958 100644 +index ea02e001f..a793bed5d 100644 --- a/python/sglang/srt/layers/attention/xpu_backend.py +++ b/python/sglang/srt/layers/attention/xpu_backend.py @@ -1,5 +1,6 @@ @@ -777,13 +963,115 @@ index ea02e001f..5c1ef4958 100644 self.skip_prefill = skip_prefill self.is_hybrid_swa = model_runner.is_hybrid_swa self.use_sliding_window_kv_pool = ( -@@ -103,6 +178,150 @@ class XPUAttentionBackend(AttentionBackend): +@@ -103,6 +178,252 @@ class XPUAttentionBackend(AttentionBackend): self.has_swa = ( self.sliding_window_size is not None and self.sliding_window_size > -1 ) + # XPU graph capture/replay state. Populated by init_cuda_graph_state. + self._graph_state: dict = {} + ++ def _build_sglang_decode_attn_inputs_eager( ++ self, forward_batch: ForwardBatch, metadata, tp_q_head_num: int, head_dim: int ++ ): ++ """Build the flat-NHD kv_indptr/kv_indices/temp_p inputs that ++ sglang_decode_attn needs, without requiring XPU device-graph capture ++ (SGL_XPU_ENABLE_GRAPH=1) to be enabled. init_forward_metadata_capture ++ only builds these when graphs are on, which left sglang_decode_attn ++ (the proven-correct kernel for GQA ratio=8 / single-KV-head configs) ++ unreachable in eager decode -- so decode silently fell back to the ++ numerically-wrong eagle_page_attn_decode / flash kernel instead. ++ ++ Optimization (P0+P1) for the disable-XPU-graph path: the scratch ++ buffers (kv_indptr / kv_indices / temp_p) are kept as a single ++ *persistent, grow-only* cache reused across decode steps instead of ++ being re-allocated every step, and the per-step ++ ``int(kv_indptr[-1].item())`` device->host sync (which stalled the ++ CPU-GPU pipeline on every full-attention layer) is removed: ++ kv_indices is sized to the upper bound ``bs * max_seq_len_k`` (a CPU ++ int already available on ``metadata``, no device readback needed). ++ create_flashinfer_kv_indices_triton only writes into the ranges ++ indexed by kv_indptr, so an over-sized buffer is safe, and ++ sglang_decode_attn only reads the kv_indptr-delimited ranges.""" ++ bs = forward_batch.batch_size ++ # kv_indptr/kv_indices/temp_p depend only on this step's cache_seqlens ++ # (batch-level), not on the attention layer, so they are identical for ++ # every full-attention layer in a decode step. ``metadata`` is a fresh ++ # object built once per step in init_forward_metadata, so memoizing the ++ # built inputs on it lets the first full-attn layer build them and the ++ # remaining layers reuse them -- removing ~9x (cumsum + kv_indptr fill + ++ # create_flashinfer_kv_indices) dispatched ops per step on the ++ # disable-XPU-graph path. ++ cached = getattr(metadata, "_eager_kv_inputs", None) ++ if cached is not None and cached[0].numel() == bs + 1: ++ return cached ++ device = metadata.cache_seqlens_int32.device ++ _SPLIT_TILE = 64 ++ _MAX_N_SPLITS = 256 ++ graph_max_seq = _SPLIT_TILE * _MAX_N_SPLITS # 16384 ++ ++ # Upper bound on total kv entries this step: sum(seq_lens) <= bs * max_seq_len_k. ++ # max_seq_len_k is a plain python int already computed on the CPU-side ++ # seq_lens_cpu in init_forward_metadata, so reading it here costs no ++ # device synchronization (unlike kv_indptr[-1].item()). ++ max_seq_len_k = getattr(metadata, "max_seq_len_k", None) ++ if not isinstance(max_seq_len_k, int) or max_seq_len_k <= 0: ++ max_seq_len_k = self.max_context_len ++ need_kv = max(bs * max_seq_len_k, 1) ++ need_temp = max(bs * tp_q_head_num * _MAX_N_SPLITS * (1 + 1 + 256), 1) ++ ++ # Persistent grow-only scratch: allocate once, reuse (and only grow) ++ # across decode steps. No per-step torch.empty, no per-bs rebuild. ++ cache = getattr(self, "_sglang_decode_eager_scratch", None) ++ if ( ++ cache is None ++ or cache["kv_indptr"].numel() < bs + 1 ++ or cache["kv_indices"].numel() < need_kv ++ or cache["temp_p"].numel() < need_temp ++ ): ++ cache = { ++ "kv_indptr": torch.zeros( ++ max(bs + 1, 1), dtype=torch.int32, device=device ++ ), ++ "kv_indices": torch.empty( ++ need_kv, dtype=torch.int32, device=device ++ ), ++ "temp_p": torch.empty( ++ need_temp, dtype=torch.float32, device=device ++ ), ++ } ++ self._sglang_decode_eager_scratch = cache ++ ++ kv_indptr = cache["kv_indptr"][: bs + 1] ++ kv_indices = cache["kv_indices"] ++ temp_p = cache["temp_p"] ++ ++ kv_indptr[0] = 0 ++ torch.cumsum( ++ metadata.cache_seqlens_int32, ++ dim=0, ++ dtype=torch.int32, ++ out=kv_indptr[1:], ++ ) ++ from sglang.srt.layers.attention.triton_ops.kv_indices import ( ++ create_flashinfer_kv_indices_triton, ++ ) ++ ++ create_flashinfer_kv_indices_triton[(bs,)]( ++ self.req_to_token, ++ forward_batch.req_pool_indices, ++ metadata.cache_seqlens_int32, ++ kv_indptr, ++ None, ++ kv_indices, ++ self.req_to_token.stride(0), ++ ) ++ result = (kv_indptr, kv_indices, temp_p, graph_max_seq) ++ try: ++ metadata._eager_kv_inputs = result ++ except Exception: ++ pass ++ return result ++ + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + """Pre-allocate stable device buffers for XPU graph capture/replay.""" + max_num_pages = (self.max_context_len + self.page_size - 1) // self.page_size @@ -928,7 +1216,7 @@ index ea02e001f..5c1ef4958 100644 def init_forward_metadata(self, forward_batch: ForwardBatch): """Initialize forward metadata hence all layers in the forward pass can reuse it.""" -@@ -590,24 +809,52 @@ class XPUAttentionBackend(AttentionBackend): +@@ -590,24 +911,52 @@ class XPUAttentionBackend(AttentionBackend): cu_seqlens_k = metadata.encoder_cu_seqlens_k window_size = (-1, -1) @@ -998,7 +1286,7 @@ index ea02e001f..5c1ef4958 100644 if use_cascade_attn: o, softmax_lse, *rest = result -@@ -865,10 +1112,31 @@ class XPUAttentionBackend(AttentionBackend): +@@ -865,10 +1214,31 @@ class XPUAttentionBackend(AttentionBackend): -1, self.page_size, layer.tp_v_head_num, layer.head_dim ) @@ -1031,7 +1319,7 @@ index ea02e001f..5c1ef4958 100644 k_cache=key_cache, v_cache=value_cache, page_table=metadata.encoder_page_table, -@@ -887,7 +1155,7 @@ class XPUAttentionBackend(AttentionBackend): +@@ -887,7 +1257,7 @@ class XPUAttentionBackend(AttentionBackend): elif use_local_attn: # Use chunked (local) attention batching for self-attention o = flash_attn_with_kvcache( @@ -1040,7 +1328,7 @@ index ea02e001f..5c1ef4958 100644 k_cache=key_cache, v_cache=value_cache, page_table=local_attn_metadata.local_block_table, -@@ -903,32 +1171,54 @@ class XPUAttentionBackend(AttentionBackend): +@@ -903,32 +1273,75 @@ class XPUAttentionBackend(AttentionBackend): v_descale=v_descale, **kwargs, ) @@ -1049,17 +1337,22 @@ index ea02e001f..5c1ef4958 100644 - layer.sliding_window_size is not None - and layer.sliding_window_size > -1 + elif ( -+ self._graph_state.get("temp_p") is not None -+ and not use_cascade_attn ++ not use_cascade_attn + and _sglang_decode_attn_fn is not None -+ and getattr(metadata, "kv_indices", None) is not None + and layer.head_dim == 256 + and layer.tp_q_head_num % layer.tp_k_head_num == 0 + and os.environ.get("SGL_XPU_DECODE_SGLANG_ATTN", "1") == "1" + ): + # Proven-correct flat-NHD token-granular decode kernel. The paged + # eagle_page_attn_decode is numerically wrong for GQA ratio=8 / -+ # single-KV-head on this stack, so route graph decode here. ++ # single-KV-head on this stack, so route decode here instead. ++ # Prefer the graph-captured stable buffers when XPU device graphs ++ # are on (SGL_XPU_ENABLE_GRAPH=1); otherwise build the same ++ # flat-NHD kv_indptr/kv_indices/temp_p inputs on the fly for ++ # eager decode (they were previously only ever built inside ++ # init_forward_metadata_capture, which made this branch ++ # unreachable without graphs and silently pushed decode into the ++ # numerically-wrong kernel below). + B = forward_batch.batch_size + k_buf = self.token_to_kv_pool.get_key_buffer(layer.layer_id) + v_buf = self.token_to_kv_pool.get_value_buffer(layer.layer_id) @@ -1070,16 +1363,32 @@ index ea02e001f..5c1ef4958 100644 + k_buf = k_buf.to(torch.float16) + v_buf = v_buf.to(torch.float16) + o_fp16 = torch.empty_like(q_fp16) ++ if self._graph_state.get("temp_p") is not None and getattr( ++ metadata, "kv_indices", None ++ ) is not None: ++ kv_indptr = metadata.kv_indptr ++ kv_indices = metadata.kv_indices ++ temp_p = self._graph_state["sglang_temp_p"] ++ graph_max_seq = self._sglang_decode_graph_max_seq ++ else: ++ ( ++ kv_indptr, ++ kv_indices, ++ temp_p, ++ graph_max_seq, ++ ) = self._build_sglang_decode_attn_inputs_eager( ++ forward_batch, metadata, layer.tp_q_head_num, layer.head_dim ++ ) + _sglang_decode_attn_fn( + q_fp16, + k_buf, + v_buf, -+ metadata.kv_indptr, -+ metadata.kv_indices, ++ kv_indptr, ++ kv_indices, + o_fp16, + float(layer.scaling), -+ self._graph_state["sglang_temp_p"], -+ self._sglang_decode_graph_max_seq, ++ temp_p, ++ graph_max_seq, ) - - page_table = metadata.page_table @@ -1149,7 +1458,7 @@ index 59ca3f9cc..e104249d2 100644 # Store in data_container with expert/shard info if not hasattr(param, "expert_data_map"): diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py -index 96e431d4e..f2ba507fc 100644 +index 96e431d4e..cd890f194 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton.py @@ -1,5 +1,6 @@ @@ -1159,12 +1468,18 @@ index 96e431d4e..f2ba507fc 100644 from dataclasses import dataclass from typing import TYPE_CHECKING, Any, List, Optional -@@ -24,6 +25,144 @@ if TYPE_CHECKING: +@@ -24,6 +25,152 @@ if TYPE_CHECKING: ) +# Lazy-loaded esimd MoE op: registers torch.ops.moe_ops.moe_forward_full_silu_routed. +# Returns the op handle on success, None on failure (so we always fall back to Triton). ++# ESIMD MoE gates: set once at launch, constant for the process. Cached to avoid an ++# os.environ read on the per-layer MoE hot path each decode step. ++_ESIMD_MOE = os.environ.get("SGL_XPU_ESIMD_MOE", "0") == "1" ++_ESIMD_MOE_PREFILL = os.environ.get("SGL_XPU_ESIMD_MOE_PREFILL", "0") == "1" ++ ++ +def _load_esimd_moe_op(fp8_variant: str = "e4m3"): + """Load the right ESIMD MoE silu-routed kernel for the given fp8 variant. + @@ -1190,6 +1505,9 @@ index 96e431d4e..f2ba507fc 100644 + return op + + ++ ++ ++ +def _try_esimd_moe_silu_routed( + runner_input: "TritonRunnerInput", + quant_info: "TritonMoeQuantInfo", @@ -1277,10 +1595,9 @@ index 96e431d4e..f2ba507fc 100644 + s13 = _per_expert_pt_scale(quant_info.w13_scale) + s2 = _per_expert_pt_scale(quant_info.w2_scale) + -+ # We use the `_sglang` kernel variant which accepts w13 directly in -+ # sglang's [E, 2*intermediate, hidden] layout — no transpose / copy -+ # required, no extra memory cost, and the Triton fallback can still read -+ # the same parameter unchanged. ++ # Both the e4m3 (`_sglang`) and e5m2 kernel variants now accept w13 directly ++ # in sglang's [E, 2*intermediate, hidden] N-major layout — no transpose / ++ # extra copy, and the Triton fallback reads the same parameter unchanged. + w13_kernel = w13 + + topk_w = runner_input.topk_weights @@ -1304,7 +1621,7 @@ index 96e431d4e..f2ba507fc 100644 @dataclass class TritonRunnerInput(RunnerInput): -@@ -81,6 +220,15 @@ class TritonRunnerCore(MoeRunnerCore): +@@ -81,6 +228,15 @@ class TritonRunnerCore(MoeRunnerCore): running_state: dict, hooks: Optional[Any] = None, ) -> TritonRunnerOutput: @@ -1312,7 +1629,7 @@ index 96e431d4e..f2ba507fc 100644 + # path (`fused_experts_none_to_triton`) is the one that actually fires + # under the default runner config; this branch is here for the future + # case where a runner_input gets routed straight into the runner_core. -+ if os.environ.get("SGL_XPU_ESIMD_MOE", "0") == "1": ++ if _ESIMD_MOE: + out = _try_esimd_moe_silu_routed(runner_input, quant_info, self.config) + if out is not None: + return TritonRunnerOutput(hidden_states=out) @@ -1320,7 +1637,7 @@ index 96e431d4e..f2ba507fc 100644 from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( _fused_moe_kernel_sequence, ) -@@ -136,6 +284,207 @@ class TritonRunnerCore(MoeRunnerCore): +@@ -136,6 +292,207 @@ class TritonRunnerCore(MoeRunnerCore): return MoeRunnerBackend.TRITON @@ -1383,8 +1700,8 @@ index 96e431d4e..f2ba507fc 100644 + s13 = _per_expert_pt_scale(quant_info.w13_scale) + s2 = _per_expert_pt_scale(quant_info.w2_scale) + -+ # Use the `_sglang` kernel variant: accepts w13 in [E, 2*inter, hidden] -+ # directly, no transpose required. ++ # Both e4m3 and e5m2 kernel variants accept w13 in [E, 2*inter, hidden] ++ # (N-major) directly — no transpose / extra copy. + w13_kernel = w13 + + if topk_weights.dtype != torch.float16: @@ -1528,7 +1845,7 @@ index 96e431d4e..f2ba507fc 100644 @register_fused_func("none", "triton") def fused_experts_none_to_triton( dispatch_output: StandardDispatchOutput, -@@ -145,6 +494,30 @@ def fused_experts_none_to_triton( +@@ -145,6 +502,30 @@ def fused_experts_none_to_triton( from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_experts from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput @@ -1540,7 +1857,7 @@ index 96e431d4e..f2ba507fc 100644 + # a single linear layer, so a single scalar per expert is too lossy. + # TODO: migrate the per-block FP8 MoE GEMM kernel from llm-scaler/vllm + # (or feed 2D scale through a future kernel variant) before re-enabling. -+ if os.environ.get("SGL_XPU_ESIMD_MOE", "0") == "1": ++ if _ESIMD_MOE: + esimd_out = _maybe_esimd_moe_silu_fused( + dispatch_output, quant_info, runner_config + ) @@ -1549,7 +1866,7 @@ index 96e431d4e..f2ba507fc 100644 + + # Large-T (prefill) M-tiled DPAS FP8 MoE kernel. Separate env gate so it can + # be enabled independently of the decode kernel. -+ if os.environ.get("SGL_XPU_ESIMD_MOE_PREFILL", "0") == "1": ++ if _ESIMD_MOE_PREFILL: + esimd_out = _maybe_esimd_moe_silu_prefill( + dispatch_output, quant_info, runner_config + ) @@ -1559,8 +1876,35 @@ index 96e431d4e..f2ba507fc 100644 output = fused_experts( hidden_states=dispatch_output.hidden_states, w1=quant_info.w13_weight, +diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py +index b104e40d7..4da941e2e 100644 +--- a/python/sglang/srt/layers/moe/topk.py ++++ b/python/sglang/srt/layers/moe/topk.py +@@ -114,6 +114,11 @@ from sglang.srt.utils import ( + from sglang.srt.utils.patch_torch import register_fake_if_exists + + _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() ++# Benchmark-only simulation flags, constant for the process lifetime. Cached at ++# import so select_experts (called once per MoE layer per decode step) does not ++# re-read os.getenv twice on every call. ++_SIMULATE_UNIFORM_EXPERTS = envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get() ++_SIMULATE_ROUND_ROBIN_EXPERTS = envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get() + + if TYPE_CHECKING: + from sglang.srt.layers.quantization import QuantizationConfig +@@ -1739,8 +1744,8 @@ def select_experts( + renormalize=renormalize, + ) + +- simulate_uniform_experts = envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get() +- simulate_round_robin_experts = envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get() ++ simulate_uniform_experts = _SIMULATE_UNIFORM_EXPERTS ++ simulate_round_robin_experts = _SIMULATE_ROUND_ROBIN_EXPERTS + if simulate_uniform_experts and simulate_round_robin_experts: + raise ValueError( + "SGLANG_SIMULATE_UNIFORM_EXPERTS and " diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py -index b4e05df01..f88e4f1cf 100644 +index b4e05df01..c003da1f1 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -116,6 +116,145 @@ _use_aiter = envs.SGLANG_USE_AITER.get() and _is_hip @@ -1741,7 +2085,23 @@ index b4e05df01..f88e4f1cf 100644 if _is_hip and _use_hip_int4: self.process_weights_hip_int4(layer) -@@ -1807,7 +1953,6 @@ class Fp8MoEMethod(FusedMoEMethodBase): +@@ -1580,6 +1726,15 @@ class Fp8MoEMethod(FusedMoEMethodBase): + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + ++ # NOTE: e5m2 fused decode kernels used to require a transposed copy ++ # of the routed gate_up weight (w13._esimd_e5m2_t = [E, hidden, ++ # 2*inter]), which cost a full extra weight replica per rank ++ # (~10GB/tile at E=256). The e5m2 up kernels now read sglang's ++ # native N-major w13 [E, 2*inter, hidden] directly ++ # (moe_up_routed_e5m2_nmajor_kernel / moe_up_merged_e5m2_nmajor_kernel), ++ # so NO transposed copy is cached. The routed DOWN weight is likewise ++ # consumed in its natural layout ([E, hidden, inter]). ++ + if _is_hip: + self.process_weights_hip_scale_padding(layer) + +@@ -1807,7 +1962,6 @@ class Fp8MoEMethod(FusedMoEMethodBase): layer: torch.nn.Module, dispatch_output: DispatchOutput, ) -> CombineInput: @@ -1749,7 +2109,7 @@ index b4e05df01..f88e4f1cf 100644 from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput x = dispatch_output.hidden_states -@@ -1860,9 +2005,13 @@ class Fp8MoEMethod(FusedMoEMethodBase): +@@ -1860,9 +2014,13 @@ class Fp8MoEMethod(FusedMoEMethodBase): topk_weights, topk_ids, _ = dispatch_output.topk_output assert layer.w13_weight.dtype == layer.w2_weight.dtype @@ -1765,7 +2125,7 @@ index b4e05df01..f88e4f1cf 100644 x, layer.w13_weight, diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py -index 30a035a23..b10f39fbd 100644 +index 30a035a23..c82c83ed6 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -40,6 +40,7 @@ from sglang.srt.utils import ( @@ -1824,11 +2184,61 @@ index 30a035a23..b10f39fbd 100644 fp8_max = torch.finfo(fp8_dtype).max fp8_min = -fp8_max +@@ -1753,6 +1779,41 @@ else: + shape = (max(num_token_padding, input.shape[0]), shape[1]) + output = torch.empty(shape, device=input.device, dtype=fp8_dtype) + ++ # The sgl_kernel C++ quant ops (sgl_per_tensor_quant_fp8 / ++ # sgl_per_token_quant_fp8) hardcode a Float8_e4m3fn output tensor and ++ # raise "output must be Float8_e4m3fn tensor" for any other fp8 dtype. ++ # When the global weight dtype is switched to e5m2 (SGLANG_FP8_DTYPE=e5m2) ++ # fall back to a native torch quantization that honours fp8_dtype. ++ if fp8_dtype != torch.float8_e4m3fn: ++ m = shape[0] ++ src = input if num_token_padding is None else input[:m] ++ if scale is None and use_per_token_if_dynamic: ++ scale = torch.empty( ++ (shape[0], 1), device=input.device, dtype=torch.float32 ++ ) ++ absmax = src.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12) ++ s = absmax / fp8_max ++ scale[:m].copy_(s) ++ output[:m].copy_( ++ (src / s).clamp(fp8_min, fp8_max).to(fp8_dtype) ++ ) ++ elif scale is None: ++ scale = torch.zeros(1, device=input.device, dtype=torch.float32) ++ absmax = src.abs().amax().clamp(min=1e-12) ++ s = absmax / fp8_max ++ scale.copy_(s.view(-1)) ++ output[:m].copy_( ++ (src / s).clamp(fp8_min, fp8_max).to(fp8_dtype) ++ ) ++ else: ++ assert ( ++ scale.numel() == 1 ++ ), f"Expected scalar scale, got numel={scale.numel()}" ++ output[:m].copy_( ++ (src / scale).clamp(fp8_min, fp8_max).to(fp8_dtype) ++ ) ++ return output, scale ++ + if scale is None: + # Dynamic scaling + if use_per_token_if_dynamic: diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py -index 6c6827c7d..35b948c62 100755 +index 6c6827c7d..33ba1e5de 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py -@@ -46,6 +46,7 @@ from sglang.srt.utils import ( +@@ -1,6 +1,7 @@ + from __future__ import annotations + + import logging ++import os + from enum import Enum + from functools import lru_cache + from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Union +@@ -46,6 +47,7 @@ from sglang.srt.utils import ( is_sm90_supported, is_sm100_supported, is_sm120_supported, @@ -1836,7 +2246,7 @@ index 6c6827c7d..35b948c62 100755 offloader, ) from sglang.srt.utils.custom_op import register_custom_op -@@ -59,6 +60,19 @@ _is_sm100_supported = is_sm100_supported() +@@ -59,6 +61,22 @@ _is_sm100_supported = is_sm100_supported() _is_sm120_supported = is_sm120_supported() _is_gfx95_supported = is_gfx95_supported() _is_musa = is_musa() @@ -1845,6 +2255,9 @@ index 6c6827c7d..35b948c62 100755 +# Lazy-loaded handle to the merged custom_esimd_kernels_sglang +# esimd_gemm_fp8_pert kernel wrapper. +# Only initialised on XPU; None elsewhere or if the package is missing. ++# Output width below which the ESIMD kernel is preferred at any M, because ++# torch._scaled_mm is not reproducible on narrow-N shapes (see apply_fp8_linear). ++_XPU_ESIMD_FP8_NARROW_N = int(os.environ.get("SGL_XPU_ESIMD_FP8_NARROW_N", "128")) +_esimd_gemm_fp8_pert = None +if _is_xpu: + try: @@ -1856,7 +2269,17 @@ index 6c6827c7d..35b948c62 100755 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _use_aiter_gfx95 = _use_aiter and _is_gfx95_supported -@@ -188,6 +202,7 @@ class Fp8GemmRunnerBackend(Enum): +@@ -121,6 +139,9 @@ if _is_cuda: + + + use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL") ++# Cached once at import: read on every apply_fp8_linear call (100+/decode step) ++# otherwise, and the value is constant for the process lifetime. ++_enable_torch_compile = get_bool_env_var("SGLANG_ENABLE_TORCH_COMPILE") + + # Input scaling factors are no longer optional in _scaled_mm starting + # from pytorch 2.5. Allocating a dummy tensor to pass as input_scale +@@ -188,6 +209,7 @@ class Fp8GemmRunnerBackend(Enum): DEEP_GEMM = "deep_gemm" TRITON = "triton" AITER = "aiter" @@ -1864,7 +2287,7 @@ index 6c6827c7d..35b948c62 100755 def is_auto(self) -> bool: return self == Fp8GemmRunnerBackend.AUTO -@@ -213,6 +228,9 @@ class Fp8GemmRunnerBackend(Enum): +@@ -213,6 +235,9 @@ class Fp8GemmRunnerBackend(Enum): def is_aiter(self) -> bool: return self == Fp8GemmRunnerBackend.AITER @@ -1874,7 +2297,7 @@ index 6c6827c7d..35b948c62 100755 FP8_GEMM_RUNNER_BACKEND: Fp8GemmRunnerBackend | None = None -@@ -439,6 +457,17 @@ def _dispatch_explicit_backend(backend: Fp8GemmRunnerBackend) -> Callable: +@@ -439,6 +464,17 @@ def _dispatch_explicit_backend(backend: Fp8GemmRunnerBackend) -> Callable: elif backend.is_triton(): return triton_w8a8_block_fp8_linear @@ -1892,7 +2315,7 @@ index 6c6827c7d..35b948c62 100755 else: raise ValueError(f"Unknown FP8 GEMM backend: {backend}") -@@ -832,6 +861,54 @@ def triton_w8a8_block_fp8_linear( +@@ -832,6 +868,55 @@ def triton_w8a8_block_fp8_linear( input_2d = input.view(-1, input.shape[-1]) output_shape = [*input.shape[:-1], weight.shape[0]] @@ -1942,12 +2365,13 @@ index 6c6827c7d..35b948c62 100755 + N = weight_nk.shape[0] + output = torch.empty(M, N, dtype=torch.float16, device=input_2d.device) + _esimd_gemm_fp8_pert(input_fp16, weight_nk, scale_pt, output) -+ return output.to(input.dtype).view(*output_shape) ++ out = output if output.dtype == input.dtype else output.to(input.dtype) ++ return out.view(*output_shape) + q_input, x_scale = per_token_group_quant_fp8( input_2d, block_size[1], column_major_scales=False ) -@@ -843,6 +920,50 @@ def triton_w8a8_block_fp8_linear( +@@ -843,6 +928,50 @@ def triton_w8a8_block_fp8_linear( return output.to(dtype=input_2d.dtype).view(*output_shape) @@ -1998,7 +2422,73 @@ index 6c6827c7d..35b948c62 100755 @lru_cache(maxsize=1) def _get_triton_mxfp8_downcast(): try: -@@ -1491,6 +1612,46 @@ def apply_fp8_linear( +@@ -1444,6 +1573,55 @@ def _apply_fallback_scaled_mm( + input_dtype, + ): + global TORCH_DEVICE_IDENTITY ++ ++ # XPU fast path: fuse the (X*W) dequant scaling into torch._scaled_mm's ++ # RowWise epilogue instead of running an identity-scaled fp32 GEMM followed ++ # by two fp32 elementwise multiplies and an fp32->fp16 cast. XPU _scaled_mm ++ # RowWise requires scale_a=(M,1) and scale_b=(1,N), both contiguous float; ++ # per-tensor weight scales are broadcast up to (1,N). Bit-exact vs the ++ # unfused path (cos=1.0) and ~3x faster on 4k prefill. Falls back on any ++ # shape/dtype mismatch. ++ if _is_xpu: ++ try: ++ N = weight.shape[1] ++ sa = x_scale ++ if sa.dim() == 1: ++ sa = sa.reshape(-1, 1) ++ if ( ++ sa.dim() == 2 ++ and sa.shape[0] == qinput.shape[0] ++ and sa.shape[1] == 1 ++ and sa.dtype == torch.float32 ++ and sa.is_contiguous() ++ ): ++ # scale_b=(1,N) is constant per weight — build once and cache on ++ # the weight_scale tensor (mirrors the _esimd_1d caching below). ++ sb = getattr(weight_scale, "_rowwise_1N", None) ++ if sb is None or sb.shape[1] != N: ++ sb = weight_scale.reshape(1, -1).to(torch.float32) ++ if sb.shape[1] == 1: ++ sb = sb.expand(1, N) ++ sb = sb.contiguous() ++ if sb.shape[1] == N: ++ try: ++ weight_scale._rowwise_1N = sb ++ except Exception: ++ pass ++ if sb.shape[1] == N: ++ output = torch._scaled_mm( ++ qinput, ++ weight, ++ scale_a=sa, ++ scale_b=sb, ++ out_dtype=input_dtype, ++ bias=bias, ++ ) ++ return _process_scaled_mm_output( ++ output, input_2d_shape, output_shape ++ ) ++ except Exception: ++ pass # fall through to the portable unfused path below ++ + if TORCH_DEVICE_IDENTITY is None: + TORCH_DEVICE_IDENTITY = torch.ones(1, dtype=torch.float32, device=weight.device) + +@@ -1482,15 +1660,62 @@ def apply_fp8_linear( + # We also don't pad when using torch.compile, + # as it breaks with dynamic shapes. + if pad_output is None: +- pad_output = not cutlass_fp8_supported and not get_bool_env_var( +- "SGLANG_ENABLE_TORCH_COMPILE" +- ) ++ pad_output = not cutlass_fp8_supported and not _enable_torch_compile + output_padding = 17 if pad_output else None + + # View input as 2D matrix for fp8 methods input_2d = input.view(-1, input.shape[-1]) output_shape = [*input.shape[:-1], weight.shape[1]] @@ -2009,6 +2499,14 @@ index 6c6827c7d..35b948c62 100755 + # Restricted to small M (M<=64) since the kernel's M>=64 weight-stationary + # path is much slower than torch._scaled_mm (verified: 14000us vs 156us at + # M=4096). Triggered for decode (M=1) and small chunked-prefill batches. ++ # ++ # Narrow-N GEMMs take the ESIMD path at any M as well: torch._scaled_mm ++ # splits such shapes along K and reduces the partials in an order that ++ # varies with GPU occupancy, so the same operands can yield two different ++ # results (observed on GDN in_proj_ba, N=32, on prefill tail chunks with ++ # M>64 -- it made prefill non-deterministic). N is tiny there, so the ++ # weight-stationary path costs almost nothing at these shapes. ++ _n_out = weight.shape[1] + if ( + _is_xpu + and _esimd_gemm_fp8_pert is not None @@ -2016,7 +2514,7 @@ index 6c6827c7d..35b948c62 100755 + and not (cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]) + and (weight_scale.numel() == 1) + and bias is None -+ and input_2d.shape[0] <= 64 ++ and (input_2d.shape[0] <= 64 or _n_out <= _XPU_ESIMD_FP8_NARROW_N) + ): + weight_nk = getattr(weight, "_esimd_t", None) + if weight_nk is None: @@ -2040,13 +2538,14 @@ index 6c6827c7d..35b948c62 100755 + N = weight_nk.shape[0] + output = torch.empty(M, N, dtype=torch.float16, device=input_2d.device) + _esimd_gemm_fp8_pert(input_fp16, weight_nk, scale_1d, output) -+ return output.to(input.dtype).view(*output_shape) ++ out = output if output.dtype == input.dtype else output.to(input.dtype) ++ return out.view(*output_shape) + if compressed_tensor_quant: # Maybe apply padding to output, see comment in __init__ num_token_padding = output_padding diff --git a/python/sglang/srt/layers/quantization/gguf.py b/python/sglang/srt/layers/quantization/gguf.py -index 5c1e7a63e..69ccc3646 100644 +index 5c1e7a63e..9ae9c0a6d 100644 --- a/python/sglang/srt/layers/quantization/gguf.py +++ b/python/sglang/srt/layers/quantization/gguf.py @@ -1,9 +1,9 @@ @@ -2160,7 +2659,7 @@ index 5c1e7a63e..69ccc3646 100644 return GGUFMoEMethod(self) return None -@@ -456,6 +534,1318 @@ class GGUFLinearMethod(LinearMethodBase): +@@ -456,6 +534,1425 @@ class GGUFLinearMethod(LinearMethodBase): return out @@ -2261,6 +2760,9 @@ index 5c1e7a63e..69ccc3646 100644 + return dequant.view(*x.shape, hidden_size) + + ++_XPU_EMB_DEBUG = os.environ.get("SGLANG_GGUF_EMB_DEBUG", "0") == "1" ++ ++ +class GGUFEmbeddingXPUMethod(GGUFLinearMethod): + """GGUF embedding for Intel XPU (PTL Xe3). + @@ -2285,18 +2787,35 @@ index 5c1e7a63e..69ccc3646 100644 + layer._xpu_emb_rep = _xpu_prepare_shard( + layer.qweight.data, int(qweight_type), self.params_dtype + ) ++ if _XPU_EMB_DEBUG: ++ rep = layer._xpu_emb_rep ++ msg = ["kind=%s" % rep[0]] ++ for i, t in enumerate(rep[1:]): ++ if not torch.is_tensor(t): ++ continue ++ if t.dtype.is_floating_point: ++ nb = int((~torch.isfinite(t)).sum().item()) ++ msg.append( ++ "t%d %s %s nonfinite=%d absmax=%.6g" ++ % (i, tuple(t.shape), t.dtype, nb, ++ float(t.abs().float().max().item())) ++ ) ++ else: ++ msg.append("t%d %s %s" % (i, tuple(t.shape), t.dtype)) ++ logger.error("[EMBDEBUG] rep health: %s", " | ".join(msg)) + if hasattr(layer, "qweight"): + del layer.qweight + + def embedding(self, layer: torch.nn.Module, x: torch.Tensor) -> torch.Tensor: -+ debug_bounds = False ++ debug_bounds = _XPU_EMB_DEBUG + + def _log_nonfinite_once(flag_name: str, stage: str, t: torch.Tensor): + if not debug_bounds or t.numel() == 0 or bool(torch.isfinite(t).all().item()): + return -+ if getattr(layer, flag_name, False): ++ n = getattr(layer, flag_name, 0) ++ if n >= 8: + return -+ setattr(layer, flag_name, True) ++ setattr(layer, flag_name, n + 1) + row_bad = ~torch.isfinite(t.reshape(t.shape[0], -1)).all(dim=1) + bad_rows = row_bad.nonzero(as_tuple=False).flatten() + bad_rows_head = bad_rows[:16].tolist() @@ -2307,16 +2826,26 @@ index 5c1e7a63e..69ccc3646 100644 + and x_flat_dbg.numel() >= int(bad_rows.max().item()) + 1 + else [] + ) ++ # Element-level (not whole-row) badness: report the exact flat ++ # positions so a partial-row corruption is distinguishable from a ++ # bad table row. ++ bad_flat = (~torch.isfinite(t)).nonzero(as_tuple=False) + logger.error( + "GGUFEmbeddingXPUMethod non-finite at %s: dtype=%s shape=%s " -+ "nan=%d inf=%d bad_rows_head=%s bad_token_ids_head=%s", ++ "nan=%d inf=%d n_bad_rows=%d bad_rows_head=%s bad_token_ids_head=%s " ++ "bad_elem_head=%s id_min=%d id_max=%d ntok=%d", + stage, + str(t.dtype), + tuple(t.shape), + int(torch.isnan(t).sum().item()), + int(torch.isinf(t).sum().item()), ++ int(bad_rows.numel()), + bad_rows_head, + bad_ids_head, ++ bad_flat[:8].tolist(), ++ int(x_flat_dbg.min().item()) if x_flat_dbg.numel() else -1, ++ int(x_flat_dbg.max().item()) if x_flat_dbg.numel() else -1, ++ int(x_flat_dbg.numel()), + ) + + # Eagle/NEXTN embed-share: set_embed_and_head may have replaced this @@ -2464,8 +2993,10 @@ index 5c1e7a63e..69ccc3646 100644 + and M == 1 + and rep[0] in ("q6_k", "q5_k", "q4_k", "q8_0", "q4_0") + ): -+ xf = x2.to(torch.float16).contiguous() -+ out = _xpu_rep_gemv(xf, rep).to(x.dtype) # [1, vocab] ++ xf = _as_fp16c(x2) ++ out = _xpu_rep_gemv(xf, rep) # [1, vocab] fp16 ++ if out.dtype != x.dtype: ++ out = out.to(x.dtype) + if bias is not None: + out = out + bias + return out.reshape(*x.shape[:-1], out.shape[-1]) @@ -2482,11 +3013,12 @@ index 5c1e7a63e..69ccc3646 100644 + and 2 <= M <= 16 + and esimd_gemv_q6_k_m is not None + ): -+ xf = x2.to(torch.float16).contiguous() ++ xf = _as_fp16c(x2) + N = rep[1].shape[0] + out = torch.empty(M, N, dtype=torch.float16, device=xf.device) + esimd_gemv_q6_k_m(xf, rep[1], rep[2], rep[3], out) -+ out = out.to(x.dtype) ++ if out.dtype != x.dtype: ++ out = out.to(x.dtype) + if bias is not None: + out = out + bias + return out.reshape(*x.shape[:-1], out.shape[-1]) @@ -2522,6 +3054,13 @@ index 5c1e7a63e..69ccc3646 100644 +_Q4_K_BYTES = 144 # half2 dm(4) + scales[12] + qs[128] +_MOE_DOWN_REPACK_CHUNK_ROWS = int( + os.environ.get("SGLANG_GGUF_XPU_MOE_REPACK_CHUNK_ROWS", "65536")) ++# Cap on the transient int32 intermediate produced per repack chunk. The q5/q6 ++# repack path materializes several [rows, K] int32 tensors at once, so a fixed ++# row count scales badly with K: the 248320x2048 embedding needs ~1.5-2 GB per ++# 65536-row chunk, which OOMs at TP=1 (the whole table lives on one card). ++# Chunk rows are derived from this budget so wide tensors are split finer. ++_REPACK_CHUNK_BUDGET_BYTES = int( ++ os.environ.get("SGLANG_GGUF_XPU_REPACK_CHUNK_BYTES", str(32 << 20))) + + +def _xpu_repack_q4_k(qweight: torch.Tensor): @@ -2606,7 +3145,7 @@ index 5c1e7a63e..69ccc3646 100644 + + +def _xpu_repack_rows_chunked(repack_fn, qweight: torch.Tensor, -+ chunk_rows: int = _MOE_DOWN_REPACK_CHUNK_ROWS, ++ chunk_rows: int = None, + **kwargs): + """Generic row-chunked wrapper for any of the _xpu_repack_q{4,5,6}_k + functions (col_perm only reorders the K/column dim, so it composes @@ -2614,8 +3153,18 @@ index 5c1e7a63e..69ccc3646 100644 + lm_head weight, vocab_size rows) where a single whole-tensor repack call + creates multi-GB int32 intermediates and was the cause of a second + TP=1-loading OOM (in _xpu_prepare_shard, downstream of -+ GGUFEmbeddingXPUMethod.process_weights_after_loading).""" ++ GGUFEmbeddingXPUMethod.process_weights_after_loading). ++ ++ chunk_rows defaults to whatever keeps one chunk's int32 intermediates ++ under _REPACK_CHUNK_BUDGET_BYTES, so narrow tensors keep using large ++ chunks while wide ones (big K) are split finer.""" + N = qweight.shape[0] ++ if chunk_rows is None: ++ # qweight rows are packed bytes; the repack blows them up to int32 ++ # elements, so estimate the per-row cost from the byte width. ++ row_bytes = max(1, int(qweight.shape[1])) * 4 ++ chunk_rows = max(1024, _REPACK_CHUNK_BUDGET_BYTES // row_bytes) ++ chunk_rows = min(chunk_rows, _MOE_DOWN_REPACK_CHUNK_ROWS) + if N <= chunk_rows: + return repack_fn(qweight, **kwargs) + outs = None @@ -3143,6 +3692,31 @@ index 5c1e7a63e..69ccc3646 100644 + return st + + ++def _as_fp16c(x: torch.Tensor) -> torch.Tensor: ++ """Return x as a contiguous fp16 tensor WITHOUT emitting a no-op ++ ``aten::to`` / ``aten::contiguous`` dispatch when x is already fp16 and ++ contiguous (the common ``--dtype float16`` case). ++ ++ At M=1 decode the GGUF path is host-dispatch bound; the per-projection ++ ``x.to(torch.float16).contiguous()`` fired ~110 no-op casts/step (plus the ++ matching output cast). ``is_contiguous()`` is a cheap C++ property check, ++ not a device op, so this guard is pure host-overhead reduction and is ++ bit-identical (returns the same storage) for fp16 inputs. bf16/fp32 ++ networks still get a real cast + contiguous copy. ++ """ ++ if x.dtype == torch.float16: ++ return x if x.is_contiguous() else x.contiguous() ++ return x.to(torch.float16).contiguous() ++ ++ ++# Per-weight cache for the fp16-resident dense shards (GDN b/a in_proj_ba): ++# id(weight) -> transposed_contiguous_weight [K, N] ++# Reps are held by the layer for the model's lifetime so id(weight) is stable ++# and this dict is bounded by the number of fp16 shards (~60). See the fp16 ++# branch of _xpu_shard_matmul. ++_fp16_wt_cache: dict = {} ++ ++ +def _xpu_shard_matmul(x: torch.Tensor, rep) -> torch.Tensor: + """x [M,K] fp16 @ shard^T -> [M,N] fp16. rep from _xpu_prepare_shard.""" + kind = rep[0] @@ -3150,7 +3724,7 @@ index 5c1e7a63e..69ccc3646 100644 + _, packed, scale = rep + N = packed.shape[0] + M = x.shape[0] -+ xf = x.to(torch.float16).contiguous() ++ xf = _as_fp16c(x) + if M == 1: + # Decode: the ESIMD GEMV is bandwidth-optimal (~3x faster than a + # dense fp16 matmul at M=1). @@ -3170,7 +3744,7 @@ index 5c1e7a63e..69ccc3646 100644 + _, qs, scale = rep + N = qs.shape[0] + M = x.shape[0] -+ xf = x.to(torch.float16).contiguous() ++ xf = _as_fp16c(x) + if M == 1: + # Decode: ESIMD q8_0 GEMV (int8 resident, bandwidth-optimal). + out = torch.empty(M, N, dtype=torch.float16, device=x.device) @@ -3202,7 +3776,7 @@ index 5c1e7a63e..69ccc3646 100644 + _, ql, scale, minv = rep + N = ql.shape[0] + M = x.shape[0] -+ xf = x.to(torch.float16).contiguous() ++ xf = _as_fp16c(x) + if M == 1: + # Decode: ESIMD q4_K GEMV (4.5-bit resident, asymmetric scale+min). + out = torch.empty(M, N, dtype=torch.float16, device=x.device) @@ -3215,7 +3789,7 @@ index 5c1e7a63e..69ccc3646 100644 + _, ql, qh, scale, minv = rep + N = ql.shape[0] + M = x.shape[0] -+ xf = x.to(torch.float16).contiguous() ++ xf = _as_fp16c(x) + if M == 1: + out = torch.empty(M, N, dtype=torch.float16, device=x.device) + esimd_gemv_q5_k(xf, ql, qh, scale, minv, out) @@ -3226,16 +3800,28 @@ index 5c1e7a63e..69ccc3646 100644 + _, ql, qh, scale = rep + N = ql.shape[0] + M = x.shape[0] -+ xf = x.to(torch.float16).contiguous() ++ xf = _as_fp16c(x) + if M == 1: + out = torch.empty(M, N, dtype=torch.float16, device=x.device) + esimd_gemv_q6_k(xf, ql, qh, scale, out) + return out + w = _xpu_dequant_q6_k(ql, qh, scale, torch.float16) # [N, K] + return xf @ w.t() -+ # fp16-resident dense weight [N, K] -+ _, w, _ = rep -+ return x.to(w.dtype) @ w.t() ++ # fp16-resident dense weight [N, K]. Hit by the GDN b/a (beta/decay) shards ++ # of in_proj_ba (each [num_v_heads, hidden] fp16, unquantized in the GGUF), ++ # ~60 GEMVs/decode-step. The naive ``x.to(w.dtype) @ w.t()`` fired ~8 host ++ # dispatches each (no-op cast + t/transpose/as_strided weight-transpose VIEW ++ # rebuilt per call + matmul-wrapper + mm + output empty/resize_). Cache the ++ # contiguous transpose [K, N] once so each call is a single ``torch.mm`` on ++ # two contiguous operands. Bit-identical (same operands, same math). A fresh ++ # output is allocated per call (no persistent-buffer aliasing risk). ++ w = rep[1] ++ xf = x if x.dtype == w.dtype else x.to(w.dtype) ++ wt = _fp16_wt_cache.get(id(w)) ++ if wt is None: ++ wt = w.t().contiguous() # [K, N], computed once per weight ++ _fp16_wt_cache[id(w)] = wt ++ return torch.mm(xf, wt) + + +def _xpu_permute_gdn_out_cols(rep, perm): @@ -3336,14 +3922,30 @@ index 5c1e7a63e..69ccc3646 100644 + """D1: if every shard in `ids` is the SAME GEMV rep kind (q8_0 or q4_0) with + the same K, build one merged rep by row-concatenating the per-shard weights, + plus the per-shard N sizes (to slice the output). Returns -+ (merged_rep, [N0, N1, ...]) or None if not mergeable (mixed kinds / fp16 / -+ k-quant — those keep the per-shard path). Bit-exact: q8_0/q4_0 rows are ++ (merged_rep, [N0, N1, ...]) or None if not mergeable (mixed kinds / ++ k-quant — those keep the per-shard path). Bit-exact: q8_0/q4_0/fp16 rows are + independent, so cat-then-GEMV == per-shard-GEMV-then-cat. + notes §10bj.""" + kinds = {reps[i][0] for i in ids} + if len(ids) < 2 or len(kinds) != 1: + return None + kind = next(iter(kinds)) ++ if kind == "fp16": ++ # D2: the GDN in_proj_ba shards (ssm_beta/ssm_alpha) are unquantized F32 ++ # -> two fp16 dense reps of N=num_v_heads/tp. Unmerged they cost 2 tiny ++ # oneDNN mm (25us host each) + a torch.cat (62us host) per layer, i.e. ++ # ~3.4ms/step at 30 GDN layers — pure launch overhead, the GEMVs are 64KB. ++ # Row-cat them into one [sumN, K] weight so apply() issues a single mm ++ # and the cat disappears. Bit-exact (dense rows are independent). ++ Ks = {reps[i][1].shape[1] for i in ids if reps[i][1].dim() == 2} ++ if len(Ks) != 1: ++ return None ++ # An unloaded/empty shard must keep the per-shard path (shape is bogus). ++ if any(reps[i][1].dim() != 2 or reps[i][1].numel() == 0 for i in ids): ++ return None ++ sizes = [reps[i][1].shape[0] for i in ids] ++ w = torch.cat([reps[i][1] for i in ids], dim=0).contiguous() ++ return (("fp16", w, None), sizes) + if kind == "q8_0": + # rep = ("q8_0", qs[N,K] int8, scale[N,K/32] f16) + Ks = {reps[i][1].shape[1] for i in ids} @@ -3455,7 +4057,8 @@ index 5c1e7a63e..69ccc3646 100644 + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: -+ x2 = x.reshape(-1, x.shape[-1]) ++ # At M=1 decode x is already 2D [1, K]; skip the no-op reshape dispatch. ++ x2 = x if x.dim() == 2 else x.reshape(-1, x.shape[-1]) + merged = getattr(layer, "_xpu_merged", None) + if merged is not None: + # D1: one big-N GEMV over the merged q8_0/q4_0 shards (§10bj). The @@ -3470,16 +4073,19 @@ index 5c1e7a63e..69ccc3646 100644 + # The q4_0 ESIMD kernels are fp16-only (PTL has no bf16 ESIMD), so a + # bf16 network would otherwise get an fp16 tensor back here. Cast the + # result to the input activation dtype to keep the graph type-consistent. -+ out = out.to(x.dtype) ++ # Skip the no-op cast on fp16 networks (elides one host dispatch/proj). ++ if out.dtype != x.dtype: ++ out = out.to(x.dtype) + if bias is not None: + out = out + bias -+ return out.reshape(*x.shape[:-1], out.shape[-1]) ++ # x2 aliased x when x was 2D, so out is already [M, N] — skip no-op reshape. ++ return out if x.dim() == 2 else out.reshape(*x.shape[:-1], out.shape[-1]) + + class GGUFMoEMethod(FusedMoEMethodBase): """MoE method for GGUF. -@@ -565,6 +1955,253 @@ class GGUFMoEMethod(FusedMoEMethodBase): +@@ -565,6 +2062,260 @@ class GGUFMoEMethod(FusedMoEMethodBase): return StandardCombineInput(hidden_states=output) @@ -3686,15 +4292,17 @@ index 5c1e7a63e..69ccc3646 100644 + x = dispatch_output.hidden_states + topk_weights, topk_ids, _ = dispatch_output.topk_output + -+ x2 = x.reshape(-1, x.shape[-1]) ++ # At M=1 decode x is already 2D [1, K]; skip the no-op reshape dispatch. ++ x2 = x if x.dim() == 2 else x.reshape(-1, x.shape[-1]) + M = x2.shape[0] -+ out = torch.zeros_like(x2) + if M == 0: -+ return StandardCombineInput(hidden_states=out.reshape_as(x)) ++ z = torch.zeros_like(x2) ++ return StandardCombineInput( ++ hidden_states=z if x.dim() == 2 else z.reshape_as(x)) + top_k = topk_ids.shape[1] + hidden, inter = self.hidden, self.intermediate + n_routed = M * top_k -+ xf = x2.to(torch.float16).contiguous() ++ xf = _as_fp16c(x2) + sel = topk_ids.reshape(-1).to(torch.int32).contiguous() + tw = topk_weights.reshape(-1).to(torch.float16).contiguous() + @@ -3709,7 +4317,10 @@ index 5c1e7a63e..69ccc3646 100644 + self.up_ql, self.up_sc, self.up_mn, + self.down_ql, self.down_qh_plain, self.down_sc, self.down_mn, + down_is_q6=self._down_is_q6) -+ return StandardCombineInput(hidden_states=out_g.to(out.dtype).reshape_as(x)) ++ if out_g.dtype != x2.dtype: ++ out_g = out_g.to(x2.dtype) ++ return StandardCombineInput( ++ hidden_states=out_g if x.dim() == 2 else out_g.reshape_as(x)) + + # Fused: 1 up launch (gate/up Q4_K + silu*up) + 1 down launch (Q5_K/Q6_K + # weighted) over ALL routed pairs, then sum the top_k partials. Replaces @@ -3726,8 +4337,10 @@ index 5c1e7a63e..69ccc3646 100644 + esimd_moe_down_q5k(inter_buf, self.down_ql, self.down_qh_plain, self.down_sc, + self.down_mn, sel, tw, out_partial, M, hidden, inter, top_k) + # sum the top_k per-route partials back to per-token output (one op). -+ out = out_partial.view(M, top_k, hidden).sum(dim=1).to(out.dtype) -+ return StandardCombineInput(hidden_states=out.reshape_as(x)) ++ summed = out_partial.view(M, top_k, hidden).sum(dim=1) ++ out = summed if summed.dtype == x2.dtype else summed.to(x2.dtype) ++ return StandardCombineInput( ++ hidden_states=out if x.dim() == 2 else out.reshape_as(x)) + + class GGUFEmbeddingMethod(GGUFLinearMethod): @@ -3924,7 +4537,7 @@ index ebbfc2b28..68891945d 100644 if self.model_runner.server_args.debug_cuda_graph: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index 3b30eb0e1..141eb11a4 100644 +index 3b30eb0e1..93c745eef 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -849,6 +849,19 @@ class ModelRunner(ModelRunnerKVCacheMixin): @@ -3947,6 +4560,90 @@ index 3b30eb0e1..141eb11a4 100644 elif current_platform.is_out_of_tree(): self.init_attention_backend() if current_platform.support_cuda_graph(): +@@ -1273,6 +1286,73 @@ class ModelRunner(ModelRunnerKVCacheMixin): + ), + ) + ++ def _install_layer_trace_hooks(self): ++ """Dump a per-module fingerprint of every forward output. ++ ++ Used to locate where prefill indeterminism first appears: run the same ++ prompt twice, then diff the two traces to find the first module whose ++ output changes. Enabled with SGLANG_XPU_LAYER_TRACE=. ++ """ ++ import json as _json ++ import zlib as _zlib ++ ++ prefix = os.environ["SGLANG_XPU_LAYER_TRACE"] ++ rank = getattr(self, "tp_rank", 0) ++ path = f"{prefix}.rank{rank}.jsonl" ++ # Truncate any trace left over from a previous server instance. ++ open(path, "w").close() ++ ++ def fingerprint(t): ++ if not isinstance(t, torch.Tensor) or t.numel() == 0: ++ return None ++ f = t.detach().to(torch.float64) ++ # A sum/abs-sum pair can collide (any permutation preserves both), ++ # so carry a byte-exact CRC of the raw values as well: "input is ++ # identical" has to be a bit-level claim, not a statistical one. ++ try: ++ b = t.detach().contiguous().cpu().numpy().tobytes() ++ crc = _zlib.crc32(b) & 0xFFFFFFFF ++ except Exception: ++ crc = None ++ return [float(f.sum().item()), float(f.abs().sum().item()), ++ list(t.shape), crc] ++ ++ def make_hook(name): ++ def hook(_mod, inp, out): ++ first = out ++ if isinstance(out, (tuple, list)): ++ first = next((o for o in out ++ if isinstance(o, torch.Tensor)), None) ++ fp = fingerprint(first) ++ if fp is None: ++ return ++ if name == state["first"]: ++ state["step"] += 1 ++ rec = {"name": name, "fp": fp, "s": state["step"]} ++ # Input fingerprint too: an output-only trace cannot tell a ++ # module that was fed different data from one that is itself ++ # non-deterministic. ++ src = inp[0] if isinstance(inp, (tuple, list)) and inp else inp ++ ifp = fingerprint(src) ++ if ifp is not None: ++ rec["ifp"] = ifp ++ with open(path, "a") as fh: ++ fh.write(_json.dumps(rec) + "\n") ++ return hook ++ ++ state = {"step": -1, "first": None} ++ ++ want = os.environ.get("SGLANG_XPU_LAYER_TRACE_FILTER", "layers.") ++ n = 0 ++ for name, mod in self.model.named_modules(): ++ if want and want not in name: ++ continue ++ if state["first"] is None: ++ state["first"] = name ++ mod.register_forward_hook(make_hook(name)) ++ n += 1 ++ logger.warning("[LAYER_TRACE] hooked %d modules -> %s", n, path) ++ + def load_model(self): + tic_total = time.perf_counter() + before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id) +@@ -1383,6 +1463,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): + pyt_hooks = PytHooks() + pyt_hooks.register_hooks(self.model, module_prefix="model") + ++ if os.environ.get("SGLANG_XPU_LAYER_TRACE"): ++ self._install_layer_trace_hooks() ++ + if self.server_args.kv_cache_dtype == "fp8_e4m3": + if self.server_args.quantization_param_path is not None: + if callable(getattr(self.model, "load_kv_cache_scales", None)): diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 32ff5bac2..6b7d1eff6 100644 --- a/python/sglang/srt/model_loader/loader.py @@ -4341,132 +5038,1742 @@ index 32ff5bac2..6b7d1eff6 100644 # Check for FLASH_RL format early # FP8 approach: BF16/FP16 model with native FP8 quantization if load_config.load_format == LoadFormat.FLASH_RL: -diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py -index 65f78e691..bf045b854 100644 ---- a/python/sglang/srt/models/qwen3_5.py -+++ b/python/sglang/srt/models/qwen3_5.py -@@ -15,6 +15,7 @@ - """Inference-only Qwen3.5 model and Qwen3.5 MoE model compatible with HuggingFace weights.""" +diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py +index 37552404b..74b209839 100644 +--- a/python/sglang/srt/models/qwen2_moe.py ++++ b/python/sglang/srt/models/qwen2_moe.py +@@ -19,6 +19,7 @@ + """Inference-only Qwen2MoE model compatible with HuggingFace weights.""" import logging +import os - from functools import lru_cache - from typing import Iterable, Optional, Set, Tuple, Union + from contextlib import nullcontext + from typing import Any, Dict, Iterable, List, Optional, Tuple, Union -@@ -93,6 +94,7 @@ from sglang.srt.utils import ( - is_gfx95_supported, - is_hip, - is_npu, -+ is_xpu, - make_layers, - set_weight_attrs, +@@ -34,6 +35,7 @@ from sglang.srt.distributed import ( + get_pp_group, + get_pp_indices, + get_tensor_model_parallel_world_size, ++ attention_tensor_model_parallel_all_reduce, + tensor_model_parallel_all_reduce, ) -@@ -102,6 +104,7 @@ logger = logging.getLogger(__name__) + from sglang.srt.distributed.parallel_state import ( +@@ -116,6 +118,39 @@ _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() + + logger = logging.getLogger(__name__) + ++# ── NaN/Inf probe (env-gated: SGLANG_NAN_PROBE=1) ─────────────────────────── ++# Isolates whether a non-finite MoE output originates in the routed experts ++# (e.g. GGUF grouped prefill GGEMV at batch>8) or the shared expert. Zero ++# overhead when SGLANG_NAN_PROBE is unset. See qwen3_5.py for the layer-level ++# probe that pins down the origin layer / prefill-vs-decode / token count. ++import os as _os_np_moe ++ ++_NAN_PROBE_MOE = _os_np_moe.environ.get("SGLANG_NAN_PROBE", "0") == "1" ++ ++ ++def _nan_probe_moe(tag, t, layer_id=None, forward_batch=None): ++ if not _NAN_PROBE_MOE or not torch.is_tensor(t) or t.numel() == 0: ++ return ++ try: ++ if not t.dtype.is_floating_point or bool(torch.isfinite(t).all()): ++ return ++ mode = "?" ++ if forward_batch is not None: ++ fm = getattr(forward_batch, "forward_mode", None) ++ mode = getattr(fm, "name", str(fm)) if fm is not None else "?" ++ logger.error( ++ "[NANPROBE] tag=%s layer=%s mode=%s ntok=%d shape=%s nan=%d inf=%d", ++ tag, ++ layer_id, ++ mode, ++ int(t.shape[0]) if t.dim() > 0 else -1, ++ tuple(t.shape), ++ int(torch.isnan(t).sum().item()), ++ int(torch.isinf(t).sum().item()), ++ ) ++ except Exception: ++ pass ++ _is_cuda = is_cuda() - _is_npu = is_npu() _is_cpu = is_cpu() -+_is_xpu = is_xpu() - _is_gfx95 = is_gfx95_supported() - _is_hip = is_hip() + _is_cpu_amx_available = cpu_has_amx_support() +@@ -123,6 +158,805 @@ _is_hip = is_hip() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip -@@ -235,6 +238,8 @@ class Qwen3_5GatedDeltaNet(nn.Module): - conv_weights = self.conv1d.weight.view( - self.conv1d.weight.size(0), self.conv1d.weight.size(2) - ) -+ # conv_weights aliases conv1d.weight's storage; loaders that swap that -+ # storage on a device move must call rebind_device_views() afterwards. - self.attn = RadixLinearAttention( - layer_id=layer_id, - num_q_heads=self.num_k_heads // self.attn_tp_size, -@@ -275,6 +280,32 @@ class Qwen3_5GatedDeltaNet(nn.Module): - tp_size=self.attn_tp_size, - prefix=add_prefix("out_proj", prefix), - ) -+ # NOTE (Qwen3.6 ratio=2 GGUF): out_proj's input (value-head) columns are -+ # stored by GGUF in [ratio, num_k] order but HF/core_attn_out expects -+ # [num_k, ratio]. This is an INPUT-dim (column) permute. It CANNOT be done -+ # per-rank in the XPU method (the older `_gguf_gdn_col_perm` path): under -+ # TP the value-head grouping crosses the RowParallel input-shard boundary -+ # (rank0's HF heads map to GGUF cols in BOTH ratio halves), so a per-rank -+ # reshape is impossible. Instead it is applied to the GLOBAL pre-shard -+ # weight in `_gguf_gdn_transform` (raw-byte, head_v_dim-granular; safe on -+ # Q8_0 whose block=32 divides head_v_dim). At ratio=1 the layouts coincide. -+ -+ def rebind_device_views(self): -+ """Re-derive tensors that alias conv1d.weight's storage. -+ -+ ``conv_weights`` (passed to RadixLinearAttention) is a *view* of -+ ``self.conv1d.weight`` captured at construction time. Loaders that -+ load on CPU and then swap ``conv1d.weight.data`` for a device tensor -+ (e.g. ``--load-format layered_fp8``) leave that view pointing at the -+ freed CPU storage, so the conv1d kernel later dereferences an invalid -+ pointer. Rebuild the view from the current weight; the lazily-built -+ ESIMD copy (``_esimd_conv_weights``) self-heals on next forward, so -+ just drop it here. -+ """ -+ w = self.conv1d.weight -+ self.attn.conv_weights = w.view(w.size(0), w.size(2)) -+ self.attn.bias = self.conv1d.bias -+ self._esimd_conv_weights = None - - @staticmethod - def _override_weight_loader(param, loader): -@@ -300,7 +331,19 @@ class Qwen3_5GatedDeltaNet(nn.Module): - - def _bind_packed_weight_loaders(self, module): - """Bind packed-checkpoint-aware loaders to all relevant params of a merged module.""" -- for attr_name in ("weight", "weight_scale_inv", "weight_scale", "input_scale"): -+ # "qweight" / "qweight_type" cover the GGUF path: its merged params are -+ # named qweight (not weight), and its native weight_loader only accepts -+ # int shard ids. The packed wrapper splits a fused checkpoint tensor -+ # (e.g. GGUF attn_qkv = q|k|v) by the tuple shard id (0,1,2) into int -+ # shards before delegating, so GGUF GDN projections load correctly. -+ for attr_name in ( -+ "weight", -+ "weight_scale_inv", -+ "weight_scale", -+ "input_scale", -+ "qweight", -+ "qweight_type", -+ ): - param = getattr(module, attr_name, None) - if param is None: - continue -@@ -345,10 +388,16 @@ class Qwen3_5GatedDeltaNet(nn.Module): - module, param, loaded_shard_id - ) -- if loaded_weight.numel() == 1: -- # Single-element tensor (scalar or [1]): -- # broadcast to each logical shard. -- chunks = [loaded_weight.view(-1)] * len(loaded_shard_id) -+ if len(loaded_weight.shape) == 0: -+ # Scalar shard payload. Two cases: -+ # - single logical shard: load as-is (original behavior). -+ # - GGUF qweight_type: one scalar quant-type for the whole -+ # fused tensor, replicated to each int shard so the GGUF -+ # loader records the type per shard slot. -+ if len(split_sizes) == 1 and split_sizes[0] == 1: -+ chunks = [loaded_weight.reshape(1)] -+ else: -+ chunks = [loaded_weight for _ in loaded_shard_id] - else: - split_dim = getattr(param, "output_dim", 0) - if _is_cpu: -@@ -472,6 +521,324 @@ class Qwen3_5GatedDeltaNet(nn.Module): - projected_states_ba, _ = self.in_proj_ba(hidden_states) - return projected_states_qkvz, projected_states_ba -+ def _forward_xpu_fast_path( -+ self, -+ projected_states_qkvz: torch.Tensor, -+ projected_states_ba: torch.Tensor, -+ forward_batch: ForwardBatch, -+ ): -+ """Run conv1d + GDN + RMSNormGated + out_proj using the vendored -+ torch.ops.sgl_kernel.gdn_attention kernel. Returns None if the op is -+ unusable (non-contiguous input, unexpected shape, ...) so the caller -+ falls back to the default Triton/PyTorch path. ++# --------------------------------------------------------------------------- ++# ESIMD MoE router (fp8) fast path. ++# ++# vLLM's BMG path runs the MoE router logits through an ESIMD fp8 GEMV ++# (`custom_esimd_kernels_sglang.moe_ops.moe_router_forward`) instead of the ++# default fp16 `aten::mm`, saving ~1 host-dispatch/layer during decode. The ++# kernel ONLY accepts fp8 e4m3/e5m2 weights (no fp16 variant), so the fp16 gate ++# weight is quantized once (per-tensor amax/448) and cached on the module. ++# ++# Gated by SGL_XPU_MOE_ROUTER_FP8=1 (default OFF: keeps the accurate fp16 gate). ++# Quantizing the router to per-tensor fp8 perturbs ~top-8 routing on a fraction ++# of tokens, so this must be gsm8k A/B validated before enabling in production. ++# --------------------------------------------------------------------------- ++_ESIMD_MOE_ROUTER_OP = "unset" ++# Set once at launch; constant for the process. Cached to avoid an os.environ ++# read on every MoE layer's router call during decode. ++_MOE_ROUTER_FP8 = os.environ.get("SGL_XPU_MOE_ROUTER_FP8", "0") == "1" ++logger.warning( ++ "[router_fp8] module loaded: file=%s SGL_XPU_MOE_ROUTER_FP8=%s (_MOE_ROUTER_FP8=%s)", ++ __file__, os.environ.get("SGL_XPU_MOE_ROUTER_FP8"), _MOE_ROUTER_FP8, ++) ++ ++ ++def _load_esimd_moe_router_op(): ++ global _ESIMD_MOE_ROUTER_OP ++ if _ESIMD_MOE_ROUTER_OP != "unset": ++ return _ESIMD_MOE_ROUTER_OP ++ op = None ++ try: ++ from custom_esimd_kernels_sglang import moe_ops as _moe_mod ++ ++ op = getattr(_moe_mod, "moe_router_forward", None) ++ except Exception: ++ op = None ++ _ESIMD_MOE_ROUTER_OP = op ++ return op ++ ++ ++_ROUTER_DEBUG = os.environ.get("SGL_XPU_ROUTER_DEBUG", "0") == "1" ++_ROUTER_DEBUG_N = 0 ++ ++ ++def _router_dbg(reason, **kw): ++ global _ROUTER_DEBUG_N ++ if not _ROUTER_DEBUG or _ROUTER_DEBUG_N >= 12: ++ return ++ _ROUTER_DEBUG_N += 1 ++ logger.warning("[router_fp8] fallback=%s %s", reason, kw) ++ ++ ++def _esimd_router_wq_scale(gate, hidden_states: torch.Tensor): ++ """Return (wq_e4m3 [E,H], scale [1] fp32) for the router gate weight, lazily ++ quantized + cached on the weight, or None on any shape/dtype/layout mismatch. ++ Shared by _esimd_router_logits (split path) and the rtfused MoE path. ++ """ ++ if not _MOE_ROUTER_FP8: ++ _router_dbg("env_off") ++ return None ++ x = hidden_states ++ if x.device.type != "xpu" or x.dim() != 2: ++ _router_dbg("device_or_dim", dev=str(x.device), dim=x.dim(), shape=tuple(x.shape)) ++ return None ++ # Kernel is tuned for decode + tiny prefill chunks (M=1 GEMV). ++ if x.shape[0] > 8: ++ return None ++ weight = getattr(gate, "weight", None) ++ if weight is None or weight.dim() != 2: ++ _router_dbg("weight_missing", gate_type=type(gate).__name__, ++ has_w=weight is not None) ++ return None ++ N, K = weight.shape # [num_experts, hidden] ++ if x.shape[1] != K: ++ _router_dbg("K_mismatch", xK=x.shape[1], wK=K, N=N) ++ return None ++ # A bias'd gate would need to be added post-GEMV; keep it simple & safe. ++ if getattr(gate, "bias", None) is not None: ++ _router_dbg("has_bias", bias_type=type(getattr(gate, "bias")).__name__) ++ return None ++ # Lazy per-tensor fp8 quant of the gate weight, cached on the parameter. ++ wq = getattr(weight, "_esimd_router_wq", None) ++ sc = getattr(weight, "_esimd_router_scale", None) ++ if wq is None or sc is None: ++ try: ++ wf = weight.detach().float() ++ amax = wf.abs().max() ++ if not torch.isfinite(amax) or amax <= 0: ++ return None ++ scale = (amax / 448.0) ++ wq = torch.clamp(wf / scale, -448.0, 448.0).to(torch.float8_e4m3fn).contiguous() ++ sc = scale.reshape(1).to(torch.float32).contiguous() ++ weight._esimd_router_wq = wq ++ weight._esimd_router_scale = sc ++ except Exception as e: ++ _router_dbg("quant_exc", err=repr(e)) ++ return None ++ return wq, sc ++ ++ ++def _esimd_router_logits(gate, hidden_states: torch.Tensor): ++ """Compute router logits via the ESIMD fp8 GEMV, or return None to fall back. ++ ++ Returns a [T, num_experts] fp16 logits tensor on success. Falls back (None) ++ on any shape/dtype/layout mismatch so the standard fp16 gate runs instead. ++ """ ++ op = _load_esimd_moe_router_op() ++ if op is None: ++ _router_dbg("op_none") ++ return None ++ wqsc = _esimd_router_wq_scale(gate, hidden_states) ++ if wqsc is None: ++ return None ++ wq, sc = wqsc ++ x = hidden_states ++ x_in = x if x.dtype == torch.float16 else x.to(torch.float16) ++ try: ++ logits = op(x_in, wq, sc) ++ except Exception as e: ++ _router_dbg("kernel_exc", err=repr(e)) ++ return None ++ if _ROUTER_DEBUG: ++ _router_dbg("SUCCESS", out_shape=tuple(logits.shape)) ++ return logits ++ ++ ++# --------------------------------------------------------------------------- ++# ESIMD full MoE fusion (e5m2) fast path. ++# ++# `moe_forward_full_v2` fuses the whole decode MoE block into ONE dispatch: ++# router topk (softmax, renorm) + routed experts (silu) + shared expert (silu) ++# + shared_expert_gate (sigmoid) + weighted accumulate. ++# This collapses the unfused decode path (routed silu kernel + shared-expert ++# gate_up/act/down linears + gate linear/sigmoid/mul + adds ≈ 300 dispatch/step) ++# down to a single op. e5m2-only (the kernel dequantises fp8_e5m2 weights); ++# activations stay fp16. Gated by SGL_XPU_ESIMD_MOE_FULL=1. ++# ++# Weight layouts the kernel expects (all validated numerically, cos≈1.0): ++# routed gate_up : [E, 2*inter, hidden] (native sglang w13, no transpose) ++# routed down : [E, hidden, inter] (native sglang w2, no transpose) ++# shared gate_up : [NS, 2*inter, hidden] (natural) -> transpose of stored [hidden,2*inter] ++# shared down : [NS, hidden, inter] (natural) -> transpose of stored [inter,hidden] ++# shared_gate : [NS, hidden] fp16 (not quantised) ++# all scales : per-expert per-tensor fp32 (dequant scale = amax/fp8_max) ++# --------------------------------------------------------------------------- ++_ESIMD_MOE_FULL = os.environ.get("SGL_XPU_ESIMD_MOE_FULL", "0") == "1" ++_ESIMD_MOE_FULL_OP = "unset" ++_MOE_FULL_DEBUG = os.environ.get("SGL_XPU_MOE_FULL_DEBUG", "0") == "1" ++_MOE_FULL_DEBUG_N = 0 ++ ++ ++def _moe_full_dbg(reason, **kw): ++ global _MOE_FULL_DEBUG_N ++ if not _MOE_FULL_DEBUG or _MOE_FULL_DEBUG_N >= 16: ++ return ++ _MOE_FULL_DEBUG_N += 1 ++ logger.warning("[moe_full] fallback=%s %s", reason, kw) ++ ++ ++def _load_esimd_moe_full_op(): ++ global _ESIMD_MOE_FULL_OP ++ if _ESIMD_MOE_FULL_OP != "unset": ++ return _ESIMD_MOE_FULL_OP ++ ops = None ++ try: ++ from custom_esimd_kernels_sglang import moe_ops as _moe_mod ++ ++ v2 = getattr(_moe_mod, "moe_forward_full_v2", None) ++ # BSZ=1 decode: moe_forward_full uses the fused down_finalize kernel ++ # (shared-expert down + gate sigmoid + accumulate in ONE dispatch), ++ # so it launches 5 internal kernels vs v2's 7 (gate_precompute + ++ # down_shared + accumulate collapse into down_finalize). Prefer it when ++ # n_tokens==1; fall back to v2 for multi-token. ++ full = getattr(_moe_mod, "moe_forward_full", None) ++ # rtfused: BSZ=1 path that also fuses router GEMV + topk into one ++ # dispatch (takes the e4m3 gate weight + scale instead of logits). ++ rtfused = getattr(_moe_mod, "moe_forward_full_rtfused", None) ++ # rtfused_norm: BSZ=1 path that ALSO folds the pre-MoE GemmaRMSNorm ++ # (resadd + rmsnorm) into the fused router kernel head and returns the ++ # new residual, removing the standalone gemma_fused_add_rmsnorm dispatch. ++ rtfused_norm = getattr(_moe_mod, "moe_forward_full_rtfused_norm", None) ++ if v2 is not None: ++ ops = {"v2": v2, "full": full, "rtfused": rtfused, ++ "rtfused_norm": rtfused_norm} ++ except Exception: ++ ops = None ++ _ESIMD_MOE_FULL_OP = ops ++ return ops ++ ++ ++def _pt_scale_1d(scale): ++ """Collapse a per-expert weight scale to a contiguous 1-D fp32 [E] tensor.""" ++ if scale is None: ++ return None ++ s = scale.to(torch.float32) ++ if s.dim() == 0: ++ s = s.reshape(1) ++ elif s.dim() > 1: ++ # [E, *block] -> per-expert scalar (these are already per-tensor scales, ++ # so any trailing dims are size-1; mean is a safe collapse). ++ s = s.reshape(s.shape[0], -1).mean(dim=-1) ++ return s.contiguous() ++ ++ ++def _gather_moe_full_weights(block, x: torch.Tensor): ++ """Collect + validate all tensors needed by the fused decode MoE ops ++ (moe_forward_full / rtfused / rtfused_norm). Returns a dict of the routed + ++ shared expert weights/scales, the gate module, and routing dims, or None to ++ fall back. Layout/dtype checks here guarantee a bad tensor never reaches the ++ kernel.""" ++ experts = getattr(block, "experts", None) ++ shared = getattr(block, "shared_expert", None) ++ sgate = getattr(block, "shared_expert_gate", None) ++ gate = getattr(block, "gate", None) ++ topk = getattr(block, "topk", None) ++ if experts is None or shared is None or sgate is None or gate is None or topk is None: ++ _moe_full_dbg("missing_submodule", experts=experts is not None, ++ shared=shared is not None, sgate=sgate is not None) ++ return None ++ ++ # Fused topk uses softmax + renorm + plain (ungrouped, no-bias) selection. ++ tc = getattr(topk, "topk_config", None) ++ if tc is None: ++ return None ++ if getattr(tc, "scoring_func", "softmax") != "softmax": ++ return None ++ if not getattr(tc, "renormalize", True): ++ return None ++ if getattr(tc, "use_grouped_topk", False): ++ return None ++ if getattr(tc, "correction_bias", None) is not None: ++ return None ++ if getattr(tc, "custom_routing_function", None) is not None: ++ return None ++ top_k = getattr(tc, "top_k", None) ++ if top_k is None: ++ return None ++ ++ # --- routed weights (e5m2) + load-time transposed caches --- ++ w13 = getattr(experts, "w13_weight", None) ++ w2 = getattr(experts, "w2_weight", None) ++ if w13 is None or w2 is None: ++ return None ++ if w13.dtype != torch.float8_e5m2 or w2.dtype != torch.float8_e5m2: ++ _moe_full_dbg("routed_not_e5m2", w13=str(w13.dtype), w2=str(w2.dtype)) ++ return None ++ gate_up_routed = w13 # native [E, 2*inter, hidden] ++ down_routed = w2 # natural [E, hidden, inter] ++ s13 = _pt_scale_1d(getattr(experts, "w13_weight_scale", None)) ++ s2 = _pt_scale_1d(getattr(experts, "w2_weight_scale", None)) ++ if s13 is None or s2 is None: ++ return None ++ ++ # --- shared expert (dense Fp8, weights stored transposed [hidden, out]) --- ++ gu_s = getattr(shared, "gate_up_proj", None) ++ dn_s = getattr(shared, "down_proj", None) ++ if gu_s is None or dn_s is None: ++ return None ++ gw = getattr(gu_s, "weight", None) # stored [hidden, 2*inter] e5m2 ++ dw = getattr(dn_s, "weight", None) # stored [inter, hidden] e5m2 ++ if gw is None or dw is None: ++ return None ++ if gw.dtype != torch.float8_e5m2 or dw.dtype != torch.float8_e5m2: ++ _moe_full_dbg("shared_not_e5m2", gw=str(gw.dtype), dw=str(dw.dtype)) ++ return None ++ shared_gate_up = getattr(gw, "_esimd_moe_full_nat", None) ++ if shared_gate_up is None: ++ try: ++ shared_gate_up = gw.t().contiguous().unsqueeze(0) # [1, 2*inter, hidden] ++ gw._esimd_moe_full_nat = shared_gate_up ++ except Exception: ++ return None ++ shared_down = getattr(dw, "_esimd_moe_full_nat", None) ++ if shared_down is None: ++ try: ++ shared_down = dw.t().contiguous().unsqueeze(0) # [1, hidden, inter] ++ dw._esimd_moe_full_nat = shared_down ++ except Exception: ++ return None ++ ss13 = _pt_scale_1d(getattr(gu_s, "weight_scale", None)) ++ ss2 = _pt_scale_1d(getattr(dn_s, "weight_scale", None)) ++ if ss13 is None or ss2 is None: ++ return None ++ ++ # --- shared_expert_gate weight (fp16 [NS, hidden], not quantised) --- ++ sgw = getattr(sgate, "weight", None) ++ if sgw is None or sgw.dim() != 2: ++ return None ++ sgw16 = (sgw if sgw.dtype == torch.float16 else sgw.to(torch.float16)).contiguous() ++ ++ return { ++ "gate": gate, ++ "top_k": int(top_k), ++ "n_routed": int(w13.shape[0]), ++ "num_shared": int(shared_gate_up.shape[0]), ++ "gate_up_routed": gate_up_routed, "s13": s13, ++ "shared_gate_up": shared_gate_up, "ss13": ss13, ++ "down_routed": down_routed, "s2": s2, ++ "shared_down": shared_down, "ss2": ss2, ++ "sgw16": sgw16, ++ } ++ ++ ++_GGUF_INTROSPECT_DONE = set() ++def _load_shared_q8_op(): ++ try: ++ import custom_esimd_kernels_sglang.custom_esimd_kernels # noqa: F401 (registers ops) ++ except Exception: ++ pass ++ ns = getattr(torch.ops, "custom_esimd_kernels_sglang", None) ++ return getattr(ns, "esimd_shared_expert_q8", None) if ns is not None else None ++ ++ ++_GGUF_MOE_SHARED = bool( ++ os.environ.get("SGL_XPU_GGUF_MOE_SHARED") ++ or os.environ.get("SGLANG_XPU_GGUF_MOE_SHARED") ++) ++_SHARED_Q8_OP = None ++ ++ ++def _gather_gguf_shared_q8(block): ++ """Collect + validate the Q8_0 shared-expert tensors for the fused decode op. ++ Returns (gu_qs, gu_sc, d_qs, d_sc, wg, inter_s) or False to disable.""" ++ se = getattr(block, "shared_expert", None) ++ sg = getattr(block, "shared_expert_gate", None) ++ if se is None or sg is None: ++ return False ++ gu = getattr(se, "gate_up_proj", None) ++ dn = getattr(se, "down_proj", None) ++ if gu is None or dn is None: ++ return False ++ ++ # gate_up: prefer the merged rep (rows gate then up); else concat per-shard. ++ merged = getattr(gu, "_xpu_merged", None) ++ if isinstance(merged, tuple) and merged and merged[0] == "q8_0": ++ gu_qs, gu_sc = merged[1].contiguous(), merged[2].contiguous() ++ else: ++ reps = getattr(gu, "_xpu_reps", None) ++ order = getattr(gu, "_xpu_shard_order", None) ++ if not isinstance(reps, dict) or not order: ++ return False ++ parts = [reps[i] for i in order] ++ if any((not isinstance(p, tuple)) or p[0] != "q8_0" for p in parts): ++ return False ++ gu_qs = torch.cat([p[1] for p in parts], dim=0).contiguous() ++ gu_sc = torch.cat([p[2] for p in parts], dim=0).contiguous() ++ ++ dreps = getattr(dn, "_xpu_reps", None) ++ drep = dreps.get("_single") if isinstance(dreps, dict) else None ++ if not (isinstance(drep, tuple) and drep[0] == "q8_0"): ++ return False ++ d_qs, d_sc = drep[1].contiguous(), drep[2].contiguous() ++ ++ w = getattr(sg, "weight", None) ++ if w is None or w.dim() != 2 or w.shape[0] != 1: ++ return False ++ wg = (w[0] if w.dtype == torch.float16 else w[0].to(torch.float16)).contiguous() ++ ++ two_inter, hidden = gu_qs.shape ++ inter_s = two_inter // 2 ++ # shape/dtype guards: a bad tensor never reaches the kernel. ++ if gu_qs.dtype != torch.int8 or d_qs.dtype != torch.int8: ++ return False ++ if (hidden % 32) or (inter_s % 32) or (two_inter % 2): ++ return False ++ if tuple(d_qs.shape) != (hidden, inter_s): ++ return False ++ if tuple(gu_sc.shape) != (two_inter, hidden // 32): ++ return False ++ if tuple(d_sc.shape) != (hidden, inter_s // 32): ++ return False ++ if wg.shape[0] != hidden: ++ return False ++ return (gu_qs, gu_sc, d_qs, d_sc, wg, int(inter_s)) ++ ++ ++def _maybe_gguf_shared_q8(block, x): ++ """One-dispatch fused GGUF Q8_0 shared expert (gate_up+silu+down+gate*sigmoid). ++ Returns the shared-expert partial [M, hidden] fp16, or None to fall back to ++ the unfused ``_forward_shared_experts`` path. Env-gated by ++ SGL_XPU_GGUF_MOE_SHARED=1.""" ++ if not _GGUF_MOE_SHARED or x.device.type != "xpu" or x.dim() != 2: ++ return None ++ global _SHARED_Q8_OP ++ if _SHARED_Q8_OP is None: ++ _SHARED_Q8_OP = _load_shared_q8_op() ++ if _SHARED_Q8_OP is None: ++ return None ++ cache = getattr(block, "_gguf_shared_q8_cache", None) ++ if cache is None: ++ cache = _gather_gguf_shared_q8(block) ++ block._gguf_shared_q8_cache = cache ++ if not cache: ++ return None ++ gu_qs, gu_sc, d_qs, d_sc, wg, inter_s = cache ++ xf = x if x.dtype == torch.float16 else x.to(torch.float16) ++ try: ++ return _SHARED_Q8_OP(xf, gu_qs, gu_sc, d_qs, d_sc, wg, inter_s) ++ except Exception: ++ block._gguf_shared_q8_cache = False ++ return None ++ ++ ++# ═══════════════ GGUF FULL MoE fusion (topk + routed + shared -> 1 op) ═══════════════ ++_GGUF_MOE_FULL = bool( ++ os.environ.get("SGL_XPU_GGUF_MOE_FULL") ++ or os.environ.get("SGLANG_XPU_GGUF_MOE_FULL") ++) ++_GGUF_FULL_OP = None ++_GGUF_FULL_NORM_OP = None ++ ++ ++def _env_int(name, default): ++ try: ++ v = os.environ.get(name) ++ return default if v is None or v == "" else int(v) ++ except Exception: ++ return default ++ ++ ++# Largest decode batch the norm-fused GGUF MoE op is allowed to serve. Every ++# kernel stage behind it (topk / up_q4k / shared_up_q8 / down / finalize) is ++# already M-generic and the C++ side sizes all scratch from x.size(0), so this ++# is purely a policy cap: it keeps the fused path away from prefill-sized M ++# where the per-token GEMV shape stops paying off. Set to 1 to A/B the fusion. ++_GGUF_MOE_FUSE_MAX_M = _env_int("SGL_XPU_GGUF_MOE_FUSE_MAX_M", 64) ++ ++ ++def _load_gguf_moe_full_norm_op(): ++ """Norm-fused GGUF MoE op: absorbs the post-attention GemmaRMSNorm and the ++ fp16 router GEMV, cutting two host dispatches per layer on the (host-bound) ++ decode path. Returns None on older kernel builds.""" ++ try: ++ import custom_esimd_kernels_sglang.custom_esimd_kernels # noqa: F401 ++ except Exception: ++ pass ++ ns = getattr(torch.ops, "custom_esimd_kernels_sglang", None) ++ if ns is None: ++ return None ++ return getattr(ns, "esimd_moe_forward_full_gguf_norm", None) ++ ++ ++def _gguf_router_weight(gate): ++ """The dense fp16 router weight [E, hidden] behind the MoE gate linear, or ++ None when it cannot be expressed that way. ++ ++ Two layouts occur for the 35B GGUF checkpoint: ``ffn_gate_inp`` is stored ++ unquantized (F32), so depending on the quant config the gate is either a ++ plain ``UnquantizedLinearMethod`` (weight on the module) or a GGUF XPU ++ linear holding a single fp16 resident rep. Handle both; the fp16 copy is ++ cached on the module so the (possible) cast happens once, not per step. ++ """ ++ w = getattr(gate, "_esimd_router_w16", None) ++ if w is not None: ++ return w ++ reps = getattr(gate, "_xpu_reps", None) ++ if reps is None: ++ qm = getattr(gate, "quant_method", None) ++ reps = getattr(qm, "_xpu_reps", None) ++ if isinstance(reps, dict) and list(reps.keys()) == ["_single"]: ++ rep = reps["_single"] ++ w = rep[1] if rep[0] == "fp16" else None ++ else: ++ w = getattr(gate, "weight", None) ++ w = getattr(w, "data", w) ++ if w is None or w.dim() != 2: ++ return None ++ if w.dtype != torch.float16: ++ w = w.to(torch.float16) ++ w = w.contiguous() ++ gate._esimd_router_w16 = w ++ return w ++ ++ ++def _load_gguf_moe_full_op(): ++ try: ++ import custom_esimd_kernels_sglang.custom_esimd_kernels # noqa: F401 ++ except Exception: ++ pass ++ ns = getattr(torch.ops, "custom_esimd_kernels_sglang", None) ++ return getattr(ns, "esimd_moe_forward_full_gguf", None) if ns is not None else None ++ ++ ++def _gather_gguf_moe_full(block): ++ """Collect + validate everything the fused GGUF full-MoE op needs: ++ routed Q4_K gate/up + Q5_K/Q6_K down reps (from experts.quant_method), ++ the Q8_0 shared-expert reps, top_k, renorm, and dims. Returns a dict or ++ False. Every layout/dtype check here keeps a bad tensor off the kernel.""" ++ experts = getattr(block, "experts", None) ++ topk = getattr(block, "topk", None) ++ gate = getattr(block, "gate", None) ++ if experts is None or topk is None or gate is None: ++ return False ++ qm = getattr(experts, "quant_method", None) ++ if qm is None or not hasattr(qm, "gate_ql") or not hasattr(qm, "down_ql"): ++ return False ++ ++ # topk must be plain softmax + renorm (matches the fused kernel selection). ++ tc = getattr(topk, "topk_config", None) ++ if tc is None: ++ return False ++ if getattr(tc, "scoring_func", "softmax") != "softmax": ++ return False ++ if getattr(tc, "use_grouped_topk", False): ++ return False ++ if getattr(tc, "correction_bias", None) is not None: ++ return False ++ if getattr(tc, "custom_routing_function", None) is not None: ++ return False ++ top_k = getattr(tc, "top_k", None) ++ if top_k is None: ++ return False ++ renorm = bool(getattr(tc, "renormalize", True)) ++ ++ shared = _gather_gguf_shared_q8(block) ++ if not shared: ++ return False ++ gu_qs, gu_sc, d_qs, d_sc, wg, inter_s = shared ++ ++ down_mn = getattr(qm, "down_mn", None) ++ down_is_q6 = bool(getattr(qm, "_down_is_q6", False)) ++ # q5k needs down_mn; q6k ignores it (pass down_sc as a valid placeholder). ++ if down_mn is None: ++ if not down_is_q6: ++ return False ++ down_mn = qm.down_sc ++ ++ hidden = int(getattr(qm, "hidden")) ++ inter = int(getattr(qm, "intermediate")) ++ E = int(getattr(qm, "E")) ++ # shared and routed share the same hidden; inter_s (shared) may differ. ++ if wg.shape[0] != hidden: ++ return False ++ return { ++ "gate": gate, "top_k": int(top_k), "renorm": renorm, ++ "E": E, "hidden": hidden, "inter": inter, "inter_s": int(inter_s), ++ "down_is_q6": down_is_q6, ++ "gate_ql": qm.gate_ql, "gate_sc": qm.gate_sc, "gate_mn": qm.gate_mn, ++ "up_ql": qm.up_ql, "up_sc": qm.up_sc, "up_mn": qm.up_mn, ++ "down_ql": qm.down_ql, "down_qh": qm.down_qh_plain, ++ "down_sc": qm.down_sc, "down_mn": down_mn, ++ "gu_qs": gu_qs, "gu_sc": gu_sc, "d_qs": d_qs, "d_sc": d_sc, "wg": wg, ++ } ++ ++ ++def _maybe_gguf_moe_full_norm(block, h, residual, nw, eps): ++ """Norm-fused GGUF MoE: GemmaRMSNorm(resadd) + router GEMV + experts in ONE ++ dispatch. Returns ``(moe_out, new_residual)`` or None to fall back. ++ ++ ``residual`` is updated in place by the kernel (same as ++ ``gemma_fused_add_rmsnorm``), so a None return must happen BEFORE the call. ++ """ ++ if not _GGUF_MOE_FULL or h.device.type != "xpu" or h.dim() != 2: ++ return None ++ if not 1 <= h.shape[0] <= _GGUF_MOE_FUSE_MAX_M: ++ return None ++ global _GGUF_FULL_NORM_OP ++ if _GGUF_FULL_NORM_OP is None: ++ _GGUF_FULL_NORM_OP = _load_gguf_moe_full_norm_op() ++ if _GGUF_FULL_NORM_OP is None: ++ return None ++ cache = getattr(block, "_gguf_moe_full_cache", None) ++ if cache is None: ++ cache = _gather_gguf_moe_full(block) ++ block._gguf_moe_full_cache = cache ++ if not cache: ++ return None ++ if getattr(block, "_gguf_moe_full_norm_off", False): ++ return None ++ rw = cache.get("router_w", None) ++ if rw is None: ++ rw = _gguf_router_weight(cache["gate"]) ++ if rw is None or rw.shape != (int(cache["E"]), int(cache["hidden"])): ++ logger.warning( ++ "[gguf_moe_full_norm] disabled: router weight %s, want (%s, %s)", ++ None if rw is None else tuple(rw.shape), ++ cache["E"], cache["hidden"]) ++ block._gguf_moe_full_norm_off = True ++ return None ++ cache["router_w"] = rw ++ hf = h if h.dtype == torch.float16 else h.to(torch.float16) ++ hf = hf if hf.is_contiguous() else hf.contiguous() ++ if residual.dtype != torch.float16 or not residual.is_contiguous(): ++ if not getattr(block, "_gguf_norm_res_warned", False): ++ block._gguf_norm_res_warned = True ++ logger.warning("[gguf_moe_full_norm] disabled: residual dtype=%s contig=%s", ++ residual.dtype, residual.is_contiguous()) ++ return None ++ try: ++ out, res_out = _GGUF_FULL_NORM_OP( ++ hf, residual, nw, float(eps), rw, ++ cache["gate_ql"], cache["gate_sc"], cache["gate_mn"], ++ cache["up_ql"], cache["up_sc"], cache["up_mn"], ++ cache["down_ql"], cache["down_qh"], cache["down_sc"], cache["down_mn"], ++ cache["gu_qs"], cache["gu_sc"], cache["d_qs"], cache["d_sc"], cache["wg"], ++ int(cache["E"]), int(cache["top_k"]), int(cache["inter"]), ++ int(cache["inter_s"]), bool(cache["down_is_q6"]), bool(cache["renorm"]), ++ ) ++ except Exception as e: ++ logger.warning("[gguf_moe_full_norm] disabled: kernel raised %r", e) ++ block._gguf_moe_full_norm_off = True ++ return None ++ if not getattr(block, "_gguf_norm_ok_logged", False): ++ block._gguf_norm_ok_logged = True ++ logger.warning("[gguf_moe_full_norm] ACTIVE (norm + router folded into MoE op)") ++ return out, res_out ++ ++ ++def _maybe_gguf_moe_full(block, x): ++ """One-dispatch fused GGUF MoE: router topk + routed experts (Q4_K/Q5_K) + ++ Q8_0 shared expert -> the final [M, hidden] fp16 PARTIAL (routed + gate*shared ++ summed; caller does the all_reduce). Returns None to fall back to the unfused ++ router/shared path. Env-gated by SGL_XPU_GGUF_MOE_FULL=1.""" ++ if not _GGUF_MOE_FULL or x.device.type != "xpu" or x.dim() != 2 or x.shape[0] > 8: ++ return None ++ global _GGUF_FULL_OP ++ if _GGUF_FULL_OP is None: ++ _GGUF_FULL_OP = _load_gguf_moe_full_op() ++ if _GGUF_FULL_OP is None: ++ _moe_full_dbg("gguf_op_none") ++ return None ++ cache = getattr(block, "_gguf_moe_full_cache", None) ++ if cache is None: ++ cache = _gather_gguf_moe_full(block) ++ block._gguf_moe_full_cache = cache ++ if not cache: ++ return None ++ ++ gate = cache["gate"] ++ logits = _esimd_router_logits(gate, x) ++ if logits is None: ++ logits, _ = gate(x) ++ if logits.dtype != torch.float16: ++ logits = logits.to(torch.float16) ++ logits = logits if logits.dim() == 2 else logits.reshape(x.shape[0], -1) ++ logits = logits.contiguous() ++ xf = x if x.dtype == torch.float16 else x.to(torch.float16) ++ try: ++ out = _GGUF_FULL_OP( ++ xf, logits, ++ cache["gate_ql"], cache["gate_sc"], cache["gate_mn"], ++ cache["up_ql"], cache["up_sc"], cache["up_mn"], ++ cache["down_ql"], cache["down_qh"], cache["down_sc"], cache["down_mn"], ++ cache["gu_qs"], cache["gu_sc"], cache["d_qs"], cache["d_sc"], cache["wg"], ++ int(cache["E"]), int(cache["top_k"]), int(cache["inter"]), ++ int(cache["inter_s"]), bool(cache["down_is_q6"]), bool(cache["renorm"]), ++ ) ++ except Exception as e: ++ _moe_full_dbg("gguf_kernel_exc", err=repr(e)) ++ block._gguf_moe_full_cache = False ++ return None ++ if _MOE_FULL_DEBUG: ++ _moe_full_dbg("SUCCESS", T=int(x.shape[0]), E=int(cache["E"]), ++ top_k=int(cache["top_k"]), out=tuple(out.shape), path="gguf_full") ++ return out ++ ++ ++def _maybe_esimd_moe_full(block, hidden_states: torch.Tensor): ++ """One-dispatch decode MoE via moe_forward_full_v2. Returns the final ++ [T, hidden] fp16 tensor (routed + gate*shared, already summed), or None to ++ fall back to the unfused router/shared path. ++ """ ++ if not _ESIMD_MOE_FULL: ++ return None ++ ops = _load_esimd_moe_full_op() ++ if ops is None: ++ _moe_full_dbg("op_none") ++ return None ++ x = hidden_states ++ if x.device.type != "xpu" or x.dim() != 2 or x.shape[0] > 8: ++ return None ++ ++ W = _gather_moe_full_weights(block, x) ++ if W is None: ++ return None ++ gate = W["gate"] ++ top_k = W["top_k"]; num_shared = W["num_shared"]; n_routed = W["n_routed"] ++ gate_up_routed = W["gate_up_routed"]; s13 = W["s13"] ++ shared_gate_up = W["shared_gate_up"]; ss13 = W["ss13"] ++ down_routed = W["down_routed"]; s2 = W["s2"] ++ shared_down = W["shared_down"]; ss2 = W["ss2"] ++ sgw16 = W["sgw16"] ++ x_in = x if x.dtype == torch.float16 else x.to(torch.float16) ++ ++ # BSZ=1 decode: prefer moe_forward_full_rtfused, which folds the router GEMV ++ # + softmax-topk into ONE dispatch (needs the e4m3-quantized gate weight). ++ # Fall back to moe_forward_full (fused down_finalize, separate router+topk) ++ # when the quant isn't available, and to moe_forward_full_v2 for multi-token. ++ rt = ops.get("rtfused") ++ if x.shape[0] == 1 and rt is not None: ++ wqsc = _esimd_router_wq_scale(gate, x) ++ if wqsc is not None: ++ wq, sc = wqsc ++ try: ++ out = rt( ++ x_in, wq, sc, ++ gate_up_routed, s13, ++ shared_gate_up, ss13, ++ down_routed, s2, ++ shared_down, ss2, ++ sgw16, ++ int(top_k), int(num_shared), int(n_routed), ++ ) ++ except Exception as e: ++ _moe_full_dbg("rtfused_exc", err=repr(e)) ++ return None ++ if _MOE_FULL_DEBUG: ++ _moe_full_dbg("SUCCESS", T=int(x.shape[0]), E=int(n_routed), ++ top_k=int(top_k), out=tuple(out.shape), path="rtfused") ++ return out ++ ++ # --- router logits [T, E] fp16 (full/v2 do topk internally) --- ++ logits = _esimd_router_logits(gate, x) ++ if logits is None: ++ logits, _ = gate(x) ++ if logits.dtype != torch.float16: ++ logits = logits.to(torch.float16) ++ ++ # BSZ=1 decode -> moe_forward_full (fused down_finalize, 5 internal kernels). ++ # Multi-token or missing op -> moe_forward_full_v2 (7 internal kernels). ++ op = ops.get("full") if (x.shape[0] == 1 and ops.get("full") is not None) else ops["v2"] ++ try: ++ out = op( ++ x_in, logits, ++ gate_up_routed, s13, ++ shared_gate_up, ss13, ++ down_routed, s2, ++ shared_down, ss2, ++ sgw16, ++ int(top_k), int(num_shared), int(n_routed), ++ ) ++ except Exception as e: ++ _moe_full_dbg("kernel_exc", err=repr(e)) ++ return None ++ if _MOE_FULL_DEBUG: ++ _moe_full_dbg("SUCCESS", T=int(x.shape[0]), E=int(n_routed), ++ top_k=int(top_k), out=tuple(out.shape), ++ path=("full" if op is ops.get("full") else "v2")) ++ return out ++ ++ ++def _maybe_esimd_moe_full_norm(block, hidden_states, residual, norm_weight_folded, eps): ++ """BSZ=1 decode MoE that ALSO folds the pre-MoE GemmaRMSNorm (residual add + ++ rmsnorm) into the fused router kernel head. ``hidden_states`` is the pre-norm ++ attention output (already all-reduced across the attn-TP group by the ++ caller), ``residual`` the residual stream, ``norm_weight_folded`` the Gemma ++ (1 + weight) in fp16, and ``eps`` the norm epsilon. ++ ++ Returns ``(moe_out, new_residual)`` where new_residual = hidden + residual, ++ or None to fall back to the standard prepare_mlp + mlp path. ++ """ ++ if not _ESIMD_MOE_FULL: ++ return None ++ ops = _load_esimd_moe_full_op() ++ if ops is None: ++ _moe_full_dbg("op_none") ++ return None ++ rtn = ops.get("rtfused_norm") ++ if rtn is None: ++ return None ++ x = hidden_states ++ if x.device.type != "xpu" or x.dim() != 2 or x.shape[0] != 1: ++ return None ++ if residual is None or residual.dim() != 2 or residual.shape != x.shape: ++ return None ++ ++ W = _gather_moe_full_weights(block, x) ++ if W is None: ++ return None ++ ++ wqsc = _esimd_router_wq_scale(W["gate"], x) ++ if wqsc is None: ++ return None ++ wq, sc = wqsc ++ ++ h_in = x if x.dtype == torch.float16 else x.to(torch.float16) ++ h_in = h_in.contiguous() ++ res = residual if residual.dtype == torch.float16 else residual.to(torch.float16) ++ res = res if res.is_contiguous() else res.contiguous() ++ try: ++ out = rtn( ++ h_in, res, norm_weight_folded, float(eps), ++ wq, sc, ++ W["gate_up_routed"], W["s13"], ++ W["shared_gate_up"], W["ss13"], ++ W["down_routed"], W["s2"], ++ W["shared_down"], W["ss2"], ++ W["sgw16"], ++ W["top_k"], W["num_shared"], W["n_routed"], ++ ) ++ except Exception as e: ++ _moe_full_dbg("rtfused_norm_exc", err=repr(e)) ++ return None ++ if not isinstance(out, (list, tuple)) or len(out) != 2: ++ return None ++ if _MOE_FULL_DEBUG: ++ _moe_full_dbg("SUCCESS", T=1, E=W["n_routed"], top_k=W["top_k"], ++ out=tuple(out[0].shape), path="rtfused_norm") ++ return out[0], out[1] ++ ++ + def can_fuse_shared_expert( + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig], +@@ -208,6 +1042,7 @@ class Qwen2MoeMLP(nn.Module): + + + class Qwen2MoeSparseMoeBlock(nn.Module): ++ _prep_mlp_skip_seen = set() + def __init__( + self, + layer_id: int, +@@ -454,7 +1289,9 @@ class Qwen2MoeSparseMoeBlock(nn.Module): + + def _forward_router_experts(self, hidden_states: torch.Tensor): + # router_logits: (num_tokens, n_experts) +- router_logits, _ = self.gate(hidden_states) ++ router_logits = _esimd_router_logits(self.gate, hidden_states) ++ if router_logits is None: ++ router_logits, _ = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + if self.enable_shared_expert_fusion and TopKOutputChecker.format_is_standard( + topk_output +@@ -522,8 +1359,31 @@ class Qwen2MoeSparseMoeBlock(nn.Module): + hidden_states + ) + else: +- shared_output = self._forward_shared_experts(hidden_states) +- final_hidden_states = self._forward_router_experts(hidden_states) ++ fused_full = _maybe_esimd_moe_full(self, hidden_states) ++ if fused_full is None: ++ fused_full = _maybe_gguf_moe_full(self, hidden_states) ++ if fused_full is not None: ++ # Fused op returns routed + gate*shared already summed → skip the ++ # separate shared path and its add below. ++ final_hidden_states = fused_full ++ shared_output = None ++ else: ++ shared_output = _maybe_gguf_shared_q8(self, hidden_states) ++ if shared_output is None: ++ shared_output = self._forward_shared_experts(hidden_states) ++ _nan_probe_moe( ++ "moe_shared_out", ++ shared_output, ++ layer_id=getattr(self, "layer_id", None), ++ forward_batch=forward_batch, ++ ) ++ final_hidden_states = self._forward_router_experts(hidden_states) ++ _nan_probe_moe( ++ "moe_routed_out", ++ final_hidden_states, ++ layer_id=getattr(self, "layer_id", None), ++ forward_batch=forward_batch, ++ ) + + if shared_output is not None: + final_hidden_states += shared_output +@@ -542,6 +1402,120 @@ class Qwen2MoeSparseMoeBlock(nn.Module): + + return final_hidden_states.view(num_tokens, hidden_dim) + ++ def esimd_prepare_mlp_moe( ++ self, ++ hidden_states: torch.Tensor, ++ residual: torch.Tensor, ++ norm_module, ++ forward_batch, ++ use_reduce_scatter: bool, ++ should_allreduce_fusion: bool, ++ ): ++ """Phase 3: fused replacement for ``prepare_mlp`` + ``mlp.forward`` on the ++ plain-TP single-token decode path. ++ ++ Baseline does, in two dispatched pieces: ++ 1. prepare_mlp: all-reduce(attn output) then GemmaRMSNorm(resadd+norm) ++ 2. mlp.forward: MoE (router+experts+shared) then all-reduce(MoE output) ++ ++ Here the GemmaRMSNorm resadd+norm is folded into the fused MoE router ++ kernel head (``moe_forward_full_rtfused_norm``), removing the standalone ++ ``gemma_fused_add_rmsnorm`` dispatch and returning the new residual so no ++ separate python add is needed. The two all-reduces are reproduced ++ explicitly (identical to the baseline collectives). ++ ++ Returns ``(hidden_out, residual_out)`` on success, or ``None`` to signal ++ the caller to run the standard ``prepare_mlp`` + ``mlp`` path. All guards ++ that could invalidate the manual collectives are checked BEFORE any ++ communication, so a ``None`` return never leaves a stray all-reduce. ++ """ ++ def _skip(why): ++ # One-shot per reason: tells us which guard blocks the fusion ++ # without spamming a 40-layer x N-token decode loop. ++ seen = Qwen2MoeSparseMoeBlock._prep_mlp_skip_seen ++ if why not in seen: ++ seen.add(why) ++ logger.warning("[esimd_prepare_mlp_moe] disabled: %s", why) ++ return None ++ ++ if not _ESIMD_MOE_FULL and not _GGUF_MOE_FULL: ++ return _skip("no MOE_FULL env") ++ ops = _load_esimd_moe_full_op() ++ # Either the fp8 ESIMD norm-fused op or the GGUF norm-fused op must be ++ # available, otherwise there is nothing to fold the norm into. ++ have_fp8 = ( ++ _ESIMD_MOE_FULL and ops is not None and ops.get("rtfused_norm") is not None ++ ) ++ have_gguf = _GGUF_MOE_FULL and _load_gguf_moe_full_norm_op() is not None ++ if not have_fp8 and not have_gguf: ++ return _skip("no norm-fused op (fp8=%s gguf=%s)" % (have_fp8, have_gguf)) ++ if not ( ++ hidden_states.device.type == "xpu" ++ and forward_batch is not None ++ and forward_batch.forward_mode.is_decode() ++ ): ++ return _skip("not xpu decode") ++ # M>1 is only serviceable by the GGUF norm-fused op; the fp8 twin ++ # (_maybe_esimd_moe_full_norm) is still single-token. Without GGUF we ++ # keep the original M==1 gate so fp8-only setups are untouched. ++ max_m = _GGUF_MOE_FUSE_MAX_M if have_gguf else 1 ++ if hidden_states.dim() != 2 or not (1 <= hidden_states.shape[0] <= max_m): ++ return _skip("shape %s" % (tuple(hidden_states.shape),)) ++ if residual is None or residual.shape != hidden_states.shape: ++ return _skip("residual mismatch") ++ # Only the plain-TP path (no input-scatter, no DP-attention, no ++ # all-reduce fusion, no reduce-scatter) matches the manual collectives. ++ if use_reduce_scatter or should_allreduce_fusion: ++ return _skip("reduce_scatter=%s allreduce_fusion=%s" ++ % (use_reduce_scatter, should_allreduce_fusion)) ++ if getattr(hidden_states, "_sglang_needs_allreduce_fusion", False): ++ return _skip("needs_allreduce_fusion flag") ++ try: ++ from sglang.srt.layers.communicator import get_attn_tp_context ++ from sglang.srt.layers.dp_attention import get_attention_dp_size ++ ++ if get_attn_tp_context().input_scattered: ++ return _skip("input_scattered") ++ if get_attention_dp_size() != 1: ++ return _skip("dp_size != 1") ++ except Exception as e: ++ return _skip("ctx probe exc %r" % (e,)) ++ # Fold GemmaRMSNorm (1 + weight) once per layer. ++ nw = getattr(norm_module, "_esimd_moe_nw", None) ++ if nw is None: ++ w = getattr(norm_module, "weight", None) ++ eps = getattr(norm_module, "variance_epsilon", None) ++ if w is None or eps is None: ++ return None ++ nw = (w.data.to(torch.float32) + 1.0).to(torch.float16).contiguous() ++ norm_module._esimd_moe_nw = nw ++ eps = float(norm_module.variance_epsilon) ++ ++ # ── Commit: reproduce prepare_mlp's attention-output all-reduce ── ++ h_ar = attention_tensor_model_parallel_all_reduce(hidden_states) ++ ++ fused = None ++ if have_gguf: ++ fused = _maybe_gguf_moe_full_norm(self, h_ar, residual, nw, eps) ++ if fused is None and have_fp8: ++ fused = _maybe_esimd_moe_full_norm(self, h_ar, residual, nw, eps) ++ if fused is not None: ++ moe_out, new_residual = fused ++ # Reproduce mlp.forward's post-experts all-reduce (kernel is per-rank). ++ if self.tp_size > 1 and not get_moe_a2a_backend().is_flashinfer(): ++ moe_out = tensor_model_parallel_all_reduce(moe_out) ++ return moe_out, new_residual ++ ++ # Fallback (kernel unavailable / guard miss inside the MoE helper): run ++ # the baseline norm on the already-reduced hidden, then the standard MoE ++ # forward (which performs its own post-experts all-reduce). This keeps ++ # correctness without duplicating the attention all-reduce above. ++ normed, new_residual = norm_module(h_ar, residual) ++ moe_out = self.forward( ++ normed, forward_batch, use_reduce_scatter, should_allreduce_fusion ++ ) ++ return moe_out, new_residual ++ + + class Qwen2MoeAttention(nn.Module): + def __init__( +diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py +index 65f78e691..513307770 100644 +--- a/python/sglang/srt/models/qwen3_5.py ++++ b/python/sglang/srt/models/qwen3_5.py +@@ -15,8 +15,9 @@ + """Inference-only Qwen3.5 model and Qwen3.5 MoE model compatible with HuggingFace weights.""" + + import logging ++import os + from functools import lru_cache +-from typing import Iterable, Optional, Set, Tuple, Union ++from typing import Iterable, List, Optional, Set, Tuple, Union + + import torch + import torch.nn as nn +@@ -93,15 +94,107 @@ from sglang.srt.utils import ( + is_gfx95_supported, + is_hip, + is_npu, ++ is_xpu, + make_layers, + set_weight_attrs, + ) + from sglang.srt.utils.hf_transformers_utils import get_processor, get_rope_config + + logger = logging.getLogger(__name__) ++ ++# ── NaN/Inf probe (env-gated: SGLANG_NAN_PROBE=1) ─────────────────────────── ++# Debug instrumentation to locate the first non-finite tensor in the forward ++# pass (symptom: "！！！" garbage generations, more frequent at batch>8/16). ++# Logs the ORIGIN stage (input clean -> output NaN/Inf) with layer id, forward ++# mode (EXTEND/DECODE), and token count so we can tell which kernel/batch size ++# introduces the NaN. Zero overhead when SGLANG_NAN_PROBE is unset. ++import os as _os_np ++ ++_NAN_PROBE = _os_np.environ.get("SGLANG_NAN_PROBE", "0") == "1" ++ ++_nan_probe_state = {"origin_found": False, "fwd": -1} ++ ++ ++def _np_bad(t): ++ """Return (n_nan, n_inf) as python ints; (0, 0) if finite / not applicable.""" ++ if t is None or not torch.is_tensor(t) or t.numel() == 0: ++ return (0, 0) ++ if not t.dtype.is_floating_point: ++ return (0, 0) ++ if bool(torch.isfinite(t).all()): ++ return (0, 0) ++ return (int(torch.isnan(t).sum().item()), int(torch.isinf(t).sum().item())) ++ ++ ++def _nan_probe(tag, t, layer_id=None, forward_batch=None, residual=None): ++ if not _NAN_PROBE: ++ return ++ try: ++ nan_h, inf_h = _np_bad(t) ++ nan_r, inf_r = _np_bad(residual) ++ if (nan_h + inf_h + nan_r + inf_r) == 0: ++ return ++ mode = "?" ++ if forward_batch is not None: ++ fm = getattr(forward_batch, "forward_mode", None) ++ mode = getattr(fm, "name", str(fm)) if fm is not None else "?" ++ ntok = int(t.shape[0]) if (torch.is_tensor(t) and t.dim() > 0) else -1 ++ first = not _nan_probe_state["origin_found"] ++ if first: ++ _nan_probe_state["origin_found"] = True ++ extra = "" ++ if tag == "embed_out" and torch.is_tensor(t) and t.dim() == 2: ++ bad = ~torch.isfinite(t) ++ rows = bad.any(dim=1).nonzero(as_tuple=False).flatten() ++ full = int(bad.all(dim=1).sum().item()) ++ ids = getattr(forward_batch, "input_ids", None) ++ idh = [] ++ if torch.is_tensor(ids) and ids.numel() >= t.shape[0]: ++ idh = ( ++ ids.flatten()[: t.shape[0]].index_select(0, rows[:16]).tolist() ++ ) ++ extra = ( ++ " | bad_rows=%d full_nan_rows=%d rows_head=%s ids_head=%s" ++ " ids_shape=%s elem_head=%s" ++ % ( ++ int(rows.numel()), ++ full, ++ rows[:16].tolist(), ++ idh, ++ (tuple(ids.shape) if torch.is_tensor(ids) else None), ++ bad.nonzero(as_tuple=False)[:8].tolist(), ++ ) ++ ) ++ logger.error( ++ "[NANPROBE]%s fwd=%d tag=%s layer=%s mode=%s ntok=%d | " ++ "hidden nan=%d inf=%d | residual nan=%d inf=%d%s", ++ " ORIGIN" if first else "", ++ _nan_probe_state["fwd"], ++ tag, ++ layer_id, ++ mode, ++ ntok, ++ nan_h, ++ inf_h, ++ nan_r, ++ inf_r, ++ extra, ++ ) ++ except Exception: ++ pass ++ ++ ++def _nan_probe_new_forward(): ++ if not _NAN_PROBE: ++ return ++ _nan_probe_state["origin_found"] = False ++ _nan_probe_state["fwd"] += 1 ++ ++ + _is_cuda = is_cuda() + _is_npu = is_npu() + _is_cpu = is_cpu() ++_is_xpu = is_xpu() + _is_gfx95 = is_gfx95_supported() + _is_hip = is_hip() + _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +@@ -113,6 +206,240 @@ _qknorm_use_alt_stream = ( + get_bool_env_var("SGLANG_QK_NORM_ALT_STREAM", "False") and _hip_use_alt_stream + ) + _is_amx_available = cpu_has_amx_support() ++# XPU ESIMD fast-path gates. These SGL_XPU_* env vars are set once at launch and ++# are constant for the process lifetime, so cache them at import instead of ++# re-reading os.environ on every decode-step call (30-60 reads/step otherwise). ++_XPU_GDN_NORM_GEMV = os.environ.get("SGL_XPU_GDN_NORM_GEMV", "0") == "1" ++_XPU_GDN_RESADD_NORM = os.environ.get("SGL_XPU_GDN_RESADD_NORM", "0") == "1" ++_XPU_GDN_FAST_PATH = os.environ.get("SGLANG_XPU_GDN_FAST_PATH", "0") == "1" ++_XPU_GDN_ESIMD = os.environ.get("SGL_XPU_GDN_ESIMD", "0") == "1" ++_XPU_FA_ESIMD_QKV = os.environ.get("SGL_XPU_FA_ESIMD_QKV", "0") == "1" ++# Fuse the GDN in_proj_qkvz + in_proj_ba GEMVs (same hidden input, both fp8) into a ++# single esimd_gemv_fp8_pert_fused2 call, removing one GEMM dispatch + its glue per ++# GDN layer (~30/step). Decode-only (M=1). Falls back to two separate linears. ++_XPU_GDN_INPROJ_FUSED2 = os.environ.get("SGL_XPU_GDN_INPROJ_FUSED2", "0") == "1" ++# Phase 5b: fuse the full-attention layers' input_layernorm (GemmaRMSNorm resadd + ++# rmsnorm) into the qkv_proj GEMV via the single deployed ++# esimd_resadd_norm_gemv2_fp8_pert op (qkv as matrix-0, a 1-row dummy as matrix-1), ++# with kernel-written new_residual. Removes the per-full-attn-layer ++# gemma_fused_add_rmsnorm dispatch (~10/step). Decode single-token, fp8 qkv only. ++_XPU_FA_RESADD_NORM = os.environ.get("SGL_XPU_FA_RESADD_NORM", "0") == "1" ++ ++# Phase 5c (GGUF): the Phase-5b resadd-norm fusions above are fp8-only, but the ++# 35B GGUF build runs its attention projections as ESIMD q8_0 GEMVs, so they ++# never fire. `esimd_resadd_norm_gemv_q8_ba` is the GGUF counterpart: it folds ++# GemmaRMSNorm(input_layernorm) plus the q8_0 in_proj/qkv GEMV plus (for GDN) ++# the unquantised fp16 in_proj_ba GEMV into ONE op call. Decode at bs=1 is ++# host-bound (~13-25us of torch dispatch per call vs ~4us of actual enqueue), so ++# collapsing 70 op calls/step is worth far more than any kernel-level tuning. ++# On by default whenever the GGUF MoE full fusion is on; SGL_XPU_GGUF_RESADD_NORM=0 ++# forces the unfused fallback. ++_XPU_GGUF_RESADD_NORM = ( ++ os.environ.get("SGL_XPU_GGUF_MOE_FULL", "0") == "1" ++ and os.environ.get("SGL_XPU_GGUF_RESADD_NORM", "1") == "1" ++) ++# DEPRECATED. First iteration of the GGUF norm+proj fusion, superseded by ++# _gguf_norm_gemv() / _XPU_GGUF_RESADD_NORM above. It drove a single-matrix ++# `esimd_resadd_norm_gemv_q8_0` op that was never landed in ++# custom-esimd-kernels, so the path always ImportErrors and falls back; the ++# shipped design instead uses `esimd_resadd_norm_gemv_q8_ba`, which folds the ++# unquantised fp16 in_proj_ba GEMV into the SAME op call (one launch instead of ++# two) and keeps the q/k/v-merged rep. Kept for reference behind its own opt-in ++# flag so it can never shadow the supported path; remove once the q8_0-only ++# variant is confirmed unnecessary. ++_XPU_GGUF_RESADD_NORM_LEGACY = ( ++ os.environ.get("SGL_XPU_GGUF_RESADD_NORM_LEGACY", "0") == "1" ++) ++ ++ ++def _env_int(name, default): ++ try: ++ return int(os.environ.get(name, default)) ++ except (TypeError, ValueError): ++ return default ++ ++ ++# Largest decode batch the GGUF resadd-norm fusion will handle. ++# ++# The kernel (Moe_norm_q8_kernel) has always been M-generic: its grid is ++# M*blocks*K_SPLIT and every access is indexed off `token = gid / blocks`. The ++# bs=1 restriction lived purely in this file, because the cached output buffers ++# were allocated as [1, N]. Real serving decode batches are almost never 1 (a ++# BFCL multi_turn run spends 87% of its decode steps at batch > 8), so that ++# restriction meant the fusion was bypassed for the large majority of steps. ++# ++# Buffers are now allocated per distinct M, so the only reason for an upper ++# bound is to stop an unbounded prefill-sized M from allocating one buffer set ++# per batch size ever seen. Set to 1 to restore the old single-token behaviour. ++_GGUF_FUSE_MAX_M = _env_int("SGL_XPU_GGUF_FUSE_MAX_M", 64) ++_GGUF_NORM_Q8_OP = None ++ ++ ++def _load_gguf_norm_q8_op(): ++ global _GGUF_NORM_Q8_OP ++ if _GGUF_NORM_Q8_OP is None: ++ try: ++ import custom_esimd_kernels_sglang # noqa: F401 (registers the lib) ++ ++ _GGUF_NORM_Q8_OP = ( ++ torch.ops.custom_esimd_kernels_sglang.esimd_resadd_norm_gemv_q8_ba ++ ) ++ except Exception as e: ++ logger.warning("[gguf_resadd_norm] op unavailable: %r", e) ++ _GGUF_NORM_Q8_OP = False ++ return _GGUF_NORM_Q8_OP or None ++ ++ ++def _gguf_xpu_rep(lin, kind): ++ """The single ESIMD weight rep of a GGUF XPU linear, if it is `kind`. ++ ++ A GGUF linear keeps one rep per loaded shard; when every shard shares a ++ quant type they are pre-merged into ``_xpu_merged`` (row-cat in output ++ order), which is exactly the ``[N, K]`` matrix the fused kernel needs. A ++ single-shard linear has no merge, so fall back to its only rep. ++ """ ++ if lin is None or getattr(lin, "bias", None) is not None: ++ return None ++ merged = getattr(lin, "_xpu_merged", None) ++ rep = merged[0] if merged is not None else None ++ if rep is None: ++ reps = getattr(lin, "_xpu_reps", None) ++ order = getattr(lin, "_xpu_shard_order", None) ++ if isinstance(reps, dict) and order is not None and len(order) == 1: ++ rep = reps[order[0]] ++ if rep is None or rep[0] != kind: ++ return None ++ return rep ++ ++ ++def _gguf_resadd_norm_guard(hidden_states, residual, forward_batch): ++ """True when the plain-TP decode path is what will run, i.e. prepare_attn ++ reduces to a bare input_layernorm and the fused op is a drop-in.""" ++ if not _XPU_GGUF_RESADD_NORM or not _is_xpu: ++ return False ++ if forward_batch is None or not forward_batch.forward_mode.is_decode(): ++ return False ++ if forward_batch.forward_mode.is_target_verify(): ++ return False ++ if hidden_states.dim() != 2: ++ return False ++ # The kernel handles any M; the bound only caps how many distinct output ++ # buffer sets we are willing to cache (see _GGUF_FUSE_MAX_M). ++ if not 1 <= hidden_states.shape[0] <= _GGUF_FUSE_MAX_M: ++ return False ++ if hidden_states.dtype != torch.float16: ++ return False ++ # Layer 0 has no residual yet. The kernel always does the add, so it is fed ++ # a cached zero buffer: ``h + 0.0`` is exact in fp16 and reproduces the ++ # plain path's ``residual = hidden_states; hidden = norm(hidden)``. Leaving ++ # layer 0 on the unfused path meant its in_proj_ba ran a oneDNN fp16 GEMV ++ # whose split-K reduction order is not reproducible under load, which made ++ # decode non-deterministic (identical operands, two possible results). ++ if residual is not None: ++ if residual.shape != hidden_states.shape: ++ return False ++ if residual.dtype != torch.float16: ++ return False ++ if getattr(hidden_states, "_sglang_needs_allreduce_fusion", False): ++ return False ++ try: ++ from sglang.srt.layers.communicator import get_attn_tp_context ++ ++ if get_attn_tp_context().input_scattered: ++ return False ++ except Exception: ++ return False ++ return True ++ ++ ++def _gguf_norm_gemv(layer, norm, lin0, lin1, hidden_states, residual, tag): ++ """Run the fused resadd-norm + q8_0 GEMV (+ optional fp16 GEMV). ++ ++ ``lin1`` is the GDN ``in_proj_ba`` (unquantised fp16) or None for the ++ full-attention layers. Returns ``(out0, out1_or_None, new_residual)`` or ++ None to fall back. Weights and norm constants are resolved once and cached ++ on the layer; output buffers are cached per distinct token count M, so the ++ steady-state cost stays exactly one op dispatch at any batch size. ++ """ ++ op = _load_gguf_norm_q8_op() ++ if op is None or getattr(layer, "_gguf_norm_q8_off", False): ++ return None ++ cache = getattr(layer, "_gguf_norm_q8_cache", None) ++ if cache is None: ++ rep0 = _gguf_xpu_rep(lin0, "q8_0") ++ if rep0 is None: ++ logger.warning("[gguf_resadd_norm] %s: proj is not a single q8_0 rep", tag) ++ layer._gguf_norm_q8_off = True ++ return None ++ qs, sc = rep0[1], rep0[2] ++ hidden = hidden_states.shape[1] ++ if qs.dim() != 2 or qs.shape[1] != hidden: ++ logger.warning("[gguf_resadd_norm] %s: q8_0 K=%s != hidden=%s", ++ tag, tuple(qs.shape), hidden) ++ layer._gguf_norm_q8_off = True ++ return None ++ dev = qs.device ++ w1 = torch.empty(0, dtype=torch.float16, device=dev) ++ if lin1 is not None: ++ rep1 = _gguf_xpu_rep(lin1, "fp16") ++ if rep1 is None or rep1[1].dim() != 2 or rep1[1].shape[1] != hidden: ++ logger.warning("[gguf_resadd_norm] %s: ba is not a single fp16 [N,%s] rep", ++ tag, hidden) ++ layer._gguf_norm_q8_off = True ++ return None ++ w1 = rep1[1].contiguous() ++ nw = ((norm.weight.data.to(torch.float32) + 1.0) ++ .to(torch.float16).contiguous()) ++ cache = { ++ "nw": nw, ++ "eps": float(norm.variance_epsilon), ++ "qs": qs, ++ "sc": sc, ++ "w1": w1, ++ "hidden": hidden, ++ "dev": dev, ++ "has_ba": lin1 is not None, ++ # Output buffers, keyed by token count. Decode batch sizes repeat, ++ # so this is allocated a handful of times and then only looked up. ++ "bufs": {}, ++ } ++ layer._gguf_norm_q8_cache = cache ++ logger.warning("[gguf_resadd_norm] %s ACTIVE (norm + proj folded, ba=%s)", ++ tag, cache["has_ba"]) ++ ++ M = int(hidden_states.shape[0]) ++ buf = cache["bufs"].get(M) ++ if buf is None: ++ dev, hidden = cache["dev"], cache["hidden"] ++ w1 = cache["w1"] ++ buf = { ++ "o0": torch.empty((M, cache["qs"].shape[0]), dtype=torch.float16, device=dev), ++ "o1": (torch.empty((M, w1.shape[0]), dtype=torch.float16, device=dev) ++ if cache["has_ba"] ++ else torch.empty(0, dtype=torch.float16, device=dev)), ++ "xn": torch.empty((M, hidden), dtype=torch.float16, device=dev), ++ # Post-add residual goes to its own buffer: the kernel's block 0 ++ # stores it while the other row-blocks still read the OLD residual. ++ # Per-layer, so it is never the same buffer the kernel reads (the ++ # incoming residual belongs to the previous layer). ++ "nr": torch.empty((M, hidden), dtype=torch.float16, device=dev), ++ # Zero residual for layer 0 (residual is None there). ++ "zr": torch.zeros((M, hidden), dtype=torch.float16, device=dev), ++ } ++ cache["bufs"][M] = buf ++ ++ h = hidden_states if hidden_states.is_contiguous() else hidden_states.contiguous() ++ res = buf["zr"] if residual is None else ( ++ residual if residual.is_contiguous() else residual.contiguous()) ++ try: ++ op(h, res, cache["nw"], cache["eps"], buf["xn"], buf["nr"], ++ cache["qs"], cache["sc"], buf["o0"], cache["w1"], buf["o1"]) ++ except Exception as e: ++ logger.warning("[gguf_resadd_norm] %s: kernel raised %r", tag, e) ++ layer._gguf_norm_q8_off = True ++ return None ++ return buf["o0"], (buf["o1"] if cache["has_ba"] else None), buf["nr"] ++ + + cached_get_processor = lru_cache(get_processor) + +@@ -123,6 +450,134 @@ def _disable_shared_experts_fusion() -> bool: + return get_global_server_args().disable_shared_experts_fusion + + ++def _esimd_fp8_weight_nk_scale(lin): ++ """Resolve an fp8 Linear's weight as row-major ``[N, K]`` plus a per-tensor ++ scalar scale, for the ESIMD GEMV fusion kernels. ++ ++ Handles the two fp8 weight layouts SGLang produces: ++ ++ * **block-quant** (``weight_scale_inv`` present): ``layer.weight`` is stored ++ ``[N, K]`` already; the block scale is collapsed to its mean. ++ * **online / dynamic per-tensor** (``weight_scale`` present): the loader ++ stores ``layer.weight`` **transposed** as ``[K, N]`` (see ++ ``Fp8LinearMethod.process_weights_after_loading``). ESIMD needs ``[N, K]``, ++ so we transpose (reusing the ``_esimd_t`` cache shared with the dense fp8 ++ fast path) and take the per-tensor scalar. ++ ++ Returns ``(weight_nk, scale_pt_fp32_1d)`` on success, or ``None`` to force a ++ safe fallback. The result is validated against the layer's known ++ ``output_size_per_partition`` / ``input_size_per_partition`` so a layout we ++ did not anticipate never reaches the kernel (which reads raw pointers and ++ would otherwise fault with UR_RESULT_ERROR_DEVICE_LOST). ++ """ ++ w = getattr(lin, "weight", None) ++ if w is None or w.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2): ++ return None ++ # Weight and its fp8 scale are static after loading, so the transposed ++ # [N, K] layout and the scalar per-tensor scale never change across decode ++ # steps. Cache the fully-resolved (w_nk, scale_pt) on the weight so we skip ++ # the per-step t()/contiguous()/to(fp32)/mean()/contiguous() device ops ++ # (~5 dispatched ops x 60 calls/step of pure redundant recompute). ++ _cached = getattr(w, "_esimd_nk_scale", None) ++ if _cached is not None: ++ return _cached ++ if w.dim() != 2: ++ return None ++ N = getattr(lin, "output_size_per_partition", None) ++ K = getattr(lin, "input_size_per_partition", None) ++ if N is None or K is None: ++ return None ++ ++ ws_inv = getattr(lin, "weight_scale_inv", None) ++ if ws_inv is not None: ++ # Block-quant: weight already [N, K]. ++ if w.shape[0] != N or w.shape[1] != K: ++ return None ++ w_nk = w ++ scale_src = ws_inv ++ else: ++ ws = getattr(lin, "weight_scale", None) ++ if ws is None: ++ return None ++ # Online per-tensor: stored [K, N] -> transpose to [N, K]. ++ if w.shape[0] != K or w.shape[1] != N: ++ return None ++ w_nk = getattr(w, "_esimd_t", None) ++ if w_nk is None: ++ w_nk = w.t().contiguous() ++ try: ++ w._esimd_t = w_nk ++ except Exception: ++ pass ++ scale_src = ws ++ ++ if w_nk.shape[0] != N or w_nk.shape[1] != K: ++ return None ++ scale_pt = ( ++ scale_src.data.to(torch.float32).reshape(-1).mean().reshape(1).contiguous() ++ ) ++ result = (w_nk, scale_pt) ++ try: ++ w._esimd_nk_scale = result ++ except Exception: ++ pass ++ return result ++ ++ ++def _esimd_q8_0_weight_scale(lin): ++ """DEPRECATED (see _XPU_GGUF_RESADD_NORM_LEGACY). ++ ++ Resolve a GGUF-q8_0 Linear's quant weight as ``(qs [N, K] int8, scale ++ [N, K/32] fp16)`` for the legacy single-matrix ESIMD q8_0 GEMV kernel. ++ ++ Superseded by ``_gguf_xpu_rep()``, which returns the packed rep tuple ++ directly and additionally handles the fp16 ``in_proj_ba`` shard so both ++ projections can be folded into one ``esimd_resadd_norm_gemv_q8_ba`` call. ++ ++ The GGUF XPU linear method stores its resident quant reps on the module: ++ * ``layer._xpu_merged`` = ``(merged_rep, sizes)`` when q/k/v shards were ++ row-concatenated into one big-N GEMV (the D1 optimization); or ++ * ``layer._xpu_reps`` / ``layer._xpu_shard_order`` for the per-shard reps. ++ A single-matrix fused GEMV needs the *whole* projection as one q8_0 matrix, ++ so this returns a value only when the projection is a lone q8_0 shard or a ++ merged q8_0 rep; anything else (multi-shard unmerged, non-q8_0) yields ++ ``None`` and the caller safely falls back. Cached on the module. ++ """ ++ cached = getattr(lin, "_esimd_q8_0_ws", None) ++ if cached is not None: ++ return cached ++ rep = None ++ merged = getattr(lin, "_xpu_merged", None) ++ if merged is not None: ++ rep = merged[0] ++ else: ++ order = getattr(lin, "_xpu_shard_order", None) ++ reps = getattr(lin, "_xpu_reps", None) ++ if order is not None and reps is not None and len(order) == 1: ++ rep = reps.get(order[0]) ++ if rep is None or rep[0] != "q8_0": ++ return None ++ _, qs, scale = rep ++ # Validate the exact [N, K] int8 + [N, K/32] fp16 layout the kernel reads ++ # from raw pointers (a mismatch would fault with DEVICE_LOST). ++ if qs.dim() != 2 or scale.dim() != 2: ++ return None ++ if qs.dtype != torch.int8 or scale.dtype != torch.float16: ++ return None ++ if qs.shape[1] % 32 != 0 or scale.shape[0] != qs.shape[0]: ++ return None ++ if scale.shape[1] != qs.shape[1] // 32: ++ return None ++ if not (qs.is_contiguous() and scale.is_contiguous()): ++ return None ++ result = (qs, scale) ++ try: ++ lin._esimd_q8_0_ws = result ++ except Exception: ++ pass ++ return result ++ ++ + if _is_npu: + from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import ( + split_qkvgate_gemma_rmsnorm_rope, +@@ -235,6 +690,8 @@ class Qwen3_5GatedDeltaNet(nn.Module): + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) ++ # conv_weights aliases conv1d.weight's storage; loaders that swap that ++ # storage on a device move must call rebind_device_views() afterwards. + self.attn = RadixLinearAttention( + layer_id=layer_id, + num_q_heads=self.num_k_heads // self.attn_tp_size, +@@ -275,6 +732,32 @@ class Qwen3_5GatedDeltaNet(nn.Module): + tp_size=self.attn_tp_size, + prefix=add_prefix("out_proj", prefix), + ) ++ # NOTE (Qwen3.6 ratio=2 GGUF): out_proj's input (value-head) columns are ++ # stored by GGUF in [ratio, num_k] order but HF/core_attn_out expects ++ # [num_k, ratio]. This is an INPUT-dim (column) permute. It CANNOT be done ++ # per-rank in the XPU method (the older `_gguf_gdn_col_perm` path): under ++ # TP the value-head grouping crosses the RowParallel input-shard boundary ++ # (rank0's HF heads map to GGUF cols in BOTH ratio halves), so a per-rank ++ # reshape is impossible. Instead it is applied to the GLOBAL pre-shard ++ # weight in `_gguf_gdn_transform` (raw-byte, head_v_dim-granular; safe on ++ # Q8_0 whose block=32 divides head_v_dim). At ratio=1 the layouts coincide. ++ ++ def rebind_device_views(self): ++ """Re-derive tensors that alias conv1d.weight's storage. ++ ++ ``conv_weights`` (passed to RadixLinearAttention) is a *view* of ++ ``self.conv1d.weight`` captured at construction time. Loaders that ++ load on CPU and then swap ``conv1d.weight.data`` for a device tensor ++ (e.g. ``--load-format layered_fp8``) leave that view pointing at the ++ freed CPU storage, so the conv1d kernel later dereferences an invalid ++ pointer. Rebuild the view from the current weight; the lazily-built ++ ESIMD copy (``_esimd_conv_weights``) self-heals on next forward, so ++ just drop it here. ++ """ ++ w = self.conv1d.weight ++ self.attn.conv_weights = w.view(w.size(0), w.size(2)) ++ self.attn.bias = self.conv1d.bias ++ self._esimd_conv_weights = None + + @staticmethod + def _override_weight_loader(param, loader): +@@ -300,7 +783,19 @@ class Qwen3_5GatedDeltaNet(nn.Module): + + def _bind_packed_weight_loaders(self, module): + """Bind packed-checkpoint-aware loaders to all relevant params of a merged module.""" +- for attr_name in ("weight", "weight_scale_inv", "weight_scale", "input_scale"): ++ # "qweight" / "qweight_type" cover the GGUF path: its merged params are ++ # named qweight (not weight), and its native weight_loader only accepts ++ # int shard ids. The packed wrapper splits a fused checkpoint tensor ++ # (e.g. GGUF attn_qkv = q|k|v) by the tuple shard id (0,1,2) into int ++ # shards before delegating, so GGUF GDN projections load correctly. ++ for attr_name in ( ++ "weight", ++ "weight_scale_inv", ++ "weight_scale", ++ "input_scale", ++ "qweight", ++ "qweight_type", ++ ): + param = getattr(module, attr_name, None) + if param is None: + continue +@@ -345,10 +840,16 @@ class Qwen3_5GatedDeltaNet(nn.Module): + module, param, loaded_shard_id + ) + +- if loaded_weight.numel() == 1: +- # Single-element tensor (scalar or [1]): +- # broadcast to each logical shard. +- chunks = [loaded_weight.view(-1)] * len(loaded_shard_id) ++ if len(loaded_weight.shape) == 0: ++ # Scalar shard payload. Two cases: ++ # - single logical shard: load as-is (original behavior). ++ # - GGUF qweight_type: one scalar quant-type for the whole ++ # fused tensor, replicated to each int shard so the GGUF ++ # loader records the type per shard slot. ++ if len(split_sizes) == 1 and split_sizes[0] == 1: ++ chunks = [loaded_weight.reshape(1)] ++ else: ++ chunks = [loaded_weight for _ in loaded_shard_id] + else: + split_dim = getattr(param, "output_dim", 0) + if _is_cpu: +@@ -444,7 +945,58 @@ class Qwen3_5GatedDeltaNet(nn.Module): + + return query, key, value, z, b, a + ++ def _esimd_fused_input_proj(self, hidden_states: torch.Tensor): ++ """Fuse in_proj_qkvz + in_proj_ba into one esimd_gemv_fp8_pert_fused2 call. ++ ++ Both projections read the same ``hidden_states`` and share the contraction ++ dim K, so the two fp8 GEMVs collapse into a single kernel dispatch. Returns ++ ``(qkvz, ba)`` on success or ``None`` to fall back to two separate linears. ++ Guarded to decode single-token (M=1), fp16 input, both weights fp8. ++ """ ++ if not _XPU_GDN_INPROJ_FUSED2: ++ return None ++ if hidden_states.dim() != 2 or hidden_states.shape[0] != 1: ++ return None ++ # A bias'd projection would need a post-GEMV add; keep it simple. ++ if getattr(self.in_proj_qkvz, "bias", None) is not None: ++ return None ++ if getattr(self.in_proj_ba, "bias", None) is not None: ++ return None ++ r0 = _esimd_fp8_weight_nk_scale(self.in_proj_qkvz) ++ r1 = _esimd_fp8_weight_nk_scale(self.in_proj_ba) ++ if r0 is None or r1 is None: ++ return None ++ w0, s0 = r0 ++ w1, s1 = r1 ++ if w0.shape[1] != w1.shape[1] or w0.shape[1] != hidden_states.shape[1]: ++ return None ++ try: ++ from custom_esimd_kernels_sglang import esimd_gemv_fp8_pert_fused2 ++ except Exception: ++ return None ++ x = ( ++ hidden_states ++ if hidden_states.dtype == torch.float16 ++ else hidden_states.to(torch.float16) ++ ) ++ x = x if x.is_contiguous() else x.contiguous() ++ scratch = getattr(self, "_esimd_inproj_scratch", None) ++ if scratch is None or scratch[0].shape[1] != w0.shape[0] or scratch[1].shape[1] != w1.shape[0]: ++ o0 = torch.empty((1, w0.shape[0]), dtype=torch.float16, device=x.device) ++ o1 = torch.empty((1, w1.shape[0]), dtype=torch.float16, device=x.device) ++ self._esimd_inproj_scratch = (o0, o1) ++ else: ++ o0, o1 = scratch ++ try: ++ esimd_gemv_fp8_pert_fused2(x, w0, s0, o0, w1, s1, o1) ++ except Exception: ++ return None ++ return o0, o1 ++ + def _forward_input_proj(self, hidden_states: torch.Tensor): ++ fused = self._esimd_fused_input_proj(hidden_states) ++ if fused is not None: ++ return fused + if ( + _is_cpu + or _is_npu +@@ -472,6 +1024,578 @@ class Qwen3_5GatedDeltaNet(nn.Module): + projected_states_ba, _ = self.in_proj_ba(hidden_states) + return projected_states_qkvz, projected_states_ba + ++ def _gdn_seq_to_interleaved( ++ self, ++ projected_states_qkvz: torch.Tensor, ++ projected_states_ba: torch.Tensor, ++ ): ++ """Reorder sglang's sequential qkvz/ba into the GQA-interleaved layout ++ the native ``sgl_kernel.gdn_attention`` kernel reads. ++ ++ Input (sglang, per-tp columns): ++ qkvz = [q_all(nk*hk) | k_all(nk*hk) | v_all(nv*hv) | z_all(nv*hv)] ++ ba = [b_all(nv) | a_all(nv)] ++ Output (kernel, per-k-head interleaved blocks): ++ qkvz = [ q(hk) k(hk) v(ratio*hv) z(ratio*hv) ] x nk ++ ba = [ b(ratio) a(ratio) ] x nk ++ where ratio = nv // nk and v-head v belongs to k-head v // ratio ++ (contiguous GQA grouping). Returns contiguous tensors. ++ """ ++ T = projected_states_qkvz.shape[0] ++ nk = self.num_k_heads // self.attn_tp_size ++ nv = self.num_v_heads // self.attn_tp_size ++ ratio = nv // nk ++ hk = self.head_k_dim ++ hv = self.head_v_dim ++ key_dim = nk * hk ++ val_dim = nv * hv ++ ++ qkvz = projected_states_qkvz ++ q = qkvz[:, 0:key_dim].reshape(T, nk, hk) ++ k = qkvz[:, key_dim : 2 * key_dim].reshape(T, nk, hk) ++ v = qkvz[:, 2 * key_dim : 2 * key_dim + val_dim].reshape(T, nk, ratio * hv) ++ z = qkvz[:, 2 * key_dim + val_dim :].reshape(T, nk, ratio * hv) ++ qkvz_il = torch.cat([q, k, v, z], dim=2).reshape(T, -1).contiguous() ++ ++ ba = projected_states_ba ++ b = ba[:, 0:nv].reshape(T, nk, ratio) ++ a = ba[:, nv : 2 * nv].reshape(T, nk, ratio) ++ ba_il = torch.cat([b, a], dim=2).reshape(T, -1).contiguous() ++ ++ return qkvz_il, ba_il ++ ++ def _forward_xpu_fast_path( ++ self, ++ projected_states_qkvz: torch.Tensor, ++ projected_states_ba: torch.Tensor, ++ forward_batch: ForwardBatch, ++ ): ++ """Run conv1d + GDN + RMSNormGated + out_proj using the vendored ++ torch.ops.sgl_kernel.gdn_attention kernel. Returns None if the op is ++ unusable (non-contiguous input, unexpected shape, ...) so the caller ++ falls back to the default Triton/PyTorch path. + + Layout assumption: num_v_heads // num_k_heads == 1 — then sglang's + MergedColumnParallelLinear sequential `[Q|K|V|Z]` layout and the @@ -4537,9 +6844,20 @@ index 65f78e691..bf045b854 100644 + + num_actual_tokens = projected_states_qkvz.shape[0] + -+ # Contiguity the kernel asserts on. -+ projected_states_qkvz = projected_states_qkvz.contiguous() -+ projected_states_ba = projected_states_ba.contiguous() ++ # Layout adapter: sglang produces projected_states_qkvz in SEQUENTIAL ++ # [q_all | k_all | v_all | z_all] order and projected_states_ba as ++ # [b_all | a_all] (see fix_query_key_value_ordering). The native ++ # gdn_attention kernel instead reads qkvz as GQA-INTERLEAVED per-k-head ++ # blocks [q(head_k), k(head_k), v(head_v*ratio), z(head_v*ratio)] and ba ++ # as per-k-head [b(ratio), a(ratio)] (see chunk_causal_conv1d_xe2.hpp: ++ # qkvz_elems_offset = k_head_id*qkvz_dim+off; chunk_reorder_zba step = ++ # (token*num_v + k_head*ratio)*2). Reorder here so the kernel sees the ++ # layout it expects. conv_weights/conv_state/ssm_state stay sequential ++ # (the kernel reads those via reordered_elems_offset), so they are ++ # untouched. All dims are per-tp (the projection is column-parallel). ++ projected_states_qkvz, projected_states_ba = self._gdn_seq_to_interleaved( ++ projected_states_qkvz, projected_states_ba ++ ) + + # Output buffers (kernel writes into these). + nv_tp = self.num_v_heads // self.attn_tp_size @@ -4607,6 +6925,178 @@ index 65f78e691..bf045b854 100644 + output, _ = self.out_proj(core_attn_out) + return output + ++ def _gguf_norm_out_proj( ++ self, core_attn_out: torch.Tensor, z_out: torch.Tensor ++ ) -> Optional[torch.Tensor]: ++ """RMSNormGated + GGUF q8_0 out_proj in one op dispatch. ++ ++ The GGUF counterpart of ``_esimd_norm_out_proj`` (which is fp8-only). ++ Removes the standalone ``gdn_rms_norm_gated`` launch plus the reshape / ++ cast glue around it, ~30 dispatches per decode step. ++ ++ The out_proj rep already carries the GGUF->HF value-head column permute ++ (baked into the global pre-shard weight, see __init__), so feeding the ++ raw rep is exactly what the unfused ``out_proj(...)`` would do. ++ """ ++ op = getattr(torch.ops.custom_esimd_kernels_sglang, ++ "esimd_norm_gemv_q8_0", None) ++ if op is None or getattr(self, "_gguf_outproj_off", False): ++ return None ++ cache = getattr(self, "_gguf_outproj_const", None) ++ if cache is None: ++ rep = _gguf_xpu_rep(self.out_proj, "q8_0") ++ if rep is None: ++ logger.warning("[gguf_norm_out_proj] out_proj is not a single q8_0 rep") ++ self._gguf_outproj_off = True ++ return None ++ qs, sc = rep[1], rep[2] ++ HV, V = int(core_attn_out.shape[1]), int(core_attn_out.shape[2]) ++ if qs.shape[1] != HV * V: ++ logger.warning("[gguf_norm_out_proj] K=%s != HV*V=%s", ++ qs.shape[1], HV * V) ++ self._gguf_outproj_off = True ++ return None ++ nw = self.norm.weight ++ nw = nw.to(torch.float16).contiguous() if nw.dtype != torch.float16 \ ++ else nw.contiguous() ++ if nw.numel() != V: ++ logger.warning("[gguf_norm_out_proj] norm weight %s != V=%s", ++ nw.numel(), V) ++ self._gguf_outproj_off = True ++ return None ++ cache = { ++ "qs": qs, "sc": sc, "nw": nw, "HV": HV, "V": V, ++ "eps": float(self.layer_norm_epsilon), ++ # Buffers keyed by token count; the kernel takes any M. ++ "bufs": {}, ++ } ++ self._gguf_outproj_const = cache ++ logger.warning("[gguf_norm_out_proj] ACTIVE (gated norm + out_proj folded)") ++ HV, V = cache["HV"], cache["V"] ++ M = int(core_attn_out.shape[0]) ++ buf = cache["bufs"].get(M) ++ if buf is None: ++ dev = cache["qs"].device ++ buf = { ++ "y": torch.empty((M * HV, V), dtype=torch.float16, device=dev), ++ "out": torch.empty((M, cache["qs"].shape[0]), ++ dtype=torch.float16, device=dev), ++ } ++ cache["bufs"][M] = buf ++ x = core_attn_out.reshape(M * HV, V) ++ z = z_out.reshape(M * HV, V) ++ if not x.is_contiguous(): ++ x = x.contiguous() ++ if not z.is_contiguous(): ++ z = z.contiguous() ++ try: ++ op(x, z, cache["nw"], buf["y"], ++ cache["qs"], cache["sc"], buf["out"], ++ HV, V, cache["eps"]) ++ except Exception as e: ++ logger.warning("[gguf_norm_out_proj] kernel raised %r", e) ++ self._gguf_outproj_off = True ++ return None ++ return buf["out"] ++ ++ def _esimd_norm_out_proj( ++ self, core_attn_out: torch.Tensor, z_out: torch.Tensor ++ ) -> Optional[torch.Tensor]: ++ """Fused RMSNormGated + FP8 out_proj GEMV for the GDN decode path. ++ ++ Replaces ``self.norm(x, z)`` (a standalone RMSNormGated launch) + ++ ``self.out_proj(...)`` (a separate fp8 GEMV) + the surrounding ++ cast/reshape/empty glue with a single ESIMD launch ++ (``esimd_norm_gemv_fp8_pert``). Gated by ``SGL_XPU_GDN_NORM_GEMV=1``. ++ ++ Semantics matched to RMSNormGated(norm_before_gate=True, ++ activation="swish"): ``out = (rmsnorm(x) * norm_weight * silu(z)) @ ++ out_proj.weight^T * scale``. The per-block out_proj weight_scale is ++ collapsed to a single per-tensor scalar (mean), mirroring the dense ++ ESIMD fp8 GEMV fast path. ++ ++ Returns the layer output ``[M, hidden]`` on success, or ``None`` to fall ++ back to the eager norm + out_proj path. The GGUF q8_0 branch handles any ++ decode batch size; the fp8 branch below is still single-token only ++ (``esimd_norm_gemv_fp8_pert`` emits ``[1, N]``). ++ """ ++ if core_attn_out.dim() != 3: ++ return None ++ if core_attn_out.shape != z_out.shape: ++ return None ++ if getattr(self.norm, "activation", "swish") not in ("swish", "silu"): ++ return None ++ if not getattr(self.norm, "norm_before_gate", True): ++ return None ++ # Phase 5c: GGUF q8_0 out_proj (the fp8 path below never matches it). ++ # M-generic: x/z are [M, HV, V] and the kernel derives M from the row ++ # count, so batched decode is still a single launch. ++ if ( ++ _XPU_GGUF_RESADD_NORM ++ and core_attn_out.dtype == torch.float16 ++ and 1 <= core_attn_out.shape[0] <= _GGUF_FUSE_MAX_M ++ ): ++ r = self._gguf_norm_out_proj(core_attn_out, z_out) ++ if r is not None: ++ return r ++ # fp8 kernel assumes a single decode token: x/z are [HV, V]. ++ if core_attn_out.shape[0] != 1: ++ return None ++ if not _XPU_GDN_NORM_GEMV: ++ return None ++ try: ++ from custom_esimd_kernels_sglang import esimd_norm_gemv_fp8_pert ++ except ImportError: ++ return None ++ resolved = _esimd_fp8_weight_nk_scale(self.out_proj) ++ if resolved is None: ++ return None ++ w_nk, scale_pt = resolved ++ HV = core_attn_out.shape[1] ++ V = core_attn_out.shape[2] ++ # Kernel contracts along HV*V; must equal out_proj in_features (K). ++ if HV * V != w_nk.shape[1]: ++ return None ++ x = core_attn_out.reshape(HV, V) ++ z = z_out.reshape(HV, V) ++ if x.dtype != torch.float16: ++ x = x.to(torch.float16) ++ if z.dtype != torch.float16: ++ z = z.to(torch.float16) ++ x = x.contiguous() ++ z = z.contiguous() ++ # Cache fp16 norm weight, collapsed per-tensor scale, and the output ++ # buffer once per layer (constants across calls; the buffer is reused so ++ # its data ptr stays stable across XPU-graph replays). ++ cache = getattr(self, "_esimd_outproj_const", None) ++ if cache is None: ++ nw = self.norm.weight ++ nw = ( ++ nw.to(torch.float16).contiguous() ++ if nw.dtype != torch.float16 ++ else nw.contiguous() ++ ) ++ out_buf = torch.empty( ++ (1, w_nk.shape[0]), dtype=torch.float16, device=w_nk.device ++ ) ++ cache = {"nw": nw, "scale": scale_pt, "w": w_nk, "out": out_buf} ++ self._esimd_outproj_const = cache ++ try: ++ esimd_norm_gemv_fp8_pert( ++ x, ++ z, ++ cache["nw"], ++ cache["w"], ++ cache["scale"], ++ cache["out"], ++ HV, ++ V, ++ float(self.layer_norm_epsilon), ++ ) ++ except Exception: ++ return None ++ return cache["out"] ++ + def _forward_xpu_esimd_gdn_decode( + self, + projected_states_qkvz: torch.Tensor, @@ -4648,24 +7138,33 @@ index 65f78e691..bf045b854 100644 + mamba_cache_params = linear_backend.req_to_token_pool.mamba2_layer_cache( + self.layer_id + ) -+ # sglang pool conv: (cache, conv_dim, W-1); kernel expects -+ # (cache, W-1, conv_dim) view (per-batch contiguous via stride). ++ # sglang pool conv native layout: (cache, conv_dim, W-1). The updated ++ # ESIMD kernel reads AND writes conv_state directly in this native ++ # layout (auto-detected in C++ by the small trailing dim W-1), so when ++ # the pool is fp16 we pass it in-place and skip both the whole-pool ++ # transpose copy and the per-layer index_copy_ writeback (~3 dispatched ++ # ops/layer * 30 GDN layers eliminated). Falls back to the legacy ++ # transposed-copy path when the pool dtype is not fp16. + pool_conv = mamba_cache_params.conv[0] + pool_ssm = mamba_cache_params.temporal -+ # The kernel reads conv state in (cache, W-1, conv_dim) layout, which -+ # the pool does not store, so we materialize a transposed-contiguous -+ # copy. Reuse a fixed buffer (copy_ rather than a fresh -+ # transpose().contiguous() each call) so its data ptr stays stable -+ # across XPU-graph replays. -+ cvcache = getattr(self, "_esimd_gdn_conv_view", None) -+ if cvcache is None or cvcache.shape != pool_conv.shape[:1] + pool_conv.shape[1:][::-1]: -+ cvcache = torch.empty( -+ (pool_conv.size(0), pool_conv.size(2), pool_conv.size(1)), -+ dtype=pool_conv.dtype, device=pool_conv.device, -+ ) -+ self._esimd_gdn_conv_view = cvcache -+ cvcache.copy_(pool_conv.transpose(-1, -2)) -+ conv_state_view = cvcache ++ _conv_native = (pool_conv.dtype == torch.float16) ++ if _conv_native: ++ conv_state_view = pool_conv ++ else: ++ # The kernel reads conv state in (cache, W-1, conv_dim) layout, which ++ # the pool does not store, so we materialize a transposed-contiguous ++ # copy. Reuse a fixed buffer (copy_ rather than a fresh ++ # transpose().contiguous() each call) so its data ptr stays stable ++ # across XPU-graph replays. ++ cvcache = getattr(self, "_esimd_gdn_conv_view", None) ++ if cvcache is None or cvcache.shape != pool_conv.shape[:1] + pool_conv.shape[1:][::-1]: ++ cvcache = torch.empty( ++ (pool_conv.size(0), pool_conv.size(2), pool_conv.size(1)), ++ dtype=pool_conv.dtype, device=pool_conv.device, ++ ) ++ self._esimd_gdn_conv_view = cvcache ++ cvcache.copy_(pool_conv.transpose(-1, -2)) ++ conv_state_view = cvcache + + # Cache the conv1d.weight view + zeros bias once per layer. + if getattr(self, "_esimd_conv_weights", None) is None: @@ -4744,22 +7243,44 @@ index 65f78e691..bf045b854 100644 + except Exception: + return None + -+ # Write conv_state back into pool: (cache, W-1, conv_dim) → (cache, conv_dim, W-1). -+ # index_copy_ writes only the touched slots; the conv_state_view above -+ # was a fresh copy of the whole pool, so untouched slots round-trip -+ # unchanged. -+ cache_indices_long = cache_indices.to(torch.long) -+ pool_conv.index_copy_( -+ 0, cache_indices_long, -+ conv_state_view.index_select(0, cache_indices_long).transpose(-1, -2).contiguous().to(pool_conv.dtype), -+ ) -+ if pool_ssm.dtype != torch.float16: ++ # conv_state writeback. In the native fp16 path the ESIMD kernel already ++ # shifted conv_state in-place into pool_conv, so no python writeback is ++ # needed. Only the legacy transposed-copy path (or a non-fp16 ssm pool) ++ # needs the index_copy_ round-trip. ++ # cache_indices is a step-level (batch) tensor identical for all 30 GDN ++ # layers; ``.to(torch.long)`` is a real dispatched copy (int32->int64). ++ # Memoize the long view on the per-step fwd_md object so only the first ++ # GDN layer that needs it pays the cast and the other layers reuse it. ++ need_conv_wb = not _conv_native ++ need_ssm_wb = pool_ssm.dtype != torch.float16 ++ if need_conv_wb or need_ssm_wb: ++ cache_indices_long = getattr(fwd_md, "_cache_indices_long", None) ++ if cache_indices_long is None or cache_indices_long.numel() != cache_indices.numel(): ++ cache_indices_long = cache_indices.to(torch.long) ++ try: ++ fwd_md._cache_indices_long = cache_indices_long ++ except Exception: ++ pass ++ if need_conv_wb: ++ # (cache, W-1, conv_dim) → (cache, conv_dim, W-1); only touched slots. ++ pool_conv.index_copy_( ++ 0, cache_indices_long, ++ conv_state_view.index_select(0, cache_indices_long).transpose(-1, -2).contiguous().to(pool_conv.dtype), ++ ) ++ if need_ssm_wb: + pool_ssm.index_copy_( + 0, cache_indices_long, + ssm_state_view.index_select(0, cache_indices_long).to(pool_ssm.dtype), + ) + + # Norm + out_proj. Mirrors the default path. ++ # Fast path: fuse RMSNormGated + fp8 out_proj into one ESIMD launch, ++ # eliminating the standalone norm kernel + separate GEMV + cast/reshape ++ # glue (decode single-token, SGL_XPU_GDN_NORM_GEMV=1). Falls back on None. ++ fused_out = self._esimd_norm_out_proj(core_attn_out, z_out) ++ if fused_out is not None: ++ return fused_out ++ + core_attn_out = core_attn_out.to(orig_dtype) + z_out = z_out.to(orig_dtype) + z_shape_og = z_out.shape @@ -4778,21 +7299,39 @@ index 65f78e691..bf045b854 100644 def forward( self, hidden_states: torch.Tensor, -@@ -487,6 +854,47 @@ class Qwen3_5GatedDeltaNet(nn.Module): +@@ -486,6 +1610,62 @@ class Qwen3_5GatedDeltaNet(nn.Module): + projected_states_qkvz, projected_states_ba = self._forward_input_proj( hidden_states ) - ++ return self._forward_from_projected( ++ projected_states_qkvz, projected_states_ba, forward_batch ++ ) ++ ++ def _forward_from_projected( ++ self, ++ projected_states_qkvz: torch.Tensor, ++ projected_states_ba: torch.Tensor, ++ forward_batch: ForwardBatch, ++ ): ++ """Core attention + output path starting from the in_proj outputs. ++ ++ Split out of ``forward`` so the fused input_layernorm + in_proj decode ++ path (esimd_resadd_norm_gemv2_fp8_pert) can supply the projected states ++ directly, skipping the standalone in_proj GEMV launches. ++ """ + # --- XPU native conv1d+GDN fast path (sgl_kernel.gdn_attention) --- + # Cherry-picked from origin/dev 7680aecdd4. Env-gated; falls back to + # the default Triton path if the kernel isn't usable on this shape. -+ _ENABLE_XPU_FAST_PATH = os.environ.get( -+ "SGLANG_XPU_GDN_FAST_PATH", "0" -+ ) == "1" ++ # Prefill(extend)-only: decode keeps the tuned ESIMD recurrent path ++ # below (esimd_gdn_conv_fused_seq); the native Xe2 kernel is a chunked ++ # (parallel) algorithm that only wins on long prefill sequences. ++ _ENABLE_XPU_FAST_PATH = _XPU_GDN_FAST_PATH + if ( + _ENABLE_XPU_FAST_PATH + and _is_xpu ++ and not forward_batch.forward_mode.is_decode() + and not forward_batch.forward_mode.is_target_verify() -+ and self.num_v_heads // self.num_k_heads == 1 ++ and self.num_v_heads % self.num_k_heads == 0 + ): + output = self._forward_xpu_fast_path( + projected_states_qkvz, @@ -4806,9 +7345,7 @@ index 65f78e691..bf045b854 100644 + # Calls the BMG-validated ESIMD kernel that fuses conv1d + the GDN + # recurrence in one launch. Sequential [q|k|v|z] layout — no gather. + # Decode-only; prefill stays on the Triton/PyTorch path. -+ _ENABLE_XPU_GDN_ESIMD = os.environ.get( -+ "SGL_XPU_GDN_ESIMD", "0" -+ ) == "1" ++ _ENABLE_XPU_GDN_ESIMD = _XPU_GDN_ESIMD + if ( + _ENABLE_XPU_GDN_ESIMD + and _is_xpu @@ -4822,24 +7359,394 @@ index 65f78e691..bf045b854 100644 + ) + if output is not None: + return output -+ + if ( self.num_v_heads // self.num_k_heads in [1, 2, 4] - and not _is_cpu -@@ -911,6 +1319,102 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): +@@ -623,56 +1803,341 @@ class Qwen3_5LinearDecoderLayer(nn.Module): + is_last_layer=(layer_id == config.num_hidden_layers - 1), + ) + +- def forward( ++ def _esimd_fused_input_norm_in_proj_q8_0( ++ self, ++ hidden_states: torch.Tensor, ++ residual: torch.Tensor, ++ qkvz_lin, ++ ba_lin, ++ r0q, ++ captured_last_layer_outputs: Optional[List[torch.Tensor]], ++ ): ++ """DEPRECATED (see _XPU_GGUF_RESADD_NORM_LEGACY). ++ ++ GGUF-q8_0 variant of the fused input_layernorm + GDN in_proj. ++ ++ ``in_proj_qkvz`` is q8_0; ``in_proj_ba`` is fp16 (GGUF keeps the tiny ++ b/a tensors unquantized). A single-matrix q8_0 kernel cannot fuse both, ++ so this fuses the GemmaRMSNorm into the q8_0 qkvz GEMV ++ (``esimd_resadd_norm_gemv_q8_0``, which also writes ``normed_out`` and a ++ separate ``new_residual``), then runs the fp16 ba projection on the ++ kernel-written normed hidden (its fp16 matmul is already ++ transpose-cached). Returns ``(qkvz, ba, new_residual)`` or ``None``. ++ Shared guards are checked by the caller. ++ ++ Superseded by ``_gguf_norm_gemv()``: ``esimd_resadd_norm_gemv_q8_ba`` ++ folds the fp16 ba GEMV into the same launch, so the shipped path costs ++ one op call here instead of two. ``esimd_resadd_norm_gemv_q8_0`` was ++ never landed, so this always ImportErrors and falls back. ++ """ ++ try: ++ from custom_esimd_kernels_sglang import esimd_resadd_norm_gemv_q8_0 ++ except ImportError: ++ return None ++ qs, scale = r0q ++ K = hidden_states.shape[1] ++ if qs.shape[1] != K: ++ return None ++ cache = getattr(self, "_esimd_resadd_q8_const", None) ++ if cache is None: ++ nw = ( ++ (self.input_layernorm.weight.data.to(torch.float32) + 1.0) ++ .to(torch.float16) ++ .contiguous() ++ ) ++ o0 = torch.empty((1, qs.shape[0]), dtype=torch.float16, device=qs.device) ++ normed = torch.empty((1, K), dtype=torch.float16, device=qs.device) ++ nr = torch.empty((1, K), dtype=torch.float16, device=qs.device) ++ cache = { ++ "nw": nw, ++ "qs": qs, ++ "scale": scale, ++ "o0": o0, ++ "normed": normed, ++ "nr": nr, ++ } ++ self._esimd_resadd_q8_const = cache ++ h = ( ++ hidden_states ++ if hidden_states.dtype == torch.float16 ++ else hidden_states.to(torch.float16) ++ ) ++ h = h.contiguous() ++ res = residual if residual.is_contiguous() else residual.contiguous() ++ try: ++ esimd_resadd_norm_gemv_q8_0( ++ h, ++ res, ++ cache["nw"], ++ cache["qs"], ++ cache["scale"], ++ cache["o0"], ++ cache["normed"], ++ cache["nr"], ++ float(self.input_layernorm.variance_epsilon), ++ ) ++ except Exception: ++ return None ++ # ba projection reads the normed hidden (matches in_proj_ba(normed) in ++ # the standard prepare_attn -> in_proj flow). fp16 GGUF matmul path. ++ try: ++ o1, _ = ba_lin(cache["normed"]) ++ except Exception: ++ return None ++ new_residual = cache["nr"] ++ if captured_last_layer_outputs is not None: ++ captured_last_layer_outputs.append(new_residual.clone()) ++ return cache["o0"], o1, new_residual ++ ++ def _esimd_fused_input_norm_in_proj( + self, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], +- **kwargs, ++ forward_batch: ForwardBatch, ++ captured_last_layer_outputs: Optional[List[torch.Tensor]], + ): +- forward_batch = kwargs.get("forward_batch", None) ++ """Fused GemmaRMSNorm(input_layernorm) + GDN in_proj (qkvz + ba). + +- hidden_states, residual = ( +- self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( ++ Replaces the prepare_attn ``input_layernorm`` launch plus the two ++ ``in_proj`` fp8 GEMV launches with a single ++ ``esimd_resadd_norm_gemv2_fp8_pert`` launch. Gated by ++ ``SGL_XPU_GDN_RESADD_NORM=1``; decode single-token only. ++ ++ Returns ``(projected_qkvz, projected_ba, new_residual)`` on success, or ++ ``None`` to fall back to the standard prepare_attn + in_proj path. Only ++ the plain TP decode path (no input-scatter, no allreduce-fusion) is ++ intercepted; every other configuration falls back. ++ """ ++ # Phase 5c: GGUF q8_0 build (the fp8 path below never matches it). ++ if _gguf_resadd_norm_guard(hidden_states, residual, forward_batch): ++ gdn = self.linear_attn ++ r = _gguf_norm_gemv( ++ self, ++ self.input_layernorm, ++ gdn.in_proj_qkvz, ++ gdn.in_proj_ba, + hidden_states, + residual, +- forward_batch, +- captured_last_layer_outputs=kwargs.get( +- "captured_last_layer_outputs", None +- ), ++ "gdn_in_proj", ++ ) ++ if r is not None: ++ o0, o1, new_residual = r ++ if captured_last_layer_outputs is not None: ++ captured_last_layer_outputs.append(new_residual.clone()) ++ return o0, o1, new_residual ++ # DEPRECATED legacy GGUF q8_0 path, opt-in only and never reached unless ++ # SGL_XPU_GGUF_RESADD_NORM_LEGACY=1 (see the flag's comment). ++ if _XPU_GGUF_RESADD_NORM_LEGACY and residual is not None and _gguf_resadd_norm_guard( ++ hidden_states, residual, forward_batch ++ ): ++ gdn = self.linear_attn ++ r0q = _esimd_q8_0_weight_scale(gdn.in_proj_qkvz) ++ if r0q is not None: ++ q8out = self._esimd_fused_input_norm_in_proj_q8_0( ++ hidden_states, ++ residual, ++ gdn.in_proj_qkvz, ++ gdn.in_proj_ba, ++ r0q, ++ captured_last_layer_outputs, ++ ) ++ if q8out is not None: ++ return q8out ++ if not _XPU_GDN_RESADD_NORM: ++ return None ++ if not (_is_xpu and forward_batch.forward_mode.is_decode()): ++ return None ++ if forward_batch.forward_mode.is_target_verify(): ++ return None ++ # Single decode token only (kernel emits [1, N]). ++ if hidden_states.dim() != 2 or hidden_states.shape[0] != 1: ++ return None ++ # First layer has no residual yet; kernel requires the residual add. ++ if residual is None: ++ return None ++ # Only interceptable when prepare_attn reduces to plain input_layernorm. ++ try: ++ from sglang.srt.layers.communicator import get_attn_tp_context ++ ++ if get_attn_tp_context().input_scattered: ++ return None ++ except Exception: ++ return None ++ if getattr(hidden_states, "_sglang_needs_allreduce_fusion", False): ++ return None ++ gdn = self.linear_attn ++ qkvz_lin = gdn.in_proj_qkvz ++ ba_lin = gdn.in_proj_ba ++ try: ++ from custom_esimd_kernels_sglang import esimd_resadd_norm_gemv2_fp8_pert ++ except ImportError: ++ return None ++ ++ # Cache row-major [N, K] fp8 weights + per-tensor scales, the ++ # (1 + weight) Gemma norm weight, and the reusable output buffers once ++ # per layer. _esimd_fp8_weight_nk_scale handles both the block-quant ++ # ([N, K]) and online per-tensor (stored [K, N], transposed) layouts and ++ # validates shapes so a bad layout never reaches the kernel. ++ cache = getattr(self, "_esimd_resadd_const", None) ++ if cache is None: ++ r0 = _esimd_fp8_weight_nk_scale(qkvz_lin) ++ r1 = _esimd_fp8_weight_nk_scale(ba_lin) ++ if r0 is None or r1 is None: ++ return None ++ w0, s0 = r0 ++ w1, s1 = r1 ++ # Contraction dim (K) must match the hidden size for both. ++ K = hidden_states.shape[1] ++ if w0.shape[1] != K or w1.shape[1] != K: ++ return None ++ # GemmaRMSNorm scales by (1 + weight); the kernel expects the ++ # pre-folded weight. ++ nw = ( ++ (self.input_layernorm.weight.data.to(torch.float32) + 1.0) ++ .to(torch.float16) ++ .contiguous() ++ ) ++ o0 = torch.empty((1, w0.shape[0]), dtype=torch.float16, device=w0.device) ++ o1 = torch.empty((1, w1.shape[0]), dtype=torch.float16, device=w1.device) ++ # Per-layer buffer for the kernel-written post-add residual ++ # (hidden + residual). Safe: each layer object owns a distinct nr, so ++ # the kernel never reads and writes the same buffer within one launch ++ # (input residual belongs to the *previous* layer's nr). ++ nr = torch.empty( ++ (1, hidden_states.shape[1]), dtype=torch.float16, device=w0.device + ) ++ cache = { ++ "nw": nw, ++ "w0": w0, ++ "s0": s0, ++ "w1": w1, ++ "s1": s1, ++ "o0": o0, ++ "o1": o1, ++ "nr": nr, ++ } ++ self._esimd_resadd_const = cache ++ ++ h = ( ++ hidden_states ++ if hidden_states.dtype == torch.float16 ++ else hidden_states.to(torch.float16) + ) ++ h = h.contiguous() ++ # ESIMD kernels read raw contiguous pointers; guard residual layout. ++ res = residual if residual.is_contiguous() else residual.contiguous() ++ try: ++ esimd_resadd_norm_gemv2_fp8_pert( ++ h, ++ res, ++ cache["nw"], ++ cache["w0"], ++ cache["s0"], ++ cache["o0"], ++ cache["w1"], ++ cache["s1"], ++ cache["o1"], ++ cache["nr"], ++ float(self.input_layernorm.variance_epsilon), ++ ) ++ except Exception: ++ return None ++ # The kernel now writes the post-add residual (hidden + residual, fp16) ++ # into cache["nr"] via its gid==0 group, eliminating the separate ++ # aten::add dispatch. ++ new_residual = cache["nr"] ++ if captured_last_layer_outputs is not None: ++ captured_last_layer_outputs.append(new_residual.clone()) ++ return cache["o0"], cache["o1"], new_residual ++ ++ def forward( ++ self, ++ hidden_states: torch.Tensor, ++ residual: Optional[torch.Tensor], ++ **kwargs, ++ ): ++ forward_batch = kwargs.get("forward_batch", None) + ++ fused = None + if not forward_batch.forward_mode.is_idle(): +- hidden_states = self.linear_attn( ++ fused = self._esimd_fused_input_norm_in_proj( + hidden_states, ++ residual, + forward_batch, ++ kwargs.get("captured_last_layer_outputs", None), + ) + +- # Fully Connected +- hidden_states, residual = self.layer_communicator.prepare_mlp( +- hidden_states, residual, forward_batch +- ) ++ if fused is not None: ++ projected_qkvz, projected_ba, residual = fused ++ hidden_states = self.linear_attn._forward_from_projected( ++ projected_qkvz, projected_ba, forward_batch ++ ) ++ else: ++ hidden_states, residual = ( ++ self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( ++ hidden_states, ++ residual, ++ forward_batch, ++ captured_last_layer_outputs=kwargs.get( ++ "captured_last_layer_outputs", None ++ ), ++ ) ++ ) ++ ++ if not forward_batch.forward_mode.is_idle(): ++ hidden_states = self.linear_attn( ++ hidden_states, ++ forward_batch, ++ ) + ++ if _NAN_PROBE: _nan_probe( ++ "gdn_attn_out", ++ hidden_states, ++ layer_id=self.layer_id, ++ forward_batch=forward_batch, ++ residual=residual, ++ ) ++ # Fully Connected + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) +- + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) ++ ++ # Phase 3: fold the post_attention_layernorm (resadd + rmsnorm) into the ++ # fused MoE router kernel, removing a per-layer dispatch. Falls back to ++ # the standard prepare_mlp + mlp path when not applicable. ++ fused_mlp = None + if isinstance(self.mlp, Qwen2MoeSparseMoeBlock): +- hidden_states = self.mlp( ++ fused_mlp = self.mlp.esimd_prepare_mlp_moe( + hidden_states, ++ residual, ++ self.post_attention_layernorm, + forward_batch, + use_reduce_scatter, + should_allreduce_fusion, + ) ++ ++ if fused_mlp is not None: ++ hidden_states, residual = fused_mlp + else: +- hidden_states = self.mlp( +- hidden_states, should_allreduce_fusion, use_reduce_scatter ++ hidden_states, residual = self.layer_communicator.prepare_mlp( ++ hidden_states, residual, forward_batch + ) ++ if isinstance(self.mlp, Qwen2MoeSparseMoeBlock): ++ hidden_states = self.mlp( ++ hidden_states, ++ forward_batch, ++ use_reduce_scatter, ++ should_allreduce_fusion, ++ ) ++ else: ++ hidden_states = self.mlp( ++ hidden_states, should_allreduce_fusion, use_reduce_scatter ++ ) + if should_allreduce_fusion: + hidden_states._sglang_needs_allreduce_fusion = True + else: +@@ -909,8 +2374,133 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): + positions: torch.Tensor, + hidden_states: torch.Tensor, forward_batch: ForwardBatch, ++ precomputed_qkv: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Full attention forward pass.""" +- """Full attention forward pass.""" ++ """Full attention forward pass. ++ ++ When ``precomputed_qkv`` is provided (Phase 5b fused input-norm + qkv ++ path), the ``qkv_proj`` GEMV is skipped and the supplied ``[1, N]`` fp16 ++ qkv tensor is fed straight into the ESIMD split/norm/rope path. ++ """ + # vllm parity: fuse split + qk_norm + rope into single ESIMD call. + # Hard-coded requirements: head_dim=256, fp16 model, GemmaRMSNorm + # weight+1.0 convention. The kernel is fp16-only; on a bf16 model the + # implicit cast loses dynamic range during RoPE and gsm8k drops from + # 0.80 to 0.40 (verified). Restrict to fp16 inputs only. + if ( -+ os.environ.get("SGL_XPU_FA_ESIMD_QKV") == "1" ++ _XPU_FA_ESIMD_QKV + and self.head_dim == 256 -+ and hidden_states.dim() == 2 -+ and hidden_states.dtype == torch.float16 ++ and ( ++ precomputed_qkv is not None ++ or (hidden_states.dim() == 2 and hidden_states.dtype == torch.float16) ++ ) + ): + try: + # Prefer the BMG sglang variant; fall back to the vllm one @@ -4853,7 +7760,10 @@ index 65f78e691..bf045b854 100644 + except ImportError: + esimd_qkv_split_norm_rope = None + if esimd_qkv_split_norm_rope is not None: -+ qkv, _ = self.qkv_proj(hidden_states) ++ if precomputed_qkv is not None: ++ qkv = precomputed_qkv ++ else: ++ qkv, _ = self.qkv_proj(hidden_states) + nTokens = qkv.shape[0] + orig_dtype = qkv.dtype + qkv_fp16 = qkv.to(torch.float16).contiguous() @@ -4904,7 +7814,21 @@ index 65f78e691..bf045b854 100644 + ) + else: + _, q_out, gate_out, k_out, v_out = scratch -+ pos_i32 = positions.to(torch.int32).contiguous() ++ # Phase 5c: positions is identical across all full-attention ++ # layers within a decode step, so the int32 conversion (a real ++ # dispatched copy) is redundant after the first layer. Memoize it ++ # on the per-step forward_batch (fresh each step -> auto-invalidates) ++ # keyed by the positions object identity. Saves ~9/step copy_. ++ pos_i32 = getattr(forward_batch, "_esimd_pos_i32", None) ++ if pos_i32 is None or getattr( ++ forward_batch, "_esimd_pos_id", None ++ ) != id(positions): ++ pos_i32 = positions.to(torch.int32).contiguous() ++ try: ++ forward_batch._esimd_pos_i32 = pos_i32 ++ forward_batch._esimd_pos_id = id(positions) ++ except Exception: ++ pass + esimd_qkv_split_norm_rope( + qkv_fp16, + q_out, gate_out, k_out, v_out, @@ -4915,10 +7839,14 @@ index 65f78e691..bf045b854 100644 + self.attn_output_gate, + cache["rotary_dim"], cache["cos_sin"], + ) -+ q = q_out.to(orig_dtype) -+ k = k_out.to(orig_dtype) -+ v = v_out.to(orig_dtype) -+ gate = gate_out.to(orig_dtype) if self.attn_output_gate else None ++ q = q_out if q_out.dtype == orig_dtype else q_out.to(orig_dtype) ++ k = k_out if k_out.dtype == orig_dtype else k_out.to(orig_dtype) ++ v = v_out if v_out.dtype == orig_dtype else v_out.to(orig_dtype) ++ gate = ( ++ (gate_out if gate_out.dtype == orig_dtype else gate_out.to(orig_dtype)) ++ if self.attn_output_gate ++ else None ++ ) + attn_output = self.attn(q, k, v, forward_batch) + if self.attn_output_gate: + # ESIMD kernel already applies sigmoid; don't re-sigmoid. @@ -4929,7 +7857,342 @@ index 65f78e691..bf045b854 100644 if ( not _is_npu or forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed() -@@ -1091,6 +1595,12 @@ class Qwen3_5ForCausalLM(nn.Module): +@@ -936,6 +2526,227 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): + output, _ = self.o_proj(attn_output) + return output + ++ def _esimd_fused_input_norm_qkv_q8_0( ++ self, ++ hidden_states: torch.Tensor, ++ residual: torch.Tensor, ++ r0q, ++ captured_last_layer_outputs: Optional[List[torch.Tensor]], ++ ): ++ """DEPRECATED (see _XPU_GGUF_RESADD_NORM_LEGACY). ++ ++ GGUF-q8_0 variant of the fused input_layernorm + qkv_proj. ++ ++ Uses the single-matrix ``esimd_resadd_norm_gemv_q8_0`` op (int8 weight + ++ per-32-block fp16 scale). ``normed_out`` is a scratch buffer here (the ++ full-attn path feeds ``o0`` straight into ``esimd_qkv_split_norm_rope`` ++ and does not reuse the normed hidden). Returns ``(qkv, new_residual)`` or ++ ``None`` to fall back. Shared guards are checked by the caller. ++ ++ Superseded by ``_gguf_norm_gemv()`` on ``esimd_resadd_norm_gemv_q8_ba``. ++ ``esimd_resadd_norm_gemv_q8_0`` was never landed in ++ custom-esimd-kernels, so this always ImportErrors and falls back. ++ """ ++ try: ++ from custom_esimd_kernels_sglang import esimd_resadd_norm_gemv_q8_0 ++ except ImportError: ++ return None ++ qs, scale = r0q ++ K = hidden_states.shape[1] ++ if qs.shape[1] != K: ++ return None ++ cache = getattr(self, "_esimd_fa_norm_q8_const", None) ++ if cache is None: ++ nw = ( ++ (self.input_layernorm.weight.data.to(torch.float32) + 1.0) ++ .to(torch.float16) ++ .contiguous() ++ ) ++ o0 = torch.empty((1, qs.shape[0]), dtype=torch.float16, device=qs.device) ++ normed = torch.empty((1, K), dtype=torch.float16, device=qs.device) ++ nr = torch.empty((1, K), dtype=torch.float16, device=qs.device) ++ cache = { ++ "nw": nw, ++ "qs": qs, ++ "scale": scale, ++ "o0": o0, ++ "normed": normed, ++ "nr": nr, ++ } ++ self._esimd_fa_norm_q8_const = cache ++ h = ( ++ hidden_states ++ if hidden_states.dtype == torch.float16 ++ else hidden_states.to(torch.float16) ++ ) ++ h = h.contiguous() ++ res = residual if residual.is_contiguous() else residual.contiguous() ++ try: ++ esimd_resadd_norm_gemv_q8_0( ++ h, ++ res, ++ cache["nw"], ++ cache["qs"], ++ cache["scale"], ++ cache["o0"], ++ cache["normed"], ++ cache["nr"], ++ float(self.input_layernorm.variance_epsilon), ++ ) ++ except Exception: ++ return None ++ new_residual = cache["nr"] ++ if captured_last_layer_outputs is not None: ++ captured_last_layer_outputs.append(new_residual.clone()) ++ return cache["o0"], new_residual ++ ++ def _esimd_fused_input_norm_qkv( ++ self, ++ hidden_states: torch.Tensor, ++ residual: Optional[torch.Tensor], ++ forward_batch: ForwardBatch, ++ captured_last_layer_outputs: Optional[List[torch.Tensor]], ++ ): ++ """Fused GemmaRMSNorm(input_layernorm) + qkv_proj (Phase 5b). ++ ++ Replaces the prepare_attn ``input_layernorm`` launch ++ (gemma_fused_add_rmsnorm) plus the ``qkv_proj`` fp8 GEMV with a single ++ ``esimd_resadd_norm_gemv2_fp8_pert`` launch (qkv as matrix-0, a 1-row ++ dummy as matrix-1), with the post-add residual written back by the ++ kernel. Gated by ``SGL_XPU_FA_RESADD_NORM=1``; decode single-token, fp8 ++ qkv only. Requires the ESIMD qkv split path (head_dim==256, fp16) since ++ the resulting qkv feeds ``esimd_qkv_split_norm_rope``. ++ ++ Returns ``(qkv, new_residual)`` on success, or ``None`` to fall back to ++ the standard prepare_attn + self_attention path. ++ """ ++ # Phase 5c: GGUF q8_0 build (the fp8 path below never matches it). ++ if _gguf_resadd_norm_guard(hidden_states, residual, forward_batch): ++ r = _gguf_norm_gemv( ++ self, ++ self.input_layernorm, ++ # self_attention is a method on this layer, not a submodule: ++ # qkv_proj hangs off the layer itself. ++ self.qkv_proj, ++ None, ++ hidden_states, ++ residual, ++ "fa_qkv", ++ ) ++ if r is not None: ++ o0, _, new_residual = r ++ if captured_last_layer_outputs is not None: ++ captured_last_layer_outputs.append(new_residual.clone()) ++ return o0, new_residual ++ # DEPRECATED legacy GGUF q8_0 path, opt-in only and never reached unless ++ # SGL_XPU_GGUF_RESADD_NORM_LEGACY=1 (see the flag's comment). ++ if ( ++ _XPU_GGUF_RESADD_NORM_LEGACY ++ and _XPU_FA_ESIMD_QKV ++ and residual is not None ++ and _gguf_resadd_norm_guard(hidden_states, residual, forward_batch) ++ ): ++ r0q = _esimd_q8_0_weight_scale(self.qkv_proj) ++ if r0q is not None: ++ q8out = self._esimd_fused_input_norm_qkv_q8_0( ++ hidden_states, residual, r0q, captured_last_layer_outputs ++ ) ++ if q8out is not None: ++ return q8out ++ if not (_XPU_FA_RESADD_NORM and _XPU_FA_ESIMD_QKV): ++ return None ++ if not (_is_xpu and forward_batch.forward_mode.is_decode()): ++ return None ++ if forward_batch.forward_mode.is_target_verify(): ++ return None ++ if self.head_dim != 256: ++ return None ++ # Single decode token only (kernel emits [1, N]). ++ if hidden_states.dim() != 2 or hidden_states.shape[0] != 1: ++ return None ++ # First layer has no residual yet; kernel requires the residual add. ++ if residual is None: ++ return None ++ # Only interceptable when prepare_attn reduces to plain input_layernorm. ++ try: ++ from sglang.srt.layers.communicator import get_attn_tp_context ++ ++ if get_attn_tp_context().input_scattered: ++ return None ++ except Exception: ++ return None ++ if getattr(hidden_states, "_sglang_needs_allreduce_fusion", False): ++ return None ++ try: ++ from custom_esimd_kernels_sglang import esimd_resadd_norm_gemv2_fp8_pert ++ except ImportError: ++ return None ++ ++ # Cache row-major [N, K] fp8 qkv weight + per-tensor scale, the ++ # (1 + weight) Gemma norm weight, the reusable output/residual buffers, ++ # and a 1-row dummy second matrix (the deployed op fuses two matrices; ++ # the dummy contributes a discarded zero column but lets us reuse the op ++ # without a new kernel). All static after load. ++ cache = getattr(self, "_esimd_fa_norm_const", None) ++ if cache is None: ++ r0 = _esimd_fp8_weight_nk_scale(self.qkv_proj) ++ if r0 is None: ++ return None ++ w0, s0 = r0 ++ K = hidden_states.shape[1] ++ if w0.shape[1] != K: ++ return None ++ nw = ( ++ (self.input_layernorm.weight.data.to(torch.float32) + 1.0) ++ .to(torch.float16) ++ .contiguous() ++ ) ++ o0 = torch.empty((1, w0.shape[0]), dtype=torch.float16, device=w0.device) ++ # 1-row dummy second matrix (zeros -> contributes nothing). ++ w1 = torch.zeros((1, K), dtype=w0.dtype, device=w0.device) ++ s1 = torch.ones((1,), dtype=torch.float32, device=w0.device) ++ o1 = torch.empty((1, 1), dtype=torch.float16, device=w0.device) ++ nr = torch.empty((1, K), dtype=torch.float16, device=w0.device) ++ cache = { ++ "nw": nw, ++ "w0": w0, ++ "s0": s0, ++ "w1": w1, ++ "s1": s1, ++ "o0": o0, ++ "o1": o1, ++ "nr": nr, ++ } ++ self._esimd_fa_norm_const = cache ++ ++ h = ( ++ hidden_states ++ if hidden_states.dtype == torch.float16 ++ else hidden_states.to(torch.float16) ++ ) ++ h = h.contiguous() ++ res = residual if residual.is_contiguous() else residual.contiguous() ++ try: ++ esimd_resadd_norm_gemv2_fp8_pert( ++ h, ++ res, ++ cache["nw"], ++ cache["w0"], ++ cache["s0"], ++ cache["o0"], ++ cache["w1"], ++ cache["s1"], ++ cache["o1"], ++ cache["nr"], ++ float(self.input_layernorm.variance_epsilon), ++ ) ++ except Exception: ++ return None ++ new_residual = cache["nr"] ++ if captured_last_layer_outputs is not None: ++ captured_last_layer_outputs.append(new_residual.clone()) ++ return cache["o0"], new_residual ++ + def forward( + self, + positions: torch.Tensor, +@@ -945,46 +2756,93 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): + captured_last_layer_outputs: Optional[list[torch.Tensor]] = None, + **kwargs, + ): +- hidden_states, residual = ( +- self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( ++ # Phase 5b: try to fuse input_layernorm (resadd + rmsnorm) into the ++ # qkv_proj GEMV, replacing prepare_attn's gemma_fused_add_rmsnorm launch. ++ # On a hit we get qkv + new_residual directly and feed qkv straight into ++ # self_attention, skipping prepare_attn and qkv_proj. Falls back to the ++ # standard prepare_attn path otherwise. ++ fused_qkv = None ++ if not forward_batch.forward_mode.is_idle(): ++ fused_qkv = self._esimd_fused_input_norm_qkv( + hidden_states, + residual, + forward_batch, +- captured_last_layer_outputs=captured_last_layer_outputs, ++ captured_last_layer_outputs, + ) +- ) + +- if not forward_batch.forward_mode.is_idle(): ++ if fused_qkv is not None: ++ qkv, residual = fused_qkv + hidden_states = self.self_attention( + positions=positions, +- hidden_states=hidden_states, ++ hidden_states=None, + forward_batch=forward_batch, ++ precomputed_qkv=qkv, ++ ) ++ else: ++ hidden_states, residual = ( ++ self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( ++ hidden_states, ++ residual, ++ forward_batch, ++ captured_last_layer_outputs=captured_last_layer_outputs, ++ ) + ) + +- # Fully Connected +- hidden_states, residual = self.layer_communicator.prepare_mlp( +- hidden_states, residual, forward_batch ++ if not forward_batch.forward_mode.is_idle(): ++ hidden_states = self.self_attention( ++ positions=positions, ++ hidden_states=hidden_states, ++ forward_batch=forward_batch, ++ ) ++ ++ if _NAN_PROBE: _nan_probe( ++ "full_attn_out", ++ hidden_states, ++ layer_id=self.layer_id, ++ forward_batch=forward_batch, ++ residual=residual, + ) ++ # Fully Connected + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) +- + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) ++ ++ # Phase 3: fold the post_attention_layernorm (resadd + rmsnorm) into the ++ # fused MoE router kernel, removing a per-layer dispatch. Falls back to ++ # the standard prepare_mlp + mlp path when not applicable. ++ fused_mlp = None + if isinstance(self.mlp, Qwen2MoeSparseMoeBlock): +- hidden_states = self.mlp( ++ fused_mlp = self.mlp.esimd_prepare_mlp_moe( + hidden_states, ++ residual, ++ self.post_attention_layernorm, + forward_batch, + use_reduce_scatter, + should_allreduce_fusion, + ) ++ ++ if fused_mlp is not None: ++ hidden_states, residual = fused_mlp + else: +- hidden_states = self.mlp( +- hidden_states, should_allreduce_fusion, use_reduce_scatter ++ hidden_states, residual = self.layer_communicator.prepare_mlp( ++ hidden_states, residual, forward_batch + ) ++ if isinstance(self.mlp, Qwen2MoeSparseMoeBlock): ++ hidden_states = self.mlp( ++ hidden_states, ++ forward_batch, ++ use_reduce_scatter, ++ should_allreduce_fusion, ++ ) ++ else: ++ hidden_states = self.mlp( ++ hidden_states, should_allreduce_fusion, use_reduce_scatter ++ ) + if should_allreduce_fusion: + hidden_states._sglang_needs_allreduce_fusion = True + else: +@@ -1091,6 +2949,12 @@ class Qwen3_5ForCausalLM(nn.Module): config.hidden_size, org_num_embeddings=config.vocab_size, enable_tp=not is_dp_attention_enabled(), @@ -4942,7 +8205,31 @@ index 65f78e691..bf045b854 100644 ) else: self.embed_tokens = PPMissingLayer() -@@ -1718,6 +2228,113 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): +@@ -1167,6 +3031,9 @@ class Qwen3_5ForCausalLM(nn.Module): + residual = pp_proxy_tensors["residual"] + + aux_hidden_states = [] ++ if _NAN_PROBE: ++ _nan_probe_new_forward() ++ _nan_probe("embed_out", hidden_states, forward_batch=forward_batch) + # Pass through decoder layers + for layer_idx in range(self.start_layer, self.end_layer): + layer = self.layers[layer_idx] +@@ -1184,6 +3051,13 @@ class Qwen3_5ForCausalLM(nn.Module): + else None + ), + ) ++ if _NAN_PROBE: _nan_probe( ++ "layer_out", ++ hidden_states, ++ layer_id=layer_idx, ++ forward_batch=forward_batch, ++ residual=residual, ++ ) + + # Process deepstack embeddings if provided + if ( +@@ -1718,6 +3592,113 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): torch.cuda.empty_cache() torch.cuda.synchronize() @@ -5056,7 +8343,7 @@ index 65f78e691..bf045b854 100644 def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ # (param_name, shard_name, shard_id) -@@ -1733,6 +2350,27 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): +@@ -1733,6 +3714,27 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): ("in_proj_ba.", "in_proj_a.", 1), ] @@ -5084,7 +8371,7 @@ index 65f78e691..bf045b854 100644 num_experts = self.config.num_experts # Params for weights, fp8 weight scales, fp8 activation scales -@@ -1832,6 +2470,14 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): +@@ -1832,6 +3834,14 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): continue if "mtp" in name: continue @@ -5099,7 +8386,7 @@ index 65f78e691..bf045b854 100644 if "language_model" in name: name = name.replace(r"model.language_model.", r"model.") if ".self_attn." in name: -@@ -1886,6 +2532,24 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): +@@ -1886,6 +3896,24 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): if "mlp.experts" in name: continue name = name.replace(weight_name, param_name) @@ -5124,7 +8411,7 @@ index 65f78e691..bf045b854 100644 # Skip loading extra parameters for GPTQ/modelopt models. if name.endswith(ignore_suffixes) and name not in params_dict: continue -@@ -2013,6 +2677,24 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): +@@ -2013,6 +4041,24 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): if name in params_dict.keys(): param = params_dict[name] @@ -5180,6 +8467,18 @@ index 6c77ff64f..78d88e51c 100644 if self.speculative_num_draft_tokens is not None: assert not self.enable_mamba_extra_buffer_lazy(), ( "extra_buffer_lazy is not yet supported with speculative decoding. " +diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py +index 4556d06b1..c5e35caf9 100644 +--- a/python/sglang/srt/utils/common.py ++++ b/python/sglang/srt/utils/common.py +@@ -329,6 +329,7 @@ def xpu_has_xmx_support(): + return False + + ++@lru_cache(maxsize=1) + def use_intel_xpu_backend(): + return get_bool_env_var("SGLANG_USE_SGL_XPU") and is_xpu() + diff --git a/python/sglang/srt/utils/hf_transformers/config.py b/python/sglang/srt/utils/hf_transformers/config.py index f66ea9a67..8e5f7beda 100644 --- a/python/sglang/srt/utils/hf_transformers/config.py diff --git a/sglang/scripts/start_qwen3_6_service.sh b/sglang/scripts/start_qwen3_6_service.sh index 9138557b..8e8ae830 100755 --- a/sglang/scripts/start_qwen3_6_service.sh +++ b/sglang/scripts/start_qwen3_6_service.sh @@ -1,9 +1,9 @@ #!/usr/bin/env bash # Launch SGLang server for Qwen3.6-35B-A3B online fp8 on Intel BMG, TP=2. # -# Golden fp8 + full-ESIMD + XPU-graph config (matches the sgl-fp8-perf setup). -# All ESIMD fast-paths + prefill fast-paths enabled. Required env knobs are -# documented inline. +# e5m2 online-fp8 + full-ESIMD config, XPU-graph DISABLED (accuracy). +# All ESIMD fast-paths + prefill fast-paths + e5m2 fused decode kernels enabled. +# Required env knobs are documented inline. set -euo pipefail @@ -31,6 +31,10 @@ export SGLANG_MAMBA_SSM_DTYPE=float16 export SGL_XPU_ESIMD_DECODE=1 # MoE silu routed kernel (replaces triton fused_moe on XPU) export SGL_XPU_ESIMD_MOE=1 +# Full decode MoE fusion: router topk + routed + shared + gate -> 1 dispatch. +# e5m2 only (SGLANG_FP8_DTYPE=e5m2 below). Reads native N-major w13 (no +# transposed weight copy). This is the main decode TPOT lever for this model. +export SGL_XPU_ESIMD_MOE_FULL=1 # MoE prefill ESIMD (M-tiled DPAS fp8 MoE prefill) export SGL_XPU_ESIMD_MOE_PREFILL=1 # Full-attention fused QKV split + RMSNorm + RoPE (Qwen3.5/3.6) @@ -46,10 +50,29 @@ export SGL_XPU_GDN_EXTEND_ESIMD=1 export SGL_XPU_PREFILL_DPAS=1 # --- XPU Graph (CUDA-graph-equivalent) --- -# Captures the decode forward graph for a TPOT speedup at BS=1. -# Safe to leave on with the kernels in this image; falls back to eager -# replay on sequences > 16384 tokens (kernel MAX_SPLITS cap). -export SGL_XPU_ENABLE_GRAPH=1 +# DISABLED: xpu-graph accuracy is unstable on this model, so decode runs eager. +# The e5m2 MoE-full fusion + resadd-norm fusions below recover the per-step +# host-dispatch cost that the graph would otherwise have hidden. +export SGL_XPU_ENABLE_GRAPH=0 + +# --- e5m2 online-quant + fused decode kernels --- +# Quantize online fp8 to e5m2 (the fused MoE-full decode kernels require e5m2). +export SGLANG_FP8_DTYPE=e5m2 +# GDN gated-RMSNorm as an ESIMD GEMV (decode). +export SGL_XPU_GDN_NORM_GEMV=1 +# Superseded by GDN_RESADD_NORM (which fuses in_proj qkvz+ba WITH input_layernorm); +# the standalone in_proj fused2 gave no e2e gain -> keep OFF. +export SGL_XPU_GDN_INPROJ_FUSED2=0 +# Fuse input_layernorm (resadd+rmsnorm) + GDN in_proj (qkvz+ba) into one ESIMD GEMV. +export SGL_XPU_GDN_RESADD_NORM=1 +# Fuse full-attention input_layernorm (resadd+rmsnorm) into qkv_proj (decode). +export SGL_XPU_FA_RESADD_NORM=1 +# MoE router as fp8 ESIMD GEMV instead of fp16 aten::mm (saves a host launch per +# MoE layer). Perturbs top-8 routing on a fraction of tokens -> gsm8k A/B before +# trusting; set to 0 to fall back to the accurate fp16 gate. +export SGL_XPU_MOE_ROUTER_FP8=1 +# Skip the per-step TP token-count sync (host overhead) on this single-node TP setup. +export SGLANG_XPU_TP_SYNC_TOKENS=0 # --load-format layered_fp8: build on CPU, load the full bf16 checkpoint into # host RAM, then move + quantize each module onto the device one at a time.