feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码

来源:
  1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
     - inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
     - inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
     - contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
     - contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
     - csrc/include/ixformer/: C++ kernel headers + cmake

  2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
     - npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
     - npu_torch/qwen3_5_gated_delta_net.cpp/.h
     - npu_torch/qwen3_next_*.cpp/.h (6 files)
     - npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
     - models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
     - models/vlm/qwen3_5.h

调用链完整性:
  ixformer_sdk/inference/functions/vllm.py
    → ops.infer.moe_topk_softmax() (C++ 层)
    → 这就是 base 镜像 libixformer.so 里的实现

  upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
    → ixformer::infer::topk_softmax() (直接 C++ 调用)
    → ixformer::infer::group_gemm() → 完整 7-step MoE pipeline

ixformer_sdk/contrib/flashinfer/__init__.py

@@ -0,0 +1,17 @@
+from .decode import BatchDecodeWithPagedKVCacheWrapper
+from .prefill import (
+    BatchPrefillWithPagedKVCacheWrapper,
+    BatchPrefillWithRaggedKVCacheWrapper,
+)
+
+
+def bmm_fp8():
+    pass
+
+
+def SegmentGEMMWrapper():
+    pass
+
+
+def bmm_fp8():
+    pass

ixformer_sdk/contrib/flashinfer/activation.py

@@ -0,0 +1,29 @@
+import ixformer.inference.functions as ops
+import torch
+
+
+def gelu_and_mul():
+    pass
+
+
+def gelu_tanh_and_mul():
+    pass
+
+
+def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
+    r"""Fused SiLU and Mul operation.
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (..., 2 * hidden_size).
+
+    out: Optional[torch.Tensor]
+        The the output tensor, if specified, the kernel will update this tensor inplace.
+
+    Returns
+    -------
+    output: torch.Tensor
+        Output tensor, shape (..., hidden_size).
+    """
+    return ops.silu_and_mul(input=input, output=out)

ixformer_sdk/contrib/flashinfer/cascade.py

@@ -0,0 +1,2 @@
+def merge_state():
+    pass

ixformer_sdk/contrib/flashinfer/decode.py

@@ -0,0 +1,101 @@
+import math
+from typing import Optional, Tuple, Union
+
+import ixformer.inference.functions as ops
+import torch
+
+
+def _grouped_size_compiled_for_decode_kernels(
+    num_qo_heads: int, num_kv_heads: int
+) -> bool:
+    return (num_qo_heads // num_kv_heads) in [1, 2, 4, 8]
+
+
+class BatchDecodeWithPagedKVCacheWrapper:
+    def __init__(
+        self,
+        float_workspace_buffer: torch.Tensor,
+        kv_layout: str = "NHD",
+        use_cuda_graph: bool = False,
+        use_tensor_cores: bool = False,
+    ) -> None:
+        pass
+
+    def plan(
+        self,
+        indptr: torch.Tensor,
+        indices: torch.Tensor,
+        last_page_len: torch.Tensor,
+        num_qo_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        page_size: int,
+        # pos_encoding_mode: str = "NONE",
+        # window_left: int = -1,
+        # logits_soft_cap: Optional[float] = None,
+        data_type: Union[str, torch.dtype] = "float16",
+        q_data_type: Optional[Union[str, torch.dtype]] = None,
+        sm_scale: Optional[float] = None,
+        # rope_scale: Optional[float] = None,
+        # rope_theta: Optional[float] = None,
+        max_seqlen_q: int = None,
+        max_seqlen_k: int = None,
+    ) -> None:
+        self.indptr = indptr
+        self.indices = indices
+        self.last_page_len = last_page_len
+        self.num_qo_heads = num_qo_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+
+        assert page_size == 1
+
+        self.cu_seqlens_q = torch.ones_like(indptr)
+        self.cu_seqlens_q[0] = 0
+        self.cu_seqlens_q = torch.cumsum(self.cu_seqlens_q, dim=0).int()
+
+        self.cu_seqlens_k = indptr
+        if sm_scale is None:
+            sm_scale = 1.0 / math.sqrt(head_dim)
+
+            self.sm_scale = sm_scale
+            self.max_seqlen_q = max_seqlen_q
+            self.max_seqlen_k = max_seqlen_k
+
+    begin_forward = plan
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+        pos_encoding_mode: str = "NONE",
+        q_scale: Optional[float] = None,
+        k_scale: Optional[float] = None,
+        v_scale: Optional[float] = None,
+        window_left: int = -1,
+        logits_soft_cap: Optional[float] = None,
+        sm_scale: Optional[float] = None,
+        rope_scale: Optional[float] = None,
+        rope_theta: Optional[float] = None,
+    ) -> torch.Tensor:
+        k_cache, v_cache = paged_kv_cache
+
+        out = torch.empty_like(q)
+
+        ops.paged_attention_flashinfer(
+            output=out,
+            query=q,
+            paged_kv_data=(k_cache.unsqueeze(1), v_cache.unsqueeze(1)),
+            paged_kv_indptr=self.indptr,
+            paged_kv_indices=self.indices,
+            paged_kv_last_page_len=self.last_page_len,
+            scale=self.sm_scale,
+            max_seq_len=self.max_seqlen_k,
+            kv_cache_format="NHD",
+        )
+
+        return out
+
+    def end_forward(self) -> None:
+        r"""Warning: this function is deprecated and has no effect."""
+        pass

ixformer_sdk/contrib/flashinfer/norm.py

@@ -0,0 +1,61 @@
+import ixformer.inference.functions as ops
+import torch
+
+
+def fused_add_rmsnorm(
+    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
+):
+    r"""Fused add root mean square normalization.
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size).
+    residual: torch.Tensor
+        Residual tensor, shape (batch_size, hidden_size).
+    weight: torch.Tensor
+        Weight tensor, shape (hidden_size,).
+    eps: float
+        Epsilon for numerical stability.
+    """
+    return ops.residual_rms_norm(
+        input=input,
+        residual=residual,
+        weight=weight,
+        eps=eps,
+    )
+
+
+def gemma_fused_add_rmsnorm():
+    pass
+
+
+def gemma_rmsnorm():
+    pass
+
+
+def rmsnorm(
+    input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
+) -> torch.Tensor:
+    r"""Root mean square normalization.
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size).
+    weight: torch.Tensor
+        Weight tensor, shape (hidden_size,).
+    eps: float
+        Epsilon for numerical stability.
+
+    Returns
+    -------
+    output: torch.Tensor
+        Normalized tensor, shape (batch_size, hidden_size).
+    """
+
+    return ops.rms_norm(
+        input=input,
+        weight=weight,
+        eps=eps,
+    )

ixformer_sdk/contrib/flashinfer/prefill.py

@@ -0,0 +1,113 @@
+import math
+from typing import Optional, Tuple, Union
+
+import ixformer._C as ops
+import torch
+
+
+class BatchPrefillWithRaggedKVCacheWrapper:
+    def __init__(
+        self,
+        float_workspace_buffer: torch.Tensor,
+        kv_layout: str = "NHD",
+    ):
+        pass
+
+    def plan(
+        self,
+        qo_indptr: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        num_qo_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        # custom_mask: Optional[torch.Tensor] = None,
+        # packed_custom_mask: Optional[torch.Tensor] = None,
+        causal: bool = True,
+        # pos_encoding_mode: str = "NONE",
+        # allow_fp16_qk_reduction: bool = False,
+        # window_left: int = -1,
+        # logits_soft_cap: Optional[float] = None,
+        sm_scale: Optional[float] = None,
+        # rope_scale: Optional[float] = None,
+        # rope_theta: Optional[float] = None,
+        # q_data_type: str = "float16",
+    ) -> None:
+        batch_size = len(qo_indptr) - 1
+        if len(kv_indptr) != batch_size + 1:
+            raise ValueError(
+                "The kv_indptr length should be equal to qk_indptr length."
+            )
+        self._causal = causal
+        self._sm_scale = sm_scale
+        if sm_scale is None:
+            sm_scale = 1.0 / math.sqrt(head_dim)
+
+        self.cu_seqlens_q = qo_indptr
+        self.cu_seqlens_k = kv_indptr
+        self.num_qo_heads = num_qo_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.max_seqlen_q = max_seqlen_q
+        self.max_seqlen_k = max_seqlen_k
+
+    begin_forward = plan
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        causal: bool = True,
+        # pos_encoding_mode: str = "NONE",
+        # allow_fp16_qk_reduction: bool = False,
+        # window_left: int = -1,
+        logits_soft_cap: Optional[float] = None,
+        sm_scale: Optional[float] = None,
+        # rope_scale: Optional[float] = None,
+        # rope_theta: Optional[float] = None,
+    ) -> torch.Tensor:
+        r"""Warning: This function is deprecated, please use :meth:`run` instead."""
+
+        q = q.view(-1, self.num_qo_heads, self.head_dim)
+        k = k.view(-1, self.num_kv_heads, self.head_dim)
+        v = v.view(-1, self.num_kv_heads, self.head_dim)
+
+        out = torch.empty_like(q)
+
+        assert causal
+        assert (
+            logits_soft_cap is None or logits_soft_cap == 0
+        ), f"logits_soft_cap not supported, but got logits_soft_cap={logits_soft_cap}"
+
+        ops.infer.ixinfer_flash_attn_unpad(
+            q,
+            k,
+            v,
+            out,
+            self.cu_seqlens_q,
+            self.cu_seqlens_k,
+            self.max_seqlen_q,
+            self.max_seqlen_k,
+            causal,
+            False,  # need_lse =False
+            sm_scale,
+            False,
+            None,
+        )
+        return out
+
+    def end_forward(self) -> None:
+        r"""Warning: this function is deprecated and has no effect."""
+        pass
+
+
+class BatchPrefillWithPagedKVCacheWrapper:
+    def __init__(
+        self,
+        float_workspace_buffer: torch.Tensor,
+        kv_layout: str = "NHD",
+        use_cuda_graph: bool = False,
+    ) -> None:
+        pass

ixformer_sdk/contrib/flashinfer/sampling.py

@@ -0,0 +1,14 @@
+def min_p_sampling_from_probs():
+    pass
+
+
+def top_k_renorm_prob():
+    pass
+
+
+def top_k_top_p_sampling_from_probs():
+    pass
+
+
+def top_p_renorm_prob():
+    pass