SGL: optimize Qwen3.6 27B and 35B A3B online fp8 quantization performance on BMG w/o XPU Graph and resolve acc issue.

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).

sglang/custom-esimd-kernels/.gitignore

@@ -1,5 +1,5 @@
 *.so
-build/
+build*/
 __pycache__/
 *.egg-info/
 *.egg-link

sglang/custom-esimd-kernels/csrc/eagle/eagle.sycl

@@ -503,15 +503,16 @@ 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<torch::Tensor, torch::Tensor> chunk_gated_delta_rule_extend(
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> chunk_gated_delta_rule_extend(
     torch::Tensor& q,
     torch::Tensor& k,
     torch::Tensor& v,
     torch::Tensor& g,
     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<torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor> 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,17 +595,17 @@ std::tuple<torch::Tensor, torch::Tensor> chunk_gated_delta_rule_extend(
           cgh.parallel_for<class chunkGdnExtendFp16Sf16>(
             range, [=](sycl::nd_item<2> ndi) SYCL_ESIMD_KERNEL {
               chunkGatedDeltaRuleExtendFp16<sycl::half>(
-                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 {
         kernel = [=](sycl::handler& cgh) {
           cgh.parallel_for<class chunkGdnExtendFp16Sf32>(
             range, [=](sycl::nd_item<2> ndi) SYCL_ESIMD_KERNEL {
               chunkGatedDeltaRuleExtendFp16<float>(
-                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,17 +617,17 @@ std::tuple<torch::Tensor, torch::Tensor> chunk_gated_delta_rule_extend(
           cgh.parallel_for<class chunkGdnExtendBf16Sf16>(
             range, [=](sycl::nd_item<2> ndi) SYCL_ESIMD_KERNEL {
               chunkGatedDeltaRuleExtendBf16<sycl::half>(
-                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 {
         kernel = [=](sycl::handler& cgh) {
           cgh.parallel_for<class chunkGdnExtendBf16Sf32>(
             range, [=](sycl::nd_item<2> ndi) SYCL_ESIMD_KERNEL {
               chunkGatedDeltaRuleExtendBf16<float>(
-                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<torch::Tensor, torch::Tensor> 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.

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<StateT, 128> hv = fp32InS_persistent.select<128, 1>(r * 128);
+        block_store<StateT, 128>(hDst + r * headDim, hv);
+      }
+      hSnapIdx += 1;
+      hNextSnapTok += hChunkSize;
+    }
+
     simd<bf16, 128> q_bf16 = block_load<bf16, 128>(qPtr + t * qkTokStride + kHeadIdx * headDim);
     simd<bf16, 128> k_bf16 = block_load<bf16, 128>(kPtr + t * qkTokStride + kHeadIdx * headDim);
     simd<bf16, 128> v_bf16 = block_load<bf16, 128>(vPtr + t * vTokStride  + headIdx  * headDim);

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<StateT, 128> hv = fp32InS_persistent.select<128, 1>(r * 128);
+        block_store<StateT, 128>(hDst + r * headDim, hv);
+      }
+      hSnapIdx += 1;
+      hNextSnapTok += hChunkSize;
+    }
+
     // --- Load q, k, v for the token's head -----------------------------
     simd<fp16, 128> q_fp16 = block_load<fp16, 128>(qPtr + t * qkTokStride + kHeadIdx * headDim);
     simd<fp16, 128> k_fp16 = block_load<fp16, 128>(kPtr + t * qkTokStride + kHeadIdx * headDim);