Qwen3.5 moe optimizations (#113)

* tilelang gated RMSNorm for Qwen3.5 GDN output

(cherry picked from commit 6dcdb3f)

* fused zba spilt

(cherry picked from commit 58d9e4a)

* support decode gdn mate

(cherry picked from commit 8b463e9)

* qwen3.5 tilelang prefill causal conv1d

(cherry picked from commit 55c085e)

* feed mate no-copy strided q/k/v views and fused gdn gating

(cherry picked from commit e5f6b40)

* del causal_conv1d d2h and gemmarmsnorm 3D to 2D

(cherry picked from commit e947cdb)
Series patch renumbered 0086 -> 0081; v0.24.0-dev already uses 0086.

* use mate output to reduce dtod and del cu_seqlens to int64

(cherry picked from commit 32b7287)

* del fused moe deepseek v4 name and fused glue deepgemm moe prefill

(cherry picked from commit 088e0fb)
Keep a default-on VLLM_MUSA_MOE_DEEPGEMM_PREFILL kill switch so the
generalized prefill path can be turned off without a code change.

* shared expert folded into the routed FP8 grouped GEMM as slot 256

(cherry picked from commit 7312722)
Widen the existing fold gate to accept FP8 [128,128] block weights in
addition to unquantized, keep it off under expert parallelism / EPLB,
and add the VLLM_MUSA_MOE_SHARED_EXPERT_FUSION switch.

* row-tiled fused silu+mul+FP8-quant kernel replace silu_and_mul_per_token_group_fp8_quant

(cherry picked from commit 44268a4)

* optional separate contiguous mamba state pool for in-place GDN decode

(cherry picked from commit 7e8c9b3)
Only the dedicated BlockPool half: the contiguous per-state striding it
also carried is already unconditional in the gpu_model_runner patch.
Regenerated the kv_cache_utils patch, whose upstream copy was truncated.

* default the separate mamba pool on

Sharing one KV tensor between a mamba layer and an attention layer is only
safe while both index it as block_id * page_size. The segregated per-state
mamba layout breaks that congruence, so without the dedicated pool the two
overlap and decode output degrades past ~9 concurrent requests.

Verified on Qwen3.5-35B-A3B-FP8, tp1, 16 capture sizes: shared pool is clean
at 4 and 8 concurrent and garbage at 16; the dedicated pool is 16/16 clean.

* per token group quant fp8 and del some redundant op

(cherry picked from commit 8eae317)
Series patch renumbered 0090 -> 0085. Also aligns the GDN decode env
default with the kv-cache one so the dedicated pool takes the in-place
path instead of falling back to gather/scatter.

* rebase two build patches onto the current series

0060 carried pre-#108 context for the PrivateUse1 registration block that
0044 now emits two lines longer, so the final hunk no longer matched.

0085 re-applied the is_musa() branch that 0012 already adds to the
per-token-group quant dispatch; it now only introduces the vec fast path.

Both were caught by running vllm_musa/patches/build_apply.py, which is what
setup.py invokes; a plain three-way apply absorbs the duplicate and hides it.

* fuse QK-RMSNorm and MRoPE for interleaved MRoPE gated attention

Ports the MUSA kernel from the qwen35_moe branch. The IR provider registered
for gated_qkv_rms_norm_rope only enters the priority list under Inductor and
excludes routed MoE, so Qwen3.5 ran the eager rotary path; this dispatches the
fused kernel ahead of the IR call and leaves it as the fallback.

Measured on Qwen3.5-35B-A3B-FP8 tp1, 2500/1500 x16 @ concurrency 64:
output throughput 509.45 -> 557.54 tok/s, mean TPOT 30.65 -> 27.99 ms.

* perf: dispatch residual Gemma RMSNorm through the fused JIT kernel

Author: zichen-kuuga

vllm_musa/jit_kernel/csrc/norm.py

@@ -150,3 +150,138 @@ def _fused_add_rmsnorm_custom_fake(
     mutates_args=["input", "residual"],
     fake_impl=_fused_add_rmsnorm_custom_fake,
 )
+
+
+@cache_once
+def _qk_mrope_module():
+    import tilelang
+
+    tilelang_dir = Path(tilelang.__file__).resolve().parent
+    return load_musa_jit(
+        "vllm_musa_norm_qk_mrope",
+        ("norm/qk_mrope.mu",),
+        extra_musa_cflags=(
+            f"-I{(tilelang_dir / 'src').resolve()}",
+            f"-I{(tilelang_dir / '3rdparty' / 'mutlass' / 'include').resolve()}",
+            "-Wno-error=address-of-temporary",
+            "-fmusa-flush-denormals-to-zero",
+            "-fno-signed-zeros",
+            "-D__MUSA_ARCH_LIST__=310",
+            "-mllvm",
+            "-mtgpu-opt-level=1",
+            "-mllvm",
+            "-mtgpu-load-store-opt=1",
+            "-mllvm",
+            "-mtgpu-fold-global-ldst=1",
+            "-mllvm",
+            "-mtgpu-load-cluster-mutation=1",
+            "-mllvm",
+            "-mtgpu-store-cluster-mutation=1",
+            "-mllvm",
+            "-mtgpu-memory-sched-mutation=1",
+            "-mllvm",
+            "-mtgpu-alloc-shared-memory-from-zero=1",
+        ),
+    )
+
+
+def fused_qk_rmsnorm_mrope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    is_interleaved: bool,
+    eps: float = 1e-6,
+    gemma: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """QK-RMSNorm + MRoPE in one kernel. q/k are (tokens, heads, head_dim)."""
+    q_out = torch.empty_like(q)
+    k_out = torch.empty_like(k)
+    torch.ops.vllm.musa_csrc_fused_qk_rmsnorm_mrope(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        positions,
+        cos_sin_cache,
+        q_out,
+        k_out,
+        bool(is_neox),
+        int(mrope_section_t),
+        int(mrope_section_h),
+        int(mrope_section_w),
+        bool(is_interleaved),
+        float(eps),
+        bool(gemma),
+    )
+    return q_out, k_out
+
+
+def _fused_qk_rmsnorm_mrope_custom(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    is_neox: bool,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    is_interleaved: bool,
+    eps: float,
+    gemma: bool,
+) -> None:
+    _qk_mrope_module().sgl_musa_fused_qk_rmsnorm_mrope(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        positions,
+        cos_sin_cache,
+        q_out,
+        k_out,
+        bool(is_neox),
+        int(mrope_section_t),
+        int(mrope_section_h),
+        int(mrope_section_w),
+        bool(is_interleaved),
+        float(eps),
+        bool(gemma),
+    )
+
+
+def _fused_qk_rmsnorm_mrope_custom_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    is_neox: bool,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    is_interleaved: bool,
+    eps: float,
+    gemma: bool,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="musa_csrc_fused_qk_rmsnorm_mrope",
+    op_func=_fused_qk_rmsnorm_mrope_custom,
+    mutates_args=["q_out", "k_out"],
+    fake_impl=_fused_qk_rmsnorm_mrope_custom_fake,
+)

vllm_musa/jit_kernel/csrc/norm/common.mu

@@ -0,0 +1,158 @@
+template <typename T>
+struct __align__(16) Vec8Storage {
+  T elem[8];
+};
+
+struct __align__(32) Float8Storage {
+  float elem[8];
+};
+
+template <typename T>
+struct __align__(16) Vec8 {
+  union {
+    Vec8Storage<T> storage;
+    T elem[8];
+  } val;
+
+  __device__ __forceinline__ Vec8() {}
+
+  template <typename Offset>
+  static __device__ __forceinline__ Vec8 load(const T* ptr, Offset idx) {
+    return *(const Vec8*)(ptr + idx);
+  }
+
+  template <typename Offset>
+  static __device__ __forceinline__ Vec8 load_byp_slc(const T* ptr, Offset idx) {
+#if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 310))
+    Vec8 dst;
+    const T* addr = ptr + idx;
+    asm volatile(
+        "LSU.LD.B128 %0, %1, _, 16, 1, 1, inner_persist=0, outer_persist=2, "
+        "chrnt=l2_l3, slc=byp, persist=0, stride_add_first=0"
+        : "=R"(dst)
+        : "R"(addr));
+    return dst;
+#else
+    return *(const Vec8*)(ptr + idx);
+#endif
+  }
+};
+
+struct __align__(32) Float8 {
+  union {
+    Float8Storage storage;
+    float elem[8];
+  } val;
+
+  __device__ __forceinline__ Float8() {}
+};
+
+__device__ __forceinline__ int mrope_24_20_20_interleaved_axis(int rot_offset) {
+  constexpr unsigned long long axis1_mask = 0x492492492492492ULL;
+  constexpr unsigned long long axis2_mask = 0x924924924924924ULL;
+  const unsigned long long bit = 1ULL << rot_offset;
+  return ((axis1_mask & bit) != 0ULL) + (((axis2_mask & bit) != 0ULL) << 1);
+}
+
+__device__ __forceinline__ int mrope_11_11_10_interleaved_axis(int rot_offset) {
+  constexpr unsigned int axis1_mask = 0x92492492U;
+  constexpr unsigned int axis2_mask = 0x24924924U;
+  const unsigned int bit = 1U << rot_offset;
+  return ((axis1_mask & bit) != 0U) + (((axis2_mask & bit) != 0U) << 1);
+}
+
+__device__ __forceinline__ float fast_rsqrt(float value) {
+#if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 310))
+  const float half_value = 0.5f * value;
+  float y = __frsqrt_rn(value);
+  y = y * (1.5f - half_value * y * y);
+  return y;
+#else
+  return rsqrtf(value);
+#endif
+}
+
+__device__ __forceinline__ float block_sum(float value, float* warp_sums) {
+  const int tid = (int)threadIdx.x;
+  const int lane = tid & 31;
+  const int warp = tid >> 5;
+  const int num_warps = ((int)blockDim.x + 31) >> 5;
+
+#pragma unroll
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    value += __shfl_down_sync(0xffffffff, value, offset, 32);
+  }
+  if (lane == 0) {
+    warp_sums[warp] = value;
+  }
+  __syncthreads_lm();
+
+  value = tid < num_warps ? warp_sums[lane] : 0.0f;
+  if (warp == 0) {
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      value += __shfl_down_sync(0xffffffff, value, offset, 32);
+    }
+    if (lane == 0) {
+      warp_sums[0] = value;
+    }
+  }
+  __syncthreads_lm();
+  return warp_sums[0];
+}
+
+__device__ __forceinline__ float block_sum_8warps(float value, float* warp_sums) {
+  const int tid = (int)threadIdx.x;
+  const int lane = tid & 31;
+  const int warp = tid >> 5;
+
+#pragma unroll
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    value += __shfl_down_sync(0xffffffff, value, offset, 32);
+  }
+  if (lane == 0) {
+    warp_sums[warp] = value;
+  }
+  __syncthreads_lm();
+
+  value = lane < 8 ? warp_sums[lane] : 0.0f;
+  if (warp == 0) {
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      value += __shfl_down_sync(0xffffffff, value, offset, 32);
+    }
+    if (lane == 0) {
+      warp_sums[0] = value;
+    }
+  }
+  __syncthreads_lm();
+  return warp_sums[0];
+}
+
+__device__ __forceinline__ float block_sum_4warps(float value, float* warp_sums) {
+  const int tid = (int)threadIdx.x;
+  const int lane = tid & 31;
+  const int warp = tid >> 5;
+
+#pragma unroll
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    value += __shfl_down_sync(0xffffffff, value, offset, 32);
+  }
+  if (lane == 0) {
+    warp_sums[warp] = value;
+  }
+  __syncthreads_lm();
+
+  value = lane < 4 ? warp_sums[lane] : 0.0f;
+  if (warp == 0) {
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      value += __shfl_down_sync(0xffffffff, value, offset, 32);
+    }
+    if (lane == 0) {
+      warp_sums[0] = value;
+    }
+  }
+  __syncthreads_lm();
+  return warp_sums[0];
+}