From 0cb4a302977dec099cbf89a7fea8fb63484e2d9c Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sat, 11 Jul 2026 19:04:04 +0000
Subject: [PATCH 01/13] feat(moe): select the fused Qwen3 MoE backend

Expose auto, legacy, and fused policies. Auto selects fused only for non-quantized CUDA FP16/BF16 qwen3_moe workloads and preserves explicit legacy behavior elsewhere.
---
 examples/test_infer.py                     |  5 +-
 python/infinilm/base_config.py             |  8 ++
 python/infinilm/llm/llm.py                 | 89 +++++++++++++++++++++-
 python/infinilm/server/inference_server.py |  5 ++
 4 files changed, 104 insertions(+), 3 deletions(-)

diff --git a/examples/test_infer.py b/examples/test_infer.py
index b0250042..334c1251 100644
--- a/examples/test_infer.py
+++ b/examples/test_infer.py
@@ -31,7 +31,8 @@ def test(
     video_num_frames=None,
     skip_load=False,
     weight_load_mode="async",
-    skip_legacy_moe=False,
+    skip_legacy_moe=None,
+    moe_backend="auto",
 ):
     model_path = os.path.expanduser(model_path)
     # ---------------------------------------------------------------------------- #
@@ -62,6 +63,7 @@ def test(
         skip_load=skip_load,
         weight_load_mode=weight_load_mode,
         skip_legacy_moe=skip_legacy_moe,
+        moe_backend=moe_backend,
     )
 
     conversations = [
@@ -150,4 +152,5 @@ def test(
         skip_load=cfg.skip_load,
         weight_load_mode=cfg.weight_load_mode,
         skip_legacy_moe=cfg.skip_legacy_moe,
+        moe_backend=cfg.moe_backend,
     )
diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py
index 3b8f6533..05ac361e 100644
--- a/python/infinilm/base_config.py
+++ b/python/infinilm/base_config.py
@@ -66,6 +66,7 @@ def __init__(self):
         self.dp = self.args.dp
         self.ep = self.args.ep
         self.moe_ep_backend = self.args.moe_ep_backend
+        self.moe_backend = self.args.moe_backend
         self.skip_legacy_moe = self.args.skip_legacy_moe
 
         self.attn = self.args.attn
@@ -209,9 +210,16 @@ def _add_common_args(self):
             default="auto",
             help=MOE_EP_BACKEND_HELP,
         )
+        self.parser.add_argument(
+            "--moe-backend",
+            choices=["auto", "legacy", "fused"],
+            default="auto",
+            help="MoE execution backend",
+        )
         self.parser.add_argument(
             "--skip-legacy-moe",
             action="store_true",
+            default=None,
             help="use the new fused MoE implementation instead of the legacy Qwen3 MoE MLP",
         )
 
diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py
index 59e9e94b..9f50bd43 100644
--- a/python/infinilm/llm/llm.py
+++ b/python/infinilm/llm/llm.py
@@ -7,6 +7,7 @@
 """
 
 import asyncio
+import json
 import logging
 import os
 import threading
@@ -33,6 +34,68 @@
 
 logger = logging.getLogger(__name__)
 
+_MOE_BACKENDS = {"auto", "legacy", "fused"}
+
+
+def _resolve_moe_backend(
+    model_path: str,
+    device: str,
+    dtype: str,
+    moe_backend: str,
+    skip_legacy_moe: Optional[bool],
+):
+    """Resolve the MoE policy while preserving the legacy boolean override."""
+    if moe_backend not in _MOE_BACKENDS:
+        raise ValueError(
+            f"Unsupported moe_backend {moe_backend!r}; expected one of "
+            f"{sorted(_MOE_BACKENDS)}"
+        )
+
+    if skip_legacy_moe is not None:
+        resolved_backend = "fused" if skip_legacy_moe else "legacy"
+        if moe_backend != "auto" and moe_backend != resolved_backend:
+            raise ValueError(
+                "Conflicting MoE options: "
+                f"moe_backend={moe_backend!r}, skip_legacy_moe={skip_legacy_moe!r}"
+            )
+        logger.info(
+            "Resolved MoE backend to %s from explicit skip_legacy_moe=%s",
+            resolved_backend,
+            skip_legacy_moe,
+        )
+        return skip_legacy_moe, resolved_backend
+
+    if moe_backend != "auto":
+        return moe_backend == "fused", moe_backend
+
+    config_path = os.path.join(model_path, "config.json")
+    with open(config_path, encoding="utf-8") as config_file:
+        hf_config = json.load(config_file)
+
+    text_config = hf_config.get("text_config") or hf_config
+    quantization_config = text_config.get("quantization_config") or {}
+    quant_method = quantization_config.get("quant_method") or text_config.get(
+        "quant_method"
+    )
+    model_type = text_config.get("model_type", "")
+    fused_supported = (
+        model_type == "qwen3_moe"
+        and device == "cuda"
+        and dtype in {"float16", "bfloat16"}
+        and not quant_method
+    )
+    resolved_backend = "fused" if fused_supported else "legacy"
+    logger.info(
+        "Resolved moe_backend=auto to %s (model_type=%s, device=%s, "
+        "dtype=%s, quant_method=%s)",
+        resolved_backend,
+        model_type or "unknown",
+        device,
+        dtype,
+        quant_method or "none",
+    )
+    return fused_supported, resolved_backend
+
 
 class LLMEngine:
     """Low-level LLM engine that handles inference execution."""
@@ -332,7 +395,8 @@ def __init__(
         use_mla: bool = False,
         weight_load_mode: str = "async",
         skip_load: bool = False,
-        skip_legacy_moe: bool = False,
+        skip_legacy_moe: Optional[bool] = None,
+        moe_backend: str = "auto",
     ):
         """Initialize LLM.
 
@@ -354,7 +418,16 @@ def __init__(
             attn_backend: Attention backend to use ('default', 'flash-attn').
             use_mla: Whether to use DeepSeek V2 MLA attention when supported.
             weight_load_mode: Weight loading mode across tensor-parallel workers.
+            skip_legacy_moe: Compatibility override for the legacy boolean API.
+            moe_backend: MoE execution policy ('auto', 'legacy', or 'fused').
         """
+        skip_legacy_moe, resolved_moe_backend = _resolve_moe_backend(
+            model_path=model_path,
+            device=device,
+            dtype=dtype,
+            moe_backend=moe_backend,
+            skip_legacy_moe=skip_legacy_moe,
+        )
         config = EngineConfig(
             model_path=model_path,
             draft_model_path=draft_model_path,
@@ -382,6 +455,7 @@ def __init__(
         )
         self.engine = LLMEngine(config)
         self.config = config
+        self.moe_backend = resolved_moe_backend
 
     def generate(
         self,
@@ -539,7 +613,8 @@ def __init__(
         kv_transfer_config: Optional[KVTransferConfig] = None,
         use_mla: bool = False,
         weight_load_mode: str = "async",
-        skip_legacy_moe: bool = False,
+        skip_legacy_moe: Optional[bool] = None,
+        moe_backend: str = "auto",
     ):
         """Initialize AsyncLLMEngine.
 
@@ -564,7 +639,16 @@ def __init__(
             kv_connector_extra_config: Extra config dict for KV connector.
             use_mla: Whether to use DeepSeek V2 MLA attention when supported.
             weight_load_mode: Weight loading mode across tensor-parallel workers.
+            skip_legacy_moe: Compatibility override for the legacy boolean API.
+            moe_backend: MoE execution policy ('auto', 'legacy', or 'fused').
         """
+        skip_legacy_moe, resolved_moe_backend = _resolve_moe_backend(
+            model_path=model_path,
+            device=device,
+            dtype=dtype,
+            moe_backend=moe_backend,
+            skip_legacy_moe=skip_legacy_moe,
+        )
         config = EngineConfig(
             model_path=model_path,
             draft_model_path=draft_model_path,
@@ -592,6 +676,7 @@ def __init__(
         )
         self.engine = LLMEngine(config)
         self.config = config
+        self.moe_backend = resolved_moe_backend
 
         self._running = False
         self._step_thread: Optional[threading.Thread] = None
diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py
index 242593e2..69d0829b 100644
--- a/python/infinilm/server/inference_server.py
+++ b/python/infinilm/server/inference_server.py
@@ -117,6 +117,7 @@ def __init__(
         weight_load_mode: str = "async",
         ignore_eos: bool = False,
         kv_transfer_config: Optional[KVTransferConfig] = None,
+        moe_backend: str = "auto",
     ):
         """Initialize inference server.
 
@@ -144,6 +145,7 @@ def __init__(
             weight_load_mode: Weight loading mode across tensor-parallel workers.
             ignore_eos: Whether to ignore EOS tokens during generation.
             kv_transfer_config: Optional configuration for the KV transfer mechanism.
+            moe_backend: MoE execution backend.
         """
         self.model_path = model_path
         # vLLM-like served model id: directory name of model_path
@@ -153,6 +155,7 @@ def __init__(
         self.tensor_parallel_size = tensor_parallel_size
         self.moe_ep_backend = moe_ep_backend
         self.moe_ep_size = moe_ep_size
+        self.moe_backend = moe_backend
         self.cache_type = cache_type
         self.max_tokens = max_tokens
         self.max_batch_size = max_batch_size
@@ -192,6 +195,7 @@ async def lifespan(app: FastAPI):
                 tensor_parallel_size=self.tensor_parallel_size,
                 moe_ep_backend=self.moe_ep_backend,
                 moe_ep_size=self.moe_ep_size,
+                moe_backend=self.moe_backend,
                 cache_type=self.cache_type,
                 max_batch_size=self.max_batch_size,
                 max_tokens=self.max_tokens,
@@ -607,6 +611,7 @@ def main():
         tensor_parallel_size=cfg.tp,
         moe_ep_backend=moe_ep_backend,
         moe_ep_size=ep,
+        moe_backend=cfg.moe_backend,
         cache_type="paged" if cfg.enable_paged_attn else "static",
         max_tokens=cfg.max_new_tokens,
         max_batch_size=cfg.max_batch_size,

From 4a537ce34030cdd120768cd54de97bb01cd9e33c Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sat, 11 Jul 2026 19:04:04 +0000
Subject: [PATCH 02/13] perf(graph): bound paged captures by batch capacity

Propagate max_batch_size through the Python and C++ cache configuration so Paged CUDA Graph captures only scheduler-reachable batch shapes.
---
 csrc/cache/kv_cache.cpp                          | 15 ++++++++++++++-
 csrc/cache/kv_cache.hpp                          |  6 ++++++
 csrc/engine/compiler/paged_compiler.cpp          | 16 ++++++++++++++++
 csrc/pybind11/cache/cache.hpp                    |  8 ++++++--
 python/infinilm/cache/cache.py                   |  8 ++++++--
 python/infinilm/llm/model_runner/model_runner.py |  4 +++-
 6 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/csrc/cache/kv_cache.cpp b/csrc/cache/kv_cache.cpp
index 88d7367e..1d30b72d 100644
--- a/csrc/cache/kv_cache.cpp
+++ b/csrc/cache/kv_cache.cpp
@@ -77,11 +77,19 @@ infinicore::Tensor create_layer_kv_cache(
 // ==========================
 // PagedKVCacheConfig
 // ==========================
+// Preserve the legacy Graph capture ceiling for two-argument C++ callers.
 PagedKVCacheConfig::PagedKVCacheConfig(
     size_t num_blocks,
     size_t block_size)
+    : PagedKVCacheConfig(num_blocks, block_size, 512) {}
+
+PagedKVCacheConfig::PagedKVCacheConfig(
+    size_t num_blocks,
+    size_t block_size,
+    size_t max_batch_size)
     : num_blocks_(num_blocks),
-      block_size_(block_size) {
+      block_size_(block_size),
+      max_batch_size_(max_batch_size) {
 }
 
 std::unique_ptr<CacheConfig>
@@ -99,6 +107,11 @@ PagedKVCacheConfig::block_size() const {
     return block_size_;
 }
 
+size_t
+PagedKVCacheConfig::max_batch_size() const {
+    return max_batch_size_;
+}
+
 namespace PagedKVCache {
 // ==========================
 // PagedKVCache
diff --git a/csrc/cache/kv_cache.hpp b/csrc/cache/kv_cache.hpp
index 4d0a9a70..7ca94a58 100644
--- a/csrc/cache/kv_cache.hpp
+++ b/csrc/cache/kv_cache.hpp
@@ -40,14 +40,20 @@ class PagedKVCacheConfig final : public CacheConfig {
     PagedKVCacheConfig(
         size_t num_blocks,
         size_t block_size = 256);
+    PagedKVCacheConfig(
+        size_t num_blocks,
+        size_t block_size,
+        size_t max_batch_size);
 
     std::unique_ptr<CacheConfig> unique_copy() const override;
     size_t num_blocks() const;
     size_t block_size() const;
+    size_t max_batch_size() const;
 
 private:
     size_t num_blocks_;
     size_t block_size_;
+    size_t max_batch_size_;
 };
 
 namespace PagedKVCache {
diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp
index f267794e..29681b78 100644
--- a/csrc/engine/compiler/paged_compiler.cpp
+++ b/csrc/engine/compiler/paged_compiler.cpp
@@ -18,6 +18,22 @@ PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBa
     for (size_t b = 256; b <= 512; b += 64) {
         decode_batch_sizes_.push_back(b);
     }
+
+    const auto *config = dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config());
+    if (config == nullptr) {
+        return;
+    }
+    const size_t max_batch_size = config->max_batch_size();
+    if (max_batch_size == 0) {
+        throw std::invalid_argument("Paged Graph max_batch_size must be greater than zero");
+    }
+    decode_batch_sizes_.erase(
+        std::remove_if(decode_batch_sizes_.begin(), decode_batch_sizes_.end(),
+                       [max_batch_size](size_t b) { return b > max_batch_size; }),
+        decode_batch_sizes_.end());
+    if (decode_batch_sizes_.empty() || decode_batch_sizes_.back() != max_batch_size) {
+        decode_batch_sizes_.push_back(max_batch_size);
+    }
 }
 
 void PagedCompiler::compile() {
diff --git a/csrc/pybind11/cache/cache.hpp b/csrc/pybind11/cache/cache.hpp
index 492f6c30..439c9da2 100644
--- a/csrc/pybind11/cache/cache.hpp
+++ b/csrc/pybind11/cache/cache.hpp
@@ -34,15 +34,19 @@ inline void bind_cache(py::module &m) {
                infinilm::cache::CacheConfig,
                std::shared_ptr<infinilm::cache::PagedKVCacheConfig>>(m, "PagedKVCacheConfig")
         .def(
-            py::init<size_t, size_t>(),
+            py::init<size_t, size_t, size_t>(),
             py::arg("num_blocks"),
-            py::arg("block_size") = 256)
+            py::arg("block_size") = 256,
+            py::arg("max_batch_size") = 512)
         .def(
             "num_blocks",
             &infinilm::cache::PagedKVCacheConfig::num_blocks)
         .def(
             "block_size",
             &infinilm::cache::PagedKVCacheConfig::block_size)
+        .def(
+            "max_batch_size",
+            &infinilm::cache::PagedKVCacheConfig::max_batch_size)
         .def("__repr__", [](const infinilm::cache::PagedKVCacheConfig &) {
             return "<PagedKVCacheConfig>";
         });
diff --git a/python/infinilm/cache/cache.py b/python/infinilm/cache/cache.py
index 4a9bcc44..b8b8e41c 100644
--- a/python/infinilm/cache/cache.py
+++ b/python/infinilm/cache/cache.py
@@ -18,5 +18,9 @@ def __init__(
 
 
 class PagedKVCacheConfig(CacheConfig, _infinilm.PagedKVCacheConfig):
-    def __init__(self, num_blocks: int, block_size: int = 256):
-        _infinilm.PagedKVCacheConfig.__init__(self, num_blocks, block_size)
+    def __init__(
+        self, num_blocks: int, block_size: int = 256, max_batch_size: int = 512
+    ):
+        _infinilm.PagedKVCacheConfig.__init__(
+            self, num_blocks, block_size, max_batch_size
+        )
diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py
index 45f1eaee..c611b1f3 100644
--- a/python/infinilm/llm/model_runner/model_runner.py
+++ b/python/infinilm/llm/model_runner/model_runner.py
@@ -59,7 +59,9 @@ def __init__(self, config: EngineConfig):
             )
         elif config.cache_type == "paged":
             cache_config = PagedKVCacheConfig(
-                num_blocks=config.num_blocks, block_size=config.block_size
+                num_blocks=config.num_blocks,
+                block_size=config.block_size,
+                max_batch_size=config.max_batch_size,
             )
             logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}")
         else:

From 4f749c2adf4fc434782664d21a52c1d1030d3336 Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sat, 11 Jul 2026 19:04:04 +0000
Subject: [PATCH 03/13] perf(graph): sparsify high-batch paged captures

Keep exact graphs for batches 1 through 16 and coarse exact buckets above 16. Missing shapes retain the existing eager fallback and the configured maximum remains captured.
---
 csrc/engine/compiler/paged_compiler.cpp | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp
index 29681b78..7176bf7d 100644
--- a/csrc/engine/compiler/paged_compiler.cpp
+++ b/csrc/engine/compiler/paged_compiler.cpp
@@ -34,6 +34,17 @@ PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBa
     if (decode_batch_sizes_.empty() || decode_batch_sizes_.back() != max_batch_size) {
         decode_batch_sizes_.push_back(max_batch_size);
     }
+
+    // Keep 1-16 dense, retain multiples of eight from the legacy bucket list,
+    // and always capture the configured maximum.
+    if (max_batch_size > 16) {
+        decode_batch_sizes_.erase(
+            std::remove_if(decode_batch_sizes_.begin(), decode_batch_sizes_.end(),
+                           [max_batch_size](size_t b) {
+                               return b > 16 && b != max_batch_size && b % 8 != 0;
+                           }),
+            decode_batch_sizes_.end());
+    }
 }
 
 void PagedCompiler::compile() {

From 010b4398fd8a89158c2326f65dbf4feb4ebbbc20 Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sat, 11 Jul 2026 19:04:04 +0000
Subject: [PATCH 04/13] perf(scheduler): bound Qwen3 MoE high-concurrency
 prefill

Use a 4096-token scheduler admission default only for high-concurrency Paged qwen3_moe serving. Explicit overrides, other models, and small capacities keep their prior behavior.
---
 python/infinilm/llm/llm.py | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py
index 9f50bd43..ddb19a98 100644
--- a/python/infinilm/llm/llm.py
+++ b/python/infinilm/llm/llm.py
@@ -35,6 +35,11 @@
 logger = logging.getLogger(__name__)
 
 _MOE_BACKENDS = {"auto", "legacy", "fused"}
+# Apply the conservative prefill budget from 32 concurrent requests onward.
+# It still admits 64 x 32-token prompts in one wave while bounding longer packed
+# prefills on 24 GiB GPUs. `INFINILM_MAX_NUM_BATCHED_TOKENS` overrides it.
+_QWEN3_MOE_HIGH_CONCURRENCY_BATCH_SIZE = 32
+_QWEN3_MOE_MAX_BATCHED_TOKENS = 4096
 
 
 def _resolve_moe_backend(
@@ -97,6 +102,17 @@ def _resolve_moe_backend(
     return fused_supported, resolved_backend
 
 
+def _default_max_num_batched_tokens(
+    llm_config: dict, max_batch_size: int, max_position_embeddings: int
+) -> int:
+    if (
+        llm_config.get("model_type") == "qwen3_moe"
+        and max_batch_size >= _QWEN3_MOE_HIGH_CONCURRENCY_BATCH_SIZE
+    ):
+        return min(max_position_embeddings, _QWEN3_MOE_MAX_BATCHED_TOKENS)
+    return max_position_embeddings
+
+
 class LLMEngine:
     """Low-level LLM engine that handles inference execution."""
 
@@ -145,10 +161,18 @@ def __init__(self, config: EngineConfig):
             )
             num_mamba_cache_blocks = max(2, config.num_blocks // 4)
 
+            # Bound packed prefill activations for high-concurrency Qwen3-MoE.
+            default_max_num_batched_tokens = _default_max_num_batched_tokens(
+                llm_config, config.max_batch_size, max_position_embeddings
+            )
             max_num_batched_tokens = int(
-                os.getenv("INFINILM_MAX_NUM_BATCHED_TOKENS", max_position_embeddings)
+                os.getenv(
+                    "INFINILM_MAX_NUM_BATCHED_TOKENS",
+                    default_max_num_batched_tokens,
+                )
             )
             assert 1024 <= max_num_batched_tokens <= max_position_embeddings
+            logger.info("Scheduler max_num_batched_tokens=%d", max_num_batched_tokens)
 
             self.scheduler = Scheduler(
                 max_batch_size=config.max_batch_size,

From fb495c1d36b7945cdd782e3f1d72c17ed4871fca Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sat, 11 Jul 2026 19:04:05 +0000
Subject: [PATCH 05/13] perf(lm-head): project only required prefill logits

Project normal single-request prefill logits only at the sampled final position. sample_all_positions callers and decode retain complete logits semantics.
---
 csrc/engine/infer_engine.cpp                  | 29 ++++++++++++++++++-
 csrc/engine/rank_worker.cpp                   |  6 ++--
 .../causal_lm_templates/text_causal_lm.hpp    | 29 ++++++++++++++++++-
 csrc/models/infinilm_model.hpp                |  4 +++
 4 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp
index a5221eb3..b47ab2bf 100644
--- a/csrc/engine/infer_engine.cpp
+++ b/csrc/engine/infer_engine.cpp
@@ -147,6 +147,11 @@ std::vector<std::string> InferEngine::state_dict_keys() {
 infinilm::InfinilmModel::Input
 InferEngine::Input::to_model_input(infinicore::Device device) const {
 
+    if (input_offsets.has_value()
+        && input_offsets.value()->device().getType() != infinicore::Device::Type::CPU) {
+        throw std::invalid_argument("`input_offsets` must reside on CPU.");
+    }
+
     auto to_device = [&](const std::optional<infinicore::Tensor> &t)
         -> std::optional<infinicore::Tensor> {
         return t.has_value() ? t.value()->to(device) : t;
@@ -163,6 +168,26 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         }
         return result;
     };
+    std::optional<std::vector<size_t>> last_token_indices;
+    if (!sample_all_positions
+        && input_offsets.has_value()
+        && input_ids.has_value()
+        && input_ids.value()->shape().size() == 2
+        && input_ids.value()->size(0) == 1
+        && input_ids.value()->size(1) > input_offsets.value()->size(0) - 1) {
+        const auto &offsets_tensor = input_offsets.value();
+        const size_t num_requests = offsets_tensor->size(0) - 1;
+        const auto *offsets = reinterpret_cast<const int32_t *>(offsets_tensor->data());
+        std::vector<size_t> indices;
+        indices.reserve(num_requests);
+        for (size_t i = 0; i < num_requests; ++i) {
+            if (offsets[i + 1] <= offsets[i]) {
+                throw std::invalid_argument("input_offsets must describe non-empty requests");
+            }
+            indices.push_back(static_cast<size_t>(offsets[i + 1] - 1));
+        }
+        last_token_indices = std::move(indices);
+    }
 
     infinilm::InfinilmModel::Input input = {
         to_device(input_ids), // @todo: on device in the future
@@ -181,7 +206,9 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         to_device_vec(image_grid_thw),
         image_req_ids,
         visual_token_ranges,
-        to_device(target_hidden_states)};
+        to_device(target_hidden_states),
+        std::move(last_token_indices),
+        sample_all_positions};
 
     infinilm::global_state::get_forward_context().attn_metadata = {
         input.past_sequence_lengths,
diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp
index 11696e1e..54c191ad 100644
--- a/csrc/engine/rank_worker.cpp
+++ b/csrc/engine/rank_worker.cpp
@@ -449,15 +449,17 @@ void RankWorker::thread_loop() {
                             int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data();
 
                             const bool sample_all_positions = local_args.sample_all_positions;
+                            const size_t logits_positions = batch_size * total_len;
+                            const bool logits_are_last_token_only = !sample_all_positions && logits_positions == n_req;
                             const size_t n_out = sample_all_positions ? static_cast<size_t>(input_offsets[n_req]) : n_req;
                             auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device)};
 
                             for (size_t i{0}; i < n_out; ++i) {
                                 size_t score_idx = i;
-                                if (!sample_all_positions) {
+                                if (!sample_all_positions && !logits_are_last_token_only) {
                                     score_idx = static_cast<size_t>(input_offsets[i + 1] - 1);
                                 }
-                                auto score{logits->view({batch_size * total_len, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})};
+                                auto score{logits->view({logits_positions, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})};
                                 auto out{output_ids->narrow({{0, i, 1}})->view({})};
                                 float random_val = std::uniform_real_distribution<float>(0, 1)(rng_);
                                 infinicore::op::random_sample_(
diff --git a/csrc/layers/causal_lm_templates/text_causal_lm.hpp b/csrc/layers/causal_lm_templates/text_causal_lm.hpp
index eb4f2b47..b536a16c 100644
--- a/csrc/layers/causal_lm_templates/text_causal_lm.hpp
+++ b/csrc/layers/causal_lm_templates/text_causal_lm.hpp
@@ -3,6 +3,8 @@
 #include "../../models/infinilm_model.hpp"
 #include "../linear/linear.hpp"
 #include "infinicore/device.hpp"
+#include "infinicore/ops/cat.hpp"
+#include <stdexcept>
 
 namespace infinilm::layers::causal_lm_templates {
 
@@ -44,7 +46,32 @@ class TextCausalLM : public InfinilmModel {
      */
     Output forward(const Input &input) const override {
         auto hidden_states = model_->forward(input);
-        auto logits = lm_head_->forward(hidden_states);
+        auto lm_head_input = hidden_states;
+        // Ordinary prefill needs only the final-position logits.
+        if (!input.sample_all_positions
+            && input.last_token_indices.has_value()
+            && !input.last_token_indices->empty()
+            && hidden_states->size(0) == 1
+            && hidden_states->size(1) > input.last_token_indices->size()) {
+            const auto &indices = input.last_token_indices.value();
+            for (const size_t index : indices) {
+                if (index >= hidden_states->size(1)) {
+                    throw std::out_of_range("last-token index exceeds packed hidden states");
+                }
+            }
+            if (indices.size() == 1) {
+                lm_head_input = hidden_states->narrow({{1, indices.front(), 1}});
+            } else {
+                std::vector<infinicore::Tensor> final_hidden_states;
+                final_hidden_states.reserve(indices.size());
+                for (const size_t index : indices) {
+                    final_hidden_states.push_back(
+                        hidden_states->narrow({{1, index, 1}}));
+                }
+                lm_head_input = infinicore::op::cat(std::move(final_hidden_states), 1);
+            }
+        }
+        auto logits = lm_head_->forward(lm_head_input);
         return {logits};
     }
 
diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp
index ac994fd6..0eb27c39 100644
--- a/csrc/models/infinilm_model.hpp
+++ b/csrc/models/infinilm_model.hpp
@@ -55,6 +55,10 @@ class InfinilmModel : public infinicore::nn::Module {
         std::optional<std::vector<size_t>> visual_token_ranges;
         /// Target model hidden states consumed by draft/MTP models.
         std::optional<infinicore::Tensor> target_hidden_states;
+        /// Packed hidden-state positions sampled by ordinary prefill.
+        std::optional<std::vector<size_t>> last_token_indices;
+        /// Preserve logits for every packed position for speculative/MTP callers.
+        bool sample_all_positions{false};
     };
 
     struct Output {

From 98565e2e3953bd939230659c1585cb2fdcc1b3b8 Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sat, 11 Jul 2026 19:04:05 +0000
Subject: [PATCH 06/13] perf(lm-head): shard the Qwen3 MoE vocabulary
 projection

Shard the non-quantized Qwen3-MoE vocabulary projection across tensor-parallel ranks. Preserve complete-logit fallbacks for raw and non-greedy callers.
---
 csrc/engine/compiler/paged_compiler.cpp       |   7 +-
 .../compiler/static_batching_compiler.cpp     |   5 +
 csrc/engine/infer_engine.cpp                  |   9 ++
 csrc/engine/rank_worker.cpp                   | 146 ++++++++++++++----
 csrc/engine/rank_worker.hpp                   |   7 +
 .../causal_lm_templates/text_causal_lm.hpp    |  55 ++++++-
 csrc/models/infinilm_model.hpp                |   5 +
 csrc/pybind11/engine/engine.hpp               |   1 +
 python/infinilm/infer_engine.py               |  47 +++---
 9 files changed, 229 insertions(+), 53 deletions(-)

diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp
index 7176bf7d..1e0b97e1 100644
--- a/csrc/engine/compiler/paged_compiler.cpp
+++ b/csrc/engine/compiler/paged_compiler.cpp
@@ -58,6 +58,7 @@ void PagedCompiler::compile() {
 
         auto make_decode_input = [&](size_t b) {
             InfinilmModel::Input input;
+            input.use_local_vocab_logits = true;
             input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice());
             input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
             input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice());
@@ -133,8 +134,10 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input &
         size_t batch_size = input.block_tables.value()->size(0);
         size_t block_per_req = input.block_tables.value()->size(1);
 
-        // only support decode only batch
-        if (batch_size != input.input_ids.value()->size(1)) {
+        // Vocab-parallel graphs are valid only for local-logit sampling.
+        if ((model_->uses_vocab_parallel_logits()
+             && !input.use_local_vocab_logits)
+            || batch_size != input.input_ids.value()->size(1)) {
             return {nullptr, nullptr};
         } else {
             auto result = compiled_map_decode_.find(batch_size);
diff --git a/csrc/engine/compiler/static_batching_compiler.cpp b/csrc/engine/compiler/static_batching_compiler.cpp
index af2f4799..76f35dcf 100644
--- a/csrc/engine/compiler/static_batching_compiler.cpp
+++ b/csrc/engine/compiler/static_batching_compiler.cpp
@@ -11,6 +11,7 @@ void StaticBatchingCompiler::compile() {
     if (model_->get_cache_config() != nullptr && dynamic_cast<const cache::StaticKVCacheConfig *>(model_->get_cache_config())) {
         size_t b = dynamic_cast<const cache::StaticKVCacheConfig *>(model_->get_cache_config())->max_batch_size();
         InfinilmModel::Input input;
+        input.use_local_vocab_logits = true;
         input.input_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I64, infinicore::context::getDevice());
         input.position_ids = infinicore::Tensor::empty({b, 1}, infinicore::DataType::I64, infinicore::context::getDevice());
         input.past_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
@@ -46,6 +47,10 @@ void StaticBatchingCompiler::compile() {
 StaticBatchingCompiler::Compiled StaticBatchingCompiler::get_compiled(
     const InfinilmModel::Input &input) {
     if (model_->get_cache_config() != nullptr && dynamic_cast<const cache::StaticKVCacheConfig *>(model_->get_cache_config())) {
+        if (model_->uses_vocab_parallel_logits()
+            && !input.use_local_vocab_logits) {
+            return std::make_tuple(nullptr, nullptr);
+        }
         size_t batch_size = input.input_ids.value()->size(0);
         size_t seqlen = input.input_ids.value()->size(1);
         auto result = compiled_map_.find(std::make_tuple(batch_size, seqlen));
diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp
index b47ab2bf..a7eded92 100644
--- a/csrc/engine/infer_engine.cpp
+++ b/csrc/engine/infer_engine.cpp
@@ -174,10 +174,15 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         && input_ids.has_value()
         && input_ids.value()->shape().size() == 2
         && input_ids.value()->size(0) == 1
+        && input_offsets.value()->size(0) > 1
         && input_ids.value()->size(1) > input_offsets.value()->size(0) - 1) {
         const auto &offsets_tensor = input_offsets.value();
         const size_t num_requests = offsets_tensor->size(0) - 1;
         const auto *offsets = reinterpret_cast<const int32_t *>(offsets_tensor->data());
+        if (offsets[0] != 0
+            || static_cast<size_t>(offsets[num_requests]) != input_ids.value()->size(1)) {
+            throw std::invalid_argument("input_offsets must cover the packed input");
+        }
         std::vector<size_t> indices;
         indices.reserve(num_requests);
         for (size_t i = 0; i < num_requests; ++i) {
@@ -188,6 +193,9 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         }
         last_token_indices = std::move(indices);
     }
+    const bool use_local_vocab_logits = allow_local_vocab_logits
+                                     && !sample_all_positions
+                                     && is_greedy_sampling();
 
     infinilm::InfinilmModel::Input input = {
         to_device(input_ids), // @todo: on device in the future
@@ -208,6 +216,7 @@ InferEngine::Input::to_model_input(infinicore::Device device) const {
         visual_token_ranges,
         to_device(target_hidden_states),
         std::move(last_token_indices),
+        use_local_vocab_logits,
         sample_all_positions};
 
     infinilm::global_state::get_forward_context().attn_metadata = {
diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp
index 54c191ad..4d2c9a47 100644
--- a/csrc/engine/rank_worker.cpp
+++ b/csrc/engine/rank_worker.cpp
@@ -1,6 +1,8 @@
 #include "rank_worker.hpp"
 #include "../models/model_factory.hpp"
 #include "infinicore/ops.hpp"
+#include "infinicore/ops/distributed/allgather.hpp"
+#include "infinicore/ops/take.hpp"
 #include <spdlog/spdlog.h>
 #include <stdexcept>
 
@@ -419,7 +421,9 @@ void RankWorker::thread_loop() {
                         infinicore::Tensor hidden_states;
                         // All-position speculative/MTP runs need eager mode because
                         // hidden states are not part of compiled graph outputs.
-                        if (!local_args.sample_all_positions && compiler_ != nullptr) {
+                        if (!local_args.sample_all_positions
+                            && local_args.allow_local_vocab_logits
+                            && compiler_ != nullptr) {
                             auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device::cpu()));
                             if (graph != nullptr && output != nullptr) {
                                 graph->run();
@@ -434,45 +438,135 @@ void RankWorker::thread_loop() {
                             hidden_states = model_output.hidden_states;
                         }
 
-                        // Random sampling (rank 0 only)
-                        if (rank_info_.tp_rank == 0) {
-                            auto temperature{local_args.temperature};
-                            auto top_p{local_args.top_p};
-                            auto top_k{local_args.top_k};
-
-                            const auto &logits_shape{logits->shape()};
-                            const auto &vocab_size{logits_shape[2]};
-                            const auto &total_len{logits_shape[1]};
-                            const auto &batch_size{logits_shape[0]};
-
-                            auto n_req = local_args.input_offsets.value()->size(0) - 1;
-                            int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data();
+                        auto temperature{local_args.temperature};
+                        auto top_p{local_args.top_p};
+                        auto top_k{local_args.top_k};
+                        const auto &logits_shape{logits->shape()};
+                        const size_t vocab_size{logits_shape[2]};
+                        const size_t total_len{logits_shape[1]};
+                        const size_t batch_size{logits_shape[0]};
+                        const size_t logits_positions = batch_size * total_len;
+                        const size_t n_req = local_args.input_offsets.value()->size(0) - 1;
+                        int32_t *input_offsets = reinterpret_cast<int32_t *>(
+                            local_args.input_offsets.value()->data());
+                        const bool sample_all_positions = local_args.sample_all_positions;
+                        const bool logits_are_last_token_only = !sample_all_positions && logits_positions == n_req;
+                        const size_t n_out = sample_all_positions
+                                               ? static_cast<size_t>(input_offsets[n_req])
+                                               : n_req;
+
+                        const size_t global_vocab_size = model_config_->get<size_t>("vocab_size");
+                        const bool distributed_greedy = !sample_all_positions
+                                                     && local_args.is_greedy_sampling()
+                                                     && rank_info_.tp_size > 1
+                                                     && vocab_size * static_cast<size_t>(rank_info_.tp_size)
+                                                            == global_vocab_size;
+
+                        if (distributed_greedy) {
+                            // Exchange one local argmax candidate per TP rank.
+                            auto local_ids = infinicore::Tensor::empty(
+                                {n_out}, infinicore::DataType::I64, rank_info_.device);
+                            auto local_values = infinicore::Tensor::empty(
+                                {n_out}, logits->dtype(), rank_info_.device);
+                            infinicore::Tensor output_ids;
+                            infinicore::Tensor winning_ranks;
+                            if (rank_info_.tp_rank == 0) {
+                                output_ids = infinicore::Tensor::empty(
+                                    {n_out}, infinicore::DataType::I64, rank_info_.device);
+                                winning_ranks = infinicore::Tensor::empty(
+                                    {n_out}, infinicore::DataType::I64, rank_info_.device);
+                            }
+                            auto flat_logits = logits->view({logits_positions, vocab_size});
+                            for (size_t i = 0; i < n_out; ++i) {
+                                size_t score_idx = i;
+                                if (!logits_are_last_token_only) {
+                                    score_idx = static_cast<size_t>(
+                                        input_offsets[i + 1] - 1);
+                                }
+                                auto score = flat_logits->narrow(
+                                                            {{0, score_idx, 1}})
+                                                 ->view({vocab_size});
+                                auto local_id = local_ids->narrow({{0, i, 1}});
+                                infinicore::op::random_sample_(
+                                    local_id->view({}), score, 0.0f,
+                                    top_p, top_k, temperature);
+                                infinicore::op::take_(
+                                    local_values->narrow({{0, i, 1}}), score, local_id);
+                            }
 
-                            const bool sample_all_positions = local_args.sample_all_positions;
-                            const size_t logits_positions = batch_size * total_len;
-                            const bool logits_are_last_token_only = !sample_all_positions && logits_positions == n_req;
-                            const size_t n_out = sample_all_positions ? static_cast<size_t>(input_offsets[n_req]) : n_req;
-                            auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device)};
+                            const size_t tp_size = static_cast<size_t>(rank_info_.tp_size);
+                            auto candidate_values = infinicore::op::distributed::allgather(
+                                local_values, tp_size, rank_info_.comm);
+                            auto candidate_ids = infinicore::op::distributed::allgather(
+                                local_ids, tp_size, rank_info_.comm);
+                            if (n_out > 1) {
+                                candidate_values = candidate_values
+                                                       ->view({tp_size, n_out})
+                                                       ->permute({1, 0})
+                                                       ->contiguous();
+                                candidate_ids = candidate_ids
+                                                    ->view({tp_size, n_out})
+                                                    ->permute({1, 0})
+                                                    ->contiguous();
+                            } else {
+                                candidate_values = candidate_values->view({1, tp_size});
+                                candidate_ids = candidate_ids->view({1, tp_size});
+                            }
 
+                            if (rank_info_.tp_rank == 0) {
+                                for (size_t i = 0; i < n_out; ++i) {
+                                    auto winning_rank = winning_ranks->narrow(
+                                        {{0, i, 1}});
+                                    auto values = candidate_values->narrow(
+                                                                      {{0, i, 1}})
+                                                      ->view({tp_size});
+                                    infinicore::op::random_sample_(
+                                        winning_rank->view({}), values,
+                                        0.0f, 0.0f, 1, 0.0f);
+                                    auto ids = candidate_ids->narrow(
+                                                                {{0, i, 1}})
+                                                   ->view({tp_size});
+                                    infinicore::op::take_(
+                                        output_ids->narrow({{0, i, 1}}), ids,
+                                        winning_rank);
+                                }
+                                auto output_ids_cpu = output_ids->to(
+                                    infinicore::Device::cpu());
+                                auto winning_ranks_cpu = winning_ranks->to(
+                                    infinicore::Device::cpu());
+                                infinicore::context::syncStream();
+                                auto *output_ptr = reinterpret_cast<int64_t *>(
+                                    output_ids_cpu->data());
+                                const auto *rank_ptr = reinterpret_cast<const int64_t *>(
+                                    winning_ranks_cpu->data());
+                                for (size_t i = 0; i < n_out; ++i) {
+                                    output_ptr[i] += rank_ptr[i]
+                                                   * static_cast<int64_t>(vocab_size);
+                                }
+                                output_ = Output{
+                                    output_ids_cpu, logits, hidden_states};
+                            }
+                        } else if (rank_info_.tp_rank == 0) {
+                            auto output_ids = infinicore::Tensor::empty(
+                                {n_out}, infinicore::DataType::I64, rank_info_.device);
+                            auto flat_logits = logits->view({logits_positions, vocab_size});
                             for (size_t i{0}; i < n_out; ++i) {
                                 size_t score_idx = i;
                                 if (!sample_all_positions && !logits_are_last_token_only) {
                                     score_idx = static_cast<size_t>(input_offsets[i + 1] - 1);
                                 }
-                                auto score{logits->view({logits_positions, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})};
-                                auto out{output_ids->narrow({{0, i, 1}})->view({})};
+                                auto score = flat_logits->narrow(
+                                                            {{0, score_idx, 1}})
+                                                 ->view({vocab_size});
+                                auto out = output_ids->narrow({{0, i, 1}})->view({});
                                 float random_val = std::uniform_real_distribution<float>(0, 1)(rng_);
                                 infinicore::op::random_sample_(
                                     out, score, random_val, top_p, top_k, temperature);
                             }
 
                             output_ids = output_ids->to(infinicore::Device::cpu());
-
                             infinicore::context::syncStream();
-
-                            auto out{Output{output_ids, logits, hidden_states}};
-
-                            output_ = std::move(out);
+                            output_ = Output{output_ids, logits, hidden_states};
                         }
 
                         job_done_ = true;
diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp
index d396ef6f..085f6aca 100644
--- a/csrc/engine/rank_worker.hpp
+++ b/csrc/engine/rank_worker.hpp
@@ -73,12 +73,19 @@ class RankWorker {
         /// Sample logits at every packed input position instead of one token per request.
         bool sample_all_positions{false};
 
+        /// Allow callers to consume vocabulary-parallel logits.
+        bool allow_local_vocab_logits{false};
+
         float temperature{1};
 
         int top_k{50};
 
         float top_p{1};
 
+        bool is_greedy_sampling() const {
+            return top_k == 1 || top_p == 0.0f || temperature == 0.0f;
+        }
+
         infinilm::InfinilmModel::Input to_model_input(infinicore::Device device) const;
     };
 
diff --git a/csrc/layers/causal_lm_templates/text_causal_lm.hpp b/csrc/layers/causal_lm_templates/text_causal_lm.hpp
index b536a16c..e1469a85 100644
--- a/csrc/layers/causal_lm_templates/text_causal_lm.hpp
+++ b/csrc/layers/causal_lm_templates/text_causal_lm.hpp
@@ -1,9 +1,12 @@
 #pragma once
 
+#include "../../global_state/global_state.hpp"
 #include "../../models/infinilm_model.hpp"
 #include "../linear/linear.hpp"
 #include "infinicore/device.hpp"
 #include "infinicore/ops/cat.hpp"
+#include "infinicore/ops/distributed/allgather.hpp"
+#include <infiniccl.h>
 #include <stdexcept>
 
 namespace infinilm::layers::causal_lm_templates {
@@ -38,7 +41,21 @@ class TextCausalLM : public InfinilmModel {
         const auto &dtype{model_config->get_dtype()};
 
         model_ = this->register_module<Model>("model", model_config, device);
-        lm_head_ = this->register_module<infinilm::layers::linear::ReplicatedLinear>("lm_head", hidden_size, vocab_size, false, dtype, device);
+        const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info();
+        tp_size_ = static_cast<size_t>(rank_info.tp_size);
+        tp_communicator_ = rank_info.comm;
+        vocab_parallel_ = tp_size_ > 1
+                       && model_config->get<std::string>("model_type") == "qwen3_moe"
+                       && vocab_size % tp_size_ == 0
+                       && model_config->get_quant_scheme() == infinilm::quantization::QuantScheme::NONE;
+        if (vocab_parallel_) {
+            lm_head_ = this->register_module<infinilm::layers::linear::ColumnParallelLinear>(
+                "lm_head", hidden_size, vocab_size, false, dtype, device,
+                rank_info.tp_rank, rank_info.tp_size);
+        } else {
+            lm_head_ = this->register_module<infinilm::layers::linear::ReplicatedLinear>(
+                "lm_head", hidden_size, vocab_size, false, dtype, device);
+        }
     }
 
     /**
@@ -72,18 +89,50 @@ class TextCausalLM : public InfinilmModel {
             }
         }
         auto logits = lm_head_->forward(lm_head_input);
+        if (!input.use_local_vocab_logits) {
+            logits = gather_vocab_logits_(std::move(logits));
+        }
         return {logits};
     }
 
+    /// Return full-vocabulary logits, gathering rank-local shards when needed.
     infinicore::Tensor logits_from_hidden(const infinicore::Tensor &hidden_states) const {
-        return lm_head_->forward(const_cast<infinicore::Tensor &>(hidden_states));
+        return gather_vocab_logits_(
+            lm_head_->forward(const_cast<infinicore::Tensor &>(hidden_states)));
     }
 
     Model &model() { return *model_; }
 
+    bool uses_vocab_parallel_logits() const override {
+        return vocab_parallel_;
+    }
+
 protected:
     INFINICORE_NN_MODULE(Model, model);
-    INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head);
+    infinicore::Tensor gather_vocab_logits_(infinicore::Tensor logits) const {
+        if (!vocab_parallel_) {
+            return logits;
+        }
+        auto output_shape = logits->shape();
+        const size_t local_vocab_size = output_shape.back();
+        const size_t positions = logits->numel() / local_vocab_size;
+        auto gathered = infinicore::op::distributed::allgather(
+            logits, tp_size_, tp_communicator_);
+        output_shape.back() = local_vocab_size * tp_size_;
+        if (positions == 1) {
+            return gathered->view(output_shape);
+        }
+        return gathered->view({tp_size_, positions, local_vocab_size})
+            ->permute({1, 0, 2})
+            ->contiguous()
+            ->view(output_shape);
+    }
+
+    // The base type holds either a replicated or column-parallel registered head.
+    INFINICORE_NN_MODULE(infinilm::layers::linear::BaseLinear, lm_head);
+    bool vocab_parallel_{false};
+    size_t tp_size_{1};
+    infinicclComm_t tp_communicator_{nullptr};
 };
 
 } // namespace infinilm::layers::causal_lm_templates
diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp
index 0eb27c39..609f83f0 100644
--- a/csrc/models/infinilm_model.hpp
+++ b/csrc/models/infinilm_model.hpp
@@ -57,6 +57,8 @@ class InfinilmModel : public infinicore::nn::Module {
         std::optional<infinicore::Tensor> target_hidden_states;
         /// Packed hidden-state positions sampled by ordinary prefill.
         std::optional<std::vector<size_t>> last_token_indices;
+        /// Keep vocabulary-parallel logits sharded for distributed greedy sampling.
+        bool use_local_vocab_logits{false};
         /// Preserve logits for every packed position for speculative/MTP callers.
         bool sample_all_positions{false};
     };
@@ -70,6 +72,9 @@ class InfinilmModel : public infinicore::nn::Module {
 
     virtual ~InfinilmModel() = default;
     virtual Output forward(const Input &input) const = 0;
+    virtual bool uses_vocab_parallel_logits() const {
+        return false;
+    }
     virtual void reset_cache(const cache::CacheConfig *cache_config);
     virtual const cache::CacheConfig *get_cache_config() const {
         return cache_config_.get();
diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp
index b73ce06a..b3f5dac5 100644
--- a/csrc/pybind11/engine/engine.hpp
+++ b/csrc/pybind11/engine/engine.hpp
@@ -239,6 +239,7 @@ inline void bind_infer_engine(py::module &m) {
         .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges)
         .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states)
         .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions)
+        .def_readwrite("allow_local_vocab_logits", &InferEngine::Input::allow_local_vocab_logits)
         .def_readwrite("temperature", &InferEngine::Input::temperature)
         .def_readwrite("top_k", &InferEngine::Input::top_k)
         .def_readwrite("top_p", &InferEngine::Input::top_p);
diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py
index 3e874046..4d1937c2 100644
--- a/python/infinilm/infer_engine.py
+++ b/python/infinilm/infer_engine.py
@@ -230,6 +230,7 @@ def _build_input(
         visual_token_ranges=None,
         target_hidden_states=None,
         sample_all_positions=False,
+        allow_local_vocab_logits=False,
         temperature=None,
         top_k=None,
         top_p=None,
@@ -269,7 +270,7 @@ def convert_tensor_list(tensor_list_):
         top_k = 1 if top_k is None else top_k
         top_p = 1.0 if top_p is None else top_p
 
-        return super().Input(
+        model_input = super().Input(
             input_ids,
             position_ids=position_ids,
             past_sequence_lengths=past_kv_lengths,
@@ -292,6 +293,8 @@ def convert_tensor_list(tensor_list_):
             top_k=top_k,
             top_p=top_p,
         )
+        model_input.allow_local_vocab_logits = allow_local_vocab_logits
+        return model_input
 
     def forward(
         self,
@@ -385,6 +388,7 @@ def convert_tensor_list(tensor_list_):
                         image_req_ids=image_req_ids,
                         visual_token_ranges=visual_token_ranges,
                         target_hidden_states=target_hidden_states,
+                        allow_local_vocab_logits=True,
                         temperature=temperature,
                         top_k=top_k,
                         top_p=top_p,
@@ -419,28 +423,27 @@ def forward_raw(
         top_p=None,
     ):
         try:
-            output = super().forward(
-                self._build_input(
-                    input_ids,
-                    position_ids=position_ids,
-                    past_kv_lengths=past_kv_lengths,
-                    total_kv_lengths=total_kv_lengths,
-                    input_offsets=input_offsets,
-                    cu_seqlens=cu_seqlens,
-                    block_tables=block_tables,
-                    slot_mapping=slot_mapping,
-                    pixel_values=pixel_values,
-                    image_bound=image_bound,
-                    tgt_sizes=tgt_sizes,
-                    image_req_ids=image_req_ids,
-                    visual_token_ranges=visual_token_ranges,
-                    target_hidden_states=target_hidden_states,
-                    sample_all_positions=sample_all_positions,
-                    temperature=temperature,
-                    top_k=top_k,
-                    top_p=top_p,
-                )
+            model_input = self._build_input(
+                input_ids,
+                position_ids=position_ids,
+                past_kv_lengths=past_kv_lengths,
+                total_kv_lengths=total_kv_lengths,
+                input_offsets=input_offsets,
+                cu_seqlens=cu_seqlens,
+                block_tables=block_tables,
+                slot_mapping=slot_mapping,
+                pixel_values=pixel_values,
+                image_bound=image_bound,
+                tgt_sizes=tgt_sizes,
+                image_req_ids=image_req_ids,
+                visual_token_ranges=visual_token_ranges,
+                target_hidden_states=target_hidden_states,
+                sample_all_positions=sample_all_positions,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
             )
+            output = super().forward(model_input)
             return {
                 "output_ids": infinicore.Tensor(output.output_ids),
                 "logits": infinicore.Tensor(output.logits),

From a756f898d7b8746cf188b87659b29c3c83c37d8f Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sat, 11 Jul 2026 19:04:05 +0000
Subject: [PATCH 07/13] perf(graph): skip empty non-quant runtime resets

Skip recursive runtime-state reset only when the model is explicitly non-quantized. Quantized Marlin lock and workspace resets remain unchanged.
---
 csrc/models/infinilm_model.cpp | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp
index c2cc6876..a16951f8 100644
--- a/csrc/models/infinilm_model.cpp
+++ b/csrc/models/infinilm_model.cpp
@@ -94,6 +94,10 @@ void InfinilmModel::process_weights_after_loading() {
 }
 
 void InfinilmModel::reset_runtime_state() const {
+    // Only quantized kernels currently keep resettable runtime state.
+    if (model_config_ && model_config_->get_quant_scheme() == quantization::QuantScheme::NONE) {
+        return;
+    }
     for (const auto &[_, sub] : children()) {
         reset_runtime_state_recursive_(sub.get());
     }

From bf807bc7dace7dddc48af6d7b10aea415f9c163f Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sun, 12 Jul 2026 03:36:53 +0800
Subject: [PATCH 08/13] fix(graph): rebuild paged capture buckets after cache
 reset

Derive the decode capture set from the current PagedKVCacheConfig on every compile so cache resets cannot retain stale batch limits or permanently discard larger buckets.
---
 csrc/engine/compiler/paged_compiler.cpp | 175 ++++++++++++------------
 1 file changed, 86 insertions(+), 89 deletions(-)

diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp
index 1e0b97e1..c170a96b 100644
--- a/csrc/engine/compiler/paged_compiler.cpp
+++ b/csrc/engine/compiler/paged_compiler.cpp
@@ -5,7 +5,20 @@
 namespace infinilm::engine {
 
 PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBarrier *barrier)
-    : GraphCompiler(model, barrier) {
+    : GraphCompiler(model, barrier) {}
+
+void PagedCompiler::compile() {
+    compiled_map_decode_.clear();
+    const auto *config = dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config());
+    if (config == nullptr) {
+        return;
+    }
+    const size_t max_batch_size = config->max_batch_size();
+    if (max_batch_size == 0) {
+        throw std::invalid_argument("Paged Graph max_batch_size must be greater than zero");
+    }
+
+    decode_batch_sizes_.clear();
     for (size_t b = 1; b < 64; ++b) {
         decode_batch_sizes_.push_back(b);
     }
@@ -18,15 +31,6 @@ PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBa
     for (size_t b = 256; b <= 512; b += 64) {
         decode_batch_sizes_.push_back(b);
     }
-
-    const auto *config = dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config());
-    if (config == nullptr) {
-        return;
-    }
-    const size_t max_batch_size = config->max_batch_size();
-    if (max_batch_size == 0) {
-        throw std::invalid_argument("Paged Graph max_batch_size must be greater than zero");
-    }
     decode_batch_sizes_.erase(
         std::remove_if(decode_batch_sizes_.begin(), decode_batch_sizes_.end(),
                        [max_batch_size](size_t b) { return b > max_batch_size; }),
@@ -45,87 +49,80 @@ PagedCompiler::PagedCompiler(const std::shared_ptr<InfinilmModel> &model, RankBa
                            }),
             decode_batch_sizes_.end());
     }
-}
-
-void PagedCompiler::compile() {
-    if (model_->get_cache_config() != nullptr && dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config())) {
-        size_t nblocks = dynamic_cast<const cache::PagedKVCacheConfig *>(model_->get_cache_config())->num_blocks();
-        size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end());
-        compiled_map_decode_.clear();
-        block_tables_holder_ = infinicore::Tensor::empty(
-            {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice());
-        set_zeros(block_tables_holder_);
-
-        auto make_decode_input = [&](size_t b) {
-            InfinilmModel::Input input;
-            input.use_local_vocab_logits = true;
-            input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice());
-            input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
-            input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice());
-            set_zeros(input.input_ids.value());
-            set_zeros(input.position_ids.value());
-            set_zeros(input.total_sequence_lengths.value());
-            std::vector<int32_t> total_sequence_lengths_vec(b, 1);
-            infinicore::context::memcpyH2D(input.total_sequence_lengths.value()->data(), total_sequence_lengths_vec.data(), b * sizeof(int32_t), false);
-            input.input_offsets = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice());
-            std::vector<int32_t> input_offsets_vec(b + 1, 0);
-            for (size_t i = 0; i <= b; i++) {
-                input_offsets_vec[i] = i;
-            }
-            infinicore::context::memcpyH2D(input.input_offsets.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false);
-            input.cu_seqlens = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice());
-            infinicore::context::memcpyH2D(input.cu_seqlens.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false);
-            const size_t block_per_req = nblocks;
-            input.block_tables = block_tables_holder_->as_strided({b, block_per_req}, {(ptrdiff_t)block_per_req, 1});
-            input.slot_mapping = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
-            set_zeros(input.slot_mapping.value());
-
-            // Attention reads attn_metadata from thread-local forward context.
-            infinilm::global_state::get_forward_context().attn_metadata = {
-                input.past_sequence_lengths,
-                input.total_sequence_lengths,
-                input.input_offsets,
-                input.cu_seqlens,
-                input.block_tables,
-                input.slot_mapping,
-            };
-            return input;
-        };
-
-        {
-            const size_t warmup_batch_size = std::min(max_batch_size, static_cast<size_t>(64));
-            auto input = make_decode_input(warmup_batch_size);
-            model_->forward(input);
-            infinicore::context::syncStream();
-            // Warmup runs the eager Marlin path and may leave per-layer lock
-            // workspaces dirty. Reset before CUDA graph capture so capture
-            // starts from the same all-zero lock state as normal execution.
-            model_->reset_runtime_state();
-            infinicore::context::syncStream();
+    size_t nblocks = config->num_blocks();
+    block_tables_holder_ = infinicore::Tensor::empty(
+        {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice());
+    set_zeros(block_tables_holder_);
+
+    auto make_decode_input = [&](size_t b) {
+        InfinilmModel::Input input;
+        input.use_local_vocab_logits = true;
+        input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice());
+        input.position_ids = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
+        input.total_sequence_lengths = infinicore::Tensor::empty({b}, infinicore::DataType::I32, infinicore::context::getDevice());
+        set_zeros(input.input_ids.value());
+        set_zeros(input.position_ids.value());
+        set_zeros(input.total_sequence_lengths.value());
+        std::vector<int32_t> total_sequence_lengths_vec(b, 1);
+        infinicore::context::memcpyH2D(input.total_sequence_lengths.value()->data(), total_sequence_lengths_vec.data(), b * sizeof(int32_t), false);
+        input.input_offsets = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice());
+        std::vector<int32_t> input_offsets_vec(b + 1, 0);
+        for (size_t i = 0; i <= b; i++) {
+            input_offsets_vec[i] = i;
         }
+        infinicore::context::memcpyH2D(input.input_offsets.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false);
+        input.cu_seqlens = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice());
+        infinicore::context::memcpyH2D(input.cu_seqlens.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false);
+        const size_t block_per_req = nblocks;
+        input.block_tables = block_tables_holder_->as_strided({b, block_per_req}, {(ptrdiff_t)block_per_req, 1});
+        input.slot_mapping = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice());
+        set_zeros(input.slot_mapping.value());
+
+        // Attention reads attn_metadata from thread-local forward context.
+        infinilm::global_state::get_forward_context().attn_metadata = {
+            input.past_sequence_lengths,
+            input.total_sequence_lengths,
+            input.input_offsets,
+            input.cu_seqlens,
+            input.block_tables,
+            input.slot_mapping,
+        };
+        return input;
+    };
+
+    {
+        const size_t warmup_batch_size = std::min(max_batch_size, static_cast<size_t>(64));
+        auto input = make_decode_input(warmup_batch_size);
+        model_->forward(input);
+        infinicore::context::syncStream();
+        // Warmup runs the eager Marlin path and may leave per-layer lock
+        // workspaces dirty. Reset before CUDA graph capture so capture
+        // starts from the same all-zero lock state as normal execution.
+        model_->reset_runtime_state();
+        infinicore::context::syncStream();
+    }
 
-        for (size_t b : decode_batch_sizes_) {
-            auto input = make_decode_input(b);
-
-            barrier_->wait();
-            (void)model_->forward(input);
-            infinicore::context::syncStream();
-            // Capture must not start with stale Marlin locks from previous
-            // warmup/capture attempts. This reset is intentionally outside
-            // graph capture; the current implementation still pays a memset
-            // before every graph replay in get_compiled().
-            model_->reset_runtime_state();
-            infinicore::context::syncStream();
-            infinicore::context::startGraphRecording();
-            auto output = model_->forward(input);
-            auto graph = infinicore::context::stopGraphRecording();
-            barrier_->wait();
-
-            auto shared_output = std::shared_ptr<InfinilmModel::Output>(
-                new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)});
-
-            compiled_map_decode_[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)};
-        }
+    for (size_t b : decode_batch_sizes_) {
+        auto input = make_decode_input(b);
+
+        barrier_->wait();
+        (void)model_->forward(input);
+        infinicore::context::syncStream();
+        // Capture must not start with stale Marlin locks from previous
+        // warmup/capture attempts. This reset is intentionally outside
+        // graph capture; the current implementation still pays a memset
+        // before every graph replay in get_compiled().
+        model_->reset_runtime_state();
+        infinicore::context::syncStream();
+        infinicore::context::startGraphRecording();
+        auto output = model_->forward(input);
+        auto graph = infinicore::context::stopGraphRecording();
+        barrier_->wait();
+
+        auto shared_output = std::shared_ptr<InfinilmModel::Output>(
+            new InfinilmModel::Output{infinicore::graph::GraphTensor(output.logits)});
+
+        compiled_map_decode_[b] = CompiledResult{std::move(input), std::make_tuple(graph, shared_output)};
     }
 }
 

From 8a10da5b8a5d365d6eab7677c0a06a76053af41d Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sun, 12 Jul 2026 03:36:59 +0800
Subject: [PATCH 09/13] test(moe): cover serving policy selection

Exercise automatic backend gating, explicit compatibility overrides, quantization fallback, and the high-concurrency prefill budget boundaries.
---
 test/llm/test_moe_policy.py | 70 +++++++++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)
 create mode 100644 test/llm/test_moe_policy.py

diff --git a/test/llm/test_moe_policy.py b/test/llm/test_moe_policy.py
new file mode 100644
index 00000000..e1af4dd2
--- /dev/null
+++ b/test/llm/test_moe_policy.py
@@ -0,0 +1,70 @@
+import json
+
+import pytest
+from infinilm.llm.llm import (
+    _default_max_num_batched_tokens,
+    _resolve_moe_backend,
+)
+
+
+def _model_dir(tmp_path, config):
+    (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8")
+    return str(tmp_path)
+
+
+def test_auto_selects_fused_qwen3_moe(tmp_path):
+    model_path = _model_dir(tmp_path, {"model_type": "qwen3_moe"})
+
+    assert _resolve_moe_backend(model_path, "cuda", "bfloat16", "auto", None) == (
+        True,
+        "fused",
+    )
+
+
+@pytest.mark.parametrize(
+    ("config", "device", "dtype"),
+    [
+        ({"model_type": "qwen3_moe"}, "cpu", "bfloat16"),
+        ({"model_type": "qwen3_moe"}, "cuda", "float32"),
+        ({"model_type": "llama"}, "cuda", "bfloat16"),
+        (
+            {
+                "model_type": "qwen3_moe",
+                "quantization_config": {"quant_method": "gptq"},
+            },
+            "cuda",
+            "bfloat16",
+        ),
+    ],
+)
+def test_auto_preserves_legacy_for_unsupported_configs(tmp_path, config, device, dtype):
+    model_path = _model_dir(tmp_path, config)
+
+    assert _resolve_moe_backend(model_path, device, dtype, "auto", None) == (
+        False,
+        "legacy",
+    )
+
+
+def test_explicit_backend_and_legacy_override(tmp_path):
+    model_path = _model_dir(tmp_path, {"model_type": "llama"})
+
+    assert _resolve_moe_backend(model_path, "cpu", "float32", "fused", None) == (
+        True,
+        "fused",
+    )
+    assert _resolve_moe_backend(model_path, "cpu", "float32", "auto", False) == (
+        False,
+        "legacy",
+    )
+    with pytest.raises(ValueError, match="Conflicting MoE options"):
+        _resolve_moe_backend(model_path, "cpu", "float32", "legacy", True)
+
+
+def test_qwen3_moe_prefill_budget():
+    qwen_config = {"model_type": "qwen3_moe"}
+
+    assert _default_max_num_batched_tokens(qwen_config, 31, 32768) == 32768
+    assert _default_max_num_batched_tokens(qwen_config, 32, 32768) == 4096
+    assert _default_max_num_batched_tokens(qwen_config, 64, 2048) == 2048
+    assert _default_max_num_batched_tokens({"model_type": "llama"}, 64, 32768) == 32768

From 3cd5ed17c801f085c35aea507288b67087997d88 Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sun, 12 Jul 2026 03:45:02 +0800
Subject: [PATCH 10/13] docs(moe): document backend selection

---
 README.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/README.md b/README.md
index 4c1eaf1e..bac23981 100644
--- a/README.md
+++ b/README.md
@@ -76,6 +76,8 @@
       ```bash
       CUDA_VISIBLE_DEVICES=0,1,2,3 python python/infinilm/server/inference_server.py --device nvidia --model=/models/9G7B_MHA/ --enable-paged-attn --attn=flash-attn --enable-graph
       ```
+
+    - Qwen3-MoE 执行后端可通过 `--moe-backend=auto|legacy|fused` 选择；默认 `auto` 仅对支持的 CUDA FP16/BF16 非量化模型启用 fused 路径。
     
     - 测试推理服务性能：
       ```bash

From bda732b87846c8e864ff1598344f4f1ba84451ac Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sun, 12 Jul 2026 05:42:38 +0800
Subject: [PATCH 11/13] test(moe): add contest serving reproduction script

---
 scripts/run_qwen3_moe_contest.sh       | 327 +++++++++++++++++++++++++
 scripts/run_vllm_bench_serve_client.py |  43 ++++
 2 files changed, 370 insertions(+)
 create mode 100755 scripts/run_qwen3_moe_contest.sh
 create mode 100755 scripts/run_vllm_bench_serve_client.py

diff --git a/scripts/run_qwen3_moe_contest.sh b/scripts/run_qwen3_moe_contest.sh
new file mode 100755
index 00000000..716762b2
--- /dev/null
+++ b/scripts/run_qwen3_moe_contest.sh
@@ -0,0 +1,327 @@
+#!/usr/bin/env bash
+# Run the T2-1-3 Qwen3-MoE correctness or serving benchmark profiles.
+set -euo pipefail
+
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+REPO_DIR=$(cd "$SCRIPT_DIR/.." && pwd)
+
+usage() {
+    cat <<'EOF'
+Usage: scripts/run_qwen3_moe_contest.sh PROFILE
+
+Profiles:
+  check       Run the focused policy tests without starting a server.
+  smoke       Start a small server and issue one real chat request.
+  c1-short    Benchmark concurrency=1, input=32, output=256.
+  c64-short   Benchmark concurrency=64, input=32, output=256.
+  c64-long    Benchmark concurrency=64, input=256, output=1024.
+
+Environment:
+  MODEL_DIR        Qwen3-30B-A3B-Thinking-2507 directory (required except for check)
+  INFINICORE_DIR   Optional InfiniCore checkout; omit when installed in Python
+  INFINI_ROOT      Installed InfiniCore prefix (default: $HOME/.infini)
+  PYTHON           Server Python executable (default: python3)
+  BENCH_PYTHON     Python with vLLM 0.9.2 client modules (default: PYTHON)
+  RESULT_DIR       Log/result directory (default: current-directory benchmark_results/)
+  CUDA_VISIBLE_DEVICES  Four TP devices (default: 0,1,2,3)
+  PORT             Local service port (default: 8102)
+  BUILD=1          Rebuild/install InfiniLM before running the profile
+  DRY_RUN=1        Print the resolved commands without executing them
+  SHOW_COMMANDS=1  Print full commands during a real run
+
+Examples:
+  MODEL_DIR=/path/to/Qwen3-30B-A3B-Thinking-2507 \
+    scripts/run_qwen3_moe_contest.sh smoke
+
+  MODEL_DIR=/path/to/Qwen3-30B-A3B-Thinking-2507 \
+    BENCH_PYTHON=/path/to/vllm-client/bin/python \
+    scripts/run_qwen3_moe_contest.sh c64-long
+EOF
+}
+
+PROFILE=${1:-}
+if [ -z "$PROFILE" ] || [ "$PROFILE" = "-h" ] || [ "$PROFILE" = "--help" ]; then
+    usage
+    exit 0
+fi
+
+MODEL_DIR=${MODEL_DIR:-}
+INFINICORE_DIR=${INFINICORE_DIR:-}
+INFINI_ROOT=${INFINI_ROOT:-${HOME:-$REPO_DIR}/.infini}
+PYTHON=${PYTHON:-python3}
+BENCH_PYTHON=${BENCH_PYTHON:-$PYTHON}
+RESULT_DIR=${RESULT_DIR:-benchmark_results/qwen3_moe}
+CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}
+HOST=${HOST:-127.0.0.1}
+PORT=${PORT:-8102}
+TP=${TP:-4}
+BUILD=${BUILD:-0}
+DRY_RUN=${DRY_RUN:-0}
+SHOW_COMMANDS=${SHOW_COMMANDS:-0}
+SERVER_START_TIMEOUT=${SERVER_START_TIMEOUT:-1800}
+
+MAX_BATCH_SIZE=1
+NUM_BLOCKS=18
+NUM_PROMPTS=1
+MAX_CONCURRENCY=1
+INPUT_LEN=32
+OUTPUT_LEN=64
+RUN_BENCH=0
+
+case "$PROFILE" in
+    check)
+        ;;
+    smoke)
+        ;;
+    c1-short)
+        OUTPUT_LEN=256
+        RUN_BENCH=1
+        ;;
+    c64-short)
+        MAX_BATCH_SIZE=64
+        NUM_BLOCKS=128
+        NUM_PROMPTS=64
+        MAX_CONCURRENCY=64
+        OUTPUT_LEN=256
+        RUN_BENCH=1
+        ;;
+    c64-long)
+        MAX_BATCH_SIZE=64
+        NUM_BLOCKS=120
+        NUM_PROMPTS=64
+        MAX_CONCURRENCY=64
+        INPUT_LEN=256
+        OUTPUT_LEN=1024
+        RUN_BENCH=1
+        ;;
+    *)
+        echo "Unknown profile: $PROFILE" >&2
+        usage >&2
+        exit 2
+        ;;
+esac
+
+if [ "$PROFILE" != "check" ] && [ -z "$MODEL_DIR" ]; then
+    echo "MODEL_DIR is required for profile $PROFILE" >&2
+    exit 2
+fi
+
+if [ -n "$MODEL_DIR" ]; then
+    MODEL_NAME=${MODEL_NAME:-$(basename "$MODEL_DIR")}
+else
+    MODEL_NAME=${MODEL_NAME:-Qwen3-30B-A3B-Thinking-2507}
+fi
+TIMESTAMP=$(date -u +%Y%m%dT%H%M%SZ)
+RUN_NAME="${PROFILE}-${TIMESTAMP}"
+SERVER_LOG="$RESULT_DIR/${RUN_NAME}-server.log"
+CLIENT_LOG="$RESULT_DIR/${RUN_NAME}-client.log"
+RESULT_FILE="$RUN_NAME.json"
+
+export INFINI_ROOT
+export CUDA_VISIBLE_DEVICES
+export LD_LIBRARY_PATH="$INFINI_ROOT/lib:${LD_LIBRARY_PATH:-}"
+if [ -n "$INFINICORE_DIR" ]; then
+    export PYTHONPATH="$INFINICORE_DIR/python:$REPO_DIR/python:${PYTHONPATH:-}"
+else
+    export PYTHONPATH="$REPO_DIR/python:${PYTHONPATH:-}"
+fi
+if [ "$MAX_BATCH_SIZE" -ge 32 ]; then
+    export INFINILM_MAX_NUM_BATCHED_TOKENS=4096
+else
+    unset INFINILM_MAX_NUM_BATCHED_TOKENS || true
+fi
+
+SERVER_COMMAND=(
+    "$PYTHON" "$REPO_DIR/python/infinilm/server/inference_server.py"
+    --device nvidia
+    --model "$MODEL_DIR"
+    --dtype bfloat16
+    --tp "$TP"
+    --dp 1
+    --ep 1
+    --moe-ep-backend disabled
+    --host "$HOST"
+    --port "$PORT"
+    --max-batch-size "$MAX_BATCH_SIZE"
+    --max-new-tokens 4096
+    --num-blocks "$NUM_BLOCKS"
+    --block-size 256
+    --temperature 1.0
+    --top-p 0.8
+    --top-k 1
+    --enable-paged-attn
+    --attn paged-attn
+    --enable-graph
+    --moe-backend auto
+    --ignore-eos
+)
+
+BENCH_COMMAND=(
+    "$BENCH_PYTHON" "$SCRIPT_DIR/run_vllm_bench_serve_client.py"
+    --backend openai-chat
+    --endpoint-type openai-chat
+    --model "$MODEL_NAME"
+    --tokenizer "$MODEL_DIR"
+    --endpoint /v1/chat/completions
+    --host "$HOST"
+    --port "$PORT"
+    --dataset-name random
+    --request-rate inf
+    --seed 714
+    --num-prompts "$NUM_PROMPTS"
+    --max-concurrency "$MAX_CONCURRENCY"
+    --random-input-len "$INPUT_LEN"
+    --random-output-len "$OUTPUT_LEN"
+    --temperature 1.0
+    --top-p 0.8
+    --top-k 1
+    --extra-body "{\"max_tokens\": $OUTPUT_LEN}"
+    --save-result
+    --save-detailed
+    --result-dir "$RESULT_DIR"
+    --result-filename "$RESULT_FILE"
+)
+
+print_command() {
+    printf '%q ' "$@"
+    printf '\n'
+}
+
+echo "Profile: $PROFILE"
+if [ -n "$MODEL_DIR" ]; then
+    echo "Model: $MODEL_NAME"
+fi
+
+if [ "$PROFILE" = "check" ]; then
+    CHECK_COMMAND=(
+        "$PYTHON" -m pytest -q "$REPO_DIR/test/llm/test_moe_policy.py"
+    )
+    if [ "$DRY_RUN" = "1" ]; then
+        print_command "${CHECK_COMMAND[@]}"
+        exit 0
+    fi
+    "${CHECK_COMMAND[@]}"
+    exit 0
+fi
+
+if [ "$DRY_RUN" = "1" ]; then
+    echo "Server command:"
+    print_command "${SERVER_COMMAND[@]}"
+    if [ "$RUN_BENCH" = "1" ]; then
+        echo "Client command:"
+        print_command "${BENCH_COMMAND[@]}"
+    fi
+    exit 0
+fi
+
+if [ ! -f "$MODEL_DIR/config.json" ]; then
+    echo "Model config not found: $MODEL_DIR/config.json" >&2
+    exit 3
+fi
+if [ -n "$INFINICORE_DIR" ] && [ ! -d "$INFINICORE_DIR/python/infinicore" ]; then
+    echo "InfiniCore Python package not found: $INFINICORE_DIR/python/infinicore" >&2
+    exit 3
+fi
+if ! command -v curl >/dev/null 2>&1; then
+    echo "curl is required" >&2
+    exit 3
+fi
+
+mkdir -p "$RESULT_DIR"
+
+if [ "$BUILD" = "1" ]; then
+    echo "Building and installing InfiniLM..."
+    (
+        cd "$REPO_DIR"
+        "$PYTHON" -m pip install -e . --no-build-isolation
+    )
+fi
+
+"$PYTHON" -c \
+    'import infinicore; from infinilm.lib import _infinilm; print("InfiniLM extension import: OK")'
+
+GPU_COUNT=$(
+    "$PYTHON" -c 'import torch; print(torch.cuda.device_count())'
+)
+if [ "$GPU_COUNT" -lt "$TP" ]; then
+    echo "Expected at least $TP visible GPUs, found $GPU_COUNT" >&2
+    exit 4
+fi
+
+if [ "$RUN_BENCH" = "1" ]; then
+    "$BENCH_PYTHON" -c 'import vllm; print(f"vLLM client {vllm.__version__}")'
+fi
+
+if curl -fsS "http://$HOST:$PORT/health" >/dev/null 2>&1; then
+    echo "A service is already listening at $HOST:$PORT; choose another PORT" >&2
+    exit 5
+fi
+
+SERVER_PID=
+cleanup() {
+    if [ -n "$SERVER_PID" ] && kill -0 "$SERVER_PID" >/dev/null 2>&1; then
+        echo "Stopping server process $SERVER_PID..."
+        kill "$SERVER_PID" >/dev/null 2>&1 || true
+        wait "$SERVER_PID" 2>/dev/null || true
+    fi
+}
+trap cleanup EXIT INT TERM
+
+echo "Starting InfiniLM server..."
+if [ "$SHOW_COMMANDS" = "1" ]; then
+    print_command "${SERVER_COMMAND[@]}"
+fi
+"${SERVER_COMMAND[@]}" >"$SERVER_LOG" 2>&1 &
+SERVER_PID=$!
+
+READY=0
+ATTEMPT=0
+START_SECONDS=$SECONDS
+START_DEADLINE=$((START_SECONDS + SERVER_START_TIMEOUT))
+while [ "$SECONDS" -lt "$START_DEADLINE" ]; do
+    if curl -fsS "http://$HOST:$PORT/health" >/dev/null 2>&1; then
+        READY=1
+        break
+    fi
+    if ! kill -0 "$SERVER_PID" >/dev/null 2>&1; then
+        echo "Server exited before becoming ready" >&2
+        tail -n 100 "$SERVER_LOG" >&2 || true
+        exit 6
+    fi
+    ATTEMPT=$((ATTEMPT + 1))
+    if [ $((ATTEMPT % 6)) -eq 0 ]; then
+        echo "Waiting for model load and graph capture ($((SECONDS - START_SECONDS))s)..."
+    fi
+    sleep 5
+done
+
+if [ "$READY" != "1" ]; then
+    echo "Server did not become ready within ${SERVER_START_TIMEOUT}s" >&2
+    tail -n 100 "$SERVER_LOG" >&2 || true
+    exit 7
+fi
+
+echo "Server is ready."
+unset http_proxy https_proxy all_proxy ALL_PROXY || true
+
+if [ "$RUN_BENCH" = "1" ]; then
+    echo "Running the contest-style vLLM client..."
+    if [ "$SHOW_COMMANDS" = "1" ]; then
+        print_command "${BENCH_COMMAND[@]}"
+    fi
+    "${BENCH_COMMAND[@]}" 2>&1 | tee "$CLIENT_LOG"
+else
+    PAYLOAD=$(printf \
+        '{"model":"%s","messages":[{"role":"user","content":"What is 15 + 27? Answer briefly."}],"max_tokens":%d,"temperature":0.0}' \
+        "$MODEL_NAME" "$OUTPUT_LEN")
+    curl -fsS "http://$HOST:$PORT/v1/chat/completions" \
+        -H 'Content-Type: application/json' \
+        --data-binary "$PAYLOAD" | tee "$RESULT_DIR/${RUN_NAME}-smoke.json"
+    printf '\n'
+fi
+
+echo "Server log: $SERVER_LOG"
+if [ "$RUN_BENCH" = "1" ]; then
+    echo "Client log: $CLIENT_LOG"
+    echo "Result JSON: $RESULT_DIR/$RESULT_FILE"
+fi
diff --git a/scripts/run_vllm_bench_serve_client.py b/scripts/run_vllm_bench_serve_client.py
new file mode 100755
index 00000000..0cb5d11c
--- /dev/null
+++ b/scripts/run_vllm_bench_serve_client.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+"""Run vLLM 0.9.2's official serving benchmark client.
+
+This launcher imports only the serving benchmark modules, so the client does
+not need to initialize a vLLM inference engine. It also supplies the generic
+``--extra-body`` option that was added after vLLM 0.9.2.
+"""
+
+import argparse
+import json
+
+from vllm.benchmarks import endpoint_request_func
+from vllm.benchmarks.serve import add_cli_args, main
+
+
+def json_object(value: str) -> dict:
+    """Parse a JSON object for ``--extra-body``."""
+    parsed = json.loads(value)
+    if not isinstance(parsed, dict):
+        raise argparse.ArgumentTypeError("--extra-body must be a JSON object")
+    return parsed
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="vLLM serving benchmark client")
+    add_cli_args(parser)
+    parser.add_argument("--extra-body", type=json_object, default=None)
+    args = parser.parse_args()
+
+    if args.extra_body:
+        original = endpoint_request_func.ASYNC_REQUEST_FUNCS["openai-chat"]
+
+        async def openai_chat_with_extra_body(request_func_input, pbar=None):
+            merged = dict(request_func_input.extra_body or {})
+            merged.update(args.extra_body)
+            request_func_input.extra_body = merged
+            return await original(request_func_input, pbar)
+
+        endpoint_request_func.ASYNC_REQUEST_FUNCS["openai-chat"] = (
+            openai_chat_with_extra_body
+        )
+
+    main(args)

From d33c518d0084fbfe770d8542f16ea9b8d2287b49 Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sun, 12 Jul 2026 17:56:48 +0800
Subject: [PATCH 12/13] fix(graph): preserve runtime resets for other models

---
 csrc/models/infinilm_model.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp
index a16951f8..2a77be12 100644
--- a/csrc/models/infinilm_model.cpp
+++ b/csrc/models/infinilm_model.cpp
@@ -94,8 +94,10 @@ void InfinilmModel::process_weights_after_loading() {
 }
 
 void InfinilmModel::reset_runtime_state() const {
-    // Only quantized kernels currently keep resettable runtime state.
-    if (model_config_ && model_config_->get_quant_scheme() == quantization::QuantScheme::NONE) {
+    // Non-quantized Qwen3-MoE currently has no resettable module state.
+    if (model_config_
+        && model_config_->get_or<std::string>("model_type", "") == "qwen3_moe"
+        && model_config_->get_quant_scheme() == quantization::QuantScheme::NONE) {
         return;
     }
     for (const auto &[_, sub] : children()) {

From 4a44bd363033f41afcebc4d9e3ef155e426db3d3 Mon Sep 17 00:00:00 2001
From: surprisely <46776020+surprisely@users.noreply.github.com>
Date: Sun, 12 Jul 2026 17:56:51 +0800
Subject: [PATCH 13/13] test(moe): validate contest smoke responses

---
 scripts/run_qwen3_moe_contest.sh | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/scripts/run_qwen3_moe_contest.sh b/scripts/run_qwen3_moe_contest.sh
index 716762b2..ad7b3084 100755
--- a/scripts/run_qwen3_moe_contest.sh
+++ b/scripts/run_qwen3_moe_contest.sh
@@ -65,7 +65,7 @@ NUM_BLOCKS=18
 NUM_PROMPTS=1
 MAX_CONCURRENCY=1
 INPUT_LEN=32
-OUTPUT_LEN=64
+OUTPUT_LEN=128
 RUN_BENCH=0
 
 case "$PROFILE" in
@@ -311,13 +311,40 @@ if [ "$RUN_BENCH" = "1" ]; then
     fi
     "${BENCH_COMMAND[@]}" 2>&1 | tee "$CLIENT_LOG"
 else
+    SMOKE_FILE="$RESULT_DIR/${RUN_NAME}-smoke.json"
     PAYLOAD=$(printf \
         '{"model":"%s","messages":[{"role":"user","content":"What is 15 + 27? Answer briefly."}],"max_tokens":%d,"temperature":0.0}' \
         "$MODEL_NAME" "$OUTPUT_LEN")
     curl -fsS "http://$HOST:$PORT/v1/chat/completions" \
         -H 'Content-Type: application/json' \
-        --data-binary "$PAYLOAD" | tee "$RESULT_DIR/${RUN_NAME}-smoke.json"
+        --data-binary "$PAYLOAD" | tee "$SMOKE_FILE"
     printf '\n'
+    "$PYTHON" - "$SMOKE_FILE" <<'PY'
+import json
+import sys
+
+with open(sys.argv[1], encoding="utf-8") as response_file:
+    response = json.load(response_file)
+
+try:
+    choice = response["choices"][0]
+    content = choice["message"]["content"]
+    completion_tokens = response["usage"]["completion_tokens"]
+except (KeyError, IndexError, TypeError) as error:
+    raise SystemExit(f"Smoke validation failed: malformed response: {error}") from error
+
+if not isinstance(content, str) or not content.strip():
+    raise SystemExit("Smoke validation failed: empty completion")
+if not isinstance(completion_tokens, int) or completion_tokens <= 0:
+    raise SystemExit("Smoke validation failed: no completion tokens")
+if "42" not in content:
+    raise SystemExit("Smoke validation failed: arithmetic answer 42 not found")
+
+print(
+    "Smoke validation: PASS "
+    f"(completion_tokens={completion_tokens}, answer_contains_42=yes)"
+)
+PY
 fi
 
 echo "Server log: $SERVER_LOG"