ref(upstream): add Deep-Spark/vllm latest + xllm ILU kernel sources

Cloned from GitHub:
  - Deep-Spark/vllm (latest): qwen3_5.py with multimodal support,
    transformers configs, multimodal registry, model registry
  - jd-opensource/xllm (latest): ILU kernel implementations
    (attention, fused_moe, group_gemm, activation, norm, rope, matmul)
    + GatedDeltaNet layer for Qwen3.5

These are the REAL upstream implementations that the base Docker image
is compiled from. Our dlopen modules should match these interfaces:
  - ilu_ops_api.h: 14 functions in xllm::kernel::ilu namespace
  - ixformer.h: 15 functions in ixformer::infer namespace

Key interface signatures for dlopen targets:
  batch_prefill()  → ixinfer_flash_attn_unpad_with_block_tables
  batch_decode()   → xllm_paged_attention
  moe_active_topk()→ topk_softmax
  moe_gen_idx()    → moe_compute_token_index_api
  group_gemm()     → moe_w16a16_group_gemm
  silu_and_mul()   → silu_and_mul
  rms_norm()       → rms_norm + residual_rms_norm

upstream_ref/ds_vllm_latest/vllm/model_executor/models/qwen3_5.py

@@ -0,0 +1,819 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2025 The vLLM team.
+# Copyright 2025 The Qwen Team.
+# Copyright 2025 The HuggingFace Inc. team.
+# All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inference-only Qwen3.5 Series compatible with HuggingFace weights."""
+
+import typing
+from collections.abc import Callable, Iterable
+
+import torch
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed import (
+    get_pp_group,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import (
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.layernorm import (
+    GemmaRMSNorm as Qwen3_5RMSNorm,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+    QwenGatedDeltaNetAttention,
+)
+from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateCopyFunc,
+    MambaStateCopyFuncCalculator,
+    MambaStateDtypeCalculator,
+    MambaStateShapeCalculator,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.qwen3_5 import (
+    Qwen3_5Config,
+    Qwen3_5TextConfig,
+)
+from vllm.transformers_utils.configs.qwen3_5_moe import (
+    Qwen3_5MoeConfig,
+    Qwen3_5MoeTextConfig,
+)
+
+from .interfaces import (
+    HasInnerState,
+    IsHybrid,
+    MixtureOfExperts,
+    MultiModalEmbeddings,
+    SupportsEagle3,
+    SupportsLoRA,
+    SupportsPP,
+    _require_is_multimodal,
+)
+from .qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
+from .qwen3_next import (
+    Qwen3NextAttention,
+    Qwen3NextDecoderLayer,
+    Qwen3NextModel,
+    Qwen3NextSparseMoeBlock,
+    QwenNextMixtureOfExperts,
+)
+from .qwen3_vl import (
+    Qwen3_VisionTransformer,
+    Qwen3VLDummyInputsBuilder,
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLMultiModalProcessor,
+    Qwen3VLProcessingInfo,
+)
+from .utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    _merge_multimodal_embeddings,
+    extract_layer_index,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+
+logger = init_logger(__name__)
+
+
+class Qwen3_5ProcessingInfo(Qwen3VLProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.get_hf_config(Qwen3_5Config)
+
+
+class Qwen3_5MoeProcessingInfo(Qwen3VLProcessingInfo):
+    def get_hf_config(self):
+        return self.ctx.get_hf_config(Qwen3_5MoeConfig)
+
+
+class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        layer_type: str,
+        prefix: str = "",
+    ) -> None:
+        super(Qwen3NextDecoderLayer, self).__init__()
+
+        config = vllm_config.model_config.hf_text_config
+        model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+
+        self.layer_type = layer_type
+        self.layer_idx = extract_layer_index(prefix)
+
+        if self.layer_type == "linear_attention":
+            self.linear_attn = QwenGatedDeltaNetAttention(
+                config=config,
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.linear_attn",
+                gqa_interleaved_layout=False,
+            )
+        elif self.layer_type == "full_attention":
+            self.self_attn = Qwen3NextAttention(
+                config,
+                model_config=model_config,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.self_attn",
+            )
+        else:
+            raise ValueError(f"Invalid layer_type {self.layer_type}")
+
+        # NOTE: Determine the MLP type based on the model type
+        # Qwen3.5 use all layers for MLP / Qwen3.5-MoE use sparse MoE blocks
+        if config.model_type == "qwen3_5_moe_text":
+            self.mlp = Qwen3NextSparseMoeBlock(
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.mlp",
+            )
+        elif config.model_type == "qwen3_5_text":
+            self.mlp = Qwen3NextMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+        else:
+            raise ValueError(f"Invalid model_type {config.model_type}")
+
+        self.input_layernorm = Qwen3_5RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm = Qwen3_5RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+        self.layer_scale = getattr(config, "layer_scale", False)
+        if self.layer_scale:
+            self.attn_layer_scale = torch.nn.Parameter(
+                torch.zeros(
+                    1,
+                    1,
+                    config.hidden_size,
+                ),
+            )
+            self.ffn_layer_scale = torch.nn.Parameter(
+                torch.zeros(
+                    1,
+                    1,
+                    config.hidden_size,
+                ),
+            )
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
+        # otherwise (seq_len, ).
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+    }
+)
+class Qwen3_5Model(Qwen3NextModel):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super(Qwen3NextModel, self).__init__()
+
+        config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig = (
+            vllm_config.model_config.hf_text_config
+        )
+        parallel_config = vllm_config.parallel_config
+
+        eplb_config = parallel_config.eplb_config
+        self.num_redundant_experts = eplb_config.num_redundant_experts
+
+        self.config = config
+
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+        )
+
+        def get_layer(prefix: str):
+            return Qwen3_5DecoderLayer(
+                vllm_config,
+                layer_type=config.layer_types[extract_layer_index(prefix)],
+                prefix=prefix,
+            )
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
+        )
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
+
+        if get_pp_group().is_last_rank:
+            self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+        self.aux_hidden_state_layers: tuple[int, ...] = ()
+
+    def load_fused_expert_weights(
+        self,
+        name: str,
+        params_dict: dict,
+        loaded_weight: torch.Tensor,
+        shard_id: str,
+        num_experts: int,
+    ) -> bool:
+        param = params_dict[name]
+        weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
+        loaded_local_expert = False
+        for expert_id in range(num_experts):
+            curr_expert_weight = loaded_weight[expert_id]
+            success = weight_loader(
+                param,
+                curr_expert_weight,
+                name,
+                shard_id=shard_id,
+                expert_id=expert_id,
+                return_success=True,
+            )
+            if success:
+                loaded_local_expert = True
+
+        return loaded_local_expert
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            # GDN
+            ("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)),
+            ("in_proj_qkvz", "in_proj_z", 3),
+            # self attention
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            # mlp
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+            ("in_proj_ba", "in_proj_b", 0),
+            ("in_proj_ba", "in_proj_a", 1),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        expert_params_mapping = self.get_expert_mapping()
+        is_fused_expert = False
+        fused_expert_params_mapping: list[tuple[str, str, int, str]] = []
+        for param_name, ckpt_name, _, shard_id in fused_moe_make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="gate_up_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="gate_up_proj",
+            num_experts=1,
+        ):
+            if shard_id == "w3":
+                continue
+            parts = ckpt_name.split(".")
+            fused_expert_params_mapping.append(
+                (f"{param_name}weight", f"{parts[0]}.{parts[2]}", 0, shard_id)
+            )
+        num_experts = (
+            self.config.num_experts if hasattr(self.config, "num_experts") else 0
+        )
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+
+            if name.startswith("mtp."):
+                continue
+
+            # Remapping the name of FP8 kv-scale.
+            if name.endswith("scale"):
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if "experts.gate_up_proj" in name or "experts.down_proj" in name:
+                    is_fused_expert = True
+                    expert_params_mapping = fused_expert_params_mapping
+
+                if weight_name not in name:
+                    continue
+
+                if "mlp.experts" in name:
+                    continue
+
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                # Skip layers on other devices.
+                if is_pp_missing_parameter(name, self):
+                    continue
+                # name = apply_attn_prefix(name, params_dict)
+                if name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                is_expert_weight = False
+                for mapping in expert_params_mapping:
+                    param_name, weight_name, expert_id, shard_id = mapping
+                    if weight_name not in name:
+                        continue
+                    is_expert_weight = True
+                    name_mapped = name.replace(weight_name, param_name)
+                    # Skip layers on other devices.
+                    if is_pp_missing_parameter(name_mapped, self):
+                        continue
+                    if is_fused_expert:
+                        # qwen3.5 no need to transpose
+                        # loaded_weight = loaded_weight.transpose(-1, -2)
+                        if "experts.gate_up_proj" in name:
+                            loaded_weight = loaded_weight.chunk(2, dim=-2)
+                            success_w1 = self.load_fused_expert_weights(
+                                name_mapped,
+                                params_dict,
+                                loaded_weight[0],
+                                "w1",
+                                num_experts,
+                            )
+                            success_w3 = self.load_fused_expert_weights(
+                                name_mapped,
+                                params_dict,
+                                loaded_weight[1],
+                                "w3",
+                                num_experts,
+                            )
+                            success = success_w1 and success_w3
+                        else:
+                            # down_proj
+                            success = self.load_fused_expert_weights(
+                                name_mapped,
+                                params_dict,
+                                loaded_weight,
+                                shard_id,
+                                num_experts,
+                            )
+                        if success:
+                            name = name_mapped
+                            break
+                    else:
+                        # Skip loading extra bias for GPTQ models.
+                        if (
+                            name_mapped.endswith(".bias")
+                            or name_mapped.endswith("_bias")
+                        ) and name_mapped not in params_dict:
+                            continue
+                        param = params_dict[name_mapped]
+                        weight_loader = param.weight_loader
+                        success = weight_loader(
+                            param,
+                            loaded_weight,
+                            name_mapped,
+                            shard_id=shard_id,
+                            expert_id=expert_id,
+                            return_success=True,
+                        )
+                    if success:
+                        name = name_mapped
+                        break
+                else:
+                    if is_expert_weight:
+                        # We've checked that this is an expert weight
+                        # However it's not mapped locally to this rank
+                        # So we simply skip it
+                        continue
+                    # Skip loading extra bias for GPTQ models.
+                    if name.endswith(".bias") and name not in params_dict:
+                        continue
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    if name not in params_dict:
+                        logger.warning_once(
+                            f"Parameter {name} not found in params_dict, skip loading"
+                        )
+                        continue
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class Qwen3_5ForCausalLMBase(
+    nn.Module,
+    HasInnerState,
+    SupportsEagle3,
+    SupportsLoRA,
+    SupportsPP,
+):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+        # GDN fused projections.
+        "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+        "in_proj_ba": ["in_proj_b", "in_proj_a"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        config = vllm_config.model_config.hf_text_config
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
+
+        scheduler_config = vllm_config.scheduler_config
+        if cache_config.mamba_cache_mode == "all":
+            raise NotImplementedError(
+                "Qwen3.5 currently does not support 'all' prefix caching, "
+                "please use '--mamba-cache-mode=align' instead"
+            )
+        self.quant_config = vllm_config.quant_config
+
+        super().__init__()
+        self.config = config
+        self.scheduler_config = scheduler_config
+        self.model = Qwen3_5Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+
+        if get_pp_group().is_last_rank:
+            if config.tie_word_embeddings:
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=self.quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
+        self.model.aux_hidden_state_layers = layers
+
+    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
+        num_layers = len(self.model.layers)
+        return (2, num_layers // 2, num_layers - 3)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ):
+        hidden_states = self.model(
+            input_ids, positions, intermediate_tensors, inputs_embeds
+        )
+
+        return hidden_states
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=["mtp."],
+        )
+        return loader.load_weights(weights)
+
+
+class Qwen3_5ForCausalLM(Qwen3_5ForCausalLMBase):
+    pass
+
+
+class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+
+        # set MoE hyperparameters
+        self.set_moe_parameters()
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        return self.model.get_expert_mapping()
+
+
+########################################################
+# Qwen3_5-Dense
+########################################################
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=Qwen3_5ProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid):
+    # Qwen3.5 does not support multimodal pruning (EVS).
+    supports_multimodal_pruning = False
+
+    packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | {
+        "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+        "in_proj_ba": ["in_proj_b", "in_proj_a"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
+        # protocols have not __init__ method, so we need to use nn.Module.__init__
+        nn.Module.__init__(self)
+        config: Qwen3_5Config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+
+        self.config = config
+        self.model_config = vllm_config.model_config
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        # Qwen3.5 does not support multimodal pruning (EVS).
+        self.is_multimodal_pruning_enabled = False
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual = Qwen3_VisionTransformer(
+                config.vision_config,
+                norm_eps=getattr(config, "rms_norm_eps", 1e-6),
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "visual"),
+            )
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = Qwen3_5ForCausalLM(
+                vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model")
+            )
+
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self._embed_text_input_ids(
+            input_ids,
+            self.language_model.embed_input_ids,
+            is_multimodal=is_multimodal,
+        )
+
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+
+        is_multimodal = _require_is_multimodal(is_multimodal)
+
+        inputs_embeds = _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
+
+        return inputs_embeds
+
+    def recompute_mrope_positions(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Qwen3.5 does not support multimodal pruning (EVS). "
+            "recompute_mrope_positions should never be called."
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        """Run forward pass for Qwen3.5.
+
+        Args:
+            input_ids: Flattened (concatenated) input_ids corresponding to a
+                batch.
+            positions: Flattened (concatenated) position ids corresponding to a
+                batch.
+                **NOTE**: If mrope is enabled (default setting for Qwen3VL
+                opensource models), the shape will be `(3, seq_len)`,
+                otherwise it will be `(seq_len,).
+            intermediate_tensors: Intermediate tensors from previous pipeline
+                stages.
+            inputs_embeds: Pre-computed input embeddings.
+            **kwargs: Additional keyword arguments including:
+                - pixel_values: Pixel values to be fed to a model.
+                    `None` if no images are passed.
+                - image_grid_thw: Tensor `(n_images, 3)` of image 3D grid in
+                    LLM. `None` if no images are passed.
+                - pixel_values_videos: Pixel values of videos to be fed to a
+                    model. `None` if no videos are passed.
+                - video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in
+                    LLM. `None` if no videos are passed.
+        """
+
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+
+        hidden_states = self.language_model.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=["mtp."],
+        )
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: "VllmConfig",
+    ) -> tuple[torch.dtype, torch.dtype]:
+        return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
+            vllm_config.model_config.dtype,
+            vllm_config.cache_config.mamba_cache_dtype,
+            vllm_config.cache_config.mamba_ssm_cache_dtype,
+        )
+
+    @classmethod
+    def get_mamba_state_shape_from_config(
+        cls, vllm_config: "VllmConfig"
+    ) -> tuple[tuple[int, int], tuple[int, int]]:
+        parallel_config = vllm_config.parallel_config
+        hf_config = vllm_config.model_config.hf_text_config
+        tp_size = parallel_config.tensor_parallel_size
+        num_spec = (
+            vllm_config.speculative_config.num_speculative_tokens
+            if vllm_config.speculative_config
+            else 0
+        )
+        return MambaStateShapeCalculator.gated_delta_net_state_shape(
+            tp_size,
+            hf_config.linear_num_key_heads,
+            hf_config.linear_num_value_heads,
+            hf_config.linear_key_head_dim,
+            hf_config.linear_value_head_dim,
+            hf_config.linear_conv_kernel_dim,
+            num_spec,
+        )
+
+    @classmethod
+    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
+        return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
+
+
+########################################################
+# Qwen3_5-MoE
+########################################################
+
+
+class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts):
+    def update_physical_experts_metadata(
+        self,
+        num_physical_experts: int,
+        num_local_physical_experts: int,
+    ) -> None:
+        assert self.num_local_physical_experts == num_local_physical_experts
+        self.num_physical_experts = num_physical_experts
+        self.num_local_physical_experts = num_local_physical_experts
+        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+        for layer in self.language_model.model.layers:
+            if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
+                moe = layer.mlp
+                moe.n_local_physical_experts = num_local_physical_experts
+                moe.n_physical_experts = num_physical_experts
+                moe.n_redundant_experts = self.num_redundant_experts
+                moe.experts.update_expert_map()
+
+    def set_moe_parameters(self):
+        self.expert_weights = []
+
+        self.moe_layers = []
+        example_moe = None
+        for layer in self.language_model.model.layers:
+            if isinstance(layer, Qwen3_5DecoderLayer) and isinstance(
+                layer.mlp, Qwen3NextSparseMoeBlock
+            ):
+                example_moe = layer.mlp
+                self.moe_layers.append(layer.mlp.experts)
+
+        if example_moe is None:
+            raise RuntimeError(
+                "No Qwen3_5 layer found in the language_model.model.layers."
+            )
+
+        # Set MoE hyperparameters
+        self.num_moe_layers = len(self.moe_layers)
+        self.num_expert_groups = 1
+        self.num_shared_experts = 0
+        self.num_logical_experts = example_moe.n_logical_experts
+        self.num_physical_experts = example_moe.n_physical_experts
+        self.num_local_physical_experts = example_moe.n_local_physical_experts
+        self.num_routed_experts = example_moe.n_routed_experts
+        self.num_redundant_experts = example_moe.n_redundant_experts
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=Qwen3_5MoeProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class Qwen3_5MoeForConditionalGeneration(
+    Qwen3_5ForConditionalGeneration, Qwen3_5_MoeMixtureOfExperts
+):
+    # For MoE LoRA weights loading
+    is_3d_moe_weight: bool = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
+        # protocols have not __init__ method, so we need to use nn.Module.__init__
+        nn.Module.__init__(self)
+        config: Qwen3_5MoeConfig = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+
+        self.config = config
+        self.model_config = vllm_config.model_config
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        # Qwen3.5 does not support multimodal pruning (EVS).
+        self.is_multimodal_pruning_enabled = False
+
+        with self._mark_tower_model(vllm_config, {"image", "video"}):
+            self.visual = Qwen3_VisionTransformer(
+                config.vision_config,
+                norm_eps=getattr(config, "rms_norm_eps", 1e-6),
+                quant_config=quant_config,
+                prefix=maybe_prefix(prefix, "visual"),
+            )
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = Qwen3_5MoeForCausalLM(
+                vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model")
+            )
+
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
+
+        # set MoE hyperparameters
+        self.set_moe_parameters()

upstream_ref/ds_vllm_latest/vllm/model_executor/models/qwen3_5_mtp.py

@@ -0,0 +1,466 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only Qwen3_5 MTP model."""
+
+import typing
+from collections.abc import Callable, Iterable
+
+import torch
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import (
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.linear import ColumnParallelLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.interfaces import LocalArgmaxMixin
+from vllm.model_executor.models.qwen3_5 import Qwen3_5DecoderLayer, Qwen3_5RMSNorm
+from vllm.model_executor.models.qwen3_next import QwenNextMixtureOfExperts
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5TextConfig
+from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig
+
+from .interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    _require_is_multimodal,
+)
+from .utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    _merge_multimodal_embeddings,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    maybe_prefix,
+)
+
+logger = init_logger(__name__)
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
+        # otherwise (seq_len, ).
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+        "hidden_states": 0,
+    }
+)
+class Qwen3_5MultiTokenPredictor(nn.Module):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        model_config = vllm_config.model_config
+        quant_config = vllm_config.quant_config
+
+        config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig = model_config.hf_text_config
+
+        self.config = config
+
+        self.vocab_size = config.vocab_size
+
+        self.mtp_start_layer_idx = config.num_hidden_layers
+        self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1)
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+        )
+
+        # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is
+        # missing from hf_quant_config.json exclude_modules. Force unquantized.
+        # Ref: https://github.com/vllm-project/vllm/pull/38650
+        # Ref: https://github.com/NVIDIA/Model-Optimizer/pull/1124
+        fc_quant = (
+            None
+            if (quant_config and quant_config.get_name() == "modelopt_fp4")
+            else quant_config
+        )
+        self.fc = ColumnParallelLinear(
+            self.config.hidden_size * 2,
+            self.config.hidden_size,
+            gather_output=True,
+            bias=False,
+            return_bias=False,
+            quant_config=fc_quant,
+            prefix=f"{prefix}.fc",
+        )
+
+        self.layers = torch.nn.ModuleList(
+            Qwen3_5DecoderLayer(
+                vllm_config,
+                layer_type="full_attention",
+                prefix=f"{prefix}.layers.{idx}",
+            )
+            for idx in range(self.num_mtp_layers)
+        )
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
+
+        self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.pre_fc_norm_hidden = Qwen3_5RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.pre_fc_norm_embedding = Qwen3_5RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is None:
+                inputs_embeds = self.embed_input_ids(input_ids)
+            assert hidden_states.shape[-1] == inputs_embeds.shape[-1]
+            inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds)
+            hidden_states = self.pre_fc_norm_hidden(hidden_states)
+            hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1)
+            hidden_states = self.fc(hidden_states)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+
+        current_step_idx = spec_step_idx % self.num_mtp_layers
+        hidden_states, residual = self.layers[current_step_idx](
+            positions=positions,
+            hidden_states=hidden_states,
+            residual=residual,
+        )
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+    def load_fused_expert_weights(
+        self,
+        name: str,
+        params_dict: dict,
+        loaded_weight: torch.Tensor,
+        shard_id: str,
+        num_experts: int,
+    ) -> bool:
+        param = params_dict[name]
+        weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
+        loaded_local_expert = False
+        for expert_id in range(num_experts):
+            curr_expert_weight = loaded_weight[expert_id]
+            success = weight_loader(
+                param,
+                curr_expert_weight,
+                name,
+                shard_id=shard_id,
+                expert_id=expert_id,
+                return_success=True,
+            )
+            if success:
+                loaded_local_expert = True
+
+        return loaded_local_expert
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
+        # Params for weights, fp8 weight scales, fp8 activation scales
+        # (param_name, weight_name, expert_id, shard_id)
+        expert_params_mapping = fused_moe_make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.num_experts
+            if hasattr(self.config, "num_experts")
+            else 0,
+        )
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        is_fused_expert = False
+        fused_expert_params_mapping: list[tuple[str, str, int, str]] = []
+        for param_name, ckpt_name, _, shard_id in fused_moe_make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="gate_up_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="gate_up_proj",
+            num_experts=1,
+        ):
+            if shard_id == "w3":
+                continue
+            parts = ckpt_name.split(".")
+            fused_expert_params_mapping.append(
+                (f"{param_name}weight", f"{parts[0]}.{parts[2]}", 0, shard_id)
+            )
+        num_experts = (
+            self.config.num_experts if hasattr(self.config, "num_experts") else 0
+        )
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if "experts.gate_up_proj" in name or "experts.down_proj" in name:
+                    is_fused_expert = True
+                    expert_params_mapping = fused_expert_params_mapping
+
+                if weight_name not in name:
+                    continue
+
+                if "mlp.experts" in name:
+                    continue
+
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                # Skip layers on other devices.
+                if is_pp_missing_parameter(name, self):
+                    continue
+                if name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                is_expert_weight = False
+                for mapping in expert_params_mapping:
+                    param_name, weight_name, expert_id, shard_id = mapping
+                    if weight_name not in name:
+                        continue
+                    is_expert_weight = True
+                    name_mapped = name.replace(weight_name, param_name)
+                    # Skip layers on other devices.
+                    if is_pp_missing_parameter(name_mapped, self):
+                        continue
+                    if is_fused_expert:
+                        # qwen3.5 no need to transpose
+                        # loaded_weight = loaded_weight.transpose(-1, -2)
+                        if "experts.gate_up_proj" in name:
+                            loaded_weight = loaded_weight.chunk(2, dim=-2)
+                            success_w1 = self.load_fused_expert_weights(
+                                name_mapped,
+                                params_dict,
+                                loaded_weight[0],
+                                "w1",
+                                num_experts,
+                            )
+                            success_w3 = self.load_fused_expert_weights(
+                                name_mapped,
+                                params_dict,
+                                loaded_weight[1],
+                                "w3",
+                                num_experts,
+                            )
+                            success = success_w1 and success_w3
+                        else:
+                            # down_proj
+                            success = self.load_fused_expert_weights(
+                                name_mapped,
+                                params_dict,
+                                loaded_weight,
+                                shard_id,
+                                num_experts,
+                            )
+                        if success:
+                            name = name_mapped
+                            break
+                    else:
+                        # Skip loading extra bias for GPTQ models.
+                        if (
+                            name_mapped.endswith(".bias")
+                            or name_mapped.endswith("_bias")
+                        ) and name_mapped not in params_dict:
+                            continue
+                        param = params_dict[name_mapped]
+                        weight_loader = param.weight_loader
+                        success = weight_loader(
+                            param,
+                            loaded_weight,
+                            name_mapped,
+                            shard_id=shard_id,
+                            expert_id=expert_id,
+                            return_success=True,
+                        )
+                    if success:
+                        name = name_mapped
+                        break
+                else:
+                    if is_expert_weight:
+                        # We've checked that this is an expert weight
+                        # However it's not mapped locally to this rank
+                        # So we simply skip it
+                        continue
+                    # Skip loading extra bias for GPTQ models.
+                    if name.endswith(".bias") and name not in params_dict:
+                        continue
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    if name not in params_dict:
+                        logger.warning_once(
+                            f"Parameter {name} not found in params_dict, skip loading"
+                        )
+                        continue
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
+        # otherwise (seq_len, ).
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+        "hidden_states": 0,
+    }
+)
+class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal):
+    packed_modules_mapping = {
+        "qkv_proj": [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+        ],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        config = vllm_config.model_config.hf_text_config
+        self.vllm_config = vllm_config
+        cache_config = vllm_config.cache_config
+        if cache_config.mamba_cache_mode == "all":
+            raise NotImplementedError(
+                "Qwen3_5MTP currently does not support 'all' prefix caching, "
+                "please use '--mamba-cache-mode=align' instead"
+            )
+
+        self.quant_config = vllm_config.quant_config
+
+        super().__init__()
+        self.config = config
+        self.model = Qwen3_5MultiTokenPredictor(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp")
+        )
+
+        if get_pp_group().is_last_rank:
+            if config.tie_word_embeddings:
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=self.quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self._embed_text_input_ids(
+            input_ids,
+            self.model.embed_input_ids,
+            is_multimodal=is_multimodal,
+        )
+
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+
+        is_multimodal = _require_is_multimodal(is_multimodal)
+
+        inputs_embeds = _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
+
+        return inputs_embeds
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ):
+        hidden_states = self.model(
+            input_ids, positions, hidden_states, intermediate_tensors, inputs_embeds
+        )
+        return hidden_states
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor | None:
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        def remap_weight_names(weights):
+            for name, weight in weights:
+                if name.startswith("mtp."):
+                    name = name.replace("mtp.", "model.")
+                elif any(key in name for key in ["embed_tokens", "lm_head"]):
+                    if "embed_tokens" in name:
+                        name = name.replace("language_model.", "")
+                else:
+                    continue
+                yield name, weight
+
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(remap_weight_names(weights))
+
+
+class Qwen3_5MoeMTP(Qwen3_5MTP, QwenNextMixtureOfExperts):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        self.set_moe_parameters()

upstream_ref/ds_vllm_latest/vllm/multimodal/__init__.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from .hasher import MultiModalHasher
+from .inputs import BatchedTensorInputs, MultiModalKwargsItems, NestedTensors
+from .registry import MultiModalRegistry
+
+MULTIMODAL_REGISTRY = MultiModalRegistry()
+"""
+The global [`MultiModalRegistry`][vllm.multimodal.registry.MultiModalRegistry]
+is used by model runners to dispatch data processing according to the target
+model.
+
+Info:
+    [mm_processing](../../../design/mm_processing.md)
+"""
+
+__all__ = [
+    "BatchedTensorInputs",
+    "MultiModalHasher",
+    "MultiModalKwargsItems",
+    "NestedTensors",
+    "MULTIMODAL_REGISTRY",
+    "MultiModalRegistry",
+]

upstream_ref/ds_vllm_latest/vllm/multimodal/registry.py

@@ -0,0 +1,378 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import threading
+from collections import defaultdict
+from collections.abc import Mapping
+from dataclasses import dataclass
+from multiprocessing.synchronize import Lock as LockType
+from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypeVar, cast
+
+from vllm.inputs import MultiModalInput
+from vllm.logger import init_logger
+from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config
+
+from .cache import (
+    BaseMultiModalProcessorCache,
+    BaseMultiModalReceiverCache,
+    MultiModalProcessorOnlyCache,
+    MultiModalProcessorSenderCache,
+    MultiModalReceiverCache,
+    ShmObjectStoreReceiverCache,
+    ShmObjectStoreSenderCache,
+)
+from .processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    InputProcessingContext,
+    TimingContext,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import ModelConfig, ObservabilityConfig, VllmConfig
+    from vllm.model_executor.models.interfaces import SupportsMultiModal
+
+logger = init_logger(__name__)
+
+N = TypeVar("N", bound=type["SupportsMultiModal"])
+_I = TypeVar("_I", bound=BaseProcessingInfo)
+_I_co = TypeVar("_I_co", bound=BaseProcessingInfo, covariant=True)
+
+
+class ProcessingInfoFactory(Protocol[_I_co]):
+    """
+    Constructs a
+    [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor]
+    instance from the context.
+    """
+
+    def __call__(
+        self,
+        ctx: InputProcessingContext,
+    ) -> _I_co: ...
+
+
+class DummyInputsBuilderFactory(Protocol[_I]):  # type: ignore[misc]
+    """
+    Constructs a
+    [`BaseDummyInputsBuilder`][vllm.multimodal.processing.BaseDummyInputsBuilder]
+    instance from the context.
+    """
+
+    def __call__(self, info: _I) -> BaseDummyInputsBuilder[_I]: ...
+
+
+class MultiModalProcessorFactory(Protocol[_I]):  # type: ignore[misc]
+    """
+    Constructs a
+    [`BaseMultiModalProcessor`][vllm.multimodal.processing.BaseMultiModalProcessor]
+    instance from the context.
+    """
+
+    def __call__(
+        self,
+        info: _I,
+        dummy_inputs: BaseDummyInputsBuilder[_I],
+        *,
+        cache: BaseMultiModalProcessorCache | None = None,
+    ) -> BaseMultiModalProcessor[_I]: ...
+
+
+@dataclass(frozen=True)
+class _ProcessorFactories(Generic[_I]):
+    info: ProcessingInfoFactory[_I]
+    processor: MultiModalProcessorFactory[_I]
+    dummy_inputs: DummyInputsBuilderFactory[_I]
+
+    def build_processor(
+        self,
+        ctx: InputProcessingContext,
+        *,
+        cache: BaseMultiModalProcessorCache | None = None,
+    ):
+        info = self.info(ctx)
+        dummy_inputs_builder = self.dummy_inputs(info)
+        return self.processor(info, dummy_inputs_builder, cache=cache)
+
+
+class MultiModalRegistry:
+    """
+    A registry that dispatches data processing according to the model.
+    """
+
+    def supports_multimodal_inputs(self, model_config: "ModelConfig") -> bool:
+        """
+        Checks if the model supports multimodal inputs.
+        Returns True if the model is multimodal with any non-zero supported
+        modalities, otherwise returns False, effectively running in
+        text-only mode.
+        """
+        if not model_config.is_multimodal_model:
+            return False
+
+        mm_config = model_config.get_multimodal_config()
+        try:
+            info = self._create_processing_info(model_config, tokenizer=None)
+        except ValueError:
+            logger.warning_once(
+                "Model %s is treated as multimodal but has no registered "
+                "multimodal processor; running in text-only mode.",
+                model_config.model,
+            )
+            return False
+
+        # Check if all supported modalities have limit == 0
+        if all(
+            mm_config.get_limit_per_prompt(modality) == 0
+            for modality in info.supported_mm_limits
+        ):
+            # If enable_mm_embeds is True, we still need MM infrastructure
+            # to process pre-computed embeddings even though encoder won't run
+            if mm_config.enable_mm_embeds:
+                return True
+
+            logger.info_once(
+                "All limits of multimodal modalities supported by the model "
+                "are set to 0, running in text-only mode."
+            )
+            return False
+
+        return True
+
+    def register_processor(
+        self,
+        processor: MultiModalProcessorFactory[_I],
+        *,
+        info: ProcessingInfoFactory[_I],
+        dummy_inputs: DummyInputsBuilderFactory[_I],
+    ):
+        """
+        Register a multi-modal processor to a model class. The processor
+        is constructed lazily, hence a factory method should be passed.
+
+        When the model receives multi-modal data, the provided function is
+        invoked to transform the data into a dictionary of model inputs.
+        """
+
+        def wrapper(model_cls: N) -> N:
+            if "_processor_factory" in model_cls.__dict__:
+                logger.warning(
+                    "Model class %s already has a multi-modal processor "
+                    "registered to %s. It is overwritten by the new one.",
+                    model_cls,
+                    self,
+                )
+
+            model_cls._processor_factory = _ProcessorFactories(
+                info=info,
+                dummy_inputs=dummy_inputs,
+                processor=processor,
+            )
+
+            return model_cls
+
+        return wrapper
+
+    def _get_model_cls(self, model_config: "ModelConfig") -> "SupportsMultiModal":
+        # Avoid circular import
+        from vllm.model_executor.model_loader import get_model_architecture
+
+        model_cls, _ = get_model_architecture(model_config)
+        if not hasattr(model_cls, "_processor_factory"):
+            raise ValueError(
+                f"Model class {model_cls.__name__} has no registered "
+                "multimodal processor"
+            )
+        return cast("SupportsMultiModal", model_cls)
+
+    def _create_processing_ctx(
+        self,
+        model_config: "ModelConfig",
+        tokenizer: TokenizerLike | None = None,
+    ) -> InputProcessingContext:
+        if tokenizer is None:
+            tokenizer = cached_tokenizer_from_config(model_config)
+
+        return InputProcessingContext(model_config, tokenizer)
+
+    def _create_processing_info(
+        self,
+        model_config: "ModelConfig",
+        tokenizer: TokenizerLike | None = None,
+    ) -> BaseProcessingInfo:
+        model_cls = self._get_model_cls(model_config)
+        factories = model_cls._processor_factory
+        ctx = self._create_processing_ctx(model_config, tokenizer)
+        return factories.info(ctx)
+
+    def get_processing_info(self, model_config: "ModelConfig") -> BaseProcessingInfo:
+        return self._create_processing_info(model_config, tokenizer=None)
+
+    def create_processor(
+        self,
+        model_config: "ModelConfig",
+        *,
+        tokenizer: TokenizerLike | None = None,
+        cache: BaseMultiModalProcessorCache | None = None,
+    ) -> BaseMultiModalProcessor[BaseProcessingInfo]:
+        """
+        Create a multi-modal processor for a specific model and tokenizer.
+        """
+        if not model_config.is_multimodal_model:
+            model_name = model_config.served_model_name or model_config.model
+            raise ValueError(f"{model_name} is not a multimodal model")
+
+        model_cls = self._get_model_cls(model_config)
+        factories = model_cls._processor_factory
+
+        ctx = self._create_processing_ctx(model_config, tokenizer)
+
+        return factories.build_processor(ctx, cache=cache)
+
+    def get_dummy_mm_inputs(
+        self,
+        model_config: "ModelConfig",
+        mm_counts: Mapping[str, int],
+        *,
+        cache: BaseMultiModalProcessorCache | None = None,
+        processor: BaseMultiModalProcessor | None = None,
+    ) -> MultiModalInput:
+        """
+        Create dummy data for profiling the memory usage of a model.
+
+        The model is identified by `model_config`.
+        """
+        seq_len = model_config.max_model_len
+
+        if processor is None:
+            processor = self.create_processor(model_config, cache=cache)
+
+        mm_config = model_config.get_multimodal_config()
+        processor_inputs = processor.dummy_inputs.get_dummy_processor_inputs(
+            seq_len=seq_len,
+            mm_counts=mm_counts,
+            mm_options=mm_config.limit_per_prompt,
+        )
+        mm_inputs = processor.apply(
+            processor_inputs,
+            timing_ctx=TimingContext(enabled=False),
+        )
+
+        prompt_token_ids = mm_inputs["prompt_token_ids"]
+        total_len = len(prompt_token_ids)
+        if total_len < seq_len:
+            prompt_token_ids.extend([0] * (seq_len - total_len))
+
+        return mm_inputs
+
+    def _get_cache_type(
+        self,
+        vllm_config: "VllmConfig",
+    ) -> Literal[None, "processor_only", "lru", "shm"]:
+        model_config = vllm_config.model_config
+        if not self.supports_multimodal_inputs(model_config):
+            return None
+
+        # Check if the cache is disabled.
+        mm_config = model_config.get_multimodal_config()
+        if mm_config.mm_processor_cache_gb <= 0:
+            return None
+
+        # Check if IPC caching is supported.
+        parallel_config = vllm_config.parallel_config
+        is_ipc_supported = parallel_config._api_process_count == 1 and (
+            parallel_config.data_parallel_size == 1
+            or parallel_config.data_parallel_external_lb
+        )
+
+        if not is_ipc_supported:
+            return "processor_only"
+
+        mm_config = model_config.get_multimodal_config()
+        return mm_config.mm_processor_cache_type
+
+    def processor_cache_from_config(
+        self,
+        vllm_config: "VllmConfig",
+    ) -> BaseMultiModalProcessorCache | None:
+        """Return a `BaseMultiModalProcessorCache`, if enabled."""
+        cache_type = self._get_cache_type(vllm_config)
+        if cache_type is None:
+            return None
+        elif cache_type == "processor_only":
+            return MultiModalProcessorOnlyCache(vllm_config.model_config)
+        elif cache_type == "lru":
+            return MultiModalProcessorSenderCache(vllm_config.model_config)
+        elif cache_type == "shm":
+            return ShmObjectStoreSenderCache(vllm_config)
+        else:
+            raise ValueError(f"Unknown cache type: {cache_type!r}")
+
+    def processor_only_cache_from_config(
+        self,
+        vllm_config: "VllmConfig",
+    ) -> MultiModalProcessorOnlyCache | None:
+        """Return a `MultiModalProcessorOnlyCache`, if enabled."""
+        cache_type = self._get_cache_type(vllm_config)
+        if cache_type is None:
+            return None
+
+        return MultiModalProcessorOnlyCache(vllm_config.model_config)
+
+    def engine_receiver_cache_from_config(
+        self,
+        vllm_config: "VllmConfig",
+    ) -> BaseMultiModalReceiverCache | None:
+        """Return a `BaseMultiModalReceiverCache` for the engine process."""
+        cache_type = self._get_cache_type(vllm_config)
+        if cache_type in (None, "processor_only", "shm"):
+            return None
+        elif cache_type == "lru":
+            return MultiModalReceiverCache(vllm_config.model_config)
+        else:
+            raise ValueError(f"Unknown cache type: {cache_type!r}")
+
+    def worker_receiver_cache_from_config(
+        self,
+        vllm_config: "VllmConfig",
+        shared_worker_lock: LockType,
+    ) -> BaseMultiModalReceiverCache | None:
+        """Return a `BaseMultiModalReceiverCache` for the worker process."""
+        cache_type = self._get_cache_type(vllm_config)
+        if cache_type in (None, "processor_only", "lru"):
+            return None
+        elif cache_type == "shm":
+            return ShmObjectStoreReceiverCache(vllm_config, shared_worker_lock)
+        else:
+            raise ValueError(f"Unknown cache type: {cache_type!r}")
+
+
+class MultiModalTimingRegistry:
+    def __init__(self, observability_config: "ObservabilityConfig | None") -> None:
+        super().__init__()
+
+        if observability_config and observability_config.enable_mm_processor_stats:
+            self._lock = threading.Lock()
+            self._ctx_by_request_id = defaultdict[str, TimingContext](TimingContext)
+            self._enabled = True
+        else:
+            self._enabled = False
+
+    def get(self, request_id: str) -> TimingContext:
+        if not self._enabled:
+            return TimingContext(enabled=False)
+
+        with self._lock:
+            return self._ctx_by_request_id[request_id]
+
+    def stat(self) -> dict[str, dict[str, float]]:
+        if not self._enabled:
+            return {}
+
+        with self._lock:
+            stats = {
+                req_id: ctx.get_stats_dict()
+                for req_id, ctx in self._ctx_by_request_id.items()
+            }
+            self._ctx_by_request_id.clear()
+            return stats

upstream_ref/ds_vllm_latest/vllm/transformers_utils/configs/qwen3_5.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Qwen3.5 model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class Qwen3_5TextConfig(PretrainedConfig):
+    model_type = "qwen3_5_text"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+    base_config_key = "text_config"
+
+    def __init__(
+        self,
+        vocab_size=248320,
+        hidden_size=4096,
+        intermediate_size=12288,
+        num_hidden_layers=32,
+        num_attention_heads=16,
+        num_key_value_heads=4,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_parameters=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        head_dim=256,
+        linear_conv_kernel_dim=4,
+        linear_key_head_dim=128,
+        linear_value_head_dim=128,
+        linear_num_key_heads=16,
+        linear_num_value_heads=32,
+        layer_types=None,
+        pad_token_id=None,
+        bos_token_id=None,
+        eos_token_id=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.head_dim = head_dim
+        self.rope_parameters = rope_parameters
+        kwargs.setdefault("partial_rotary_factor", 0.25)
+
+        self.layer_types = layer_types
+        if self.layer_types is None:
+            interval_pattern = kwargs.get("full_attention_interval", 4)
+            self.layer_types = [
+                "linear_attention"
+                if bool((i + 1) % interval_pattern)
+                else "full_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+        kwargs["ignore_keys_at_rope_validation"] = {
+            "mrope_section",
+            "mrope_interleaved",
+        }
+        self.validate_layer_type()
+
+        # linear attention part
+        self.linear_conv_kernel_dim = linear_conv_kernel_dim
+        self.linear_key_head_dim = linear_key_head_dim
+        self.linear_value_head_dim = linear_value_head_dim
+        self.linear_num_key_heads = linear_num_key_heads
+        self.linear_num_value_heads = linear_num_value_heads
+        super().__init__(**kwargs)
+        # Set these AFTER super().__init__() because transformers v4's
+        # PretrainedConfig.__init__ has these as explicit params with different
+        # defaults (e.g. tie_word_embeddings=True) that would overwrite our values.
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.tie_word_embeddings = tie_word_embeddings
+
+
+class Qwen3_5VisionConfig(PretrainedConfig):
+    model_type = "qwen3_5"
+    base_config_key = "vision_config"
+
+    def __init__(
+        self,
+        depth=27,
+        hidden_size=1152,
+        hidden_act="gelu_pytorch_tanh",
+        intermediate_size=4304,
+        num_heads=16,
+        in_channels=3,
+        patch_size=16,
+        spatial_merge_size=2,
+        temporal_patch_size=2,
+        out_hidden_size=3584,
+        num_position_embeddings=2304,
+        initializer_range=0.02,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.depth = depth
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.num_heads = num_heads
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.spatial_merge_size = spatial_merge_size
+        self.temporal_patch_size = temporal_patch_size
+        self.out_hidden_size = out_hidden_size
+        self.num_position_embeddings = num_position_embeddings
+        self.initializer_range = initializer_range
+
+
+class Qwen3_5Config(PretrainedConfig):
+    model_type = "qwen3_5"
+    sub_configs = {
+        "vision_config": Qwen3_5VisionConfig,
+        "text_config": Qwen3_5TextConfig,
+    }
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        text_config=None,
+        vision_config=None,
+        image_token_id=248056,
+        video_token_id=248057,
+        vision_start_token_id=248053,
+        vision_end_token_id=248054,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        if isinstance(vision_config, dict):
+            self.vision_config = self.sub_configs["vision_config"](**vision_config)
+        elif vision_config is None:
+            self.vision_config = self.sub_configs["vision_config"]()
+
+        if isinstance(text_config, dict):
+            self.text_config = self.sub_configs["text_config"](**text_config)
+        elif text_config is None:
+            self.text_config = self.sub_configs["text_config"]()
+
+        self.image_token_id = image_token_id
+        self.video_token_id = video_token_id
+        self.vision_start_token_id = vision_start_token_id
+        self.vision_end_token_id = vision_end_token_id
+        super().__init__(**kwargs)
+        # Set after super().__init__() to avoid v4 PretrainedConfig overwrite
+        self.tie_word_embeddings = tie_word_embeddings
+
+
+__all__ = ["Qwen3_5Config", "Qwen3_5TextConfig"]

upstream_ref/ds_vllm_latest/vllm/transformers_utils/configs/qwen3_5_moe.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Qwen3.5-MoE model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class Qwen3_5MoeTextConfig(PretrainedConfig):
+    model_type = "qwen3_5_moe_text"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
+        "layers.*.mlp.experts.down_proj": "rowwise",
+        "layers.*.mlp.shared_expert.gate_proj": "colwise",
+        "layers.*.mlp.shared_expert.up_proj": "colwise",
+        "layers.*.mlp.shared_expert.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+    base_config_key = "text_config"
+
+    def __init__(
+        self,
+        vocab_size=248320,
+        hidden_size=2048,
+        num_hidden_layers=40,
+        num_attention_heads=16,
+        num_key_value_heads=2,
+        hidden_act="silu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=False,
+        rope_parameters=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        head_dim=256,
+        linear_conv_kernel_dim=4,
+        linear_key_head_dim=128,
+        linear_value_head_dim=128,
+        linear_num_key_heads=16,
+        linear_num_value_heads=32,
+        moe_intermediate_size=512,
+        shared_expert_intermediate_size=512,
+        num_experts_per_tok=8,
+        num_experts=256,
+        output_router_logits=False,
+        router_aux_loss_coef=0.001,
+        layer_types=None,
+        pad_token_id=None,
+        bos_token_id=None,
+        eos_token_id=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.head_dim = head_dim
+        self.rope_parameters = rope_parameters
+        kwargs.setdefault("partial_rotary_factor", 0.25)
+
+        self.layer_types = layer_types
+        if self.layer_types is None:
+            interval_pattern = kwargs.get("full_attention_interval", 4)
+            self.layer_types = [
+                "linear_attention"
+                if bool((i + 1) % interval_pattern)
+                else "full_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+        kwargs["ignore_keys_at_rope_validation"] = {
+            "mrope_section",
+            "mrope_interleaved",
+        }
+        self.validate_layer_type()
+
+        # linear attention part
+        self.linear_conv_kernel_dim = linear_conv_kernel_dim
+        self.linear_key_head_dim = linear_key_head_dim
+        self.linear_value_head_dim = linear_value_head_dim
+        self.linear_num_key_heads = linear_num_key_heads
+        self.linear_num_value_heads = linear_num_value_heads
+        self.moe_intermediate_size = moe_intermediate_size
+        self.shared_expert_intermediate_size = shared_expert_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.num_experts = num_experts
+        self.output_router_logits = output_router_logits
+        self.router_aux_loss_coef = router_aux_loss_coef
+        super().__init__(**kwargs)
+        # Set these AFTER super().__init__() because transformers v4's
+        # PretrainedConfig.__init__ has these as explicit params with different
+        # defaults (e.g. tie_word_embeddings=True) that would overwrite our values.
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.tie_word_embeddings = tie_word_embeddings
+
+
+class Qwen3_5MoeVisionConfig(PretrainedConfig):
+    model_type = "qwen3_5_moe"
+    base_config_key = "vision_config"
+
+    def __init__(
+        self,
+        depth=27,
+        hidden_size=1152,
+        hidden_act="gelu_pytorch_tanh",
+        intermediate_size=4304,
+        num_heads=16,
+        in_channels=3,
+        patch_size=16,
+        spatial_merge_size=2,
+        temporal_patch_size=2,
+        out_hidden_size=3584,
+        num_position_embeddings=2304,
+        initializer_range=0.02,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.depth = depth
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.num_heads = num_heads
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.spatial_merge_size = spatial_merge_size
+        self.temporal_patch_size = temporal_patch_size
+        self.out_hidden_size = out_hidden_size
+        self.num_position_embeddings = num_position_embeddings
+        self.initializer_range = initializer_range
+
+
+class Qwen3_5MoeConfig(PretrainedConfig):
+    model_type = "qwen3_5_moe"
+    sub_configs = {
+        "vision_config": Qwen3_5MoeVisionConfig,
+        "text_config": Qwen3_5MoeTextConfig,
+    }
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        text_config=None,
+        vision_config=None,
+        image_token_id=248056,
+        video_token_id=248057,
+        vision_start_token_id=248053,
+        vision_end_token_id=248054,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        if isinstance(vision_config, dict):
+            self.vision_config = self.sub_configs["vision_config"](**vision_config)
+        elif vision_config is None:
+            self.vision_config = self.sub_configs["vision_config"]()
+
+        if isinstance(text_config, dict):
+            self.text_config = self.sub_configs["text_config"](**text_config)
+        elif text_config is None:
+            self.text_config = self.sub_configs["text_config"]()
+
+        self.image_token_id = image_token_id
+        self.video_token_id = video_token_id
+        self.vision_start_token_id = vision_start_token_id
+        self.vision_end_token_id = vision_end_token_id
+        super().__init__(**kwargs)
+        # Set after super().__init__() to avoid v4 PretrainedConfig overwrite
+        self.tie_word_embeddings = tie_word_embeddings
+
+
+__all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"]

upstream_ref/xllm_latest/core/kernels/ilu/activation.cpp

@@ -0,0 +1,32 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "ilu_ops_api.h"
+
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void act_and_mul(torch::Tensor out,
+                 torch::Tensor input,
+                 const std::string& act_mode) {
+  if (act_mode == "silu") {
+    infer::silu_and_mul(input, out);
+  } else {
+    LOG(FATAL) << "Unsupported act mode: " << act_mode
+               << ", only support silu, gelu, gelu_tanh";
+  }
+}
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm_latest/core/kernels/ilu/attention.cpp

@@ -0,0 +1,163 @@
+
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "ilu_ops_api.h"
+#include "ixinfer.h"
+#include "utils.h"
+
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void reshape_paged_cache(torch::Tensor& key,
+                         std::optional<torch::Tensor>& value,
+                         torch::Tensor& key_cache,
+                         std::optional<torch::Tensor>& value_cache,
+                         torch::Tensor& slot_mapping) {
+  auto value_ = value.value_or(torch::Tensor());
+  auto value_cache_ = value_cache.value_or(torch::Tensor());
+
+  int64_t key_token_stride = key.stride(0);
+  int64_t value_token_stride = 0;
+  if (value_.defined()) {
+    value_token_stride = value_.stride(0);
+  }
+  slot_mapping = slot_mapping.to(at::kLong);
+  infer::xllm_reshape_and_cache(key,
+                                value_,
+                                key_cache,
+                                value_cache_,
+                                slot_mapping,
+                                key_token_stride,
+                                value_token_stride);
+}
+
+void batch_prefill(torch::Tensor& query,
+                   const torch::Tensor& key,
+                   const std::optional<torch::Tensor>& value,
+                   torch::Tensor& output,
+                   std::optional<torch::Tensor>& output_lse,
+                   const std::optional<torch::Tensor>& q_cu_seq_lens,
+                   const std::optional<torch::Tensor>& kv_cu_seq_lens,
+                   const std::optional<torch::Tensor>& alibi_slope,
+                   const std::optional<torch::Tensor>& attn_bias,
+                   const std::optional<torch::Tensor>& q_quant_scale,
+                   const std::optional<torch::Tensor>& k_quant_scale,
+                   const std::optional<torch::Tensor>& v_quant_scale,
+                   const torch::Tensor& block_tables,
+                   int64_t max_query_len,
+                   int64_t max_seq_len,
+                   float scale,
+                   bool is_causal,
+                   int64_t window_size_left,
+                   int64_t window_size_right,
+                   const std::string& compute_dtype,
+                   bool return_lse) {
+  double softcap = 0.0;
+  bool sqrt_alibi = false;
+  auto q_cu_seq_lens_ = q_cu_seq_lens.value_or(torch::Tensor());
+  auto kv_cu_seq_lens_ = kv_cu_seq_lens.value_or(torch::Tensor());
+  auto q_quant_scale_ = q_quant_scale.value_or(torch::Tensor());
+  auto k_quant_scale_ = k_quant_scale.value_or(torch::Tensor());
+  auto v_quant_scale_ = v_quant_scale.value_or(torch::Tensor());
+  auto block_tables_ = block_tables;
+  auto key_ = key;
+  auto value_ = value.value();
+  infer::ixinfer_flash_attn_unpad_with_block_tables(query,
+                                                    key_,
+                                                    value_,
+                                                    output,
+                                                    block_tables_,
+                                                    q_cu_seq_lens_,
+                                                    kv_cu_seq_lens_,
+                                                    max_query_len,
+                                                    max_seq_len,
+                                                    is_causal,
+                                                    window_size_left,
+                                                    window_size_right,
+                                                    static_cast<double>(scale),
+                                                    softcap,
+                                                    sqrt_alibi,
+                                                    alibi_slope,
+                                                    c10::nullopt,
+                                                    output_lse);
+}
+
+void batch_decode(torch::Tensor& query,
+                  const torch::Tensor& k_cache,
+                  torch::Tensor& output,
+                  const torch::Tensor& block_table,
+                  const torch::Tensor& seq_lens,
+                  const std::optional<torch::Tensor>& v_cache,
+                  std::optional<torch::Tensor>& output_lse,
+                  const std::optional<torch::Tensor>& q_quant_scale,
+                  const std::optional<torch::Tensor>& k_cache_quant_scale,
+                  const std::optional<torch::Tensor>& v_cache_quant_scale,
+                  const std::optional<torch::Tensor>& out_quant_scale,
+                  const std::optional<torch::Tensor>& alibi_slope,
+                  const std::optional<torch::Tensor>& mask,
+                  const std::string& compute_dtype,
+                  int64_t max_seq_len,
+                  int64_t window_size_left,
+                  int64_t window_size_right,
+                  float scale,
+                  bool return_lse,
+                  bool is_causal,
+                  int64_t kv_cache_quant_bit_size) {
+  if (query.dim() == 4) {
+    query =
+        query
+            .view({query.size(0) * query.size(1), query.size(2), query.size(3)})
+            .contiguous();
+  }
+  if (output.dim() == 4) {
+    output = output
+                 .view({output.size(0) * output.size(1),
+                        output.size(2),
+                        output.size(3)})
+                 .contiguous();
+    ;
+  }
+  auto v_cache_ = v_cache.value_or(torch::Tensor());
+  int64_t num_kv_heads = k_cache.size(1);
+  int64_t page_block_size = k_cache.size(2);
+  double softcap = 0.0;
+  bool enable_cuda_graph = false;
+  bool use_sqrt_alibi = false;
+  auto block_table_ = block_table;
+  auto k_cache_ = k_cache;
+  auto seq_lens_ = seq_lens;
+  infer::xllm_paged_attention(output,
+                              query,
+                              k_cache_,
+                              v_cache_,
+                              num_kv_heads,
+                              scale,
+                              block_table_,
+                              seq_lens_,
+                              page_block_size,
+                              max_seq_len,
+                              alibi_slope,
+                              is_causal,
+                              (int32_t)window_size_left,
+                              (int32_t)window_size_right,
+                              softcap,
+                              enable_cuda_graph,
+                              use_sqrt_alibi,
+                              c10::nullopt);
+}
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp

@@ -0,0 +1,99 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <glog/logging.h>
+
+#include "ilu_ops_api.h"
+
+namespace xllm::kernel::ilu {
+
+std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
+    const torch::Tensor& input,
+    int64_t topk,
+    int64_t num_expert_group,
+    int64_t topk_group,
+    bool normalize,
+    const std::optional<torch::Tensor>& mask,
+    const std::string& normed_by,
+    const std::string& scoring_func,
+    double route_scale,
+    const std::optional<torch::Tensor>& e_score_correction_bias) {
+  torch::Tensor input_ = input.to(torch::kFloat32);
+  auto reduce_weight =
+      torch::empty({input.size(0), topk},
+                   torch::dtype(torch::kFloat).device(input.device()));
+  auto topk_indices =
+      torch::empty({input.size(0), topk},
+                   torch::dtype(torch::kInt32).device(input.device()));
+  auto token_expert_indices =
+      torch::empty({input.size(0), topk},
+                   torch::dtype(torch::kInt32).device(input.device()));
+
+  infer::topk_softmax(
+      reduce_weight, topk_indices, token_expert_indices, input_, false);
+
+  auto tt = reduce_weight.sum(-1);
+  if (normalize) {
+    reduce_weight = reduce_weight / reduce_weight.sum(-1).unsqueeze(-1);
+  }
+  return std::make_tuple(reduce_weight, topk_indices);
+}
+
+std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
+                                       int64_t expert_num) {
+  auto src_dst = expert_id.new_empty({expert_id.numel()});
+  auto dst_src = torch::empty_like(src_dst);
+  auto expert_sizes_gpu = expert_id.new_empty({expert_num});
+  auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1});
+  infer::moe_compute_token_index_api(expert_id,
+                                     src_dst,
+                                     dst_src,
+                                     expert_sizes_gpu,
+                                     /*expert_mask=*/std::nullopt,
+                                     /*expert_sizes_cpu*/ std::nullopt,
+                                     /*expert_sizes_gpu*/ std::nullopt,
+                                     0,
+                                     expert_num,
+                                     expert_num);
+
+  expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1);
+  return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum};
+}
+
+torch::Tensor moe_expand_input(const torch::Tensor& input,
+                               const torch::Tensor& gather_index,
+                               const torch::Tensor& combine_idx,
+                               int64_t topk) {
+  int64_t dst_tokens = input.size(0) * topk;
+  auto output = input.new_empty({dst_tokens, input.size(1)});
+  infer::moe_expand_input(
+      output, input, combine_idx, gather_index, dst_tokens, topk);
+
+  return output;
+}
+
+torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight) {
+  input = input.view({-1, weight.size(1), input.size(1)});
+  auto output = input.new_empty({input.size(0), input.size(2)});
+  infer::moe_output_reduce_sum(output,
+                               input,
+                               weight,
+                               /*mask=*/std::nullopt,
+                               /*extra_residual*/ std::nullopt,
+                               /*scaling_factor=*/1.0);
+  return output;
+}
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm_latest/core/kernels/ilu/group_gemm.cpp

@@ -0,0 +1,39 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "ilu_ops_api.h"
+
+namespace xllm::kernel::ilu {
+
+torch::Tensor group_gemm(torch::Tensor& input,
+                         torch::Tensor& weight,
+                         torch::Tensor& tokens_per_experts,
+                         const std::optional<torch::Tensor>& dst_to_src,
+                         torch::Tensor& output) {
+  infer::moe_w16a16_group_gemm(
+      output,
+      input,
+      weight,
+      tokens_per_experts,
+      dst_to_src,
+      /*bias=*/std::nullopt,
+      /*format=*/"TN",
+      /*persistent=*/0,
+      /*output_n=*/tokens_per_experts.sum().item<int64_t>());
+
+  return output;
+}
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm_latest/core/kernels/ilu/ilu_ops_api.h

@@ -0,0 +1,153 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <ATen/DynamicLibrary.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <cuda_runtime.h>
+#include <glog/logging.h>
+#include <torch/all.h>
+
+#include <optional>
+
+#include "ATen/Tensor.h"
+#include "ATen/cuda/CUDAEvent.h"
+#include "c10/core/Device.h"
+#include "c10/core/DeviceGuard.h"
+#include "c10/core/GradMode.h"
+#include "c10/core/InferenceMode.h"
+#include "c10/core/MemoryFormat.h"
+#include "c10/core/ScalarType.h"
+#include "c10/core/TensorOptions.h"
+#include "c10/cuda/CUDAFunctions.h"
+#include "c10/cuda/CUDAGuard.h"
+#include "c10/cuda/CUDAStream.h"
+#include "ixformer.h"
+#include "kernels/kernels.h"
+
+// #include "utils.h"
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
+                                      torch::Tensor& key,
+                                      torch::Tensor& cos_sin_cache,
+                                      torch::Tensor& positions,
+                                      bool interleave);
+
+// act_mode only support silu, gelu, gelu_tanh
+void act_and_mul(torch::Tensor out,
+                 torch::Tensor input,
+                 const std::string& act_mode);
+
+void reshape_paged_cache(
+    torch::Tensor& key,                   //  (num_tokens, num_heads, head_size)
+    std::optional<torch::Tensor>& value,  // (num_tokens, num_heads, head_size)
+    torch::Tensor& key_cache,  // (num_blocks, num_heads, block_size, head_size)
+    std::optional<torch::Tensor>&
+        value_cache,  // (num_blocks, num_heads, block_size, head_size)
+    torch::Tensor& slot_mapping);  //(num_tokens)
+
+void batch_prefill(torch::Tensor& query,
+                   const torch::Tensor& key,
+                   const std::optional<torch::Tensor>& value,
+                   torch::Tensor& output,
+                   std::optional<torch::Tensor>& output_lse,
+                   const std::optional<torch::Tensor>& q_cu_seq_lens,
+                   const std::optional<torch::Tensor>& kv_cu_seq_lens,
+                   const std::optional<torch::Tensor>& alibi_slope,
+                   const std::optional<torch::Tensor>& attn_bias,
+                   const std::optional<torch::Tensor>& q_quant_scale,
+                   const std::optional<torch::Tensor>& k_quant_scale,
+                   const std::optional<torch::Tensor>& v_quant_scale,
+                   const torch::Tensor& block_tables,
+                   int64_t max_query_len,
+                   int64_t max_seq_len,
+                   float scale,
+                   bool is_causal,
+                   int64_t window_size_left,
+                   int64_t window_size_right,
+                   const std::string& compute_dtype,
+                   bool return_lse);
+
+void batch_decode(torch::Tensor& query,
+                  const torch::Tensor& k_cache,
+                  torch::Tensor& output,
+                  const torch::Tensor& block_table,
+                  const torch::Tensor& seq_lens,
+                  const std::optional<torch::Tensor>& v_cache,
+                  std::optional<torch::Tensor>& output_lse,
+                  const std::optional<torch::Tensor>& q_quant_scale,
+                  const std::optional<torch::Tensor>& k_cache_quant_scale,
+                  const std::optional<torch::Tensor>& v_cache_quant_scale,
+                  const std::optional<torch::Tensor>& out_quant_scale,
+                  const std::optional<torch::Tensor>& alibi_slope,
+                  const std::optional<torch::Tensor>& mask,
+                  const std::string& compute_dtype,
+                  int64_t max_seq_len,
+                  int64_t window_size_left,
+                  int64_t window_size_right,
+                  float scale,
+                  bool return_lse,
+                  bool is_causal,
+                  int64_t kv_cache_quant_bit_size);
+
+void residual_layer_norm(torch::Tensor& input,
+                         torch::Tensor& output,
+                         std::optional<torch::Tensor>& residual,
+                         torch::Tensor& weight,
+                         std::optional<torch::Tensor>& bias,
+                         std::optional<torch::Tensor>& residual_out,
+                         double eps);
+
+void rms_norm(torch::Tensor& output,
+              torch::Tensor& input,
+              torch::Tensor& weight,
+              double eps);
+
+torch::Tensor matmul(torch::Tensor a,
+                     torch::Tensor b,
+                     std::optional<torch::Tensor> bias);
+
+std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
+    const torch::Tensor& input,
+    int64_t topk,
+    int64_t num_expert_group,
+    int64_t topk_group,
+    bool normalize,
+    const std::optional<torch::Tensor>& mask,
+    const std::string& normed_by,
+    const std::string& scoring_func,
+    double route_scale,
+    const std::optional<torch::Tensor>& e_score_correction_bias);
+
+std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
+                                       int64_t expert_num);
+
+torch::Tensor moe_expand_input(const torch::Tensor& input,
+                               const torch::Tensor& gather_index,
+                               const torch::Tensor& combine_idx,
+                               int64_t topk);
+
+torch::Tensor group_gemm(torch::Tensor& input,
+                         torch::Tensor& weight,
+                         torch::Tensor& tokens_per_experts,
+                         const std::optional<torch::Tensor>& dst_to_src,
+                         torch::Tensor& output);
+
+torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight);
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h

@@ -0,0 +1,147 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <torch/all.h>
+
+#include "ATen/Tensor.h"
+#include "utils.h"
+
+namespace ixformer::infer {
+torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
+    torch::Tensor& query,
+    torch::Tensor& key_cache,
+    torch::Tensor& value_cache,
+    torch::Tensor& out,
+    torch::Tensor& block_tables,
+    torch::Tensor& cu_seq_q,
+    torch::Tensor& cu_seq_k,
+    int64_t max_seq_q,
+    int64_t max_seq_k,
+    bool is_causal,
+    int64_t window_left,
+    int64_t window_right,
+    double scale,
+    double softcap,
+    bool sqrt_alibi,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    const std::optional<torch::Tensor>& sinks,
+    std::optional<torch::Tensor>& lse);
+
+void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
+
+torch::Tensor xllm_paged_attention(
+    torch::Tensor& out,
+    torch::Tensor& query,
+    torch::Tensor& key_cache,
+    torch::Tensor& value_cache,
+    int64_t num_kv_heads,
+    double scale,
+    torch::Tensor& block_tables,
+    torch::Tensor& context_lens,
+    int64_t block_size,
+    int64_t max_context_len,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    bool causal,
+    int32_t window_left,
+    int32_t window_right,
+    double softcap,
+    bool enable_cuda_graph,
+    bool use_sqrt_alibi,
+    const std::optional<torch::Tensor>& sinks);
+
+torch::Tensor ixformer_linear(torch::Tensor& input,
+                              torch::Tensor& weight,
+                              int64_t act_type,
+                              const std::optional<torch::Tensor>& bias,
+                              const std::optional<torch::Tensor>& out,
+                              const std::optional<bool> persistent);
+
+torch::Tensor ixformer_linear_ex(torch::Tensor& input,
+                                 torch::Tensor& weight,
+                                 const c10::optional<torch::Tensor>& bias,
+                                 const c10::optional<torch::Tensor>& out);
+
+void xllm_reshape_and_cache(torch::Tensor& key,
+                            torch::Tensor& value,
+                            torch::Tensor& key_cache,
+                            torch::Tensor& value_cache,
+                            torch::Tensor& slot_mapping,
+                            int64_t key_token_stride,
+                            int64_t value_token_stride);
+
+void xllm_rotary_embedding(torch::Tensor& positions,
+                           torch::Tensor& query,
+                           torch::Tensor& key,
+                           int64_t head_size,
+                           torch::Tensor& cos_sin_cache,
+                           bool is_neox);
+
+void residual_rms_norm(torch::Tensor& input,
+                       torch::Tensor& residual,
+                       torch::Tensor& weight,
+                       torch::Tensor& output,
+                       torch::Tensor& residual_output,
+                       const std::optional<torch::Tensor>& fused_bias,
+                       double alpha,
+                       double eps,
+                       bool is_post);
+
+void rms_norm(torch::Tensor& input,
+              torch::Tensor& weight,
+              torch::Tensor& output,
+              const std::optional<torch::Tensor>& fused_bias,
+              double eps);
+
+void topk_softmax(torch::Tensor& topk_weights,
+                  torch::Tensor& topk_indices,
+                  torch::Tensor& token_expert_indices,
+                  torch::Tensor& gating_output,
+                  bool renormalize);
+
+void moe_compute_token_index_api(
+    torch::Tensor& topk_ids,
+    torch::Tensor& src_dst,
+    torch::Tensor& dst_src,
+    torch::Tensor& expert_sizes_gpu,
+    const c10::optional<torch::Tensor>& expert_mask,
+    const c10::optional<torch::Tensor>& expert_sizes_cpu,
+    const c10::optional<torch::Tensor>& expand_tokens_gpu,
+    int64_t start_expert_id,
+    int64_t end_expert_id,
+    int64_t num_experts);
+
+void moe_expand_input(torch::Tensor outputs,
+                      torch::Tensor inputs,
+                      torch::Tensor dst_to_src,
+                      const c10::optional<torch::Tensor>& src_to_dst,
+                      int64_t dst_tokens,
+                      int64_t expand_factor);
+
+void moe_w16a16_group_gemm(torch::Tensor output,
+                           torch::Tensor inputs,
+                           torch::Tensor weights,
+                           torch::Tensor tokens_per_experts,
+                           const c10::optional<torch::Tensor>& dst_to_src,
+                           const c10::optional<torch::Tensor>& bias,
+                           std::string format,
+                           int64_t persistent,
+                           int64_t output_n);
+
+void moe_output_reduce_sum(torch::Tensor outputs,
+                           torch::Tensor inputs,
+                           const c10::optional<torch::Tensor>& mul_weight,
+                           const c10::optional<torch::Tensor>& mask,
+                           const c10::optional<torch::Tensor>& extra_residual,
+                           double scaling_factor);
+}  // namespace ixformer::infer

upstream_ref/xllm_latest/core/kernels/ilu/matmul.cpp

@@ -0,0 +1,73 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "ilu_ops_api.h"
+#include "util/env_var.h"
+
+namespace xllm::kernel::ilu {
+
+bool gemv_conditions(const torch::Tensor& input,
+                     const torch::Tensor& weight,
+                     const torch::Tensor& bias,
+                     int64_t gemv_max_batch) {
+  // gemv input:[m,k] weight:[n,k]
+  // 1. m <= gemv_max_batch
+  // 2. k % 32 == 0 && n % 2 == 0
+  // 3. bias is None
+
+  torch::Tensor input_view = input.view({-1, input.size(-1)});
+  torch::Tensor weight_view = weight.view({-1, weight.size(-1)});
+
+  int64_t m = input_view.size(0);
+  int64_t k = input_view.size(1);
+  int64_t n = weight_view.size(0);
+
+  if (bias.defined() == false && m <= gemv_max_batch && k % 32 == 0 &&
+      n % 2 == 0) {
+    return true;
+  }
+  return false;
+}
+
+torch::Tensor matmul(torch::Tensor a,
+                     torch::Tensor b,
+                     std::optional<torch::Tensor> bias) {
+  int64_t act_type = -1;
+  bool persistent = false;
+  std::vector<int64_t> output_shape = a.sizes().vec();
+  if (!output_shape.empty()) {
+    output_shape[output_shape.size() - 1] = b.size(0);
+  }
+  torch::Tensor output = a.new_empty(output_shape);
+
+  bool use_gemv = true;
+  const int64_t gemv_max_batch = 1;
+  const bool disable_infer_gemm_ex =
+      xllm::util::get_bool_env("DISABLE_INFER_GEMM_EX", false);
+
+  use_gemv =
+      use_gemv &&
+      gemv_conditions(a, b, bias.value_or(at::Tensor()), gemv_max_batch) &&
+      !disable_infer_gemm_ex && (act_type == -1);
+
+  if (use_gemv) {
+    output = infer::ixformer_linear_ex(a, b, bias, output);
+  } else {
+    output = infer::ixformer_linear(a, b, act_type, bias, output, persistent);
+  }
+  return output;
+}
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm_latest/core/kernels/ilu/norm.cpp

@@ -0,0 +1,51 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "ilu_ops_api.h"
+#include "utils.h"
+
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void residual_layer_norm(torch::Tensor& input,
+                         torch::Tensor& output,
+                         std::optional<torch::Tensor>& residual,
+                         torch::Tensor& weight,
+                         std::optional<torch::Tensor>& bias,
+                         std::optional<torch::Tensor>& residual_out,
+                         double eps) {
+  auto residual_ = residual.value_or(torch::zeros_like(input));
+  torch::Tensor residual_out_ = residual_out.value_or(torch::zeros_like(input));
+  infer::residual_rms_norm(input,
+                           residual_,
+                           weight,
+                           output,
+                           residual_out_,
+                           bias,
+                           /*alpha=*/1.0,
+                           eps,
+                           false);
+}
+
+void rms_norm(torch::Tensor& output,
+              torch::Tensor& input,
+              torch::Tensor& weight,
+              double eps) {
+  std::optional<torch::Tensor> fused_bias = std::nullopt;
+  infer::rms_norm(input, weight, output, fused_bias, eps);
+}
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm_latest/core/kernels/ilu/rope.cpp

@@ -0,0 +1,31 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "ilu_ops_api.h"
+#include "utils.h"
+
+namespace xllm::kernel::ilu {
+
+void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
+                                      torch::Tensor& key,
+                                      torch::Tensor& cos_sin_cache,
+                                      torch::Tensor& positions,
+                                      bool interleave) {
+  const int64_t head_size = cos_sin_cache.size(-1);
+  infer::xllm_rotary_embedding(
+      positions, query, key, head_size, cos_sin_cache, !interleave);
+}
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm_latest/core/kernels/ilu/utils.h

@@ -0,0 +1,63 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#pragma once
+namespace xllm::kernel::ilu {
+#undef check_tensor_contiguous
+#define check_tensor_contiguous(x, type) \
+  TORCH_CHECK(x.scalar_type() == type);  \
+  TORCH_CHECK(x.is_cuda());              \
+  TORCH_CHECK(x.is_contiguous());
+
+#undef check_tensor_half_bf_float
+#define check_tensor_half_bf_float(x)                       \
+  TORCH_CHECK(x.scalar_type() == at::ScalarType::Half ||    \
+              x.scalar_type() == at::ScalarType::Float ||   \
+              x.scalar_type() == at::ScalarType::BFloat16); \
+  TORCH_CHECK(x.is_cuda());
+
+// from torchCheckMsgImpl
+inline const char* ixformer_check_msg_impl(const char* msg) { return msg; }
+// // If there is just 1 user-provided C-string argument, use it.
+
+#define IXFORMER_CHECK_MSG(cond, type, ...)                 \
+  (ixformer_check_msg_impl(                                 \
+      "Expected " #cond                                     \
+      " to be true, but got false.  "                       \
+      "(Could this error message be improved?  If so, "     \
+      "please report an enhancement request to ixformer.)", \
+      ##__VA_ARGS__))
+
+#define IXFORMER_CHECK(cond, ...)                                            \
+  {                                                                          \
+    if (!(cond)) {                                                           \
+      std::cerr << __FILE__ << " (" << __LINE__ << ")"                       \
+                << "-" << __FUNCTION__ << " : "                              \
+                << IXFORMER_CHECK_MSG(cond, "", ##__VA_ARGS__) << std::endl; \
+      throw std::runtime_error("IXFORMER_CHECK ERROR");                      \
+    }                                                                        \
+  }
+
+#undef CUINFER_CHECK
+#define CUINFER_CHECK(func)                                                \
+  do {                                                                     \
+    cuinferStatus_t status = (func);                                       \
+    if (status != CUINFER_STATUS_SUCCESS) {                                \
+      std::cerr << "Error in file " << __FILE__ << " on line " << __LINE__ \
+                << ": " << cuinferGetErrorString(status) << std::endl;     \
+      throw std::runtime_error("CUINFER_CHECK ERROR");                     \
+    }                                                                      \
+  } while (0)
+
+}  // namespace xllm::kernel::ilu

upstream_ref/xllm_latest/core/layers/ilu/attention.cpp

@@ -0,0 +1,189 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "attention.h"
+
+#include "kernels/ilu/ilu_ops_api.h"
+#include "kernels/ops_api.h"
+
+namespace xllm {
+namespace layer {
+AttentionImpl::AttentionImpl(int64_t num_heads,
+                             int64_t head_size,
+                             float scale,
+                             int64_t num_kv_heads,
+                             int64_t sliding_window)
+    : num_heads_(num_heads),
+      head_size_(head_size),
+      scale_(scale),
+      num_kv_heads_(num_kv_heads),
+      v_head_dim_(head_size),
+      use_fused_mla_qkv_(false),
+      enable_lighting_indexer_(false),
+      enable_mla_(false),
+      sliding_window_(sliding_window) {
+  if (sliding_window_ > -1) {
+    sliding_window_ = sliding_window_ - 1;
+  }
+}
+
+AttentionImpl::AttentionImpl(int64_t num_heads,
+                             int64_t head_size,
+                             int64_t num_kv_heads,
+                             int64_t v_head_dim,
+                             int64_t sliding_window,
+                             float scale,
+                             bool use_fused_mla_qkv,
+                             bool enable_lighting_indexer,
+                             bool enable_mla)
+    : num_heads_(num_heads),
+      head_size_(head_size),
+      scale_(scale),
+      num_kv_heads_(num_kv_heads),
+      v_head_dim_(v_head_dim),
+      use_fused_mla_qkv_(use_fused_mla_qkv),
+      enable_lighting_indexer_(enable_lighting_indexer),
+      enable_mla_(enable_mla),
+      sliding_window_(sliding_window) {
+  if (sliding_window_ > -1) {
+    sliding_window_ = sliding_window_ - 1;
+  }
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> AttentionImpl::forward(
+    const AttentionMetadata& attn_metadata,
+    torch::Tensor& query,
+    torch::Tensor& key,
+    torch::Tensor& value,
+    KVCache& kv_cache) {
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+  torch::Tensor output;
+  if (enable_mla_) {
+    output = torch::empty({query.size(0), num_heads_ * v_head_dim_},
+                          query.options());
+  } else {
+    output = torch::empty_like(query);
+  }
+  if (attn_metadata.is_dummy) {
+    return std::make_tuple(output, output_lse);
+  }
+
+  bool only_prefill =
+      attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
+  int64_t num_kv_heads = (enable_mla_ && !only_prefill) ? 1 : num_kv_heads_;
+  torch::Tensor k_cache = kv_cache.get_k_cache();
+  std::optional<torch::Tensor> v_cache;
+  std::optional<torch::Tensor> v;
+  if (!enable_mla_) {
+    v = value.view({-1, num_kv_heads, head_size_});
+    v_cache = kv_cache.get_v_cache();
+  }
+
+  bool skip_process_cache = enable_mla_ && (only_prefill || use_fused_mla_qkv_);
+  if (!skip_process_cache) {
+    xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params;
+    reshape_paged_cache_params.key = key.view({-1, num_kv_heads, head_size_});
+    reshape_paged_cache_params.value = v;
+    reshape_paged_cache_params.k_cache = k_cache;
+    reshape_paged_cache_params.v_cache = v_cache;
+    reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping;
+    xllm::kernel::reshape_paged_cache(reshape_paged_cache_params);
+  }
+
+  if (enable_lighting_indexer_ || !only_prefill) {
+    decoder_forward(query, output, k_cache, v_cache, attn_metadata);
+  } else {
+    prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata);
+  }
+
+  int64_t head_size = enable_mla_ ? v_head_dim_ : head_size_;
+  output = output.view({-1, num_heads_ * head_size});
+  return {output, output_lse};
+}
+
+void AttentionImpl::prefill_forward(torch::Tensor& query,
+                                    torch::Tensor& key,
+                                    torch::Tensor& value,
+                                    torch::Tensor& output,
+                                    const torch::Tensor& k_cache,
+                                    const std::optional<torch::Tensor>& v_cache,
+                                    const AttentionMetadata& attn_metadata) {
+  int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+  query = query.view({-1, num_heads_, head_size_});
+  output = output.view({-1, num_heads_, head_size_v});
+  // torch::Tensor k_cache_ = k_cache;
+  // torch::Tensor v_cache_ = v_cache.value();
+  xllm::kernel::ilu::batch_prefill(query,
+                                   k_cache,
+                                   v_cache,
+                                   output,
+                                   output_lse,
+                                   attn_metadata.q_cu_seq_lens,
+                                   attn_metadata.kv_cu_seq_lens,
+                                   /*alibi_slope=*/std::nullopt,
+                                   /*attn_bias=*/std::nullopt,
+                                   /*q_quant_scale=*/std::nullopt,
+                                   /*k_quant_scale=*/std::nullopt,
+                                   /*v_quant_scale=*/std::nullopt,
+                                   attn_metadata.block_table,
+                                   attn_metadata.max_query_len,
+                                   attn_metadata.max_seq_len,
+                                   scale_,
+                                   attn_metadata.is_causal,
+                                   sliding_window_,
+                                   /*window_size_right=*/-1,
+                                   attn_metadata.compute_dtype,
+                                   /*return_lse=*/false);
+}
+
+void AttentionImpl::decoder_forward(torch::Tensor& query,
+                                    torch::Tensor& output,
+                                    const torch::Tensor& k_cache,
+                                    const std::optional<torch::Tensor>& v_cache,
+                                    const AttentionMetadata& attn_metadata) {
+  int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
+  query = query.view({-1, 1, num_heads_, head_size_});
+  output = output.view({-1, 1, num_heads_, head_size_v});
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+
+  int64_t block_aligned_max_seq_len =
+      attn_metadata.block_table.size(-1) * k_cache.size(2);
+
+  xllm::kernel::ilu::batch_decode(query,
+                                  k_cache,
+                                  output,
+                                  attn_metadata.block_table,
+                                  attn_metadata.kv_seq_lens,
+                                  v_cache,
+                                  output_lse,
+                                  /*q_quant_scale=*/std::nullopt,
+                                  /*k_quant_scale=*/std::nullopt,
+                                  /*v_quant_scale=*/std::nullopt,
+                                  /*out_quant_scale=*/std::nullopt,
+                                  /*alibi_slope=*/std::nullopt,
+                                  attn_metadata.attn_mask,
+                                  attn_metadata.compute_dtype,
+                                  block_aligned_max_seq_len,
+                                  sliding_window_,
+                                  /*window_size_right=*/-1,
+                                  scale_,
+                                  /*return_lse=*/false,
+                                  attn_metadata.is_causal,
+                                  /*kv_cache_quant_bit_size=*/-1);
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm_latest/core/layers/ilu/attention.h

@@ -0,0 +1,82 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include <tuple>
+
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_input_params.h"
+#include "layers/common/attention_metadata.h"
+
+namespace xllm {
+namespace layer {
+class AttentionImpl : public torch::nn::Module {
+ public:
+  AttentionImpl() = default;
+
+  AttentionImpl(int64_t num_heads,
+                int64_t head_size,
+                float scale,
+                int64_t num_kv_heads,
+                int64_t sliding_window);
+  AttentionImpl(int64_t num_heads,
+                int64_t head_size,
+                int64_t num_kv_heads,
+                int64_t v_head_dim,
+                int64_t sliding_window,
+                float scale,
+                bool use_fused_mla_qkv,
+                bool enable_lighting_indexer,
+                bool enable_mla);
+
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
+      const AttentionMetadata& attn_metadata,
+      torch::Tensor& query,
+      torch::Tensor& key,
+      torch::Tensor& value,
+      KVCache& kv_cache);
+
+  void prefill_forward(torch::Tensor& query,
+                       torch::Tensor& key,
+                       torch::Tensor& value,
+                       torch::Tensor& output,
+                       const torch::Tensor& k_cache,
+                       const std::optional<torch::Tensor>& v_cache,
+                       const AttentionMetadata& attn_metadata);
+
+  void decoder_forward(torch::Tensor& query,
+                       torch::Tensor& output,
+                       const torch::Tensor& k_cache,
+                       const std::optional<torch::Tensor>& v_cache,
+                       const AttentionMetadata& attn_metadata);
+
+ private:
+  int64_t num_heads_;
+  int64_t head_size_;
+  float scale_;
+  int64_t num_kv_heads_;
+  int64_t v_head_dim_;
+  bool use_fused_mla_qkv_;
+  bool enable_lighting_indexer_;
+  bool enable_mla_;
+  int64_t sliding_window_;
+};
+TORCH_MODULE(Attention);
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp

@@ -0,0 +1,806 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "fused_moe.h"
+
+#include <glog/logging.h>
+
+#include <iomanip>
+
+#include "common/global_flags.h"
+#include "core/framework/config/eplb_config.h"
+#include "core/framework/config/scheduler_config.h"
+#include "core/framework/config/speculative_config.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "kernels/ops_api.h"
+#include "layers/common/dp_utils.h"
+#include "util/utils.h"
+
+namespace {
+
+int32_t get_dtype_size(torch::ScalarType dtype) {
+  return static_cast<int32_t>(torch::elementSize(dtype));
+}
+
+}  // namespace
+
+namespace xllm {
+namespace layer {
+
+FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args,
+                           const FusedMoEArgs& moe_args,
+                           const QuantArgs& quant_args,
+                           const ParallelArgs& parallel_args,
+                           const torch::TensorOptions& options)
+    : num_total_experts_(static_cast<int64_t>(model_args.n_routed_experts())),
+      topk_(model_args.num_experts_per_tok()),
+      num_expert_group_(model_args.n_group()),
+      topk_group_(model_args.topk_group()),
+      route_scale_(model_args.routed_scaling_factor()),
+      hidden_size_(model_args.hidden_size()),
+      n_shared_experts_(model_args.n_shared_experts()),
+      is_gated_(moe_args.is_gated),
+      renormalize_(model_args.norm_topk_prob() ? 1 : 0),
+      hidden_act_(model_args.hidden_act()),
+      scoring_func_(model_args.scoring_func()),
+      quant_args_(quant_args),
+      parallel_args_(parallel_args),
+      options_(options),
+      device_(options.device()) {
+  const int64_t num_experts = num_total_experts_;
+  const int64_t intermediate_size =
+      static_cast<int64_t>(model_args.moe_intermediate_size());
+  const std::string& topk_method = model_args.topk_method();
+  int64_t ep_size = parallel_args.ep_size();
+  int64_t ep_rank = 0;
+  tp_pg_ = parallel_args.tp_group_;
+  if (ep_size > 1) {
+    ep_rank = parallel_args.moe_ep_group_->rank();
+    tp_pg_ = parallel_args.moe_tp_group_;
+  }
+
+  // smoothquant check: If quant_method is not empty, only w8a8 smoothquant is
+  // supported
+  if (!quant_args.quant_method().empty()) {
+    if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 ||
+        !quant_args.activation_dynamic()) {
+      LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when "
+                    "quant_method is set. "
+                 << "Got quant_method=" << quant_args.quant_method()
+                 << ", bits=" << quant_args.bits()
+                 << ", activation_dynamic=" << quant_args.activation_dynamic();
+    }
+    // If confirmed as smoothquant w8a8, set is_smoothquant_ to true
+    is_smoothquant_ = true;
+  } else {
+    is_smoothquant_ = false;
+  }
+
+  // Deep EP initialization check
+  enable_deep_ep_ =
+      ::xllm::EPLBConfig::get_instance().expert_parallel_degree() == 2 &&
+      ep_size > 1;
+  if (enable_deep_ep_) {
+    // for now, we only implement the deep ep for decode stage.
+    // so we will assume the max_token_num is limited to max_batch_size * (1+K)
+    // K is the number of speculative tokens.
+    int64_t dispatch_token_size;
+    if (quant_args.quant_method() == "smoothquant") {
+      // float32 is for the scale of the quantized input
+      dispatch_token_size = hidden_size_ * get_dtype_size(torch::kInt8) +
+                            get_dtype_size(torch::kFloat32);
+    } else {
+      dispatch_token_size =
+          hidden_size_ * get_dtype_size(options_.dtype().toScalarType());
+    }
+    torch::ScalarType combine_dtype = options_.dtype().toScalarType();
+    int64_t combine_token_size = hidden_size_ * get_dtype_size(combine_dtype);
+    // Ensure calculation base is at least ep_size
+    int64_t effective_seqs = std::max(
+        (int64_t)::xllm::SchedulerConfig::get_instance().max_seqs_per_batch(),
+        (int64_t)ep_size);
+    // NOTE: ::xllm::SchedulerConfig::get_instance().max_seqs_per_batch()
+    // represents the maximum total batch size, regardless of the dp size. To
+    // ensure robust scheduling and account for the worst-case scenario, we must
+    // guarantee that each rank is capable of handling the maximum possible
+    // number of tokens. Therefore, we define max_num_tokens_per_rank as the
+    // full maximum value, without dividing by either the rank count or the dp
+    // size.
+    int64_t max_num_tokens_per_rank =
+        (1 +
+         ::xllm::SpeculativeConfig::get_instance().num_speculative_tokens()) *
+        effective_seqs * topk_;
+
+    // make sure that all layers share the same deep ep instance
+    //  so that the memory footprint is minimized
+    deep_ep_ = DeepEPManager::get_instance(dispatch_token_size,
+                                           combine_token_size,
+                                           max_num_tokens_per_rank,
+                                           num_experts,
+                                           parallel_args,
+                                           options_);
+
+    // obtain the buffer and parameters of deep ep
+    deep_ep_buffer_ = deep_ep_->get_buffer();
+    deep_ep_params_ = deep_ep_->get_params();
+
+    // intermediate buffer that can be initialized once
+    // we place these tensor here in order to speed up forward pass
+    int64_t n_tokens_recv = deep_ep_params_.max_num_tokens_recv;
+    int64_t token_bytes = is_smoothquant_
+                              ? get_dtype_size(torch::kInt8)
+                              : get_dtype_size(options_.dtype().toScalarType());
+    token_bytes = token_bytes * hidden_size_;
+    int64_t head_size = n_tokens_recv * token_bytes;
+    dispatch_recv_token_tensor_head_ =
+        deep_ep_buffer_.combine_send_token_tensor.narrow(0, 0, head_size)
+            .view({n_tokens_recv, token_bytes});
+    // input scale in smoothquant
+    if (is_smoothquant_) {
+      int64_t tail_size = n_tokens_recv * get_dtype_size(torch::kFloat32);
+      dispatch_recv_token_tensor_tail_ =
+          deep_ep_buffer_.combine_send_token_tensor
+              .narrow(0, head_size, tail_size)
+              .view({n_tokens_recv, -1});
+    }
+  }
+
+  // calculate the number of experts per rank
+  num_experts_per_rank_ = num_experts / ep_size;
+  start_expert_id_ = ep_rank * num_experts_per_rank_;
+
+  if (topk_method == "noaux_tc") {
+    e_score_correction_bias_ = register_parameter(
+        "e_score_correction_bias", torch::empty({num_experts}, options), false);
+  }
+
+  gate_ = register_module(
+      "gate_proj",
+      ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options));
+  if (n_shared_experts_ > 0) {
+    ProcessGroup* shared_expert_pg;
+    if (parallel_args_.ep_size() > 1) {
+      // we use tp=1 for shared experts computation in deep ep mode
+      CHECK(parallel_args_.ep_size() == parallel_args_.world_size())
+          << "Models with shared experts only support ep_size equal to "
+             "world size for now.";
+      shared_expert_pg = parallel_args.moe_tp_group_;
+    } else {
+      shared_expert_pg = parallel_args.process_group_;
+    }
+    // The shared experts computation can proceed in parallel with the
+    // final communication step during the MoE computation, as long as it
+    // remains independent of any communication operations. For optimal
+    // performance, ensure that the shared experts layer on each rank always
+    // maintains its own unique weights.
+    shared_experts_ =
+        register_module("shared_experts",
+                        DenseMLP(hidden_size_,
+                                 intermediate_size * n_shared_experts_,
+                                 is_gated_,
+                                 false,
+                                 hidden_act_,
+                                 /*enable_result_reduction=*/true,
+                                 quant_args,
+                                 shared_expert_pg,
+                                 options));
+  }
+
+  // create weight buffer
+  const int64_t world_size = tp_pg_->world_size();
+  int64_t local_intermediate_size = intermediate_size / world_size;
+  if (is_smoothquant_) {
+    auto quant_option = options_.dtype(torch::kInt8);
+    auto fp_option = options_.dtype(torch::kFloat32);
+    w13_ = register_parameter(
+        "w13",
+        torch::empty(
+            {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
+            quant_option),
+        false);
+    w13_scale_ = register_parameter(
+        "w13_scale",
+        torch::empty({num_experts_per_rank_, local_intermediate_size * 2},
+                     fp_option),
+        false);
+    // Note: We do not check enable_deep_ep_ here, since smooth quantization
+    // information may be needed even when deep EP mode is disabled. This allows
+    // retrieving quantization parameters for any subset of experts as required.
+    input_smooth_ = register_parameter(
+        "input_smooth",
+        torch::empty({num_total_experts_, hidden_size_}, fp_option),
+        false);
+    w2_ = register_parameter(
+        "w2",
+        torch::empty(
+            {num_experts_per_rank_, hidden_size_, local_intermediate_size},
+            quant_option),
+        false);
+    w2_scale_ = register_parameter(
+        "w2_scale",
+        torch::empty({num_experts_per_rank_, hidden_size_}, fp_option),
+        false);
+    act_smooth_ = register_parameter(
+        "act_smooth",
+        torch::empty({num_experts_per_rank_, local_intermediate_size},
+                     fp_option),
+        false);
+
+  } else {
+    w13_ = register_parameter(
+        "w13",
+        torch::empty(
+            {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
+            options_),
+        false);
+    w2_ = register_parameter(
+        "w2",
+        torch::empty(
+            {num_experts_per_rank_, hidden_size_, local_intermediate_size},
+            options_),
+        false);
+  }
+}
+
+torch::Tensor FusedMoEImpl::create_group_gemm_output(
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& group_list,
+    torch::ScalarType dtype,
+    torch::Tensor& workspace) {
+  // unify shape logic: define the target shape once.
+  bool is_3d_weight = (b.dim() != 2);
+  int64_t num_tokens = a.size(0);
+  int64_t out_dim = is_3d_weight ? b.size(1) : b.size(0);
+
+  std::vector<int64_t> output_shape;
+  int64_t required_elements = num_tokens * out_dim;
+
+  if (is_3d_weight) {
+    output_shape = {num_tokens, out_dim};
+  } else {
+    output_shape = {group_list.size(0), num_tokens, out_dim};
+    required_elements *= group_list.size(0);
+  }
+
+  auto options = a.options().dtype(dtype);
+
+  // non-smoothquant: direct allocation
+  if (!is_smoothquant_) {
+    return torch::empty(output_shape, options);
+  }
+
+  // smoothquant: managed workspace logic
+  if (!workspace.defined()) {
+    // Lazy initialization: allocate max buffer for the lifecycle
+    // Note: accessing class members w13_ and w2_ directly for context
+    int64_t max_width = std::max(w13_.size(1), w2_.size(1));
+    workspace = torch::empty({num_tokens * max_width}, options);
+  }
+
+  // view construction
+  CHECK(workspace.numel() >= required_elements)
+      << "FusedMoE Workspace too small! Alloc: " << workspace.numel()
+      << ", Req: " << required_elements;
+
+  // utilize the pre-calculated output_shape
+  return workspace.slice(0, 0, required_elements).view(output_shape);
+}
+
+torch::Tensor FusedMoEImpl::select_experts(
+    const torch::Tensor& hidden_states_2d,
+    const torch::Tensor& router_logits_2d,
+    SelectedExpertInfo& selected_expert_info,
+    bool enable_all2all_communication) {
+  // prepare the parameters for select_experts
+  std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
+  if (e_score_correction_bias_.defined()) {
+    e_score_correction_bias = e_score_correction_bias_;
+  }
+  int64_t expert_size = w13_.size(0);
+
+  // Step 1: apply softmax topk or sigmoid topk / routing logic
+  torch::Tensor reduce_weight;
+  torch::Tensor expert_id;
+  {
+    xllm::kernel::MoeFusedTopkParams moe_active_topk_params;
+    moe_active_topk_params.input = router_logits_2d;
+    moe_active_topk_params.topk = topk_;
+    moe_active_topk_params.num_expert_group = num_expert_group_;
+    moe_active_topk_params.topk_group = topk_group_;
+    moe_active_topk_params.normalize = renormalize_;
+    moe_active_topk_params.normed_by = "topk_logit";
+    moe_active_topk_params.scoring_func = scoring_func_;
+    moe_active_topk_params.route_scale = route_scale_;
+    moe_active_topk_params.e_score_correction_bias = e_score_correction_bias;
+    std::tie(reduce_weight, expert_id) =
+        xllm::kernel::moe_active_topk(moe_active_topk_params);
+  }
+
+  // Step 2: generate expert ids
+  torch::Tensor gather_idx;
+  torch::Tensor combine_idx;
+  torch::Tensor token_count;
+  std::optional<torch::Tensor> cusum_token_count;
+  {
+    xllm::kernel::MoeGenIdxParams moe_gen_idx_params;
+    moe_gen_idx_params.expert_id = expert_id;
+    moe_gen_idx_params.expert_num = num_total_experts_;
+    std::vector<torch::Tensor> output_vec =
+        xllm::kernel::moe_gen_idx(moe_gen_idx_params);
+    gather_idx = output_vec[0];
+    combine_idx = output_vec[1];
+    token_count = output_vec[2];
+    // during all2all communication, we do not need cusum_token_count in the
+    // following computation
+    if (enable_all2all_communication) {
+      cusum_token_count = std::nullopt;
+    } else {
+      cusum_token_count = output_vec[3];
+    }
+  }
+
+  // Step 3: expand and quantize input if needed
+  torch::Tensor expand_hidden_states;
+  torch::Tensor hidden_states_scale;
+  torch::Tensor token_count_slice;
+  // all2all related variables
+  torch::Tensor dispatch_send_token_tensor;
+  // in all2all, the input is scattered, so there is no need to slice the token
+  // count, and we can use the dispatch buffer directly
+  if (enable_all2all_communication) {
+    token_count_slice = token_count;
+    int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
+    int64_t dispatch_bytes =
+        num_token_expand * deep_ep_params_.dispatch_token_size;
+    dispatch_send_token_tensor =
+        deep_ep_buffer_.dispatch_send_token_tensor.slice(0, 0, dispatch_bytes)
+            .view({num_token_expand, deep_ep_params_.dispatch_token_size});
+  } else {
+    token_count_slice =
+        token_count.slice(0, start_expert_id_, start_expert_id_ + expert_size);
+  }
+
+  if (is_smoothquant_) {
+    xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
+    scaled_quantize_params.x = hidden_states_2d;
+    // use dispatch_send_token_tensor buffer for input
+    //  to reduce memory footprint
+    if (enable_all2all_communication) {
+      scaled_quantize_params.smooth = input_smooth_;
+      scaled_quantize_params.output =
+          dispatch_send_token_tensor.slice(1, 0, hidden_size_);
+    } else {
+      scaled_quantize_params.smooth = input_smooth_.slice(
+          0, start_expert_id_, start_expert_id_ + expert_size);
+      scaled_quantize_params.gather_index_start_position =
+          cusum_token_count.value().index({start_expert_id_}).unsqueeze(0);
+    }
+    scaled_quantize_params.token_count = token_count_slice;
+    scaled_quantize_params.gather_index = gather_idx;
+    scaled_quantize_params.act_mode = "none";
+    scaled_quantize_params.active_coef = 1.0;
+    scaled_quantize_params.is_gated = false;
+    scaled_quantize_params.quant_type = torch::kChar;
+    std::tie(expand_hidden_states, hidden_states_scale) =
+        xllm::kernel::scaled_quantize(scaled_quantize_params);
+    if (enable_all2all_communication) {
+      // since view_as_dtype has not supported stride yet,
+      //  we need to copy the scale output to the dispatch buffer
+      torch::Tensor dispatch_scale_slice =
+          dispatch_send_token_tensor.slice(1, hidden_size_);
+      torch::Tensor hidden_states_scale_bytes =
+          view_as_dtype(hidden_states_scale, torch::kInt8)
+              .view_as(dispatch_scale_slice);
+      dispatch_scale_slice.copy_(hidden_states_scale_bytes);
+    }
+  } else {
+    xllm::kernel::MoeExpandInputParams moe_expand_input_params;
+    moe_expand_input_params.input = hidden_states_2d;
+    moe_expand_input_params.gather_index = gather_idx;
+    moe_expand_input_params.combine_idx = combine_idx;
+    moe_expand_input_params.topk = topk_;
+    expand_hidden_states =
+        xllm::kernel::moe_expand_input(moe_expand_input_params);
+    if (enable_all2all_communication) {
+      // use copy to place the output inside the dispatch buffer
+      torch::Tensor dispatch_tensor =
+          view_as_dtype(expand_hidden_states, torch::kChar);
+      dispatch_send_token_tensor.copy_(dispatch_tensor);
+    }
+  }
+
+  // collect the selected tensor
+  selected_expert_info.reduce_weight = reduce_weight;
+  selected_expert_info.combine_idx = combine_idx;
+  selected_expert_info.token_count_slice = token_count_slice;
+  selected_expert_info.cusum_token_count = cusum_token_count;
+  if (is_smoothquant_) {
+    selected_expert_info.input_scale = hidden_states_scale;
+  }
+
+  return expand_hidden_states;
+}
+
+torch::Tensor FusedMoEImpl::forward_experts(const torch::Tensor& hidden_states,
+                                            const torch::Tensor& router_logits,
+                                            bool enable_all2all_communication) {
+  if (!stream_initialized_) {
+    // update device record
+    device_ = xllm::Device(hidden_states.device());
+
+    // acquire streams from the pool again
+    routed_stream_ = device_.get_stream_from_pool();
+    shared_stream_ = device_.get_stream_from_pool();
+    stream_initialized_ = true;
+  }
+
+  std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
+  if (e_score_correction_bias_.defined()) {
+    e_score_correction_bias = e_score_correction_bias_;
+  }
+
+  // prepare the parameters for MoE computation
+  torch::Tensor shared_expert_output;
+  torch::IntArrayRef hidden_states_shape = hidden_states.sizes();
+  torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType();
+  torch::Tensor hidden_states_2d =
+      hidden_states.reshape({-1, hidden_states.size(-1)});
+  torch::Tensor router_logits_2d =
+      router_logits.reshape({-1, router_logits.size(-1)});
+  int64_t group_gemm_max_dim = enable_all2all_communication
+                                   ? deep_ep_params_.max_num_tokens_recv / topk_
+                                   : hidden_states_2d.size(0);
+  int64_t expert_size = w13_.size(0);
+
+  // Step 1-3: select experts
+  SelectedExpertInfo selected_expert_info;
+  torch::Tensor expand_hidden_states =
+      select_experts(hidden_states_2d,
+                     router_logits_2d,
+                     selected_expert_info,
+                     enable_all2all_communication);
+
+  // Communciation Step 1: Dipatch
+  // intermediate outputs that are used both in dispatch and combine
+  torch::Tensor gather_by_rank_index;
+  torch::Tensor token_sum;
+  if (enable_all2all_communication) {
+    int64_t dispatch_token_num = hidden_states_2d.size(0) * topk_;
+
+    // 1. Dispatch Step: Generate layout and send data
+    deep_ep_->dispatch_step(dispatch_token_num,
+                            selected_expert_info.token_count_slice);
+
+    // 2. Process Result: Generate indices and unpack to computation buffer
+    // use the buffer during initialization for the output
+    expand_hidden_states = dispatch_recv_token_tensor_head_;
+    std::optional<torch::Tensor> output_tail = std::nullopt;
+    if (is_smoothquant_) {
+      output_tail = dispatch_recv_token_tensor_tail_;
+      // update selected_expert_info with the tail (input scale)
+      selected_expert_info.input_scale = output_tail;
+    }
+
+    DeepEPMetaResult deep_ep_meta = deep_ep_->process_dispatch_result(
+        num_experts_per_rank_, expand_hidden_states, output_tail);
+
+    // Extract metadata for subsequent steps
+    gather_by_rank_index = deep_ep_meta.gather_rank_index;
+    selected_expert_info.token_count_slice = deep_ep_meta.token_count_slice;
+    token_sum = deep_ep_meta.token_sum;
+  }
+
+  // common gemm workspace for reduce memory footprint
+  torch::Tensor gemm_workspace;
+
+  // Step 4: group gemm 1
+  torch::Tensor gemm1_out =
+      create_group_gemm_output(expand_hidden_states,
+                               w13_,
+                               selected_expert_info.token_count_slice,
+                               hidden_states_dtype,
+                               gemm_workspace);
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::GroupGemmParams group_gemm_params;
+    torch::ScalarType a_dtype =
+        is_smoothquant_ ? torch::kInt8 : hidden_states_dtype;
+    group_gemm_params.a =
+        view_as_dtype(expand_hidden_states, a_dtype).view({-1, hidden_size_});
+    group_gemm_params.b = w13_;
+    group_gemm_params.token_count =
+        selected_expert_info.token_count_slice.to("cpu");
+    if (is_smoothquant_) {
+      torch::Tensor a_scale =
+          selected_expert_info.input_scale.value().flatten();
+      selected_expert_info.input_scale =
+          view_as_dtype(a_scale, torch::kFloat32);
+      group_gemm_params.a_scale = selected_expert_info.input_scale;
+      group_gemm_params.b_scale = w13_scale_;
+    }
+    group_gemm_params.max_dim = group_gemm_max_dim;
+    group_gemm_params.trans_a = false;
+    group_gemm_params.trans_b = true;
+    group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
+    group_gemm_params.output = gemm1_out;
+    group_gemm_params.combine_idx = std::nullopt;
+    gemm1_out = xllm::kernel::group_gemm(group_gemm_params);
+  }
+
+  // Step 5: activation or scaled quantization(fused with activation)
+  torch::Tensor act_out;
+  torch::Tensor act_out_scale;
+  if (is_smoothquant_) {
+    int64_t slice_dim = gemm1_out.size(1);
+    if (is_gated_) slice_dim /= 2;
+    // slice operation is a view, does not take up extra memory, but points to
+    // the same memory
+    act_out = expand_hidden_states.slice(1, 0, slice_dim);
+    act_out_scale =
+        selected_expert_info.input_scale.value().slice(0, 0, gemm1_out.size(0));
+    // call scaled quantization kernel (also fused with activation)
+    xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
+    scaled_quantize_params.x = gemm1_out;
+    scaled_quantize_params.smooth = act_smooth_;
+    scaled_quantize_params.token_count = selected_expert_info.token_count_slice;
+    scaled_quantize_params.output = act_out;
+    scaled_quantize_params.output_scale = act_out_scale;
+    scaled_quantize_params.act_mode = hidden_act_;
+    scaled_quantize_params.active_coef = 1.0;
+    scaled_quantize_params.is_gated = is_gated_;
+    scaled_quantize_params.quant_type = torch::kChar;
+    std::tie(act_out, act_out_scale) =
+        xllm::kernel::scaled_quantize(scaled_quantize_params);
+  } else {
+    act_out = is_gated_
+                  ? gemm1_out.slice(1, 0, gemm1_out.size(1) / 2).contiguous()
+                  : gemm1_out;
+    // call activation kernel
+    xllm::kernel::ActivationParams activation_params;
+    activation_params.input = gemm1_out;
+    activation_params.output = act_out;
+    activation_params.cusum_token_count =
+        selected_expert_info.cusum_token_count;
+    activation_params.act_mode = hidden_act_;
+    activation_params.is_gated = is_gated_;
+    activation_params.start_expert_id = start_expert_id_;
+    activation_params.expert_size = expert_size;
+    xllm::kernel::active(activation_params);
+  }
+
+  // Step 6: group gemm 2
+  torch::Tensor gemm2_out =
+      create_group_gemm_output(act_out,
+                               w2_,
+                               selected_expert_info.token_count_slice,
+                               hidden_states_dtype,
+                               gemm_workspace);
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::GroupGemmParams group_gemm_params;
+    group_gemm_params.a = act_out;
+    group_gemm_params.b = w2_;
+    group_gemm_params.token_count =
+        selected_expert_info.token_count_slice.to("cpu");
+    if (is_smoothquant_) {
+      group_gemm_params.a_scale = act_out_scale;
+      group_gemm_params.b_scale = w2_scale_;
+    }
+    group_gemm_params.max_dim = group_gemm_max_dim;
+    group_gemm_params.trans_a = false;
+    group_gemm_params.trans_b = true;
+    group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
+    group_gemm_params.output = gemm2_out;
+    group_gemm_params.combine_idx = selected_expert_info.combine_idx;
+    gemm2_out = xllm::kernel::group_gemm(group_gemm_params);
+  }
+
+  // Communciation Step 2: Combine
+  if (enable_all2all_communication) {
+    int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
+    // Delegate pack, layout generation and combine to DeepEP
+    torch::Tensor combine_send_layout =
+        deep_ep_->combine_step_pack(gemm2_out,
+                                    gather_by_rank_index,
+                                    token_sum,
+                                    hidden_size_,
+                                    hidden_states_dtype);
+
+    // create a wait event for the current stream to finish computation
+    auto current_stream = device_.current_stream();
+    routed_stream_->wait_stream(*current_stream);
+    // pure communciation kernel: dispatch
+    {
+      torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
+      gemm2_out = deep_ep_->combine_step_comm(combine_send_layout,
+                                              num_token_expand,
+                                              hidden_size_,
+                                              hidden_states_dtype);
+    }
+
+    // pure computation kernel: shared experts
+    if (n_shared_experts_ > 0) {
+      shared_stream_->wait_stream(*current_stream);
+      torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
+      shared_expert_output = shared_experts_(hidden_states);
+    }
+
+    // join for parallelization
+    current_stream->wait_stream(*routed_stream_);
+    if (n_shared_experts_ > 0) {
+      current_stream->wait_stream(*shared_stream_);
+    }
+  }
+
+  // After group gemm is finished, some tensors are no
+  // longer needed. We must explicitly release the memory.
+  expand_hidden_states = torch::Tensor();
+  selected_expert_info.input_scale = std::nullopt;
+  act_out = torch::Tensor();
+
+  // Step 7: combine the intermediate results and get the final hidden states
+  torch::Tensor final_hidden_states;
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::MoeCombineResultParams moe_combine_result_params;
+    moe_combine_result_params.input = gemm2_out;
+    moe_combine_result_params.reduce_weight =
+        selected_expert_info.reduce_weight;
+    moe_combine_result_params.gather_ids = selected_expert_info.combine_idx;
+    moe_combine_result_params.cusum_token_count =
+        selected_expert_info.cusum_token_count;
+    moe_combine_result_params.start_expert_id = start_expert_id_;
+    moe_combine_result_params.expert_size = expert_size;
+    moe_combine_result_params.bias = std::nullopt;
+    // if all2all communication is enabled and shared output is provided,
+    //  we will fused the add up to combine result
+    if (enable_all2all_communication && n_shared_experts_ > 0) {
+      moe_combine_result_params.residual =
+          shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
+    }
+
+    final_hidden_states =
+        xllm::kernel::moe_combine_result(moe_combine_result_params);
+  }
+
+  // reshape the final hidden states to the original shape
+  final_hidden_states = final_hidden_states.reshape(hidden_states_shape);
+
+  if (enable_all2all_communication) {
+    return final_hidden_states;
+  }
+
+  // Communciation Step 3: AllReduce for non-all2all communication
+  // shared experts can be parallelized with the final communication step
+  // during moe computation.
+  auto current_stream = device_.current_stream();
+  routed_stream_->wait_stream(*current_stream);
+  {
+    torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
+    if (tp_pg_->world_size() > 1) {
+      final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_);
+    }
+    if (parallel_args_.ep_size() > 1) {
+      final_hidden_states = parallel_state::reduce(
+          final_hidden_states, parallel_args_.moe_ep_group_);
+    }
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_stream_->wait_stream(*current_stream);
+    torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
+    // for non all2all, we compute the shared experts parallelized with the
+    // final communication step
+    shared_expert_output = shared_experts_(hidden_states);
+    shared_expert_output =
+        shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
+  }
+
+  // join for parallelization
+  current_stream->wait_stream(*routed_stream_);
+  if (n_shared_experts_ > 0) {
+    current_stream->wait_stream(*shared_stream_);
+    final_hidden_states += shared_expert_output;
+  }
+
+  return final_hidden_states;
+}
+
+torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states,
+                                    const ModelInputParams& input_params) {
+  // we only support all2all communication for decode stage for now
+  bool enable_all2all_communication =
+      enable_deep_ep_ && std::all_of(input_params.parallel.dp_is_decode.begin(),
+                                     input_params.parallel.dp_is_decode.end(),
+                                     [](int32_t val) { return val == 1; });
+
+  bool is_dp_ep_parallel =
+      parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1;
+  // during all2all communication, the output has been
+  //  gathered and sliced by dispatch and combine steps,
+  //  so we do not need to gather input and slice output again
+  bool need_gather_and_slice =
+      is_dp_ep_parallel && !enable_all2all_communication;
+
+  auto input = hidden_states;
+  if (need_gather_and_slice) {
+    input = parallel_state::gather(input,
+                                   parallel_args_.dp_local_process_group_,
+                                   input_params.parallel.dp_global_token_nums);
+  }
+  // MoE Gate
+  auto router_logits = gate_(input);
+
+  // MoE Experts
+  auto output =
+      forward_experts(input, router_logits, enable_all2all_communication);
+
+  if (need_gather_and_slice) {
+    output = get_dp_local_slice(output, input_params, parallel_args_);
+  }
+
+  return output;
+}
+
+void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) {
+  if (e_score_correction_bias_.defined() &&
+      !e_score_correction_bias_is_loaded_) {
+    LOAD_WEIGHT(e_score_correction_bias);
+  }
+}
+
+void FusedMoEImpl::load_experts(const StateDict& state_dict) {
+  const int64_t rank = tp_pg_->rank();
+  const int64_t world_size = tp_pg_->world_size();
+  const int64_t start_expert_id = start_expert_id_;
+  const int64_t num_experts_per_rank = num_experts_per_rank_;
+  const int64_t num_total_experts = num_total_experts_;
+  std::vector<std::string> prefixes = {"gate_proj.", "up_proj."};
+  if (is_smoothquant_) {
+    LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13);
+    LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale);
+    // When supporting DeepEP All2All mode,
+    // we need to load the complete set of expert weights corresponding to
+    // "up_proj.smooth". Note that even if deep EP mode is not enabled, it
+    // remains possible to retrieve the smooth quantization information for a
+    // subset of experts. Therefore, we intentionally do not check whether
+    // deep_ep_ is enabled in this case.
+    LOAD_MOE_ALL_EXPERT_WEIGHT("up_proj.", "smooth", input_smooth, -1);
+    LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1);
+    LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1);
+    LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0);
+  } else {
+    LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13);
+    LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1);
+  }
+}
+
+void FusedMoEImpl::load_state_dict(const StateDict& state_dict) {
+  if (state_dict.size() == 0) {
+    return;
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_experts_->load_state_dict(
+        state_dict.get_dict_with_prefix("shared_experts."));
+  }
+  gate_->load_state_dict(state_dict.get_dict_with_prefix("gate."));
+  load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate."));
+  load_experts(state_dict.get_dict_with_prefix("experts."));
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm_latest/core/layers/ilu/fused_moe.h

@@ -0,0 +1,131 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include "framework/model/model_args.h"
+#include "framework/model/model_input_params.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+#include "layers/common/deep_ep.h"
+#include "layers/common/dense_mlp.h"
+#include "layers/common/fused_moe_base.h"
+#include "layers/common/linear.h"
+#include "platform/device.h"
+#include "util/tensor_helper.h"
+
+namespace xllm {
+namespace layer {
+
+class FusedMoEImpl : public torch::nn::Module {
+ public:
+  FusedMoEImpl() = default;
+  FusedMoEImpl(const ModelArgs& model_args,
+               const FusedMoEArgs& moe_args,
+               const QuantArgs& quant_args,
+               const ParallelArgs& parallel_args,
+               const torch::TensorOptions& options);
+
+  torch::Tensor forward_experts(const torch::Tensor& hidden_states,
+                                const torch::Tensor& router_logits,
+                                bool enable_all2all_communication);
+  torch::Tensor forward(const torch::Tensor& hidden_states,
+                        const ModelInputParams& input_params);
+  void load_state_dict(const StateDict& state_dict);
+
+ private:
+  // struct to store the selected expert info
+  struct SelectedExpertInfo {
+    torch::Tensor reduce_weight;
+    torch::Tensor combine_idx;
+    torch::Tensor token_count_slice;
+    std::optional<torch::Tensor> cusum_token_count;
+    std::optional<torch::Tensor> input_scale;
+  };
+
+  // initial steps for MoE computation, select the experts for each token
+  torch::Tensor select_experts(const torch::Tensor& hidden_states_2d,
+                               const torch::Tensor& router_logits_2d,
+                               SelectedExpertInfo& selected_expert_info,
+                               bool enable_all2all_communication);
+
+ private:
+  int64_t num_total_experts_;
+  int64_t topk_;
+  int64_t num_expert_group_;
+  int64_t topk_group_;
+  double route_scale_;
+  int64_t hidden_size_;
+  int64_t n_shared_experts_;
+  bool is_gated_;
+  int64_t renormalize_;
+  std::string hidden_act_;
+  std::string scoring_func_;
+  bool is_smoothquant_;
+
+  int64_t num_experts_per_rank_;
+  int64_t start_expert_id_;
+
+  // Deep EP related parameters
+  bool enable_deep_ep_;
+  DeepEPBuffer deep_ep_buffer_;
+  DeepEPParams deep_ep_params_;
+  torch::Tensor dispatch_recv_token_tensor_head_;
+  torch::Tensor dispatch_recv_token_tensor_tail_;
+
+  // steams for parallel shared experts
+  std::unique_ptr<Stream> shared_stream_;
+  std::unique_ptr<Stream> routed_stream_;
+  xllm::Device device_;
+  bool stream_initialized_ = false;
+
+  ReplicatedLinear gate_{nullptr};
+  DenseMLP shared_experts_{nullptr};
+  DeepEP deep_ep_{nullptr};
+
+  QuantArgs quant_args_;
+  ParallelArgs parallel_args_;
+  torch::TensorOptions options_;
+  ProcessGroup* tp_pg_;
+
+  DEFINE_WEIGHT(w13);
+  DEFINE_FUSED_WEIGHT(w1);
+  DEFINE_FUSED_WEIGHT(w3);
+  DEFINE_FUSED_WEIGHT(w2);
+  DEFINE_WEIGHT(e_score_correction_bias);
+  DEFINE_WEIGHT(w13_scale);
+  DEFINE_FUSED_WEIGHT(w1_scale);
+  DEFINE_FUSED_WEIGHT(w3_scale);
+  DEFINE_FUSED_WEIGHT(w2_scale);
+  DEFINE_FUSED_WEIGHT(input_smooth);
+  DEFINE_FUSED_WEIGHT(act_smooth);
+
+  void load_e_score_correction_bias(const StateDict& state_dict);
+  void load_experts(const StateDict& state_dict);
+  // create the group gemm output tensor with the workspace
+  torch::Tensor create_group_gemm_output(const torch::Tensor& a,
+                                         const torch::Tensor& b,
+                                         const torch::Tensor& group_list,
+                                         torch::ScalarType dtype,
+                                         torch::Tensor& workspace);
+};
+TORCH_MODULE(FusedMoE);
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp

@@ -0,0 +1,219 @@
+/* Copyright 2025-2026 The xLLM Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "qwen3_5_gated_delta_net.h"
+
+#include <glog/logging.h>
+
+namespace xllm {
+namespace layer {
+
+Qwen3_5GatedDeltaNetImpl::Qwen3_5GatedDeltaNetImpl(
+    const ModelArgs& args,
+    const QuantArgs& quant_args,
+    const ParallelArgs& parallel_args,
+    const torch::TensorOptions& options)
+    : Qwen3NextGatedDeltaNetImpl(args,
+                                 quant_args,
+                                 parallel_args,
+                                 options,
+                                 /*init_projections=*/false) {
+  in_proj_qkv_ = register_module("in_proj_qkv",
+                                 ColumnParallelLinear(args.hidden_size(),
+                                                      k_size_ * 2 + v_size_,
+                                                      /*bias=*/false,
+                                                      /*gather_output=*/false,
+                                                      quant_args,
+                                                      parallel_args.tp_group_,
+                                                      options));
+  in_proj_z_ = register_module("in_proj_z",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    v_size_,
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/false,
+                                                    quant_args,
+                                                    parallel_args.tp_group_,
+                                                    options));
+  in_proj_b_ = register_module("in_proj_b",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    num_v_heads_,
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/false,
+                                                    quant_args,
+                                                    parallel_args.tp_group_,
+                                                    options));
+  in_proj_a_ = register_module("in_proj_a",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    num_v_heads_,
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/false,
+                                                    quant_args,
+                                                    parallel_args.tp_group_,
+                                                    options));
+}
+
+torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_qkvz_from_split_activations(
+    const torch::Tensor& qkv,
+    const torch::Tensor& z) const {
+  CHECK_EQ(qkv.dim(), 3) << "Expected qkv activation to be 3D, got "
+                         << qkv.sizes();
+  CHECK_EQ(z.dim(), 3) << "Expected z activation to be 3D, got " << z.sizes();
+  CHECK_EQ(qkv.size(0), z.size(0)) << "qkv/z batch size mismatch.";
+  CHECK_EQ(qkv.size(1), z.size(1)) << "qkv/z sequence size mismatch.";
+  CHECK_EQ(qkv.size(2), (2 * k_size_ + v_size_) / tp_size_)
+      << "Unexpected qkv hidden size for Qwen3.5.";
+  CHECK_EQ(z.size(2), v_size_ / tp_size_)
+      << "Unexpected z hidden size for Qwen3.5.";
+  CHECK_GT(num_k_heads_, 0) << "linear_num_key_heads must be positive.";
+  CHECK_EQ(num_v_heads_ % num_k_heads_, 0)
+      << "linear_num_value_heads must be divisible by linear_num_key_heads.";
+
+  const int64_t bs = qkv.size(0);
+  const int64_t seqlen = qkv.size(1);
+  const int64_t local_k_heads = num_k_heads_ / tp_size_;
+  const int64_t local_v_heads = num_v_heads_ / tp_size_;
+  const int64_t num_v_heads_per_k = num_v_heads_ / num_k_heads_;
+
+  auto qkv_split = torch::split(
+      qkv, {k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_}, 2);
+  auto q = qkv_split[0].view({bs, seqlen, local_k_heads, head_k_dim_});
+  auto k = qkv_split[1].view({bs, seqlen, local_k_heads, head_k_dim_});
+  auto v = qkv_split[2].view({bs, seqlen, local_v_heads, head_v_dim_});
+  auto z_view = z.view({bs, seqlen, local_v_heads, head_v_dim_});
+
+  v = v.view({bs, seqlen, local_k_heads, num_v_heads_per_k * head_v_dim_});
+  z_view =
+      z_view.view({bs, seqlen, local_k_heads, num_v_heads_per_k * head_v_dim_});
+
+  return torch::cat({q, k, v, z_view}, -1).view({bs, seqlen, -1}).contiguous();
+}
+
+torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations(
+    const torch::Tensor& b,
+    const torch::Tensor& a) const {
+  CHECK_EQ(b.dim(), 3) << "Expected b activation to be 3D, got " << b.sizes();
+  CHECK_EQ(a.dim(), 3) << "Expected a activation to be 3D, got " << a.sizes();
+  CHECK_EQ(b.size(0), a.size(0)) << "b/a batch size mismatch.";
+  CHECK_EQ(b.size(1), a.size(1)) << "b/a sequence size mismatch.";
+  CHECK_EQ(b.size(2), num_v_heads_ / tp_size_)
+      << "Unexpected b hidden size for Qwen3.5.";
+  CHECK_EQ(a.size(2), num_v_heads_ / tp_size_)
+      << "Unexpected a hidden size for Qwen3.5.";
+  CHECK_GT(num_k_heads_, 0) << "linear_num_key_heads must be positive.";
+  CHECK_EQ(num_v_heads_ % num_k_heads_, 0)
+      << "linear_num_value_heads must be divisible by linear_num_key_heads.";
+
+  const int64_t bs = b.size(0);
+  const int64_t seqlen = b.size(1);
+  const int64_t local_k_heads = num_k_heads_ / tp_size_;
+  const int64_t num_v_heads_per_k = num_v_heads_ / num_k_heads_;
+
+  auto b_view = b.view({bs, seqlen, local_k_heads, num_v_heads_per_k});
+  auto a_view = a.view({bs, seqlen, local_k_heads, num_v_heads_per_k});
+  return torch::cat({b_view, a_view}, -1).view({bs, seqlen, -1}).contiguous();
+}
+
+std::pair<torch::Tensor, torch::Tensor>
+Qwen3_5GatedDeltaNetImpl::project_decode_inputs(
+    const torch::Tensor& hidden_states) {
+  const auto reshape_projection = [](const torch::Tensor& projection) {
+    return projection.view({projection.size(0), -1, projection.size(-1)});
+  };
+  auto qkv = reshape_projection(in_proj_qkv_->forward(hidden_states));
+  auto z_proj = reshape_projection(in_proj_z_->forward(hidden_states));
+  auto b_proj = reshape_projection(in_proj_b_->forward(hidden_states));
+  auto a_proj = reshape_projection(in_proj_a_->forward(hidden_states));
+  return {merge_qkvz_from_split_activations(qkv, z_proj),
+          merge_ba_from_split_activations(b_proj, a_proj)};
+}
+
+std::pair<torch::Tensor, torch::Tensor>
+Qwen3_5GatedDeltaNetImpl::project_flat_inputs(
+    const torch::Tensor& hidden_states) {
+  auto qkv = in_proj_qkv_->forward(hidden_states).unsqueeze(0);
+  auto z_proj = in_proj_z_->forward(hidden_states).unsqueeze(0);
+  auto b_proj = in_proj_b_->forward(hidden_states).unsqueeze(0);
+  auto a_proj = in_proj_a_->forward(hidden_states).unsqueeze(0);
+  auto qkvz = merge_qkvz_from_split_activations(qkv, z_proj);
+  auto ba = merge_ba_from_split_activations(b_proj, a_proj);
+  return {qkvz.view({hidden_states.size(0), qkvz.size(-1)}).contiguous(),
+          ba.view({hidden_states.size(0), ba.size(-1)}).contiguous()};
+}
+
+std::optional<
+    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+Qwen3_5GatedDeltaNetImpl::project_split_inputs(
+    const torch::Tensor& hidden_states,
+    const AttentionMetadata& attn_metadata) {
+  auto qkv = reshape_projected_tokens_with_pad(
+      attn_metadata, in_proj_qkv_->forward(hidden_states));
+  auto z_proj = reshape_projected_tokens_with_pad(
+      attn_metadata, in_proj_z_->forward(hidden_states));
+  auto b_proj = reshape_projected_tokens_with_pad(
+      attn_metadata, in_proj_b_->forward(hidden_states));
+  auto a_proj = reshape_projected_tokens_with_pad(
+      attn_metadata, in_proj_a_->forward(hidden_states));
+
+  const int64_t batch_size = qkv.size(0);
+  const int64_t seq_len = qkv.size(1);
+  auto z =
+      z_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
+  auto b = b_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_});
+  auto a = a_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_});
+  return std::make_tuple(qkv, z, b, a);
+}
+
+void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict(
+    const StateDict& state_dict) {
+  auto in_proj_qkv_state_dict = state_dict.get_dict_with_prefix("in_proj_qkv.");
+  if (in_proj_qkv_state_dict.size() > 0 && !in_proj_qkv_->is_weight_loaded()) {
+    in_proj_qkv_->load_state_dict(
+        in_proj_qkv_state_dict,
+        /*shard_tensor_count=*/3,
+        /*shard_sizes=*/
+        {k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_});
+  }
+
+  auto in_proj_z_state_dict = state_dict.get_dict_with_prefix("in_proj_z.");
+  if (in_proj_z_state_dict.size() > 0 && !in_proj_z_->is_weight_loaded()) {
+    in_proj_z_->load_state_dict(in_proj_z_state_dict);
+  }
+
+  auto in_proj_b_state_dict = state_dict.get_dict_with_prefix("in_proj_b.");
+  if (in_proj_b_state_dict.size() > 0 && !in_proj_b_->is_weight_loaded()) {
+    in_proj_b_->load_state_dict(in_proj_b_state_dict);
+  }
+
+  auto in_proj_a_state_dict = state_dict.get_dict_with_prefix("in_proj_a.");
+  if (in_proj_a_state_dict.size() > 0 && !in_proj_a_->is_weight_loaded()) {
+    in_proj_a_->load_state_dict(in_proj_a_state_dict);
+  }
+}
+
+void Qwen3_5GatedDeltaNetImpl::verify_projection_weights(
+    const std::string& prefix) const {
+  CHECK(in_proj_qkv_ && in_proj_qkv_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_qkv.weight";
+  CHECK(in_proj_z_ && in_proj_z_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_z.weight";
+  CHECK(in_proj_b_ && in_proj_b_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_b.weight";
+  CHECK(in_proj_a_ && in_proj_a_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_a.weight";
+}
+
+}  // namespace layer
+}  // namespace xllm

upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h

@@ -0,0 +1,66 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include <optional>
+#include <string>
+#include <tuple>
+#include <utility>
+
+#include "qwen3_next_gated_delta_net.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl {
+ public:
+  Qwen3_5GatedDeltaNetImpl() = default;
+  Qwen3_5GatedDeltaNetImpl(const ModelArgs& args,
+                           const QuantArgs& quant_args,
+                           const ParallelArgs& parallel_args,
+                           const torch::TensorOptions& options);
+
+ protected:
+  std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
+      const torch::Tensor& hidden_states) override;
+  std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
+      const torch::Tensor& hidden_states) override;
+  std::optional<
+      std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+  project_split_inputs(const torch::Tensor& hidden_states,
+                       const AttentionMetadata& attn_metadata) override;
+  bool use_fla_ssm_state_layout() const override { return true; }
+
+  void load_projection_state_dict(const StateDict& state_dict) override;
+  void verify_projection_weights(const std::string& prefix) const override;
+
+ private:
+  torch::Tensor merge_qkvz_from_split_activations(const torch::Tensor& qkv,
+                                                  const torch::Tensor& z) const;
+  torch::Tensor merge_ba_from_split_activations(const torch::Tensor& b,
+                                                const torch::Tensor& a) const;
+
+  ColumnParallelLinear in_proj_qkv_{nullptr};
+  ColumnParallelLinear in_proj_z_{nullptr};
+  ColumnParallelLinear in_proj_b_{nullptr};
+  ColumnParallelLinear in_proj_a_{nullptr};
+};
+TORCH_MODULE(Qwen3_5GatedDeltaNet);
+
+}  // namespace layer
+}  // namespace xllm