diff --git a/.gitignore b/.gitignore
index 5b5cc3a84..16d9bea36 100644
--- a/.gitignore
+++ b/.gitignore
@@ -29,3 +29,8 @@ mstar/worker/ASYNC_REDESIGN.md
 # local AI-assistant context (kept local, not published — cf. vllm-omni)
 CLAUDE.md
 AGENTS.md
+mstar_traces/
+.claude/skills/
+
+# local golden-extraction harness (Kimi-K2.7 port dev tooling; not for main)
+tools/kimi_goldens/
diff --git a/configs/kimi_k2_7.yaml b/configs/kimi_k2_7.yaml
new file mode 100644
index 000000000..c4cac0e7a
--- /dev/null
+++ b/configs/kimi_k2_7.yaml
@@ -0,0 +1,8 @@
+model: "kimi_k2_7"
+# Single-rank Kimi-K2.7 text config; use TP8 configs for the real 1T checkpoint.
+max_seq_len: 262144
+node_groups:
+  - node_names: [LLM]
+    ranks: [0]
+    tp_size: 1
+    graph_walks: [prefill, decode]
diff --git a/configs/kimi_k2_7_code_tp8.yaml b/configs/kimi_k2_7_code_tp8.yaml
new file mode 100644
index 000000000..e05bd5f34
--- /dev/null
+++ b/configs/kimi_k2_7_code_tp8.yaml
@@ -0,0 +1,21 @@
+model: "kimi_k2_7"
+# Fallback TP8 config for local Kimi-K2.7-Code snapshot; prefer the /dev/shm variant.
+max_seq_len: 8192
+# See kimi_k2_7_code_tp8_shm.yaml — a 1T INT4 load at TP8 outruns PyTorch's
+# default process-group timeout; this path is slower still (disk, not tmpfs).
+dist_timeout_s: 7200
+# Point at your checkpoint at launch — this config hardcodes no path:
+#   mstar-serve --config configs/kimi_k2_7_code_tp8.yaml \
+#               --model-path /path/to/Kimi-K2.7-Code
+# Omit --model-path to pull moonshotai/Kimi-K2.7-Code from HuggingFace.
+model_kwargs:
+  config_variant: k27_code
+  tokenizer_mode: hf
+kv_cache:
+  max_num_pages: 512
+  page_size: 128
+node_groups:
+  - node_names: [LLM]
+    ranks: [0, 1, 2, 3, 4, 5, 6, 7]
+    tp_size: 8
+    graph_walks: [prefill, decode]
diff --git a/configs/kimi_k2_7_code_tp8_shm.yaml b/configs/kimi_k2_7_code_tp8_shm.yaml
new file mode 100644
index 000000000..3298aa027
--- /dev/null
+++ b/configs/kimi_k2_7_code_tp8_shm.yaml
@@ -0,0 +1,21 @@
+model: "kimi_k2_7"
+# TP8 Kimi-K2.7-Code config using a RAM-backed /dev/shm checkpoint copy.
+max_seq_len: 8192
+# Loading ~600GB of INT4 experts across 8 ranks
+# MSTAR_DIST_TIMEOUT_S overrides this.
+dist_timeout_s: 7200
+# Copy the checkpoint into /dev/shm first, then point at it at launch:
+#   mstar serve kimi_k2_7 --model-path /dev/shm/kimi_k2_7_code
+#   mstar-serve --config configs/kimi_k2_7_code_tp8_shm.yaml \
+#               --model-path /dev/shm/kimi_k2_7_code
+model_kwargs:
+  config_variant: k27_code
+  tokenizer_mode: hf
+kv_cache:
+  max_num_pages: 512
+  page_size: 128
+node_groups:
+  - node_names: [LLM]
+    ranks: [0, 1, 2, 3, 4, 5, 6, 7]
+    tp_size: 8
+    graph_walks: [prefill, decode]
diff --git a/configs/synthetic/kimi_k2_7_repro.yaml b/configs/synthetic/kimi_k2_7_repro.yaml
new file mode 100644
index 000000000..bc3345d35
--- /dev/null
+++ b/configs/synthetic/kimi_k2_7_repro.yaml
@@ -0,0 +1,17 @@
+model: "kimi_k2_7"
+# Reduced synthetic serve config — random weights, NOT a production deployment.
+# Generate the checkpoint first, then point at it:
+#   mstar-serve --config configs/synthetic/kimi_k2_7_repro.yaml \
+#               --model-path tools/kimi_goldens/repro/checkpoint
+max_seq_len: 512
+model_kwargs:
+  config_variant: reduced
+  tokenizer_mode: byte
+kv_cache:
+  max_num_pages: 256
+  page_size: 128
+node_groups:
+  - node_names: [LLM]
+    ranks: [0]
+    tp_size: 1
+    graph_walks: [prefill, decode]
diff --git a/configs/synthetic/kimi_k2_7_tp2.yaml b/configs/synthetic/kimi_k2_7_tp2.yaml
new file mode 100644
index 000000000..17bd7c82f
--- /dev/null
+++ b/configs/synthetic/kimi_k2_7_tp2.yaml
@@ -0,0 +1,17 @@
+model: "kimi_k2_7"
+# Reduced synthetic TP=2 config — random weights, NOT a production deployment.
+# Generate the checkpoint first, then point at it:
+#   mstar-serve --config configs/synthetic/kimi_k2_7_tp2.yaml \
+#               --model-path tools/kimi_goldens/repro/checkpoint
+max_seq_len: 512
+model_kwargs:
+  config_variant: reduced
+  tokenizer_mode: byte
+kv_cache:
+  max_num_pages: 256
+  page_size: 128
+node_groups:
+  - node_names: [LLM]
+    ranks: [0, 1]
+    tp_size: 2
+    graph_walks: [prefill, decode]
diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst
index 72430b664..6eeed3d57 100644
--- a/docs/environment_variables.rst
+++ b/docs/environment_variables.rst
@@ -35,6 +35,15 @@ Communication
      - ``19000``
      - Base of the deterministic entity-id → TCP port map (``api_server``
        = base, ``conductor`` = base+1, ``worker_<rank>`` = base+100+rank).
+   * - ``MSTAR_DIST_TIMEOUT_S``
+     - config's ``dist_timeout_s``
+     - Timeout in seconds for the NCCL world group and its parallel
+       subgroups (:func:`mstar.distributed.communication.resolve_dist_timeout`).
+       Overrides the deployment config's ``dist_timeout_s``; with neither set,
+       PyTorch's default applies. Raise it only where weight load, JIT or
+       CUDA-graph capture can exceed that default (a 1T MoE at TP8 does) —
+       a hung collective takes correspondingly longer to abort. Must be set
+       before the conductor spawns workers, which inherit it.
    * - ``MSTAR_SHM_ARENA``
      - ``0``
      - SHM tensor-transport implementation. ``0``: per-uuid files.
diff --git a/docs/serving.rst b/docs/serving.rst
index 5b01c42df..b63d9af6a 100644
--- a/docs/serving.rst
+++ b/docs/serving.rst
@@ -44,6 +44,10 @@ mstar serve
    * - ``--cache-dir``
      - HF default
      - HuggingFace weight cache directory.
+   * - ``--model-path``
+     - registry default
+     - Local checkpoint directory or HF repo id to load weights from. Use this
+       instead of hardcoding a path in the config YAML.
    * - ``--tensor-comm-protocol``
      - ``SHM``
      - Tensor transport: ``SHM`` (safe single-node default), ``TCP``, or ``RDMA``.
@@ -122,6 +126,16 @@ mstar-serve
    * - ``--cache-dir``
      - HF default
      - HuggingFace weight cache directory.
+   * - ``--model-path``
+     - registry default
+     - Local checkpoint directory or HF repo id to load weights from. Applies to
+       any model and overrides the registry default, so a config YAML never has
+       to hardcode one machine's filesystem layout:
+
+       .. code-block:: bash
+
+          mstar-serve --config configs/kimi_k2_7_code_tp8_shm.yaml \
+                      --model-path /dev/shm/kimi_k2_7_code
    * - ``--socket-path-prefix``
      - ``/tmp/mstar``
      - ZMQ IPC socket prefix (shared with conductor/workers).
@@ -178,6 +192,9 @@ A config maps the model's computation-graph nodes to physical GPU ranks. The key
        scoped to specific ``graph_walks`` and/or sharded with ``tp_size``.
    * - ``model_kwargs``
      - *(optional)* Server-init model parameters (see below).
+   * - ``dist_timeout_s``
+     - *(optional)* Timeout in seconds for the NCCL world group and its parallel
+       subgroups. Unset keeps PyTorch's default.
 
 Node names are model-specific — they are the keys of the model's
 ``get_node_engine_types`` (e.g. BAGEL's ``vit_encoder`` / ``vae_encoder`` / ``LLM``,
@@ -208,6 +225,11 @@ Because placement is config-only, the *same* model code runs single-GPU or fully
 disaggregated. ``configs/`` ships several layouts per model (``*_single_gpu``,
 ``*_colocated``, ``*_pd_disaggregated``, ``*_cfg_parallel``, …).
 
+``configs/synthetic/`` holds reduced, randomly-initialized deployments used to
+exercise the serving path without real weights (shape/plumbing checks, TP
+sanity). They are **not** production configs — they load a generated checkpoint
+and emit meaningless tokens.
+
 **Tensor parallelism.** Shard a node across GPUs with ``tp_size`` and that many ``ranks``:
 
 .. code-block:: yaml
diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py
index 3fc97ad4d..160aa7f1a 100644
--- a/mstar/api_server/entrypoint.py
+++ b/mstar/api_server/entrypoint.py
@@ -52,6 +52,19 @@ def _detect_modality(filename: str) -> str:
 # Conductor process target (top-level for picklability with spawn)
 # ------------------------------------------------------------------
 
+def _resolve_model_path(model_name: str, model_path: str | None) -> str:
+    """Where to load weights from: ``--model-path`` if given, else the registry default.
+
+    ``--model-path`` takes a local directory or an HF repo id and applies to any
+    model, so a deployment config never has to hardcode one machine's filesystem
+    layout. Both the API-server-side model instance and the conductor-side one
+    resolve through here, so they cannot diverge.
+    """
+    if model_path:
+        return model_path
+    return HF_MODELS.get(model_name, {}).get("model_path_hf", "")
+
+
 def _conductor_process_target(
     model_name: str,
     config_path: str,
@@ -61,7 +74,8 @@ def _conductor_process_target(
     log_level: str = "INFO",
     cache_dir: str | None = None,
     tensor_comm_protocol=CommProtocol.RDMA,
-    tcp_transfer_device=""
+    tcp_transfer_device="",
+    model_path: str | None = None,
 ):
     """Runs DummyConductor.run() in a spawned process."""
     logging.basicConfig(
@@ -93,7 +107,7 @@ def _conductor_process_target(
         )
 
     model = get_model_class(model_name)(
-        model_path_hf=HF_MODELS.get(model_name, {}).get("model_path_hf", ""),
+        model_path_hf=_resolve_model_path(model_name, model_path),
         cache_dir=cache_dir,
         **yaml_model_kwargs,
     )
@@ -866,6 +880,12 @@ def main(argv: list[str] | None = None):
         "--cache-dir", type=str, default=None,
         help="Directory for caching downloaded HuggingFace model files",
     )
+    parser.add_argument(
+        "--model-path", type=str, default=None,
+        help="Where to load weights from — a local checkpoint directory or an HF "
+             "repo id. Overrides the model's registry default so deployment "
+             "configs need not hardcode a filesystem path.",
+    )
     parser.add_argument(
         "--log-level", type=str, default="INFO",
         choices=["DEBUG", "INFO", "WARNING", "ERROR"],
@@ -900,7 +920,7 @@ def main(argv: list[str] | None = None):
     # (tokenization only — no GPU weights needed)
     from mstar.model.registry import get_model_class
     model = get_model_class(model_name)(
-        model_path_hf=HF_MODELS.get(model_name, {}).get("model_path_hf", ""),
+        model_path_hf=_resolve_model_path(model_name, args.model_path),
         cache_dir=args.cache_dir,
         **yaml_model_kwargs,
     )
@@ -932,7 +952,8 @@ def main(argv: list[str] | None = None):
             args.log_level,
             args.cache_dir,
             CommProtocol(args.tensor_comm_protocol),
-            args.tcp_transfer_device
+            args.tcp_transfer_device,
+            args.model_path,
         ),
     )
     conductor_proc.start()
diff --git a/mstar/cli/main.py b/mstar/cli/main.py
index 22a59acd9..45d734663 100644
--- a/mstar/cli/main.py
+++ b/mstar/cli/main.py
@@ -38,6 +38,7 @@
     "whisper_large": "whisper_large.yaml",
     "higgs_audio": "higgs_audio.yaml",
     "wan22": "wan22.yaml",
+    "kimi_k2_7": "kimi_k2_7_code_tp8_shm.yaml",
 }
 
 
@@ -149,6 +150,8 @@ def _serve(args: argparse.Namespace) -> None:
     ]
     if args.cache_dir:
         argv += ["--cache-dir", args.cache_dir]
+    if args.model_path:
+        argv += ["--model-path", args.model_path]
     if args.log_stats:
         argv += ["--log-stats"]
     if args.log_stats_file:
@@ -172,6 +175,11 @@ def build_parser() -> argparse.ArgumentParser:
     serve.add_argument("--gpus", default=None, help="CUDA_VISIBLE_DEVICES, e.g. '0' or '0,1,2'")
     serve.add_argument("--config", default=None, help="override the default config (path to YAML)")
     serve.add_argument("--cache-dir", default=None, help="HuggingFace weight cache directory")
+    serve.add_argument(
+        "--model-path", default=None,
+        help="local checkpoint directory or HF repo id to load weights from "
+             "(overrides the model's registry default)",
+    )
     serve.add_argument("--socket-path-prefix", default=None, help="ZMQ IPC socket prefix")
     serve.add_argument("--upload-dir", default=None, help="temp dir for uploaded media")
     serve.add_argument(
diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py
index fe00e3c5d..5604f80f9 100644
--- a/mstar/conductor/conductor.py
+++ b/mstar/conductor/conductor.py
@@ -414,6 +414,7 @@ def _derive_worker_info(self):
         self.parallel_config = GlobalParallelConfig(
             worker_graphs=self.worker_graphs,
             worker_ids=self.worker_ids,
+            dist_timeout_s=self.model_config.get("dist_timeout_s"),
         )
 
     def _launch_workers(self):
diff --git a/mstar/distributed/communication.py b/mstar/distributed/communication.py
index 2f2572f26..ae30b37a0 100644
--- a/mstar/distributed/communication.py
+++ b/mstar/distributed/communication.py
@@ -1,9 +1,30 @@
+import os
 from dataclasses import dataclass, field
+from datetime import timedelta
 from typing import Any
 
 import torch
 import torch.distributed as dist
 
+DIST_TIMEOUT_ENV = "MSTAR_DIST_TIMEOUT_S"
+
+
+def resolve_dist_timeout(dist_timeout_s: float | None = None) -> dict[str, timedelta]:
+    """Build the ``timeout`` kwarg for ``init_process_group`` / ``new_group``."""
+    raw = os.environ.get(DIST_TIMEOUT_ENV, "").strip()
+    if raw:
+        try:
+            dist_timeout_s = float(raw)
+        except ValueError as exc:
+            raise ValueError(
+                f"{DIST_TIMEOUT_ENV} must be a number of seconds, got {raw!r}"
+            ) from exc
+    if dist_timeout_s is None:
+        return {}
+    if dist_timeout_s <= 0:
+        raise ValueError(f"Distributed timeout must be positive, got {dist_timeout_s}")
+    return {"timeout": timedelta(seconds=float(dist_timeout_s))}
+
 
 class CommGroup:
     """A communication group over one axis of the worker device mesh.
@@ -192,6 +213,8 @@ class WorkerParallelGroups:
     # projections).
     node_to_tp_group: dict[str, CommGroup] = field(default_factory=dict)
     node_to_sp_group: dict[str, CommGroup] = field(default_factory=dict)
+    # Process-group timeout in seconds, from the deployment config's
+    dist_timeout_s: float | None = None
 
     def add(self, node: str, comm_group: CommGroup):
         # disallow colocation of multiple comm groups on the same node
@@ -233,11 +256,13 @@ def init_dist(
         if not self.any_parallelism:
             return
 
+        timeout_kwargs = resolve_dist_timeout(self.dist_timeout_s)
         dist.init_process_group(
             backend="nccl",
             init_method=init_method,
             world_size=self.num_workers,
             rank=self.global_rank,
+            **timeout_kwargs,
         )
 
         # One subgroup per distinct rank tuple across BOTH mesh axes —
@@ -247,7 +272,9 @@ def init_dist(
         # an SP group (degenerate meshes) maps to one subgroup.
         rank_tuple_to_pg: dict[tuple[int, ...], "dist.ProcessGroup"] = {}
         for rank_tuple in self.world_parallel_groups:
-            rank_tuple_to_pg[rank_tuple] = dist.new_group(ranks=list(rank_tuple))
+            rank_tuple_to_pg[rank_tuple] = dist.new_group(
+                ranks=list(rank_tuple), **timeout_kwargs
+            )
 
         seen: set[int] = set()
         for comm_group in (
@@ -309,7 +336,8 @@ class GlobalParallelConfig:
     def __init__(
         # leaving type annotation as Any due to circular import
         self, worker_graphs: dict[str, Any],
-        worker_ids: list[str]
+        worker_ids: list[str],
+        dist_timeout_s: float | None = None,
     ):
         self.num_workers = len(worker_ids)
         any_parallelism = any(
@@ -337,6 +365,7 @@ def __init__(
                 global_rank=i, num_workers=self.num_workers,
                 any_parallelism=any_parallelism,
                 world_parallel_groups=world_parallel_groups,
+                dist_timeout_s=dist_timeout_s,
             ) for i, wid in enumerate(worker_ids)
         }
 
@@ -372,4 +401,3 @@ def __init__(
                         self.per_worker_config[worker_ids[rank]].add_sp(
                             node, self.sp_comm_groups[key]
                         )
-
diff --git a/mstar/engine/cache_manager.py b/mstar/engine/cache_manager.py
index 557083e24..5d586cc73 100644
--- a/mstar/engine/cache_manager.py
+++ b/mstar/engine/cache_manager.py
@@ -12,7 +12,11 @@
     KVRequestState,
     PagedAllocationManager,
 )
-from mstar.utils.flashinfer_utils import FlashInferDecodeWrapper, FlashInferPrefillWrapper
+from mstar.utils.flashinfer_utils import (
+    FlashInferDecodeWrapper,
+    FlashInferMLAWrapper,
+    FlashInferPrefillWrapper,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -74,6 +78,36 @@ class PlanCacheKey(NamedTuple):
     dtype: torch.dtype
 
 
+@dataclass
+class MlaRequestSlice:
+    """One request's slice of an absorbed-MLA SDPA plan.
+
+    ``q_start``/``seq_len`` locate this request's rows in the packed query batch;
+    ``total_len`` is its context length after this step (so the causal mask knows
+    how many cached tokens precede the new ones); ``page_indices`` gathers its
+    latent pages.
+    """
+    q_start: int
+    seq_len: int
+    total_len: int
+    page_indices: torch.Tensor
+
+
+@dataclass
+class MlaSdpaPlan:
+    """Absorbed-MLA fallback plan: latent scatter indices + per-request gather layout.
+
+    Built only when the FlashInfer MLA kernel cannot serve the configured latent
+    dims (see :func:`mla_kernel_available_for`); the kernel path plans a
+    :class:`FlashInferMLAWrapper` instead. For an ``mla_absorb`` label exactly one
+    of ``_PlanState.wrapper`` / ``_PlanState.mla`` is set — ``run_attention_mla``
+    uses that to choose a path, so leaving a stale value in either is a bug.
+    """
+    token_to_page: torch.Tensor
+    token_to_cache: torch.Tensor
+    requests: list[MlaRequestSlice]
+
+
 @dataclass
 class _PlanState:
     """Pre-computed state from plan_attention/plan_rope for a single cache label.
@@ -98,7 +132,7 @@ class _PlanState:
     actually consumes this — the model's inner ``advance_seq_lens(pos_id_ns=...)``
     runs at capture time only and is not replayed.
     """
-    wrapper: FlashInferPrefillWrapper | FlashInferDecodeWrapper | None = None
+    wrapper: FlashInferPrefillWrapper | FlashInferDecodeWrapper | FlashInferMLAWrapper | None = None
     pos_ids: torch.Tensor | None = None
     seq_lens: list[int] | None = None
     write_store: bool = True
@@ -123,6 +157,9 @@ class _PlanState:
     # segment over its contiguous frozen prefix. None on paged plans, which
     # keep the FlashInfer path. See DenseGenCacheManager._build_dense_gen_plan.
     dense_gen: dict | None = None
+    # MLA absorb fallback plan: latent scatter indices and per-request gather
+    # layout. Mutually exclusive with ``wrapper`` (see MlaSdpaPlan).
+    mla: "MlaSdpaPlan | None" = None
 
 
 class WorkspaceBufferManager:
@@ -1488,10 +1525,278 @@ def _run_dense_gen(
         return out[0] if isinstance(out, tuple) else out
 
 
+@functools.cache
+def _mla_kernel_available(ckv: int, kpe: int, sm_major: int) -> bool:
+    """Whether the FlashInfer MLA kernel fast path can serve these latent dims.
+
+    Gated conservatively: ``flashinfer.mla.BatchMLAPagedAttentionWrapper`` is
+    hard-locked to the real Kimi dims (ckv=512, kpe=64). Off-dim calls trigger an
+    *uncatchable* illegal memory access (it corrupts the CUDA context — not a
+    catchable exception), so this decision MUST be made BEFORE any kernel
+    construction/call. Requires the flashinfer MLA module, exactly ckv=512/kpe=64,
+    and a Hopper (sm90) GPU (``backend="auto"`` -> fa3). Everything else — reduced
+    configs, pre-sm90, Blackwell (sm100, which wants the trtllm MLA path), or a
+    build without flashinfer — returns False and the manager uses the all-dims
+    SDPA fallback.
+    """
+    if not (ckv == 512 and kpe == 64):
+        return False
+    if sm_major != 9:
+        return False
+    try:
+        import flashinfer.mla  # noqa: F401
+    except Exception:  # noqa: BLE001
+        return False
+    return True
+
+
+def mla_kernel_available_for(kv_cache_config: KVCacheConfig, device) -> bool:
+    """Whether ``kv_cache_config``'s latent dims can use the FlashInfer MLA kernel.
+
+    The single predicate behind both the runtime plan (:class:`MlaAbsorbCacheManager`)
+    and the capture decision (``cuda_graph_runner.mla_absorb_capture_blocked``). They
+    must not drift: a capture that assumes the kernel while the plan takes SDPA builds
+    a paged wrapper the 4-D latent cache cannot serve.
+
+    Total by construction — a non-CUDA device returns False rather than raising, so the
+    absorbed-SDPA fallback is reachable on CPU. ``_mla_kernel_available`` explains why
+    this must be decided *before* any kernel is constructed.
+    """
+    ckv = kv_cache_config.mla_ckv_dim
+    if ckv is None:
+        return False
+    if not torch.cuda.is_available() or torch.device(device).type != "cuda":
+        return False
+    sm_major = torch.cuda.get_device_capability(device)[0]
+    return _mla_kernel_available(ckv, kv_cache_config.head_dim - ckv, sm_major)
+
+
+class MlaAbsorbCacheManager(FlashInferCacheManager):
+    """Paged-cache backend for weight-absorbed MLA.
+
+    Uses a 4D latent cache ``[layers, pages, page_size, ckv + kpe]``. Real Kimi
+    dims on sm90 use FlashInfer MLA; other dims fall back to eager SDPA.
+    """
+
+    def plan_attention(
+        self,
+        seq_lens: list[int] | None = None,
+        dtype: torch.dtype | None = None,
+        is_causal=True,
+        write_store: bool=True,
+        label: str | None = None,
+        **kwargs,
+    ):
+        """Allocate pages and plan the FlashInfer MLA or eager SDPA path."""
+        self._batched_cfg_info = None
+
+        effective_label = label if label is not None else self._active_label()
+        # Always re-plan; this backend does not support the overlap skip yet.
+        self._pre_planned_labels.discard(effective_label)
+        if effective_label not in self._plan_states:
+            self._plan_states[effective_label] = _PlanState()
+        ps = self._plan_states[effective_label]
+
+        cfg = self.kv_cache_config
+        page_size = cfg.page_size
+        ckv = cfg.mla_ckv_dim
+        kpe = (cfg.head_dim - ckv) if ckv is not None else None
+        use_kernel = mla_kernel_available_for(cfg, self.device)
+
+        if use_kernel:
+            qo_indptr_list = [0]
+            kv_indptr_list = [0]
+            all_page_indices: list[int] = []
+            kv_len_list: list[int] = []
+            for i, rid in enumerate(self.request_ids):
+                state = self._get_state(rid, effective_label)
+                sl = seq_lens[i]
+                total_len = state.seq_len + sl
+                self.alloc_manager.alloc(rid, label=effective_label, seq_len=total_len)
+                page_indices = state.page_indices
+                qo_indptr_list.append(qo_indptr_list[-1] + sl)
+                all_page_indices.extend(page_indices)
+                kv_indptr_list.append(kv_indptr_list[-1] + len(page_indices))
+                kv_len_list.append(total_len)
+
+            qo_indptr = torch.tensor(qo_indptr_list, dtype=torch.int32, device=self.device)
+            kv_indptr = torch.tensor(kv_indptr_list, dtype=torch.int32, device=self.device)
+            kv_indices = torch.tensor(all_page_indices, dtype=torch.int32, device=self.device)
+            kv_len_arr = torch.tensor(kv_len_list, dtype=torch.int32, device=self.device)
+
+            if dtype is None:
+                dtype = self.kv_cache.dtype
+            if ps.wrapper is None:
+                # Eager gets a fresh wrapper; CUDA graph injects a persistent one.
+                ps.wrapper = FlashInferMLAWrapper(
+                    workspace_buffer=self.buffer_manager.get(effective_label),
+                    num_heads=cfg.num_qo_heads,
+                    head_dim_ckv=ckv,
+                    head_dim_kpe=kpe,
+                    page_size=page_size,
+                    sm_scale=cfg.softmax_scale,
+                    device=self.device,
+                    enable_nvtx=self.enable_nvtx,
+                )
+            ps.wrapper.plan(
+                qo_indptr=qo_indptr,
+                kv_indptr=kv_indptr,
+                kv_indices=kv_indices,
+                kv_len_arr=kv_len_arr,
+                causal=is_causal,
+                dtype=dtype,
+            )
+            ps.mla = None
+        else:
+            token_to_page: list[int] = []
+            token_to_cache: list[int] = []
+            requests: list[MlaRequestSlice] = []
+            q_start = 0
+            for i, rid in enumerate(self.request_ids):
+                state = self._get_state(rid, effective_label)
+                sl = seq_lens[i]
+                old_len = state.seq_len
+                total_len = old_len + sl
+
+                self.alloc_manager.alloc(
+                    rid, label=effective_label, seq_len=total_len
+                )
+                page_indices = state.page_indices
+
+                # Same page/offset mapping as FlashInfer's token_to_page plan.
+                for j in range(sl):
+                    g = old_len + j
+                    token_to_page.append(page_indices[g // page_size])
+                    token_to_cache.append(g % page_size)
+
+                requests.append(MlaRequestSlice(
+                    q_start=q_start,
+                    seq_len=sl,
+                    total_len=total_len,
+                    page_indices=torch.tensor(
+                        page_indices, dtype=torch.long, device=self.device
+                    ),
+                ))
+                q_start += sl
+
+            ps.mla = MlaSdpaPlan(
+                token_to_page=torch.tensor(
+                    token_to_page, dtype=torch.long, device=self.device
+                ),
+                token_to_cache=torch.tensor(
+                    token_to_cache, dtype=torch.long, device=self.device
+                ),
+                requests=requests,
+            )
+            # Clear any wrapper a CUDA-graph capture injected: run_attention_mla
+            # picks its path with ``ps.wrapper is not None``, and a paged wrapper
+            # cannot read the 4-D latent cache. Mirrors how the paged path clears
+            # ``dense_gen`` in DenseGenCacheManager.
+            ps.wrapper = None
+
+        # Keep advance_seq_lens / flush_to_store aligned with the base manager.
+        ps.seq_lens = seq_lens
+        ps.write_store = write_store
+        ps.dense_gen = None
+
+    @torch.compiler.disable
+    def run_attention_mla(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        kv_c: torch.Tensor,
+        k_pe: torch.Tensor,
+        layer_idx: int | None = None,
+    ) -> torch.Tensor:
+        """Compressed-latent MLA attention over the paged latent cache.
+
+        Args:
+            q_nope: [T, H, L]        query, no-rope part (L = kv_lora_rank).
+            q_pe:   [T, H, Drope]    query, rope part.
+            kv_c:   [T, 1, L]        compressed KV latent (single MQA head).
+            k_pe:   [T, 1, Drope]    rope key (single MQA head).
+            layer_idx: transformer layer; defaults to self.layer_idx.
+        Returns:
+            [T, H, L] attention output (the ``value`` = ``kv_c`` slice width).
+        """
+        if layer_idx is None:
+            layer_idx = self.layer_idx
+
+        label = self._active_label()
+        ps = self._plan_states[label]
+        assert self.kv_cache is not None
+
+        latent_cache = self.kv_cache[layer_idx]
+        latent = torch.cat([kv_c, k_pe], dim=-1).squeeze(1)
+
+        if ps.wrapper is not None:
+            ps.wrapper.set_latent(latent_cache, latent)
+            L = q_nope.shape[-1]  # ckv width (post-w_kc absorption)
+            ckv_cache = latent_cache[..., :L]
+            kpe_cache = latent_cache[..., L:]
+            return ps.wrapper.run(q_nope, q_pe, ckv_cache, kpe_cache).to(q_nope.dtype)
+
+        mla = ps.mla
+        assert mla is not None
+
+        latent_cache[mla.token_to_page, mla.token_to_cache] = latent.to(
+            latent_cache.dtype
+        )
+
+        T, H, L = q_nope.shape
+        scale = self.kv_cache_config.softmax_scale
+        query_all = torch.cat([q_nope, q_pe], dim=-1)
+        out = torch.empty(T, H, L, dtype=q_nope.dtype, device=q_nope.device)
+
+        for req in mla.requests:
+            q_start = req.q_start
+            sl = req.seq_len
+            total_len = req.total_len
+
+            # Mirror the dense-gen page gather for the fallback.
+            gathered = latent_cache[req.page_indices].reshape(
+                -1, latent_cache.shape[-1]
+            )[:total_len]
+            key = gathered
+            value = gathered[:, :L]
+
+            q_req = query_all[q_start:q_start + sl]
+            out[q_start:q_start + sl] = self._sdpa_mla(
+                q_req, key, value, old_len=total_len - sl, scale=scale
+            )
+        return out
+
+    @staticmethod
+    def _sdpa_mla(
+        q: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        old_len: int,
+        scale: float,
+    ) -> torch.Tensor:
+        """Causal SDPA fallback for one request over the shared latent head."""
+        sl = q.shape[0]
+        total = key.shape[0]
+        qt = q.transpose(0, 1).float()
+        scores = torch.einsum("hqd,kd->hqk", qt, key.float()) * scale
+        q_pos = old_len + torch.arange(sl, device=q.device)
+        k_pos = torch.arange(total, device=q.device)
+        mask = torch.where(
+            k_pos[None, :] <= q_pos[:, None],
+            0.0,
+            torch.tensor(float("-inf"), device=q.device),
+        )
+        scores = scores + mask
+        attn = scores.softmax(-1)
+        out = torch.einsum("hqk,kd->hqd", attn, value.float())
+        return out.transpose(0, 1).to(q.dtype)
+
+
 # Backend registry: KVCacheConfig.attention_backend names one of these.
 ATTENTION_BACKENDS: dict[str, type[BatchedCacheManager]] = {
     "flashinfer": FlashInferCacheManager,
     "dense_gen": DenseGenCacheManager,
+    "mla_absorb": MlaAbsorbCacheManager,
 }
 
 
diff --git a/mstar/engine/cuda_graph_runner.py b/mstar/engine/cuda_graph_runner.py
index a9252e3b6..ed16b3579 100644
--- a/mstar/engine/cuda_graph_runner.py
+++ b/mstar/engine/cuda_graph_runner.py
@@ -30,6 +30,7 @@
     BatchedCacheManager,
     WorkspaceBufferManager,
     create_cache_manager,
+    mla_kernel_available_for,
 )
 from mstar.engine.cuda_graph_config import (
     BasicBatchedCudaGraphConfig,
@@ -56,6 +57,21 @@
 DEFAULT_AR_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16, 32, 64]
 
 
+def mla_absorb_capture_blocked(
+    kv_cache_config: KVCacheConfig, device
+) -> str | None:
+    if kv_cache_config.attention_backend != "mla_absorb":
+        return None
+    if mla_kernel_available_for(kv_cache_config, device):
+        return None
+    ckv = kv_cache_config.mla_ckv_dim
+    return (
+        "attention_backend='mla_absorb' without the FlashInfer MLA kernel "
+        f"(mla_ckv_dim={ckv}, kpe={None if ckv is None else kv_cache_config.head_dim - ckv}, "
+        f"device={device}); the absorbed-SDPA fallback runs eager and plans no wrapper"
+    )
+
+
 @dataclass
 class DummyCaptureInput:
     tensors: dict[str, list[torch.Tensor]]  # {tensor_name: [tensor(s)]}
@@ -258,6 +274,14 @@ def warmup_and_capture(self) -> None:
                            self.submodule_name)
             return
 
+        blocked = mla_absorb_capture_blocked(self.kv_cache_config, self.device)
+        if blocked is not None:
+            logger.info(
+                "Skipping CUDA graph capture for %s: %s. This node serves eagerly.",
+                self.submodule_name, blocked,
+            )
+            return
+
         if not hasattr(self.submodule, 'forward_batched'):
             logger.info("Submodule %s does not support batched forward, "
                         "skipping CUDA graph capture", self.submodule_name)
@@ -329,6 +353,7 @@ def _create_persistent_wrappers(
         from mstar.engine.cache_manager import _PlanState
         from mstar.utils.flashinfer_utils import (
             FlashInferDecodeWrapper,
+            FlashInferMLAWrapper,
             FlashInferPrefillWrapper,
         )
 
@@ -336,6 +361,18 @@ def _create_persistent_wrappers(
 
         cfg = self.kv_cache_config
 
+        # Only the FlashInfer MLA path is capturable; absorbed SDPA stays eager.
+        # warmup_and_capture already cancelled capture in that case, so reaching
+        # here with mla_absorb-and-no-kernel means the two decisions drifted —
+        # refuse rather than fall through and build a paged wrapper that cannot
+        # read the 4-D latent cache.
+        blocked = mla_absorb_capture_blocked(cfg, self.device)
+        if blocked is not None:
+            raise RuntimeError(
+                f"_create_persistent_wrappers called for an uncapturable config: {blocked}"
+            )
+        use_mla_kernel = cfg.attention_backend == "mla_absorb"
+
         # Allocate workspace buffer for CUDA graph wrappers.
         # Each (label, slot) gets its own workspace — slots must NOT share
         # workspace because plan() writes scheduling state there and the
@@ -344,7 +381,22 @@ def _create_persistent_wrappers(
         plan_states = {}
         for label in config.labels:
             ws_label = f"{label}_cugraph_slot{slot_idx}"
-            if is_decode:
+            if use_mla_kernel:
+                wrapper = FlashInferMLAWrapper(
+                    workspace_buffer=self.buffer_manager.get(ws_label),
+                    num_heads=cfg.num_qo_heads,
+                    head_dim_ckv=cfg.mla_ckv_dim,
+                    head_dim_kpe=cfg.head_dim - cfg.mla_ckv_dim,
+                    page_size=cfg.page_size,
+                    sm_scale=cfg.softmax_scale,
+                    batch_size=bs,
+                    max_num_pages=cfg.max_num_pages,
+                    max_total_tokens=total_tokens,
+                    device=self.device,
+                    use_cuda_graph=True,
+                    enable_nvtx=self.enable_nvtx,
+                )
+            elif is_decode:
                 wrapper = FlashInferDecodeWrapper(
                     workspace_buffer=self.buffer_manager.get(ws_label),
                     num_qo_heads=cfg.num_qo_heads,
diff --git a/mstar/engine/kv_cache_engine.py b/mstar/engine/kv_cache_engine.py
index df85289b5..920289859 100644
--- a/mstar/engine/kv_cache_engine.py
+++ b/mstar/engine/kv_cache_engine.py
@@ -228,11 +228,18 @@ def load_model(
             )
             num_kv_heads = cfg.num_kv_heads
 
-            kv_cache = torch.zeros(
-                num_layers, max_num_pages, 2,
-                page_size, num_kv_heads, head_dim,
-                dtype=kv_cache_type, device=device,
-            ).contiguous()
+            if cfg.attention_backend == "mla_absorb":
+                # One compressed latent per token: no K/V axis and no KV-head axis.
+                kv_cache = torch.zeros(
+                    num_layers, max_num_pages, page_size, head_dim,
+                    dtype=kv_cache_type, device=device,
+                ).contiguous()
+            else:
+                kv_cache = torch.zeros(
+                    num_layers, max_num_pages, 2,
+                    page_size, num_kv_heads, head_dim,
+                    dtype=kv_cache_type, device=device,
+                ).contiguous()
 
             cpu_page_pool = None
             if cfg.cpu_offload_pages > 0:
diff --git a/mstar/engine/kv_store.py b/mstar/engine/kv_store.py
index 3db7606ad..f9f2c6ff9 100644
--- a/mstar/engine/kv_store.py
+++ b/mstar/engine/kv_store.py
@@ -122,6 +122,10 @@ class KVCacheConfig:
     # FA3 on Hopper; models can pin ``fa2`` when their deployment toolchain
     # cannot compile the Hopper JIT kernels.
     flashinfer_backend: str = "auto"
+    # For "mla_absorb", whose scale is based on qk_head_dim rather than latent width.
+    softmax_scale: float | None = None
+    # For "mla_absorb": split combined latent head_dim into ckv + kpe.
+    mla_ckv_dim: int | None = None
 
     def __post_init__(self):
         if self.num_qo_heads is None:
@@ -728,4 +732,3 @@ def reload_request(self, request_id: str, cpu_pool) -> None:
                 state.seq_len = seq_len
                 state.position_id_start = pos_id
         cpu_pool.sync()
-
diff --git a/mstar/model/components/quantization/__init__.py b/mstar/model/components/quantization/__init__.py
new file mode 100644
index 000000000..2a39f24db
--- /dev/null
+++ b/mstar/model/components/quantization/__init__.py
@@ -0,0 +1,28 @@
+"""Model-agnostic quantization backends for mstar."""
+from mstar.model.components.quantization.base import (
+    FusedMoEQuantizeMethod,
+    process_weights_after_loading,
+)
+from mstar.model.components.quantization.compressed_tensors import (
+    CompressedTensorsQuantConfig,
+    dequant_compressed_tensors_stream,
+    dequantize_weight,
+    pack_int32,
+    unpack_int32,
+)
+from mstar.model.components.quantization.marlin_moe import MarlinMoEMethod
+from mstar.utils.quantization import QuantizationData, QuantizationType, W4A16Data
+
+__all__ = [
+    "CompressedTensorsQuantConfig",
+    "FusedMoEQuantizeMethod",
+    "MarlinMoEMethod",
+    "QuantizationData",
+    "QuantizationType",
+    "W4A16Data",
+    "dequant_compressed_tensors_stream",
+    "dequantize_weight",
+    "pack_int32",
+    "process_weights_after_loading",
+    "unpack_int32",
+]
diff --git a/mstar/model/components/quantization/base.py b/mstar/model/components/quantization/base.py
new file mode 100644
index 000000000..1ff017962
--- /dev/null
+++ b/mstar/model/components/quantization/base.py
@@ -0,0 +1,61 @@
+"""Model-agnostic quantization hooks for post-load kernel layout fixes."""
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+import torch
+from torch import nn
+
+
+@runtime_checkable
+class FusedMoEQuantizeMethod(Protocol):
+    """Kernel backend for a routed-expert (fused-MoE) W4A16 GEMM.
+
+    Implementations own their kernel-format weights after :meth:`prepare` and are
+    otherwise stateless. The MoE block that holds the method is responsible for
+    freeing the source packed params once :meth:`prepare` has consumed them.
+    """
+
+    def prepare(
+        self,
+        w13_packed: torch.Tensor,
+        w13_scale: torch.Tensor,
+        w2_packed: torch.Tensor,
+        w2_scale: torch.Tensor,
+        device: torch.device,
+    ) -> None:
+        """Transform the loaded compressed-tensors packed expert weights into the
+        backend's runtime layout (e.g. Marlin repack), storing them internally.
+        """
+        ...
+
+    def apply(
+        self,
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        *,
+        activation: str = "silu",
+        reduce_results: bool = True,
+    ) -> torch.Tensor:
+        """Run the routed-expert GEMM on the prepared weights.
+
+        Mirrors :func:`mstar.utils.fused_moe.fused_experts`: returns
+        ``(tokens, hidden)`` when ``reduce_results`` else the per-slot
+        ``(tokens, top_k, hidden)`` tensor the TP path all-reduces before folding.
+        """
+        ...
+
+
+def process_weights_after_loading(root: nn.Module, device: torch.device) -> None:
+    """Finalize kernel layouts across a freshly-loaded module tree.
+
+    Call once after ``load_weights`` and before ``eval()``/CUDA-graph capture. Any
+    submodule exposing a ``process_weights_after_loading(device)`` method gets it
+    invoked (e.g. a Marlin MoE block repacks its packed experts + allocates a
+    workspace).
+    """
+    for module in root.modules():
+        hook = getattr(module, "process_weights_after_loading", None)
+        if callable(hook):
+            hook(device)
diff --git a/mstar/model/components/quantization/compressed_tensors.py b/mstar/model/components/quantization/compressed_tensors.py
new file mode 100644
index 000000000..ad96c5ca9
--- /dev/null
+++ b/mstar/model/components/quantization/compressed_tensors.py
@@ -0,0 +1,222 @@
+"""Reader for the compressed-tensors checkpoint format (neuralmagic / vLLM).
+
+Model-agnostic: this is a *serialization format*, not one model's quirk. Kimi-K2.7
+is simply the first mstar model whose checkpoint ships in it.
+
+Packed values are stored low-order-first in int32 containers. Symmetric INT4 is
+offset-binary, so dequant subtracts 8 to match vLLM's ``uint4b8`` layout.
+
+:class:`CompressedTensorsQuantConfig` describes how a checkpoint was *written*;
+:meth:`CompressedTensorsQuantConfig.moe_quant_data` translates that into the
+:class:`~mstar.utils.quantization.QuantizationData` a specific kernel call needs.
+:attr:`~CompressedTensorsQuantConfig.quant_type` is the only place ``num_bits`` is
+mapped to a supported scheme; :meth:`~CompressedTensorsQuantConfig.ensure_kernel_support`
+is how callers reject a width no kernel implements.
+"""
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Iterator
+from dataclasses import dataclass, field
+
+import torch
+
+from mstar.utils.quantization import QuantizationData, QuantizationType, W4A16Data
+
+_PACKED = ".weight_packed"
+_SCALE = ".weight_scale"
+_ZERO_POINT = ".weight_zero_point"
+_SHAPE = ".weight_shape"
+_QUANT_SUFFIXES = (_PACKED, _SCALE, _ZERO_POINT, _SHAPE)
+
+
+@dataclass(frozen=True)
+class CompressedTensorsQuantConfig:
+    """Subset of the compressed-tensors config used by Kimi-K2.7."""
+
+    num_bits: int = 4
+    group_size: int = 32
+    symmetric: bool = True
+    strategy: str = "group"  # "group" | "channel"
+    quant_format: str = "pack-quantized"
+    quant_method: str = "compressed-tensors"
+    ignore: tuple[str, ...] = field(default_factory=tuple)
+
+    @property
+    def pack_factor(self) -> int:
+        return 32 // self.num_bits
+
+    @property
+    def quant_type(self) -> QuantizationType | None:
+        """The scheme mstar's kernels implement for this checkpoint, else ``None``.
+
+        The single ``num_bits`` -> scheme decision. ``None`` means "no kernel for
+        this width"; call :meth:`ensure_kernel_support` to turn that into an error.
+        """
+        return QuantizationType.W4A16 if self.num_bits == 4 else None
+
+    def ensure_kernel_support(self) -> QuantizationType:
+        """Return this checkpoint's scheme, raising if mstar has no kernel for it.
+
+        Call before allocating packed parameters or selecting a backend. The
+        packed-expert kernels hardcode 4-bit nibble extraction, so an INT8
+        checkpoint would otherwise allocate self-consistent shapes, load without
+        complaint, and return wrong numbers — this turns that into a load-time
+        error naming the width the checkpoint declared.
+        """
+        quant_type = self.quant_type
+        if quant_type is None:
+            raise ValueError(
+                f"compressed-tensors checkpoint declares num_bits={self.num_bits}, which "
+                "mstar does not implement (supported: 4-bit / W4A16). The fused-MoE and "
+                "Marlin kernels are INT4-only."
+            )
+        return quant_type
+
+    def moe_quant_data(
+        self,
+        w1_scale: torch.Tensor,
+        w2_scale: torch.Tensor,
+        w1_zp: torch.Tensor | None = None,
+        w2_zp: torch.Tensor | None = None,
+    ) -> QuantizationData:
+        """Companion data for one stacked routed-expert GEMM under this config.
+
+        Call per dispatch rather than caching on the module: the scales are
+        ``nn.Parameter`` s that ``Module._apply`` rebuilds on ``.to(device)``, so
+        binding them at call time keeps the returned object from ever holding a
+        stale tensor.
+        """
+        quant_type = self.ensure_kernel_support()
+        if quant_type is QuantizationType.W4A16:
+            return W4A16Data(
+                w1_scale=w1_scale,
+                w2_scale=w2_scale,
+                group_size=self.group_size,
+                w1_zp=None if self.symmetric else w1_zp,
+                w2_zp=None if self.symmetric else w2_zp,
+            )
+        raise ValueError(f"No routed-expert quantization data for {quant_type}")
+
+    @classmethod
+    def from_hf_config_dict(
+        cls, quant: dict | None
+    ) -> "CompressedTensorsQuantConfig | None":
+        if not quant:
+            return None
+        groups = quant.get("config_groups") or {}
+        weights: dict = {}
+        if groups:
+            first = next(iter(groups.values()))
+            weights = (first or {}).get("weights") or {}
+        strategy = weights.get("strategy", "group")
+        group_size = weights.get("group_size")
+        if group_size is None:
+            group_size = -1 if strategy == "channel" else 32
+        return cls(
+            num_bits=int(weights.get("num_bits", 4)),
+            group_size=int(group_size),
+            symmetric=bool(weights.get("symmetric", True)),
+            strategy=str(strategy),
+            quant_format=str(quant.get("format", "pack-quantized")),
+            quant_method=str(quant.get("quant_method", "compressed-tensors")),
+            ignore=tuple(quant.get("ignore", []) or []),
+        )
+
+def pack_int32(values_unsigned: torch.Tensor, num_bits: int) -> torch.Tensor:
+    """Pack unsigned values along the last axis into low-order-first int32."""
+    pack_factor = 32 // num_bits
+    *lead, n = values_unsigned.shape
+    if n % pack_factor != 0:
+        raise ValueError(f"last dim {n} not divisible by pack_factor {pack_factor}")
+    q = values_unsigned.to(torch.int64).reshape(*lead, n // pack_factor, pack_factor)
+    shifts = torch.arange(pack_factor, device=q.device, dtype=torch.int64) * num_bits
+    packed = (q << shifts).sum(dim=-1)
+    return (packed & 0xFFFFFFFF).to(torch.int32)
+
+
+def unpack_int32(packed: torch.Tensor, num_bits: int) -> torch.Tensor:
+    """Inverse of :func:`pack_int32`; reads int32 as an unsigned bit pattern."""
+    pack_factor = 32 // num_bits
+    mask = (1 << num_bits) - 1
+    p = packed.to(torch.int64) & 0xFFFFFFFF
+    shifts = torch.arange(pack_factor, device=p.device, dtype=torch.int64) * num_bits
+    unpacked = (p.unsqueeze(-1) >> shifts) & mask  # (..., last, pack_factor)
+    *lead, m, _ = unpacked.shape
+    return unpacked.reshape(*lead, m * pack_factor)
+
+
+def dequantize_weight(
+    packed: torch.Tensor,
+    scale: torch.Tensor,
+    *,
+    num_bits: int,
+    group_size: int,
+    symmetric: bool,
+    zero_point: torch.Tensor | None = None,
+    out_dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+    """Dequantize one packed compressed-tensors weight."""
+    nibbles = unpack_int32(packed, num_bits).to(torch.float32)  # (out, in) unsigned
+    out_f, in_f = nibbles.shape
+    gs = in_f if group_size in (-1, None) else group_size
+    if in_f % gs != 0:
+        raise ValueError(f"in_features {in_f} not divisible by group_size {gs}")
+
+    if symmetric:
+        nibbles -= float(1 << (num_bits - 1))  # offset-binary: nibble - bias
+    else:
+        if zero_point is None:
+            raise ValueError("asymmetric quantization requires a zero_point")
+        zp = zero_point.to(torch.float32)
+        if zp.shape[-1] != in_f:  # per-group -> broadcast to per-column
+            zp = zp.repeat_interleave(gs, dim=-1)
+        nibbles -= zp
+
+    s = scale.to(torch.float32)
+    if s.shape[-1] != in_f:  # per-group -> broadcast to per-column
+        s = s.repeat_interleave(gs, dim=-1)
+    return (nibbles * s).to(out_dtype)
+def dequant_compressed_tensors_stream(
+    weights: Iterable[tuple[str, torch.Tensor]],
+    quant_config: CompressedTensorsQuantConfig,
+    out_dtype: torch.dtype = torch.bfloat16,
+    keep_packed: Callable[[str], bool] | None = None,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """Convert quantized sub-key streams to bf16 weights unless ``keep_packed`` matches."""
+    buffers: dict[str, dict[str, torch.Tensor]] = {}
+
+    for name, tensor in weights:
+        suffix = next((s for s in _QUANT_SUFFIXES if name.endswith(s)), None)
+        if suffix is None:
+            yield name, tensor
+            continue
+
+        base = name[: -len(suffix)]
+        if keep_packed is not None and keep_packed(base):
+            yield name, tensor
+            continue
+
+        slot = buffers.setdefault(base, {})
+        slot[suffix] = tensor
+
+        have_core = _PACKED in slot and _SCALE in slot
+        have_zp = quant_config.symmetric or _ZERO_POINT in slot
+        if have_core and have_zp:
+            weight = dequantize_weight(
+                slot[_PACKED],
+                slot[_SCALE],
+                num_bits=quant_config.num_bits,
+                group_size=quant_config.group_size,
+                symmetric=quant_config.symmetric,
+                zero_point=slot.get(_ZERO_POINT),
+                out_dtype=out_dtype,
+            )
+            del buffers[base]
+            yield base + ".weight", weight
+
+    if buffers:
+        missing = {b: sorted(slot) for b, slot in buffers.items()}
+        raise ValueError(
+            f"compressed-tensors stream ended with incomplete quantized tensors "
+            f"(missing weight_packed and/or weight_scale): {missing}"
+        )
diff --git a/mstar/model/components/quantization/marlin_moe.py b/mstar/model/components/quantization/marlin_moe.py
new file mode 100644
index 000000000..fd707d75c
--- /dev/null
+++ b/mstar/model/components/quantization/marlin_moe.py
@@ -0,0 +1,129 @@
+"""Marlin W4A16 backend for routed-expert MoE GEMMs."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from mstar.utils.marlin import ops as marlin_ops
+from mstar.utils.quantization import QuantizationType
+
+if TYPE_CHECKING:
+    from mstar.model.components.quantization.compressed_tensors import (
+        CompressedTensorsQuantConfig,
+    )
+
+
+class MarlinMoEMethod:
+    """Marlin routed-expert GEMM backend (symmetric INT4, group-wise).
+
+    A kernel *backend*, not a quantization descriptor: stateful, and it owns
+    Marlin-layout weights after :meth:`prepare` repacks the loaded packed experts
+    (the source packed params can then be freed by the owning block).
+    :meth:`apply` runs the two Marlin GEMMs.
+    """
+
+    def __init__(self, *, num_bits: int = 4, group_size: int = 32) -> None:
+        if num_bits != 4:
+            raise ValueError(f"MarlinMoEMethod supports INT4 only, got num_bits={num_bits}")
+        self.num_bits = num_bits
+        self.group_size = group_size
+        self.pack_factor = 32 // num_bits
+        # Populated by prepare():
+        self.w13_qweight: torch.Tensor | None = None
+        self.w2_qweight: torch.Tensor | None = None
+        self.w13_scale: torch.Tensor | None = None
+        self.w2_scale: torch.Tensor | None = None
+        self.workspace: torch.Tensor | None = None
+
+    @classmethod
+    def from_quant_config(
+        cls, quant_config: "CompressedTensorsQuantConfig"
+    ) -> "MarlinMoEMethod":
+        """Build from the checkpoint descriptor, so the bit width is read forward
+        from ``num_bits`` rather than reconstructed backwards from a pack factor."""
+        quant_type = quant_config.ensure_kernel_support()
+        if quant_type is not QuantizationType.W4A16:
+            raise ValueError(
+                f"MarlinMoEMethod supports {QuantizationType.W4A16} only, got {quant_type}"
+            )
+        return cls(num_bits=quant_config.num_bits, group_size=quant_config.group_size)
+
+    @property
+    def quant_type(self) -> QuantizationType:
+        return QuantizationType.W4A16
+
+    def prepare(
+        self,
+        w13_packed: torch.Tensor,
+        w13_scale: torch.Tensor,
+        w2_packed: torch.Tensor,
+        w2_scale: torch.Tensor,
+        device: torch.device,
+    ) -> None:
+        pf, gs = self.pack_factor, self.group_size
+        E, two_inter, hidden_over_pack = w13_packed.shape
+        hidden = hidden_over_pack * pf
+        _, w2_hidden, inter_over_pack = w2_packed.shape
+        inter = inter_over_pack * pf
+        assert w2_hidden == hidden, f"w2 dim1 {w2_hidden} != hidden {hidden}"
+
+        # gate_up: (E, 2*inter, hidden/pack) -> (E, hidden/pack, 2*inter) -> marlin.
+        w13_t = w13_packed.transpose(1, 2).contiguous()
+        self.w13_qweight = marlin_ops.gptq_marlin_moe_repack(
+            w13_t, size_k=hidden, size_n=two_inter, num_bits=self.num_bits
+        )
+        w13_s = w13_scale.transpose(1, 2).contiguous()  # (E, hidden/gs, 2*inter)
+        self.w13_scale = marlin_ops.marlin_moe_permute_scales(
+            w13_s, size_k=hidden, size_n=two_inter, group_size=gs
+        )
+
+        # down: (E, hidden, inter/pack) -> (E, inter/pack, hidden) -> marlin.
+        w2_t = w2_packed.transpose(1, 2).contiguous()
+        self.w2_qweight = marlin_ops.gptq_marlin_moe_repack(
+            w2_t, size_k=inter, size_n=hidden, num_bits=self.num_bits
+        )
+        w2_s = w2_scale.transpose(1, 2).contiguous()  # (E, inter/gs, hidden)
+        self.w2_scale = marlin_ops.marlin_moe_permute_scales(
+            w2_s, size_k=inter, size_n=hidden, group_size=gs
+        )
+
+        self.workspace = marlin_ops.marlin_make_workspace(torch.device(device))
+
+    def apply(
+        self,
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        *,
+        activation: str = "silu",
+        reduce_results: bool = True,
+    ) -> torch.Tensor:
+        assert self.w13_qweight is not None, "MarlinMoEMethod.apply before prepare()"
+        return marlin_ops.fused_marlin_moe(
+            x,
+            self.w13_qweight,
+            self.w2_qweight,
+            self.w13_scale,
+            self.w2_scale,
+            topk_weights,
+            topk_ids,
+            self.workspace,
+            activation=activation,
+            reduce_results=reduce_results,
+        )
+
+    @staticmethod
+    def shapes_are_legal(hidden: int, shard_inter: int, group_size: int) -> bool:
+        """Marlin needs n%64 and k%128 for both expert GEMMs."""
+        if group_size not in (-1, 32, 64, 128):
+            return False
+        checks = [
+            hidden % 128 == 0,            # gate_up K
+            (2 * shard_inter) % 64 == 0,  # gate_up N
+            shard_inter % 128 == 0,       # down K
+            hidden % 64 == 0,             # down N
+        ]
+        if group_size != -1:
+            checks += [hidden % group_size == 0, shard_inter % group_size == 0]
+        return all(checks)
diff --git a/mstar/model/kimi_k2_7/__init__.py b/mstar/model/kimi_k2_7/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/mstar/model/kimi_k2_7/_testing.py b/mstar/model/kimi_k2_7/_testing.py
new file mode 100644
index 000000000..bcc741a0d
--- /dev/null
+++ b/mstar/model/kimi_k2_7/_testing.py
@@ -0,0 +1,51 @@
+"""Test-support helpers for the Kimi-K2.7 compressed-tensors path.
+
+NOT part of the serving path. These helpers exist only to let the test suite
+fabricate a synthetic quantized checkpoint and its exact bf16 reference, so a
+golden can assert the real load path (:mod:`mstar.model.components.quantization`
++ ``weight_loader.py``) reproduces the reference bit-for-bit. Nothing here is
+imported by the model or the loader at serve time.
+
+The real load-path primitives (``pack_int32`` / ``unpack_int32`` /
+``dequantize_weight`` / ``dequant_compressed_tensors_stream`` /
+``CompressedTensorsQuantConfig``) live in
+:mod:`mstar.model.components.quantization.compressed_tensors`; this module builds
+on them.
+"""
+from __future__ import annotations
+
+import torch
+
+from mstar.model.components.quantization import dequantize_weight, pack_int32
+
+
+def fake_quantize_weight(
+    weight: torch.Tensor,
+    *,
+    num_bits: int,
+    group_size: int,
+    symmetric: bool = True,
+    scale_dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Return packed weights, stored scales, and the exact dequantized reference."""
+    if not symmetric:
+        raise NotImplementedError("fake_quantize_weight: only symmetric is implemented")
+    out_f, in_f = weight.shape
+    gs = in_f if group_size in (-1, None) else group_size
+    if in_f % gs != 0:
+        raise ValueError(f"in_features {in_f} not divisible by group_size {gs}")
+
+    qmax = (1 << (num_bits - 1)) - 1  # 7 for INT4
+    qmin = -(1 << (num_bits - 1))     # -8
+    w = weight.to(torch.float32).reshape(out_f, in_f // gs, gs)
+    scale = w.abs().amax(dim=-1, keepdim=True) / qmax
+    scale = torch.where(scale == 0, torch.ones_like(scale), scale)
+    q = torch.round(w / scale).clamp(qmin, qmax)  # signed [qmin, qmax], fp32 scale
+
+    q_unsigned = (q + float(1 << (num_bits - 1))).reshape(out_f, in_f).to(torch.int64)
+    packed = pack_int32(q_unsigned, num_bits)
+    scale_2d = scale.squeeze(-1).to(scale_dtype)
+    dequant = dequantize_weight(
+        packed, scale_2d, num_bits=num_bits, group_size=group_size, symmetric=symmetric,
+    )
+    return packed, scale_2d, dequant
diff --git a/mstar/model/kimi_k2_7/components/__init__.py b/mstar/model/kimi_k2_7/components/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/mstar/model/kimi_k2_7/components/attention.py b/mstar/model/kimi_k2_7/components/attention.py
new file mode 100644
index 000000000..e4c44c8b1
--- /dev/null
+++ b/mstar/model/kimi_k2_7/components/attention.py
@@ -0,0 +1,191 @@
+"""Kimi-K2.7 MLA attention with absorbed and naive fallback paths."""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from mstar.distributed.communication import CommGroup
+from mstar.engine.cache_manager import BatchedCacheManager
+from mstar.model.components.distributed import ColumnParallelLinear, RowParallelLinear
+from mstar.model.components.norm import RMSNorm
+from mstar.model.kimi_k2_7.components.rope import KimiYarnRotaryEmbedding, yarn_get_mscale
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+
+class KimiMLAAttention(nn.Module):
+    def __init__(self, config: KimiK2Config, comm_group: CommGroup | None = None) -> None:
+        super().__init__()
+        if comm_group is None:
+            comm_group = CommGroup.trivial()
+
+        self.tp_size = comm_group.world_size
+        self.total_num_heads = config.num_attention_heads
+        if self.total_num_heads % self.tp_size != 0:
+            raise ValueError(
+                f"num_attention_heads={self.total_num_heads} is not divisible by "
+                f"tp_size={self.tp_size}"
+            )
+        self.num_heads = self.total_num_heads // self.tp_size
+        self.qk_nope_head_dim = config.qk_nope_head_dim
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.qk_head_dim = config.qk_head_dim
+        self.v_head_dim = config.v_head_dim
+        self.kv_lora_rank = config.kv_lora_rank
+        self.padded_head_dim = config.padded_head_dim
+        h = self.total_num_heads
+
+        self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False)
+        self.q_a_layernorm = RMSNorm(config.q_lora_rank, eps=config.rms_norm_eps)
+        self.q_b_proj = ColumnParallelLinear(
+            comm_group, config.q_lora_rank, h * self.qk_head_dim, bias=False)
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            config.hidden_size, config.kv_lora_rank + config.qk_rope_head_dim, bias=False)
+        self.kv_a_layernorm = RMSNorm(config.kv_lora_rank, eps=config.rms_norm_eps)
+        self.kv_b_proj = ColumnParallelLinear(
+            comm_group, config.kv_lora_rank,
+            h * (config.qk_nope_head_dim + config.v_head_dim), bias=False)
+
+        self.o_proj = RowParallelLinear(
+            comm_group, h * config.v_head_dim, config.hidden_size,
+            bias=False, input_is_parallel=True, reduce_results=True)
+
+        rope = config.rope_scaling
+        self.rotary = KimiYarnRotaryEmbedding(
+            rotary_dim=config.qk_rope_head_dim,
+            base=config.rope_theta,
+            factor=rope["factor"],
+            original_max_position_embeddings=rope["original_max_position_embeddings"],
+            beta_fast=rope.get("beta_fast", 32),
+            beta_slow=rope.get("beta_slow", 1),
+            mscale=rope.get("mscale", 1.0),
+            mscale_all_dim=rope.get("mscale_all_dim", 0.0),
+        )
+        # run_attention uses 1/sqrt(padded_head_dim); fold DeepSeek's intended
+        # qk_head_dim**-0.5 * mscale**2 scale into q on the padded path.
+        mscale = yarn_get_mscale(rope["factor"], rope.get("mscale_all_dim", 0.0))
+        self.softmax_scale_boost = (
+            mscale * mscale * math.sqrt(self.padded_head_dim / self.qk_head_dim)
+        )
+        self.softmax_scale = self.qk_head_dim ** -0.5 * mscale * mscale
+
+        self.mla_absorb = config.mla_absorb
+        if self.mla_absorb:
+            self.register_buffer("w_kc", None, persistent=False)  # (H_local, Dnope, L)
+            self.register_buffer("w_vc", None, persistent=False)  # (H_local, Dv,    L)
+            self.register_buffer("fused_qkv_a_proj_weight", None, persistent=False)  # (q_lora+L+Drope, hidden)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_handle: BatchedCacheManager,
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.mla_absorb:
+            return self._forward_absorbed(hidden_states, cache_handle, position_ids)
+
+        # Naive/materialized MLA — the reduced-capability fallback (see the
+        # ``mla_absorb`` note in config.py), not a co-equal path. It projects the
+        # latent up to full per-head K/V and uses the standard paged MHA
+        # interface, so it works on any dims/hardware, at the cost of caching
+        # H*(Dnope+Dv) per token instead of one shared latent. Kept for parity
+        # tests and for configs the absorbed path's kernels cannot serve.
+        #
+        # Two numerical quirks follow from reusing that MHA interface: q/k/v are
+        # padded to ``padded_head_dim`` (FlashInfer paged kernels only accept
+        # 64/128/256), and because ``run_attention`` then scales by
+        # 1/sqrt(padded_head_dim), DeepSeek's intended qk_head_dim**-0.5 * mscale**2
+        # is folded into q via ``softmax_scale_boost`` below.
+        num_tokens = hidden_states.shape[0]
+        h = self.num_heads
+
+        q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
+        q = q.view(num_tokens, h, self.qk_head_dim)
+        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+        latent = self.kv_a_proj_with_mqa(hidden_states)  # (T, L + Drope)
+        kv_a, k_pe = latent.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+        kv = self.kv_b_proj(self.kv_a_layernorm(kv_a))
+        kv = kv.view(num_tokens, h, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+        k_pe = k_pe.view(num_tokens, 1, self.qk_rope_head_dim)  # shared MQA rope key
+
+        q_pe, k_pe = self.rotary(position_ids, q_pe, k_pe)
+
+        q = torch.cat([q_nope, q_pe], dim=-1)  # (T, H, Dqk)
+        k_pe = k_pe.expand(num_tokens, h, self.qk_rope_head_dim)
+        k = torch.cat([k_nope, k_pe], dim=-1)  # (T, H, Dqk)
+
+        qk_pad = self.padded_head_dim - self.qk_head_dim
+        q = F.pad(q, [0, qk_pad])  # (T, H, Dpad)
+        k = F.pad(k, [0, qk_pad])  # (T, H, Dpad)
+        v = F.pad(v, [0, self.padded_head_dim - self.v_head_dim])  # (T, H, Dpad)
+
+        q = q * self.softmax_scale_boost
+        attn = cache_handle.run_attention(q=q, k=k, v=v)  # (T, H, Dpad)
+        attn = attn[..., : self.v_head_dim].reshape(num_tokens, h * self.v_head_dim)
+        return self.o_proj(attn)
+
+    def _forward_absorbed(
+        self,
+        hidden_states: torch.Tensor,
+        cache_handle: BatchedCacheManager,
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run MLA over the compressed latent cache after folding kv_b into Q/O."""
+        if self.w_kc is None or self.w_vc is None:
+            raise RuntimeError(
+                "mla_absorb forward requires process_weights_after_loading() to "
+                "have built w_kc/w_vc from kv_b_proj first"
+            )
+        if self.fused_qkv_a_proj_weight is None:
+            raise RuntimeError(
+                "mla_absorb forward requires process_weights_after_loading() to "
+                "have built fused_qkv_a_proj_weight first"
+            )
+        num_tokens = hidden_states.shape[0]
+        h = self.num_heads
+
+        fused = F.linear(hidden_states, self.fused_qkv_a_proj_weight)
+        q_c, kv_a, k_pe = fused.split(
+            [self.q_a_proj.out_features, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+        # FlashInfer RMSNorm needs 64-byte input alignment; decode split views can
+        # be contiguous yet start at an unaligned offset, so clone kv_a.
+        q_c = q_c.contiguous()
+        kv_a = kv_a.clone(memory_format=torch.contiguous_format)
+
+        q = self.q_b_proj(self.q_a_layernorm(q_c))
+        q = q.view(num_tokens, h, self.qk_head_dim)
+        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+        kv_c = self.kv_a_layernorm(kv_a).view(num_tokens, 1, self.kv_lora_rank)  # (T,1,L)
+        k_pe = k_pe.view(num_tokens, 1, self.qk_rope_head_dim)  # (T,1,Drope) shared MQA key
+
+        q_pe, k_pe = self.rotary(position_ids, q_pe, k_pe)
+
+        q_nope = torch.einsum("thd,hdl->thl", q_nope, self.w_kc)
+
+        attn_latent = cache_handle.run_attention_mla(
+            q_nope=q_nope, q_pe=q_pe, kv_c=kv_c, k_pe=k_pe)
+
+        out = torch.einsum("thl,hdl->thd", attn_latent, self.w_vc)
+        return self.o_proj(out.reshape(num_tokens, h * self.v_head_dim))
+
+    def process_weights_after_loading(self, device: torch.device | str | None = None) -> None:
+        """Build absorbed Q/O projections from the local-head ``kv_b_proj`` shard."""
+        if not self.mla_absorb:
+            return
+        del device  # protocol arg; kv_b_proj.weight already carries the right device
+        w = self.kv_b_proj.weight  # (H_local*(Dnope+Dv), L)
+        h, d_nope, d_v, latent = (
+            self.num_heads, self.qk_nope_head_dim, self.v_head_dim, self.kv_lora_rank)
+        w = w.view(h, d_nope + d_v, latent)
+        w_kc, w_vc = w.split([d_nope, d_v], dim=1)
+        self.w_kc = w_kc.contiguous()  # (H_local, Dnope, L)
+        self.w_vc = w_vc.contiguous()  # (H_local, Dv,    L)
+
+        self.fused_qkv_a_proj_weight = torch.cat(
+            [self.q_a_proj.weight, self.kv_a_proj_with_mqa.weight], dim=0).contiguous()
diff --git a/mstar/model/kimi_k2_7/components/causal_lm.py b/mstar/model/kimi_k2_7/components/causal_lm.py
new file mode 100644
index 000000000..90e190fdd
--- /dev/null
+++ b/mstar/model/kimi_k2_7/components/causal_lm.py
@@ -0,0 +1,78 @@
+"""Assembled Kimi-K2.7 text backbone."""
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from mstar.distributed.communication import CommGroup
+from mstar.engine.cache_manager import BatchedCacheManager
+from mstar.model.kimi_k2_7.components.decoder_layer import KimiDecoderLayer
+from mstar.model.kimi_k2_7.components.language_model import (
+    build_embedding,
+    build_lm_head,
+    build_rmsnorm,
+)
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+
+class KimiLanguageModel(nn.Module):
+    def __init__(
+        self, config: KimiK2Config, comm_group: CommGroup | None = None
+    ) -> None:
+        super().__init__()
+        self.embed_tokens = build_embedding(config, comm_group=comm_group)
+        self.layers = nn.ModuleList(
+            [
+                KimiDecoderLayer(config, layer_idx, comm_group=comm_group)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = build_rmsnorm(config)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        cache_handle: BatchedCacheManager,
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden_states = self.embed_tokens(input_ids)
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            cache_handle.set_layer_idx(layer_idx)
+            hidden_states = decoder_layer(
+                hidden_states, cache_handle, position_ids
+            )
+        cache_handle.advance_seq_lens()
+        return self.norm(hidden_states)
+
+
+class KimiForCausalLM(nn.Module):
+    def __init__(
+        self, config: KimiK2Config, comm_group: CommGroup | None = None
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.model = KimiLanguageModel(config, comm_group=comm_group)
+        self.lm_head = build_lm_head(config, comm_group=comm_group)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        cache_handle: BatchedCacheManager,
+        position_ids: torch.Tensor,
+        **kwargs,
+    ) -> torch.Tensor:
+        hidden_states = self.model(input_ids, cache_handle, position_ids)
+        return self.lm_head(hidden_states)
+
+    def load_weights(self, weights, **kwargs) -> set[str]:
+        from mstar.model.kimi_k2_7.weight_loader import load_kimi_hf_weights
+
+        packed_experts = (
+            self.config.quantization_config is not None
+            and self.config.moe_in_kernel_dequant
+        )
+        return load_kimi_hf_weights(
+            self, weights, self.config.n_routed_experts,
+            quant_config=self.config.quantization_config,
+            packed_experts=packed_experts,
+        )
diff --git a/mstar/model/kimi_k2_7/components/decoder_layer.py b/mstar/model/kimi_k2_7/components/decoder_layer.py
new file mode 100644
index 000000000..9528f7860
--- /dev/null
+++ b/mstar/model/kimi_k2_7/components/decoder_layer.py
@@ -0,0 +1,45 @@
+"""Kimi-K2.7 decoder layer with MLA position ids threaded through attention."""
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from mstar.distributed.communication import CommGroup
+from mstar.engine.cache_manager import BatchedCacheManager
+from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention
+from mstar.model.kimi_k2_7.components.language_model import (
+    build_mlp_for_layer,
+    build_rmsnorm,
+)
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+
+class KimiDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: KimiK2Config,
+        layer_idx: int,
+        comm_group: CommGroup | None = None,
+    ) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.self_attn = KimiMLAAttention(config, comm_group=comm_group)
+        self.mlp = build_mlp_for_layer(config, layer_idx, comm_group=comm_group)
+        self.input_layernorm = build_rmsnorm(config)
+        self.post_attention_layernorm = build_rmsnorm(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_handle: BatchedCacheManager,
+        position_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, cache_handle, position_ids)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        return residual + hidden_states
diff --git a/mstar/model/kimi_k2_7/components/language_model.py b/mstar/model/kimi_k2_7/components/language_model.py
new file mode 100644
index 000000000..a458cfc0d
--- /dev/null
+++ b/mstar/model/kimi_k2_7/components/language_model.py
@@ -0,0 +1,72 @@
+"""Kimi-K2.7 language-model builders over existing mstar primitives."""
+from __future__ import annotations
+
+from mstar.distributed.communication import CommGroup
+from mstar.model.components import RMSNorm
+from mstar.model.components.distributed import (
+    ColumnParallelLinear,
+    ParallelGatedMLP,
+    VocabParallelEmbedding,
+)
+from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+
+def build_embedding(
+    config: KimiK2Config, comm_group: CommGroup | None = None
+) -> VocabParallelEmbedding:
+    return VocabParallelEmbedding(
+        num_embeddings=config.vocab_size,
+        embedding_dim=config.hidden_size,
+        comm_group=comm_group,
+        padding_idx=config.pad_token_id,
+    )
+
+
+def build_lm_head(
+    config: KimiK2Config, comm_group: CommGroup | None = None
+) -> ColumnParallelLinear:
+    return ColumnParallelLinear(
+        comm_group or CommGroup.trivial(),
+        input_size=config.hidden_size,
+        output_size=config.vocab_size,
+        bias=False,
+        gather_output=True,
+    )
+
+
+def build_rmsnorm(config: KimiK2Config) -> RMSNorm:
+    return RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+
+def build_dense_mlp(
+    config: KimiK2Config, comm_group: CommGroup | None = None
+) -> ParallelGatedMLP:
+    return ParallelGatedMLP(
+        hidden_size=config.hidden_size,
+        intermediate_size=config.intermediate_size,
+        comm_group=comm_group,
+        activation=config.hidden_act,
+        bias=False,
+    )
+
+
+def is_moe_layer(config: KimiK2Config, layer_idx: int) -> bool:
+    return (
+        layer_idx >= config.first_k_dense_replace
+        and layer_idx % config.moe_layer_freq == 0
+    )
+
+
+def build_moe_block(
+    config: KimiK2Config, comm_group: CommGroup | None = None
+) -> KimiSparseMoeBlock:
+    return KimiSparseMoeBlock(config, comm_group=comm_group)
+
+
+def build_mlp_for_layer(
+    config: KimiK2Config, layer_idx: int, comm_group: CommGroup | None = None
+):
+    if is_moe_layer(config, layer_idx):
+        return build_moe_block(config, comm_group=comm_group)
+    return build_dense_mlp(config, comm_group=comm_group)
diff --git a/mstar/model/kimi_k2_7/components/moe.py b/mstar/model/kimi_k2_7/components/moe.py
new file mode 100644
index 000000000..f7a195735
--- /dev/null
+++ b/mstar/model/kimi_k2_7/components/moe.py
@@ -0,0 +1,432 @@
+"""Kimi-K2.7 fine-grained MoE: sigmoid router plus ungated shared expert."""
+from __future__ import annotations
+
+import logging
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from mstar.distributed.communication import CommGroup
+from mstar.distributed.utils import divide
+from mstar.model.components.distributed import ParallelGatedMLP
+from mstar.model.components.moe import (
+    _dispatch,
+    _down_proj_weight_loader,
+    _gate_up_weight_loader,
+)
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+logger = logging.getLogger(__name__)
+
+_BACKEND_LOGGED = False
+
+
+def _gate_up_packed_loader(
+    tp_rank: int, tp_size: int, full_inter: int,
+    param: nn.Parameter, loaded_weight: torch.Tensor,
+    loaded_shard_id: str | int | None = None,
+):
+    assert loaded_shard_id is not None
+    kind, expert_str = loaded_shard_id.split(":")
+    expert_idx = int(expert_str)
+    shard_inter = divide(full_inter, tp_size)
+    start = tp_rank * shard_inter
+    tp_slice = loaded_weight[start:start + shard_inter, :]
+    if kind == "gate":
+        param.data[expert_idx, :shard_inter, :] = tp_slice
+    else:
+        param.data[expert_idx, shard_inter:, :] = tp_slice
+
+
+def _down_packed_loader(
+    tp_rank: int, tp_size: int, full_inter: int, divisor: int,
+    param: nn.Parameter, loaded_weight: torch.Tensor,
+    loaded_shard_id: str | int | None = None,
+):
+    assert loaded_shard_id is not None
+    expert_idx = int(str(loaded_shard_id).split(":")[1])
+    shard_inter = divide(full_inter, tp_size)
+    span = divide(shard_inter, divisor)
+    start = tp_rank * span
+    param.data[expert_idx, :, :] = loaded_weight[:, start:start + span]
+
+
+class KimiMoEGate(nn.Module):
+    """DeepSeek-V3 group-limited sigmoid router with selection-only bias."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        n_routed_experts: int,
+        num_experts_per_tok: int,
+        n_group: int,
+        topk_group: int,
+        routed_scaling_factor: float,
+        scoring_func: str = "sigmoid",
+        topk_method: str = "noaux_tc",
+        norm_topk_prob: bool = True,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.n_routed_experts = n_routed_experts
+        self.top_k = num_experts_per_tok
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.routed_scaling_factor = routed_scaling_factor
+        self.scoring_func = scoring_func
+        self.topk_method = topk_method
+        self.norm_topk_prob = norm_topk_prob
+
+        self.weight = nn.Parameter(torch.zeros(n_routed_experts, hidden_size))
+        if topk_method == "noaux_tc":
+            # Per-expert selection bias; fp32, added to scores for group/top-k
+            # selection but never to the combine weights.
+            self.e_score_correction_bias = nn.Parameter(
+                torch.zeros(n_routed_experts, dtype=torch.float32)
+            )
+        else:
+            self.register_parameter("e_score_correction_bias", None)
+
+    def forward(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        h = hidden_states.reshape(-1, self.hidden_size).float()
+        gating = F.linear(h, self.weight.float())  # (T, E)
+
+        if self.scoring_func == "sigmoid":
+            scores = gating.sigmoid()
+        elif self.scoring_func == "softmax":
+            scores = gating.softmax(dim=-1)
+        else:
+            raise ValueError(f"Unsupported scoring_func: {self.scoring_func!r}")
+
+        num_token = scores.shape[0]
+        if self.e_score_correction_bias is not None:
+            # noaux_tc: bias-added scores drive group + expert *selection*; the
+            # raw sigmoid scores drive the combine weights.
+            original_scores = scores
+            scores = scores + self.e_score_correction_bias.unsqueeze(0)
+            group_scores = (
+                scores.view(num_token, self.n_group, -1)
+                .topk(2, dim=-1)[0]
+                .sum(dim=-1)
+            )  # (T, n_group)
+        else:
+            original_scores = scores
+            group_scores = scores.view(num_token, self.n_group, -1).max(dim=-1).values
+
+        group_idx = torch.topk(
+            group_scores, k=self.topk_group, dim=-1, sorted=False
+        )[1]  # (T, topk_group)
+        group_mask = torch.zeros_like(group_scores)  # (T, n_group)
+        group_mask.scatter_(1, group_idx, 1)
+        score_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(num_token, self.n_group, scores.shape[-1] // self.n_group)
+            .reshape(num_token, -1)
+        )  # (T, E)
+        masked_scores = scores.masked_fill(~score_mask.bool(), float("-inf"))
+
+        if self.e_score_correction_bias is not None:
+            topk_ids = torch.topk(masked_scores, k=self.top_k, dim=-1, sorted=False)[1]
+            topk_weights = original_scores.gather(1, topk_ids)
+        else:
+            topk_weights, topk_ids = torch.topk(
+                masked_scores, k=self.top_k, dim=-1, sorted=False
+            )
+
+        if self.norm_topk_prob:
+            topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+        if self.routed_scaling_factor != 1.0:
+            topk_weights = topk_weights * self.routed_scaling_factor
+
+        return topk_weights, topk_ids
+
+
+class KimiSparseMoeBlock(nn.Module):
+    """DeepSeek-V3 MoE block: routed experts + ungated shared expert."""
+
+    def __init__(
+        self, config: KimiK2Config, comm_group: CommGroup | None = None
+    ) -> None:
+        super().__init__()
+        if comm_group is None:
+            comm_group = CommGroup.trivial()
+        self.comm_group = comm_group
+        self.tp_size = comm_group.world_size
+        self.tp_rank = comm_group.rank
+        self.hidden_size = config.hidden_size
+        self.num_experts = config.n_routed_experts
+        self.moe_intermediate_size = config.moe_intermediate_size
+        shard_inter = divide(config.moe_intermediate_size, self.tp_size)
+
+        self.packed_experts = (
+            config.quantization_config is not None and config.moe_in_kernel_dequant
+        )
+        # Hold the checkpoint descriptor rather than copying its fields out: it is
+        # the single source for group_size / pack_factor / symmetric, and it builds
+        # the per-dispatch QuantizationData.
+        self.quant_config = config.quantization_config if self.packed_experts else None
+        self.quant_kernel = getattr(config, "quant_kernel", "auto")
+        self._marlin_method = None
+        self._use_marlin = False
+
+        self.gate = KimiMoEGate(
+            hidden_size=config.hidden_size,
+            n_routed_experts=config.n_routed_experts,
+            num_experts_per_tok=config.num_experts_per_tok,
+            n_group=config.n_group,
+            topk_group=config.topk_group,
+            routed_scaling_factor=config.routed_scaling_factor,
+            scoring_func=config.scoring_func,
+            topk_method=config.topk_method,
+            norm_topk_prob=config.norm_topk_prob,
+        )
+
+        self.experts = nn.Module()
+        if self.packed_experts:
+            qc = self.quant_config
+            qc.ensure_kernel_support()
+            hidden, gs, pf = config.hidden_size, qc.group_size, qc.pack_factor
+            assert hidden % pf == 0 and hidden % gs == 0, (
+                f"hidden {hidden} must divide pack_factor {pf} and group_size {gs}"
+            )
+            assert shard_inter % pf == 0 and shard_inter % gs == 0, (
+                f"shard_inter {shard_inter} must divide pack_factor {pf} and group_size {gs}"
+            )
+            E = config.n_routed_experts
+            self.experts.gate_up_proj_packed = nn.Parameter(
+                torch.empty(E, 2 * shard_inter, hidden // pf, dtype=torch.int32),
+                requires_grad=False,
+            )
+            self.experts.gate_up_proj_scale = nn.Parameter(
+                torch.empty(E, 2 * shard_inter, hidden // gs, dtype=torch.bfloat16),
+                requires_grad=False,
+            )
+            self.experts.down_proj_packed = nn.Parameter(
+                torch.empty(E, hidden, shard_inter // pf, dtype=torch.int32),
+                requires_grad=False,
+            )
+            self.experts.down_proj_scale = nn.Parameter(
+                torch.empty(E, hidden, shard_inter // gs, dtype=torch.bfloat16),
+                requires_grad=False,
+            )
+        else:
+            self.experts.gate_up_proj = nn.Parameter(
+                torch.empty(
+                    config.n_routed_experts,
+                    2 * shard_inter,
+                    config.hidden_size,
+                )
+            )
+            self.experts.down_proj = nn.Parameter(
+                torch.empty(
+                    config.n_routed_experts,
+                    config.hidden_size,
+                    shard_inter,
+                )
+            )
+        self._attach_expert_weight_loaders()
+
+        self.shared_expert = ParallelGatedMLP(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.moe_intermediate_size * config.n_shared_experts,
+            comm_group=comm_group,
+            activation=config.hidden_act,
+            bias=False,
+        )
+
+    def _attach_expert_weight_loaders(self) -> None:
+        """Reattach per-shard loaders after ``_apply`` rebuilds parameters."""
+        from functools import partial
+
+        full_inter = self.moe_intermediate_size
+        if self.packed_experts:
+            pf, gs = self.quant_config.pack_factor, self.quant_config.group_size
+            self.experts.gate_up_proj_packed.weight_loader = partial(
+                _gate_up_packed_loader, self.tp_rank, self.tp_size, full_inter,
+            )
+            self.experts.gate_up_proj_scale.weight_loader = partial(
+                _gate_up_packed_loader, self.tp_rank, self.tp_size, full_inter,
+            )
+            self.experts.down_proj_packed.weight_loader = partial(
+                _down_packed_loader, self.tp_rank, self.tp_size, full_inter, pf,
+            )
+            self.experts.down_proj_scale.weight_loader = partial(
+                _down_packed_loader, self.tp_rank, self.tp_size, full_inter, gs,
+            )
+        else:
+            self.experts.gate_up_proj.weight_loader = partial(
+                _gate_up_weight_loader, self.tp_rank, self.tp_size, full_inter,
+            )
+            self.experts.down_proj.weight_loader = partial(
+                _down_proj_weight_loader, self.tp_rank, self.tp_size, full_inter,
+            )
+
+    def _apply(self, fn, recurse=True):
+        result = super()._apply(fn, recurse=recurse)
+        self._attach_expert_weight_loaders()
+        return result
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_shape = hidden_states.shape
+        flat = hidden_states.view(-1, self.hidden_size).contiguous()
+
+        topk_weights, topk_ids = self.gate(flat)
+        if self._use_marlin:
+            # Marlin's GEMM takes fp32 combine weights — pass them BEFORE the bf16
+            # cast the other paths need. Otherwise identical TP story.
+            routed = self._dispatch_marlin(flat, topk_weights, topk_ids)
+        else:
+            topk_weights = topk_weights.to(flat.dtype)
+            if self.packed_experts:
+                routed = self._dispatch_packed_experts(flat, topk_weights, topk_ids)
+            elif self.tp_size == 1:
+                routed = _dispatch(
+                    flat,
+                    self.experts.gate_up_proj,
+                    self.experts.down_proj,
+                    self.num_experts,
+                    topk_ids,
+                    topk_weights,
+                )
+            else:
+                routed = self._dispatch_tp(flat, topk_weights, topk_ids)
+        shared = self.shared_expert(flat)
+        return (routed + shared).view(input_shape)
+
+    def _dispatch_packed_experts(
+        self,
+        flat: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        from mstar.utils.fused_moe import fused_experts, moe_sum_reduce_triton
+
+        reduce = self.tp_size == 1
+        # Built per dispatch, not cached: Module._apply rebuilds these Parameters
+        # on .to(device), so a cached descriptor could hold a stale tensor.
+        quant = self.quant_config.moe_quant_data(
+            self.experts.gate_up_proj_scale, self.experts.down_proj_scale,
+        )
+        out = fused_experts(
+            flat,
+            self.experts.gate_up_proj_packed,
+            self.experts.down_proj_packed,
+            topk_weights,
+            topk_ids,
+            quant=quant,
+            reduce_results=reduce,
+        )
+        if reduce:
+            return out
+        self.comm_group.all_reduce(out)
+        output = torch.empty_like(flat)
+        moe_sum_reduce_triton(out, output, routed_scaling_factor=1.0)
+        return output
+
+    def process_weights_after_loading(self, device) -> None:
+        """Resolve Triton-vs-Marlin after weights land on the real device."""
+        if not self.packed_experts:
+            return
+        from mstar.model.components.quantization import MarlinMoEMethod
+        from mstar.utils.marlin import is_marlin_available
+
+        qc = self.quant_config
+        dev = torch.device(device)
+        shard_inter = divide(self.moe_intermediate_size, self.tp_size)
+        legal_shapes = MarlinMoEMethod.shapes_are_legal(
+            self.hidden_size, shard_inter, qc.group_size
+        )
+        eligible = (
+            self.quant_kernel != "triton"
+            and dev.type == "cuda"
+            and torch.cuda.get_device_capability(dev) >= (8, 0)
+            and qc.symmetric
+            and legal_shapes
+            and is_marlin_available()
+        )
+        if self.quant_kernel == "marlin" and not eligible:
+            raise RuntimeError(
+                "quant_kernel='marlin' requested but Marlin is ineligible "
+                f"(needs CUDA sm80+, symmetric INT4, legal shapes: hidden="
+                f"{self.hidden_size}, shard_inter={shard_inter}, "
+                f"group_size={qc.group_size}, legal={legal_shapes}). "
+                "Use quant_kernel='auto' to fall back to the Triton path."
+            )
+        global _BACKEND_LOGGED
+        if not eligible:
+            if not _BACKEND_LOGGED:
+                logger.info(
+                    "KimiSparseMoeBlock routed-expert backend: Triton W4A16 "
+                    "(quant_kernel=%s, marlin ineligible: legal_shapes=%s).",
+                    self.quant_kernel, legal_shapes,
+                )
+                _BACKEND_LOGGED = True
+            return
+
+        if not _BACKEND_LOGGED:
+            logger.info(
+                "KimiSparseMoeBlock routed-expert backend: Marlin W4A16 "
+                "(quant_kernel=%s, group_size=%d, tp_size=%d).",
+                self.quant_kernel, qc.group_size, self.tp_size,
+            )
+            _BACKEND_LOGGED = True
+
+        method = MarlinMoEMethod.from_quant_config(qc)
+        method.prepare(
+            self.experts.gate_up_proj_packed.data,
+            self.experts.gate_up_proj_scale.data,
+            self.experts.down_proj_packed.data,
+            self.experts.down_proj_scale.data,
+            dev,
+        )
+        for name in (
+            "gate_up_proj_packed", "gate_up_proj_scale",
+            "down_proj_packed", "down_proj_scale",
+        ):
+            p = getattr(self.experts, name)
+            p.data = torch.empty(0, dtype=p.dtype, device=dev)
+        self._marlin_method = method
+        self._use_marlin = True
+
+    def _dispatch_marlin(
+        self,
+        flat: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        from mstar.utils.fused_moe import moe_sum_reduce_triton
+
+        reduce = self.tp_size == 1
+        out = self._marlin_method.apply(
+            flat, topk_weights, topk_ids, reduce_results=reduce
+        )
+        if reduce:
+            return out
+        self.comm_group.all_reduce(out)
+        output = torch.empty_like(flat)
+        moe_sum_reduce_triton(out, output, routed_scaling_factor=1.0)
+        return output
+
+    def _dispatch_tp(
+        self,
+        flat: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        from mstar.utils.fused_moe import fused_experts, moe_sum_reduce_triton
+
+        cache3 = fused_experts(
+            flat,
+            self.experts.gate_up_proj,
+            self.experts.down_proj,
+            topk_weights,
+            topk_ids,
+            reduce_results=False,
+        )
+        self.comm_group.all_reduce(cache3)
+        output = torch.empty_like(flat)
+        moe_sum_reduce_triton(cache3, output, routed_scaling_factor=1.0)
+        return output
diff --git a/mstar/model/kimi_k2_7/components/rope.py b/mstar/model/kimi_k2_7/components/rope.py
new file mode 100644
index 000000000..7338a59fa
--- /dev/null
+++ b/mstar/model/kimi_k2_7/components/rope.py
@@ -0,0 +1,122 @@
+"""DeepSeek YARN RoPE for Kimi-K2.7 MLA."""
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+
+def yarn_get_mscale(scale: float = 1.0, mscale: float = 1.0) -> float:
+    """DeepSeek 2-arg mscale (deepseek_v2.py:428 / deepseek_scaling_rope.py:20)."""
+    if scale <= 1:
+        return 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
+
+
+def rotate_gptj(x: torch.Tensor) -> torch.Tensor:
+    """Interleaved (GPT-J) rotate: pairs ``x[..., ::2]`` / ``x[..., 1::2]``
+    (vLLM ``common.py::rotate_gptj``)."""
+    x1 = x[..., ::2]
+    x2 = x[..., 1::2]
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def _yarn_find_correction_dim(num_rotations, dim, base, max_pos) -> float:
+    return (dim * math.log(max_pos / (num_rotations * 2 * math.pi))) / (2 * math.log(base))
+
+
+def _yarn_find_correction_range(low_rot, high_rot, dim, base, max_pos) -> tuple[int, int]:
+    low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, max_pos))
+    high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, max_pos))
+    return max(low, 0), min(high, dim - 1)
+
+
+def _yarn_linear_ramp_mask(low, high, dim, dtype) -> torch.Tensor:
+    if low == high:
+        high += 0.001  # avoid singularity
+    ramp = (torch.arange(dim, dtype=dtype) - low) / (high - low)
+    return torch.clamp(ramp, 0, 1)
+
+
+class KimiYarnRotaryEmbedding(nn.Module):
+    """deepseek_yarn rotary embedding over the ``rotary_dim`` (=qk_rope_head_dim) slice."""
+
+    def __init__(
+        self,
+        rotary_dim: int,
+        base: float,
+        factor: float,
+        original_max_position_embeddings: int,
+        beta_fast: float = 32,
+        beta_slow: float = 1,
+        mscale: float = 1.0,
+        mscale_all_dim: float = 0.0,
+        extrapolation_factor: float = 1.0,
+        attn_factor: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.rotary_dim = rotary_dim
+
+        # Do not register inv_freq: meta->to_empty leaves derived buffers
+        # uninitialized, and model.to(bf16) would downcast it. Recompute fp32 lazily.
+        self._inv_freq_args = (
+            rotary_dim, base, factor, original_max_position_embeddings,
+            beta_fast, beta_slow, extrapolation_factor,
+        )
+        self._inv_freq_cache: torch.Tensor | None = None
+
+        self.mscale = float(
+            yarn_get_mscale(factor, mscale)
+            / yarn_get_mscale(factor, mscale_all_dim)
+            * attn_factor
+        )
+
+    def _get_inv_freq(self, device: torch.device) -> torch.Tensor:
+        """Return the fp32 ``inv_freq`` for ``device``, computing + caching once."""
+        cached = self._inv_freq_cache
+        if cached is None or cached.device != device:
+            cached = self._compute_inv_freq(*self._inv_freq_args).to(device=device)
+            self._inv_freq_cache = cached
+        return cached
+
+    @staticmethod
+    def _compute_inv_freq(
+        rotary_dim, base, factor, max_pos, beta_fast, beta_slow, extrapolation_factor,
+    ) -> torch.Tensor:
+        pos_freqs = base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim)
+        inv_freq_extrapolation = 1.0 / pos_freqs
+        inv_freq_interpolation = 1.0 / (factor * pos_freqs)
+
+        low, high = _yarn_find_correction_range(
+            beta_fast, beta_slow, rotary_dim, base, max_pos
+        )
+        inv_freq_mask = (
+            1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float)
+        ) * extrapolation_factor
+        return (
+            inv_freq_interpolation * (1 - inv_freq_mask)
+            + inv_freq_extrapolation * inv_freq_mask
+        )
+
+    def forward(
+        self, position_ids: torch.Tensor, q_pe: torch.Tensor, k_pe: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Rotate the pe slices.
+
+        Args:
+            position_ids: ``(tokens,)`` int positions.
+            q_pe: ``(tokens, num_heads, rotary_dim)``.
+            k_pe: ``(tokens, 1, rotary_dim)`` (shared MQA rope key).
+        Returns:
+            rotated ``(q_pe, k_pe)`` in the input dtypes.
+        """
+        inv_freq = self._get_inv_freq(position_ids.device)
+        freqs = torch.outer(position_ids.float(), inv_freq)  # (T, rotary_dim/2)
+        cos = (freqs.cos() * self.mscale).repeat_interleave(2, dim=-1).unsqueeze(-2)
+        sin = (freqs.sin() * self.mscale).repeat_interleave(2, dim=-1).unsqueeze(-2)
+
+        q32, k32 = q_pe.float(), k_pe.float()
+        q_rot = q32 * cos + rotate_gptj(q32) * sin
+        k_rot = k32 * cos + rotate_gptj(k32) * sin
+        return q_rot.to(q_pe.dtype), k_rot.to(k_pe.dtype)
diff --git a/mstar/model/kimi_k2_7/config.py b/mstar/model/kimi_k2_7/config.py
new file mode 100644
index 000000000..19d9410dc
--- /dev/null
+++ b/mstar/model/kimi_k2_7/config.py
@@ -0,0 +1,170 @@
+"""Kimi-K2.7 text config, using the DeepSeek-V3 architecture fields."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from mstar.model.components.quantization import CompressedTensorsQuantConfig
+
+
+@dataclass
+class KimiK2Config:
+    vocab_size: int = 163840
+    hidden_size: int = 7168
+    intermediate_size: int = 18432
+    num_hidden_layers: int = 61
+    num_attention_heads: int = 64
+    num_key_value_heads: int = 64  # MLA has no separate KV heads; kept for HF parity
+    rms_norm_eps: float = 1e-5
+    max_position_embeddings: int = 262144
+    tie_word_embeddings: bool = False
+    hidden_act: str = "silu"
+
+    q_lora_rank: int = 1536
+    kv_lora_rank: int = 512
+    qk_nope_head_dim: int = 128
+    qk_rope_head_dim: int = 64
+    v_head_dim: int = 128
+
+    # Default absorbed MLA stores one compressed latent KV head. The naive path is
+    # kept as the reduced-test parity fallback.
+    mla_absorb: bool = True
+
+    n_routed_experts: int = 384
+    n_shared_experts: int = 1
+    num_experts_per_tok: int = 8
+    moe_intermediate_size: int = 2048
+    n_group: int = 1
+    topk_group: int = 1
+    routed_scaling_factor: float = 2.827
+    scoring_func: str = "sigmoid"
+    topk_method: str = "noaux_tc"
+    norm_topk_prob: bool = True
+    first_k_dense_replace: int = 1
+    moe_layer_freq: int = 1
+
+    rope_theta: float = 50000.0
+    rope_scaling: dict = field(default_factory=lambda: {
+        "rope_type": "deepseek_yarn",
+        "factor": 64.0,
+        "original_max_position_embeddings": 4096,
+        "beta_fast": 32.0,
+        "beta_slow": 1.0,
+        "mscale": 1.0,
+        "mscale_all_dim": 1.0,
+    })
+
+    bos_token_id: int = 163584
+    eos_token_id: int = 163586
+    pad_token_id: int = 163839
+    temperature: float = 1.0
+    top_p: float = 1.0
+    ignore_eos: bool = False
+
+    num_nextn_predict_layers: int = 0
+
+    quantization_config: CompressedTensorsQuantConfig | None = None
+
+    # Keeps routed experts packed; non-expert quantized weights still dequantize
+    # on load.
+    moe_in_kernel_dequant: bool = False
+
+    # "auto" probes Marlin post-load on the real device; "marlin" must not silently
+    # downgrade, and "triton" keeps the packed Triton path.
+    quant_kernel: str = "auto"
+
+    prefill_token_buckets: list[int] | None = None
+    prefill_capture_batch_sizes: list[int] | None = None
+
+    @property
+    def qk_head_dim(self) -> int:
+        return self.qk_nope_head_dim + self.qk_rope_head_dim
+
+    @property
+    def padded_head_dim(self) -> int:
+        """Naive-MLA q/k/v pad target; FlashInfer paged kernels require 64/128/256."""
+        for supported in (64, 128, 256):
+            if supported >= self.qk_head_dim:
+                return supported
+        raise ValueError(
+            f"qk_head_dim={self.qk_head_dim} exceeds the largest FlashInfer SM90 "
+            "head_dim (256); the naive-MLA pad mitigation cannot cover it."
+        )
+
+    @property
+    def num_dense_layers(self) -> int:
+        return min(self.first_k_dense_replace, self.num_hidden_layers)
+
+    @classmethod
+    def reduced(cls) -> "KimiK2Config":
+        return cls(
+            mla_absorb=False,
+            vocab_size=256,
+            hidden_size=128,
+            intermediate_size=256,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            num_key_value_heads=4,
+            max_position_embeddings=512,
+            q_lora_rank=48,
+            kv_lora_rank=32,
+            qk_nope_head_dim=16,
+            qk_rope_head_dim=8,
+            v_head_dim=16,
+            n_routed_experts=4,
+            n_shared_experts=1,
+            num_experts_per_tok=2,
+            moe_intermediate_size=64,
+            n_group=1,
+            topk_group=1,
+            first_k_dense_replace=1,
+            prefill_token_buckets=[64],
+            prefill_capture_batch_sizes=[1],
+        )
+
+    @classmethod
+    def reduced_quantized(
+        cls,
+        num_bits: int = 4,
+        group_size: int = 32,
+        symmetric: bool = True,
+    ) -> "KimiK2Config":
+        cfg = cls.reduced()
+        cfg.quantization_config = CompressedTensorsQuantConfig(
+            num_bits=num_bits, group_size=group_size, symmetric=symmetric,
+        )
+        return cfg
+
+    @classmethod
+    def reduced_quantized_inkernel(
+        cls,
+        num_bits: int = 4,
+        group_size: int = 32,
+        symmetric: bool = True,
+    ) -> "KimiK2Config":
+        cfg = cls.reduced_quantized(
+            num_bits=num_bits, group_size=group_size, symmetric=symmetric,
+        )
+        cfg.moe_in_kernel_dequant = True
+        return cfg
+
+    @classmethod
+    def reduced_marlin(
+        cls,
+        num_bits: int = 4,
+        group_size: int = 32,
+        symmetric: bool = True,
+    ) -> "KimiK2Config":
+        cfg = cls.reduced_quantized_inkernel(
+            num_bits=num_bits, group_size=group_size, symmetric=symmetric,
+        )
+        cfg.hidden_size = 256
+        cfg.moe_intermediate_size = 256
+        cfg.intermediate_size = 512
+        cfg.quant_kernel = "marlin"
+        return cfg
+
+    @classmethod
+    def k27_code(cls) -> "KimiK2Config":
+        cfg = cls()
+        cfg.moe_in_kernel_dequant = True
+        return cfg
diff --git a/mstar/model/kimi_k2_7/kimi_model.py b/mstar/model/kimi_k2_7/kimi_model.py
new file mode 100644
index 000000000..57d98db0b
--- /dev/null
+++ b/mstar/model/kimi_k2_7/kimi_model.py
@@ -0,0 +1,354 @@
+"""M* model wrapper for the Kimi-K2.7 text backbone."""
+from __future__ import annotations
+
+import logging
+
+import torch
+
+from mstar.communication.tensors import NameToTensorList
+from mstar.conductor.request_info import (
+    CurrentForwardConductorMetadata,
+    StreamingConnectionState,
+)
+from mstar.engine.base import EngineType
+from mstar.engine.kv_cache_engine import KVCacheConfig
+from mstar.graph.base import GraphEdge, GraphNode, GraphSection, Loop, TensorPointerInfo
+from mstar.graph.special_destinations import EMIT_TO_CLIENT
+from mstar.model.base import ForwardPassArgs, Model
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.submodule_base import NodeSubmodule
+from mstar.utils.sampling import SamplingConfig
+
+logger = logging.getLogger(__name__)
+
+LLM_NODE = "LLM"
+DECODE_LOOP = "decode_loop"
+
+
+def _resolve_local_hf_snapshot(repo_id: str, cache_dir: str | None = None) -> str:
+    from pathlib import Path
+
+    from huggingface_hub import snapshot_download
+
+    try:
+        local_dir = snapshot_download(
+            repo_id=repo_id, cache_dir=cache_dir, local_files_only=False,
+        )
+    except Exception as e:  # noqa: BLE001
+        logger.warning("Error downloading %r from huggingface: %s", repo_id, e)
+        return repo_id
+    return str(Path(local_dir))
+
+
+class KimiK2Model(Model):
+    def __init__(
+        self,
+        model_path_hf: str,
+        cache_dir: str | None = None,
+        **kwargs,
+    ):
+        self.cache_dir = cache_dir
+        # A local checkpoint directory or an HF repo id. Deployments point this at
+        # their own copy with ``mstar serve --model-path`` / ``mstar-serve
+        # --model-path`` rather than hardcoding a path in the config yaml.
+        self.model_path_hf = model_path_hf
+        self._config_variant = kwargs.get("config_variant", "full")
+        if self._config_variant == "reduced":
+            self.config = KimiK2Config.reduced()
+        elif self._config_variant == "reduced_quantized":
+            self.config = KimiK2Config.reduced_quantized()
+        elif self._config_variant == "reduced_quantized_inkernel":
+            self.config = KimiK2Config.reduced_quantized_inkernel()
+        elif self._config_variant == "k27_code":
+            self.config = KimiK2Config.k27_code()
+        else:
+            self.config = KimiK2Config()
+        self._tokenizer_mode = kwargs.get("tokenizer_mode", "hf")
+        self._tokenizer = None
+        self._submodule_cache: dict[str, NodeSubmodule | None] = {}
+
+    @property
+    def tokenizer(self):
+        if self._tokenizer is None:
+            from transformers import AutoTokenizer
+
+            self._tokenizer = AutoTokenizer.from_pretrained(
+                self.model_path_hf,
+                cache_dir=self.cache_dir,
+                trust_remote_code=True,
+            )
+        return self._tokenizer
+
+    def get_kv_cache_config(self) -> list[KVCacheConfig]:
+        if self.config.mla_absorb:
+            from mstar.model.kimi_k2_7.components.rope import yarn_get_mscale
+            rope = self.config.rope_scaling
+            mscale = yarn_get_mscale(rope["factor"], rope.get("mscale_all_dim", 0.0))
+            softmax_scale = self.config.qk_head_dim ** -0.5 * mscale * mscale
+            return [KVCacheConfig(
+                num_layers=self.config.num_hidden_layers,
+                num_kv_heads=1,
+                head_dim=self.config.kv_lora_rank + self.config.qk_rope_head_dim,
+                max_seq_len=self.config.max_position_embeddings,
+                num_qo_heads=self.config.num_attention_heads,
+                attention_backend="mla_absorb",
+                softmax_scale=softmax_scale,
+                mla_ckv_dim=self.config.kv_lora_rank,
+            )]
+        return [KVCacheConfig(
+            num_layers=self.config.num_hidden_layers,
+            num_kv_heads=self.config.num_attention_heads,
+            head_dim=self.config.padded_head_dim,
+            max_seq_len=self.config.max_position_embeddings,
+            num_qo_heads=self.config.num_attention_heads,
+        )]
+
+    def get_node_engine_types(self) -> dict[str, EngineType]:
+        return {LLM_NODE: EngineType.KV_CACHE}
+
+    def get_graph_walk_graphs(self) -> dict[str, GraphSection]:
+        prefill = GraphNode(
+            name=LLM_NODE,
+            input_names=["text_inputs"],
+            outputs=[
+                GraphEdge(
+                    next_node=EMIT_TO_CLIENT,
+                    name="new_token",
+                    output_modality="text",
+                    conductor_new_token=True,
+                    persist=True,
+                ),
+            ],
+        )
+
+        decode = Loop(
+            name=DECODE_LOOP,
+            section=GraphNode(
+                name=LLM_NODE,
+                input_names=["text_inputs"],
+                outputs=[
+                    GraphEdge(
+                        next_node=EMIT_TO_CLIENT,
+                        name="new_token",
+                        output_modality="text",
+                        conductor_new_token=True,
+                    ),
+                    GraphEdge(
+                        next_node=LLM_NODE,
+                        name="text_inputs",
+                    ),
+                ],
+            ),
+            max_iters=self.get_max_output_tokens(),
+            outputs=[],
+        )
+
+        return dict(prefill=prefill, decode=decode)
+
+    def get_initial_forward_pass_args(
+        self,
+        partition_name: str,
+        input_modalities: list[str],
+        output_modalities: list[str],
+        input_signals: dict[str, list[TensorPointerInfo]],
+        model_kwargs: dict | None = None,
+    ) -> ForwardPassArgs:
+        full_metadata = CurrentForwardConductorMetadata(
+            input_modalities=input_modalities,
+            output_modalities=output_modalities,
+            graph_walk="prefill",
+            is_prefill=True,
+        )
+
+        graph_edge = GraphEdge(next_node=LLM_NODE, name="text_inputs")
+        graph_edge.tensor_info = input_signals.get("text_inputs", [])
+        inputs = [graph_edge]
+        unpersist_tensors = sum([inp.tensor_info for inp in inputs], start=[])
+
+        return ForwardPassArgs(
+            full_metadata=full_metadata,
+            inputs=inputs,
+            unpersist_tensors=unpersist_tensors,
+            step_metadata={"is_prefill": True},
+        )
+
+    def get_partition_forward_pass_args(
+        self,
+        partition_name: str,
+        partition_metadata: CurrentForwardConductorMetadata,
+        persist_signals: dict[str, list[TensorPointerInfo]],
+        incoming_connections: list[StreamingConnectionState] | None = None,
+    ) -> ForwardPassArgs:
+        metadata = partition_metadata
+        request_done = False
+
+        if metadata.is_prefill:
+            metadata.is_prefill = False
+            metadata.graph_walk = "decode"
+        elif metadata.graph_walk == "decode":
+            request_done = True
+            metadata.kwargs["decode_finished"] = True
+
+        if request_done:
+            return ForwardPassArgs(
+                full_metadata=metadata,
+                inputs=[],
+                unpersist_tensors=[],
+                request_done=True,
+            )
+
+        graph_edge = GraphEdge(next_node=LLM_NODE, name="text_inputs")
+        graph_edge.tensor_info = persist_signals.get("new_token", [])
+        inputs = [graph_edge]
+        unpersist_tensors = sum([inp.tensor_info for inp in inputs], start=[])
+
+        return ForwardPassArgs(
+            full_metadata=metadata,
+            inputs=inputs,
+            unpersist_tensors=unpersist_tensors,
+            step_metadata={"is_prefill": metadata.is_prefill},
+        )
+
+    def process_prompt(
+        self,
+        prompt: str | None,
+        input_modalities: list[str],
+        output_modalities: list[str],
+        tensors: NameToTensorList | None = None,
+        **kwargs,
+    ) -> NameToTensorList:
+        if prompt is None:
+            return {}
+        if self._tokenizer_mode == "byte":
+            # Reduced serve maps UTF-8 bytes directly to token ids, avoiding HF IO.
+            vocab = self.config.vocab_size
+            byte_ids = [min(b, vocab - 1) for b in prompt.encode("utf-8")] or [0]
+            input_ids = torch.tensor(byte_ids, dtype=torch.long)
+            return {"text_inputs": [input_ids]}
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids[0]
+        return {"text_inputs": [input_ids]}
+
+    def get_sampling_config(
+        self,
+        node_name: str,
+        model_kwargs: dict | None = None,
+    ) -> SamplingConfig | None:
+        model_kwargs = model_kwargs or {}
+        return SamplingConfig(
+            vocab_size=self.config.vocab_size,
+            temperature=model_kwargs.get("temperature", self.config.temperature),
+            top_p=model_kwargs.get("top_p", self.config.top_p),
+            ignore_eos=model_kwargs.get("ignore_eos", self.config.ignore_eos),
+        )
+
+    def postprocess(
+        self,
+        output: torch.Tensor,
+        modality: str,
+        request_kwargs: dict | None = None,
+    ) -> bytes:
+        if modality == "text":
+            token_ids = output.tolist() if output.numel() else []
+            if self._tokenizer_mode == "byte":
+                # Synthetic reduced models emit arbitrary byte ids; return raw bytes.
+                return bytes((t & 0xFF) for t in token_ids)
+            text = self.tokenizer.decode(token_ids, skip_special_tokens=True)
+            return text.encode("utf-8")
+        raise ValueError(f"Unsupported modality for Kimi-K2.7: {modality!r}")
+
+    def get_default_sharding_config(self):
+        from mstar.distributed.base import ShardingConfig
+
+        return ShardingConfig(groups=[], tp_enabled_nodes={LLM_NODE}, shard_dim={})
+
+    def get_submodule(
+        self,
+        node_name: str,
+        device: str = "cpu",
+        tp_group=None,
+        autocast_dtype: torch.dtype | None = None,
+        sp_group=None,
+    ) -> NodeSubmodule | None:
+        if node_name in self._submodule_cache:
+            return self._submodule_cache[node_name]
+        submodule = self._create_submodule(
+            node_name, device, tp_group=tp_group, autocast_dtype=autocast_dtype,
+        )
+        self._submodule_cache[node_name] = submodule
+        return submodule
+
+    def _create_submodule(
+        self,
+        node_name: str,
+        device: str,
+        tp_group=None,
+        autocast_dtype: torch.dtype | None = None,
+    ) -> NodeSubmodule | None:
+        if node_name != LLM_NODE:
+            return None
+
+        source = self._resolve_checkpoint()
+        if source is None:
+            logger.info(
+                "KimiK2Model: no checkpoint resolved for node %r — dummy mode (None).",
+                node_name,
+            )
+            return None
+
+        self._maybe_apply_checkpoint_quant_config(source)
+
+        from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM
+        from mstar.model.kimi_k2_7.submodules import KimiLLMSubmodule
+        from mstar.model.loader import load_weights
+
+        with torch.device("meta"):
+            language_model = KimiForCausalLM(self.config, comm_group=tp_group)
+        if autocast_dtype is not None:
+            language_model = language_model.to(autocast_dtype)
+        language_model.to_empty(device=device)
+        load_weights(language_model, source, device=device)
+        from mstar.model.components.quantization import process_weights_after_loading
+
+        process_weights_after_loading(language_model, torch.device(device))
+        language_model.eval()
+
+        logger.info("Successfully loaded Kimi-K2.7 submodule for %s", node_name)
+        return KimiLLMSubmodule(language_model=language_model, config=self.config)
+
+    def _resolve_checkpoint(self) -> str | None:
+        from pathlib import Path
+
+        path = getattr(self, "model_path_hf", None)
+        if not path:
+            return None
+        if Path(path).exists():
+            return str(path)
+        return _resolve_local_hf_snapshot(path, cache_dir=getattr(self, "cache_dir", None))
+
+    def _maybe_apply_checkpoint_quant_config(self, source: str) -> None:
+        import json
+        from pathlib import Path
+
+        from mstar.model.components.quantization import CompressedTensorsQuantConfig
+
+        if self.config.quantization_config is not None:
+            return
+        config_json = Path(source) / "config.json"
+        if not config_json.is_file():
+            return
+        try:
+            with open(config_json) as f:
+                raw = json.load(f)
+        except (OSError, ValueError) as e:  # unreadable / malformed — stay bf16
+            logger.warning("KimiK2Model: could not read %s: %s", config_json, e)
+            return
+        quant_raw = raw.get("quantization_config") or (
+            raw.get("text_config") or {}
+        ).get("quantization_config")
+        quant = CompressedTensorsQuantConfig.from_hf_config_dict(quant_raw)
+        if quant is not None:
+            logger.info(
+                "KimiK2Model: compressed-tensors checkpoint (%d-bit, group_size=%d) "
+                "— dequantizing on load.", quant.num_bits, quant.group_size,
+            )
+            self.config.quantization_config = quant
diff --git a/mstar/model/kimi_k2_7/submodules.py b/mstar/model/kimi_k2_7/submodules.py
new file mode 100644
index 000000000..907b5fba3
--- /dev/null
+++ b/mstar/model/kimi_k2_7/submodules.py
@@ -0,0 +1,206 @@
+"""AR submodule for the Kimi-K2.7 text backbone."""
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch import nn
+
+from mstar.communication.tensors import NameToTensorList
+from mstar.conductor.request_info import CurrentForwardPassInfo
+from mstar.engine.base import NodeBatch
+from mstar.engine.cache_manager import BatchedCacheManager
+from mstar.engine.cuda_graph_config import FlashInferPackedCudaGraphConfig
+from mstar.engine.cuda_graph_runner import BasicBatchedCudaGraphConfig
+from mstar.engine.kv_store import PositionInfo
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.submodule_base import (
+    ARNodeInputs,
+    ARNodeSubmodule,
+    ModelInputsFromEngine,
+    NodeInputs,
+)
+from mstar.utils.sampling import Sampler
+
+_MAIN = "main"
+
+
+class KimiLLMSubmodule(ARNodeSubmodule):
+    def __init__(self, language_model: nn.Module, config: KimiK2Config):
+        super().__init__()
+        self.language_model = language_model  # KimiForCausalLM
+        self.lm_head = language_model.lm_head
+        self.config = config
+
+    PREFILL_TOKEN_BUCKETS = [32, 64, 128, 256, 512, 1024]
+    PREFILL_CAPTURE_BATCH_SIZES = [1, 2, 4, 8, 16]
+
+    def _build_prefill_packed(
+        self, num_tokens: int, device: torch.device,
+    ) -> dict[str, torch.Tensor]:
+        return {
+            "input_ids": torch.zeros((num_tokens,), dtype=torch.long, device=device),
+            "position_ids": torch.arange(num_tokens, dtype=torch.long, device=device),
+        }
+
+    def get_cuda_graph_configs(
+        self, device: torch.device, tp_world_size: int = 1,
+    ) -> list[BasicBatchedCudaGraphConfig | FlashInferPackedCudaGraphConfig]:
+        prefill_buckets = self.config.prefill_token_buckets or self.PREFILL_TOKEN_BUCKETS
+        prefill_batch_sizes = (
+            self.config.prefill_capture_batch_sizes or self.PREFILL_CAPTURE_BATCH_SIZES
+        )
+        prefill_packed = {
+            num_tokens: self._build_prefill_packed(num_tokens, device)
+            for num_tokens in prefill_buckets
+        }
+        return [
+            BasicBatchedCudaGraphConfig(
+                capture_graph_walk="decode",
+                requires_cfg=False,
+                labels=[_MAIN],
+                single_request_inputs=ARNodeInputs(
+                    input_ids=torch.zeros(1, dtype=torch.long, device=device),
+                    input_seq_len=1,
+                ),
+            ),
+            FlashInferPackedCudaGraphConfig(
+                capture_graph_walk="prefill",
+                replay_graph_walks=["prefill"],
+                packed_seq_len_to_inputs=prefill_packed,
+                requires_cfg=False,
+                labels=[_MAIN],
+                compile=True,
+                causal_attention=True,
+                capture_batch_sizes=prefill_batch_sizes,
+            ),
+        ]
+
+    def prepare_inputs(
+        self,
+        graph_walk: str,
+        fwd_info: CurrentForwardPassInfo,
+        inputs: NameToTensorList,
+        pos_info: dict[str, PositionInfo] = {},
+        **kwargs,
+    ) -> ARNodeInputs:
+        text_inputs = inputs["text_inputs"][0]
+        return ARNodeInputs(
+            input_ids=text_inputs,
+            input_seq_len=text_inputs.shape[0],
+        )
+
+    def preprocess(
+        self,
+        graph_walk: str,
+        engine_inputs: ModelInputsFromEngine,
+        inputs: list[ARNodeInputs],
+    ) -> dict[str, torch.Tensor | Any]:
+        cache_manager = engine_inputs.cache_manager
+        seq_lens = [inp.input_seq_len for inp in inputs]
+
+        cache_manager.set_active_label(_MAIN)
+        cache_manager.plan_attention(seq_lens=seq_lens, is_causal=True, label=_MAIN)
+        cache_manager.plan_rope(seq_lens=seq_lens, pos_ids=None, label=_MAIN)
+
+        device = self.get_device()
+        pos_ids_list: list[int] = []
+        for rid, sl in zip(cache_manager.request_ids, seq_lens, strict=True):
+            start = cache_manager._get_state(rid, _MAIN).position_id_start
+            pos_ids_list.extend(range(start, start + sl))
+        position_ids = torch.tensor(pos_ids_list, dtype=torch.long, device=device)
+
+        return {
+            "input_ids": torch.cat([inp.input_ids for inp in inputs]),
+            "position_ids": position_ids,
+        }
+
+    def _hidden(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        cache_handle: BatchedCacheManager,
+    ) -> torch.Tensor:
+        return self.language_model.model(input_ids, cache_handle, position_ids)
+
+    def forward(
+        self,
+        graph_walk: str,
+        engine_inputs: ModelInputsFromEngine,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        **kwargs,
+    ) -> NameToTensorList:
+        cache_handle = engine_inputs.cache_manager
+        hidden = self._hidden(input_ids, position_ids, cache_handle)
+        logits = self.lm_head(hidden[-1:])
+        return {"logits": [logits]}
+
+    def can_batch(self, batch: NodeBatch, model_inputs: list[NodeInputs]) -> bool:
+        return True
+
+    def forward_batched(
+        self,
+        graph_walk: str,
+        engine_inputs: ModelInputsFromEngine,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        **kwargs,
+    ) -> dict[str, NameToTensorList]:
+        cache_handle = engine_inputs.cache_manager
+        sampler = engine_inputs.sampler
+        cache_handle.set_active_label(_MAIN)
+
+        hidden = self._hidden(input_ids, position_ids, cache_handle)
+
+        if graph_walk == "prefill":
+            qo_indptr_buf = cache_handle.get_qo_indptr_buf(_MAIN)
+            assert qo_indptr_buf is not None, (
+                "prefill forward_batched requires a CUDA-graph "
+                "FlashInferPrefillWrapper (qo_indptr static buffer); got None."
+            )
+            last_token_indices = (qo_indptr_buf[1:] - 1).long()
+            hidden = hidden.index_select(0, last_token_indices)
+        elif graph_walk != "decode":
+            raise ValueError(f"Batched forward not supported for graph walk: {graph_walk!r}")
+
+        logits = self.lm_head(hidden)  # (bs, vocab)
+        request_ids = cache_handle.request_ids
+        new_tokens = self._sample(sampler, request_ids, logits)
+        return {
+            rid: {"new_token": [new_tokens[i : i + 1]]}
+            for i, rid in enumerate(request_ids)
+        }
+
+    @staticmethod
+    def _sample(
+        sampler: Sampler, request_ids: list[str], logits: torch.Tensor,
+    ) -> torch.Tensor:
+        return sampler.sample(request_ids, logits, apply_penalty=True)
+
+    def postprocess(
+        self, request_id: str,
+        request_info: CurrentForwardPassInfo,
+        outputs: dict[str, list[torch.Tensor]],
+        **kwargs,
+    ):
+        if "new_token" not in outputs:
+            return
+        outputs["text_inputs"] = outputs["new_token"]
+
+    def check_stop(
+        self, request_id: str,
+        request_info: CurrentForwardPassInfo,
+        outputs: dict[str, list[torch.Tensor]],
+    ) -> set[str]:
+        if "new_token" not in outputs:
+            return set()
+        token = outputs["new_token"][0].item()
+        eos_token_id = self.config.eos_token_id
+        ignore_eos = request_info.sampling_config["LLM"].ignore_eos
+        if (not ignore_eos and eos_token_id == token) or (
+            request_info.dynamic_loop_iter_counts.get("decode_loop", 0) + 1
+            >= request_info.max_tokens
+        ):
+            return {"decode_loop"}
+        return set()
diff --git a/mstar/model/kimi_k2_7/weight_loader.py b/mstar/model/kimi_k2_7/weight_loader.py
new file mode 100644
index 000000000..fa77cc59f
--- /dev/null
+++ b/mstar/model/kimi_k2_7/weight_loader.py
@@ -0,0 +1,139 @@
+"""HF DeepSeek-V3 checkpoint loading for the Kimi-K2.7 module tree."""
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+
+from mstar.model.loader.base import StackedParamRule
+
+if TYPE_CHECKING:
+    from mstar.model.components.quantization import CompressedTensorsQuantConfig
+
+# Keep the expert index attached while remapping both bf16 and packed sub-keys.
+_EXPERT_RE = re.compile(
+    r"(.*)\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)"
+    r"\.(weight|weight_packed|weight_scale|weight_zero_point)$"
+)
+
+_EXPERT_BASE_RE = re.compile(r"\.experts\.\d+\.(gate_proj|up_proj|down_proj)$")
+
+
+def _is_routed_expert_base(base: str) -> bool:
+    return _EXPERT_BASE_RE.search(base) is not None
+
+
+def kimi_name_remapper(name: str) -> str | None:
+    if "rotary_emb" in name:
+        return None
+    if name.startswith("language_model."):
+        name = name[len("language_model."):]
+    name = name.replace(".shared_experts.", ".shared_expert.")
+    m = _EXPERT_RE.match(name)
+    if m:
+        prefix, expert_idx, proj, suffix = m.groups()
+        return f"{prefix}.experts.{proj}.__expert{expert_idx}__.{suffix}"
+    return name
+
+
+def build_kimi_stacked_params(
+    n_routed_experts: int, packed_experts: bool = False,
+) -> list[StackedParamRule]:
+    rules: list[StackedParamRule] = []
+    for i in range(n_routed_experts):
+        if packed_experts:
+            for proj, sid in (("gate_proj", f"gate:{i}"), ("up_proj", f"up:{i}")):
+                rules.append(StackedParamRule(
+                    target_suffix=".experts.gate_up_proj_packed",
+                    source_suffix=f".experts.{proj}.__expert{i}__.weight_packed",
+                    shard_id=sid,
+                ))
+                rules.append(StackedParamRule(
+                    target_suffix=".experts.gate_up_proj_scale",
+                    source_suffix=f".experts.{proj}.__expert{i}__.weight_scale",
+                    shard_id=sid,
+                ))
+            rules.append(StackedParamRule(
+                target_suffix=".experts.down_proj_packed",
+                source_suffix=f".experts.down_proj.__expert{i}__.weight_packed",
+                shard_id=f"down:{i}",
+            ))
+            rules.append(StackedParamRule(
+                target_suffix=".experts.down_proj_scale",
+                source_suffix=f".experts.down_proj.__expert{i}__.weight_scale",
+                shard_id=f"down:{i}",
+            ))
+        else:
+            rules.append(StackedParamRule(
+                target_suffix=".experts.gate_up_proj",
+                source_suffix=f".experts.gate_proj.__expert{i}__.weight",
+                shard_id=f"gate:{i}",
+            ))
+            rules.append(StackedParamRule(
+                target_suffix=".experts.gate_up_proj",
+                source_suffix=f".experts.up_proj.__expert{i}__.weight",
+                shard_id=f"up:{i}",
+            ))
+            rules.append(StackedParamRule(
+                target_suffix=".experts.down_proj",
+                source_suffix=f".experts.down_proj.__expert{i}__.weight",
+                shard_id=f"down:{i}",
+            ))
+    # Dense/shared gate-up rules must follow expert rules because matching is first-win.
+    rules.append(StackedParamRule(".gate_up_proj", ".gate_proj", 0))
+    rules.append(StackedParamRule(".gate_up_proj", ".up_proj", 1))
+    return rules
+
+
+def restore_router_bias_fp32(module: nn.Module) -> None:
+    """Force DeepSeek's router selection bias back to fp32 before loading."""
+    for sub in module.modules():
+        bias = getattr(sub, "e_score_correction_bias", None)
+        if isinstance(bias, nn.Parameter) and bias.dtype != torch.float32:
+            bias.data = bias.data.float()
+
+
+def load_kimi_hf_weights(
+    module: nn.Module,
+    weights: Iterable[tuple[str, torch.Tensor]],
+    n_routed_experts: int,
+    quant_config: "CompressedTensorsQuantConfig | None" = None,
+    packed_experts: bool = False,
+) -> set[str]:
+    from mstar.model.loader import load_hf_weights
+
+    if quant_config is not None:
+        from mstar.model.components.quantization import (
+            dequant_compressed_tensors_stream,
+        )
+
+        keep_packed = _is_routed_expert_base if packed_experts else None
+        weights = dequant_compressed_tensors_stream(
+            weights, quant_config, keep_packed=keep_packed,
+        )
+    elif packed_experts:
+        raise ValueError("packed_experts=True requires a quant_config")
+
+    restore_router_bias_fp32(module)
+    return load_hf_weights(
+        module,
+        weights,
+        stacked_params=build_kimi_stacked_params(
+            n_routed_experts, packed_experts=packed_experts,
+        ),
+        name_remapper=kimi_name_remapper,
+    )
+
+
+def load_weights(
+    module: nn.Module,
+    source: str | Path,
+    device: torch.device | str = "cpu",
+) -> set[str]:
+    from mstar.model.loader import load_weights as _driver
+
+    return _driver(module, source, device=device)
diff --git a/mstar/model/registry.py b/mstar/model/registry.py
index 87eedb2bf..8f16511b8 100644
--- a/mstar/model/registry.py
+++ b/mstar/model/registry.py
@@ -2,6 +2,7 @@
 from mstar.model.base import Model
 from mstar.model.cosmos3.cosmos3_model import Cosmos3Model
 from mstar.model.higgs_audio.higgs_audio_model import HiggsAudioModel
+from mstar.model.kimi_k2_7.kimi_model import KimiK2Model
 from mstar.model.orpheus.orpheus_model import OrpheusModel
 from mstar.model.pi05.pi05_model import Pi05Model
 from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel
@@ -16,6 +17,7 @@
     "cosmos3_droid": Cosmos3Model,
     "cosmos3_super": Cosmos3Model,
     "higgs_audio": HiggsAudioModel,
+    "kimi_k2_7": KimiK2Model,
     "orpheus": OrpheusModel,
     "pi05": Pi05Model,
     "qwen3_omni": Qwen3OmniModel,
@@ -42,6 +44,9 @@
     # Higgs-Audio v3 STT: Whisper-style audio tower + Qwen3-1.7B LLM.
     # (The v2 checkpoints are TTS/generation models, not ASR.)
     "higgs_audio": {"model_path_hf": "bosonai/higgs-audio-v3-stt"},
+    # Kimi-K2.7-Code: 1T MoE (DeepSeek-V3 text backbone + MoonViT). M0 is a
+    # text-only scaffold; real serving is TP8 / multi-node.
+    "kimi_k2_7": {"model_path_hf": "moonshotai/Kimi-K2.7-Code"},
     "orpheus": {"model_path_hf": "canopylabs/orpheus-3b-0.1-ft"},
     # Pi0.5 PyTorch port published by lerobot — single safetensors blob
     # (~14 GB). mstar/model/pi05/weight_loader.py handles the lerobot->mstar
diff --git a/mstar/utils/flashinfer_utils.py b/mstar/utils/flashinfer_utils.py
index 899bdee3c..16028fdf7 100644
--- a/mstar/utils/flashinfer_utils.py
+++ b/mstar/utils/flashinfer_utils.py
@@ -487,3 +487,171 @@ def set_kv_cache(
         positions = self.kv_cache_locations[:n, 1]
         kv_cache_layer[pages, 0, positions] = k[:n].to(self.dtype)
         kv_cache_layer[pages, 1, positions] = v[:n].to(self.dtype)
+
+
+class FlashInferMLAWrapper:
+    """FlashInfer MLA wrapper for the 4D latent cache.
+
+    The kernel only supports real Kimi dims (ckv=512, kpe=64); callers must gate
+    before construction because off-dim calls can corrupt the CUDA context. CUDA
+    graph mode updates static index/scatter buffers with ``copy_()``.
+    """
+
+    def __init__(
+        self,
+        workspace_buffer: torch.Tensor,
+        *,
+        num_heads: int,
+        head_dim_ckv: int,
+        head_dim_kpe: int,
+        page_size: int,
+        sm_scale: float,
+        batch_size: int | None = None,
+        max_num_pages: int | None = None,
+        max_total_tokens: int | None = None,
+        device: torch.device = torch.device("cuda"),
+        use_cuda_graph: bool = False,
+        backend: str = "auto",
+        enable_nvtx: bool = False,
+    ):
+        self.device = device
+        self.use_cuda_graph = use_cuda_graph
+        self.enable_nvtx = enable_nvtx
+        self.num_heads = num_heads
+        self.head_dim_ckv = head_dim_ckv
+        self.head_dim_kpe = head_dim_kpe
+        self.page_size = page_size
+        self.sm_scale = sm_scale
+        self.batch_size = batch_size
+        self.max_total_tokens = max_total_tokens
+        self.dtype = None
+        self._total_tokens = 0
+
+        import flashinfer
+
+        if self.use_cuda_graph:
+            assert batch_size is not None, "batch_size required for CUDA graph mode"
+            assert max_num_pages is not None, "max_num_pages required for CUDA graph mode"
+            assert max_total_tokens is not None, "max_total_tokens required for CUDA graph mode"
+
+            # Stable addresses for graph replay.
+            self._qo_indptr_buf = torch.zeros(
+                batch_size + 1, dtype=torch.int32, device=device
+            )
+            self._kv_indptr_buf = torch.zeros(
+                batch_size + 1, dtype=torch.int32, device=device
+            )
+            self._kv_indices_buf = torch.zeros(
+                max_num_pages, dtype=torch.int32, device=device
+            )
+            self._kv_len_arr_buf = torch.zeros(
+                batch_size, dtype=torch.int32, device=device
+            )
+
+            self.attn_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
+                workspace_buffer,
+                use_cuda_graph=True,
+                qo_indptr=self._qo_indptr_buf,
+                kv_indptr=self._kv_indptr_buf,
+                kv_indices=self._kv_indices_buf,
+                kv_len_arr=self._kv_len_arr_buf,
+                backend=backend,
+            )
+
+            # Static buffers for the vectorized latent scatter (page, offset).
+            self.token_to_page = torch.zeros(
+                max_total_tokens, dtype=torch.long, device=device
+            )
+            self.token_to_cache = torch.zeros(
+                max_total_tokens, dtype=torch.long, device=device
+            )
+        else:
+            self.attn_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
+                workspace_buffer, backend=backend,
+            )
+            self.token_to_page = None
+            self.token_to_cache = None
+
+    @torch.compiler.disable
+    def plan(
+        self,
+        qo_indptr: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        kv_indices: torch.Tensor,
+        kv_len_arr: torch.Tensor,
+        *,
+        causal: bool = True,
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        """Plan the MLA kernel and compute latent scatter indices."""
+        self.dtype = dtype
+        self.attn_wrapper.plan(
+            qo_indptr,
+            kv_indptr,
+            kv_indices,
+            kv_len_arr,
+            self.num_heads,
+            self.head_dim_ckv,
+            self.head_dim_kpe,
+            self.page_size,
+            causal,
+            self.sm_scale,
+            dtype,
+            dtype,
+        )
+
+        # Map each new token to its latent-cache page/offset from the page table.
+        n_req = qo_indptr.shape[0] - 1
+        starts = qo_indptr[:-1].to(torch.int32)
+        lens = (qo_indptr[1:] - qo_indptr[:-1]).to(torch.int32)
+        total_tokens = int(lens.sum().item())
+        self._total_tokens = total_tokens
+
+        seg = torch.repeat_interleave(
+            torch.arange(n_req, dtype=torch.int32, device=self.device), lens
+        )
+        intra = torch.arange(
+            total_tokens, dtype=torch.int32, device=self.device
+        ) - torch.repeat_interleave(starts, lens)
+
+        start_new = kv_len_arr[seg] - lens[seg]
+        g = start_new + intra
+
+        page_off = torch.div(g, self.page_size, rounding_mode="floor").to(torch.int32)
+        off_in_page = (g - page_off * self.page_size).to(torch.int32)
+        abs_page_ptr = kv_indptr[:-1][seg] + page_off
+
+        token_to_page = kv_indices[abs_page_ptr].to(torch.long)
+        token_to_cache = off_in_page.to(torch.long)
+
+        if self.use_cuda_graph:
+            self.token_to_page[:total_tokens].copy_(token_to_page)
+            self.token_to_cache[:total_tokens].copy_(token_to_cache)
+            if total_tokens < self.max_total_tokens:
+                self.token_to_page[total_tokens:] = 0
+                self.token_to_cache[total_tokens:] = 0
+        else:
+            self.token_to_page = token_to_page
+            self.token_to_cache = token_to_cache
+
+    @torch.compiler.disable
+    def set_latent(self, latent_cache_layer: torch.Tensor, latent: torch.Tensor):
+        """Scatter one compressed latent per new token into the paged cache."""
+        n = self._total_tokens
+        page_idx = self.token_to_page[:n]
+        cache_idx = self.token_to_cache[:n]
+        latent_cache_layer[page_idx, cache_idx] = latent[:n].to(latent_cache_layer.dtype)
+
+    @torch.compiler.disable
+    def run(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        ckv_cache: torch.Tensor,
+        kpe_cache: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run the planned MLA kernel."""
+        return self.attn_wrapper.run(
+            q_nope.to(self.dtype), q_pe.to(self.dtype),
+            ckv_cache, kpe_cache, return_lse=False,
+        )
diff --git a/mstar/utils/fused_moe/kernels.py b/mstar/utils/fused_moe/kernels.py
index e31c6daae..4a2d28a8f 100644
--- a/mstar/utils/fused_moe/kernels.py
+++ b/mstar/utils/fused_moe/kernels.py
@@ -138,6 +138,130 @@ def fused_moe_kernel(
     tl.store(c_ptrs, accumulator, mask=c_mask)
 
 
+@triton.jit
+def fused_moe_kernel_w4a16(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    b_scale_ptr,
+    b_zp_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Matrix dimensions (K is the LOGICAL contraction dim, not the packed width)
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    stride_am,
+    stride_ak,
+    stride_be,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_bse,
+    stride_bsk,
+    stride_bsn,
+    stride_bze,
+    stride_bzk,
+    stride_bzn,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    top_k: tl.constexpr,
+    compute_type: tl.constexpr,
+    group_size: tl.constexpr,
+    PACK_FACTOR: tl.constexpr,
+    HAS_ZP: tl.constexpr,
+    even_Ks: tl.constexpr,
+):
+    """Fused MoE tile for packed W4A16 weights.
+
+    ``K`` is logical; weights are int32-packed low-order-first along K. K tiles
+    must span whole int32s and whole quant groups.
+    """
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
+    token_mask = offs_token < num_valid_tokens
+
+    off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak)
+    # Packed B is indexed by logical K, then shifted to the right INT4 nibble.
+    b_ptrs = (
+        b_ptr + off_experts * stride_be
+        + (offs_k[:, None] // PACK_FACTOR) * stride_bk
+        + offs_bn[None, :] * stride_bn
+    )
+    b_shifter = (offs_k[:, None] % PACK_FACTOR) * 4
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    for k_start in range(0, K, BLOCK_SIZE_K):
+        # Group scales are indexed by logical K, not packed K.
+        offs_ks = (offs_k[:, None] + k_start) // group_size
+        b_scale_ptrs = b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn + offs_ks * stride_bsk
+        if even_Ks:
+            a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0)
+            b_packed = tl.load(b_ptrs)
+            b_scale = tl.load(b_scale_ptrs).to(tl.float32)
+        else:
+            k_mask = offs_k[:, None] < K - k_start
+            a = tl.load(
+                a_ptrs,
+                mask=token_mask[:, None] & (offs_k[None, :] < K - k_start),
+                other=0.0,
+            )
+            b_packed = tl.load(b_ptrs, mask=k_mask, other=0)
+            b_scale = tl.load(b_scale_ptrs, mask=k_mask, other=1.0).to(tl.float32)
+        # Mask after arithmetic shift so the top nibble is exact when bit 31 is set.
+        b_nib = ((b_packed >> b_shifter) & 0xF).to(tl.float32)
+        if HAS_ZP:
+            # Optional asymmetric zero point; Kimi uses symmetric INT4.
+            b_zp_ptrs = b_zp_ptr + off_experts * stride_bze + offs_bn[None, :] * stride_bzn + offs_ks * stride_bzk
+            if even_Ks:
+                b_zp = tl.load(b_zp_ptrs).to(tl.float32)
+            else:
+                b_zp = tl.load(b_zp_ptrs, mask=offs_k[:, None] < K - k_start, other=0.0).to(tl.float32)
+            b = ((b_nib - b_zp) * b_scale).to(compute_type)
+        else:
+            b = ((b_nib - 8.0) * b_scale).to(compute_type)
+        accumulator += tl.dot(a, b)
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += (BLOCK_SIZE_K // PACK_FACTOR) * stride_bk
+
+    if MUL_ROUTED_WEIGHT:
+        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
+        accumulator = accumulator * moe_weight[:, None]
+
+    accumulator = accumulator.to(compute_type)
+
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, accumulator, mask=c_mask)
+
+
 def invoke_fused_moe_kernel(
     A: torch.Tensor,
     B: torch.Tensor,
@@ -218,6 +342,83 @@ def grid(META):
     )
 
 
+def invoke_fused_moe_kernel_w4a16(
+    A: torch.Tensor,
+    B_packed: torch.Tensor,
+    C: torch.Tensor,
+    B_scale: torch.Tensor,
+    B_zp: torch.Tensor | None,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    mul_routed_weight: bool,
+    top_k: int,
+    config: Dict[str, Any],
+    compute_type: tl.dtype,
+    K: int,
+    pack_factor: int,
+    group_size: int,
+) -> None:
+    """Launch the packed-W4A16 MoE kernel.
+
+    ``K`` is logical, while ``B_packed.shape[2]`` is ``K // pack_factor``. When
+    ``B_zp`` is absent, ``HAS_ZP`` gates all loads from the stand-in tensor.
+    """
+    assert topk_weights.stride(1) == 1
+    assert sorted_token_ids.stride(0) == 1
+
+    N = B_packed.shape[1]
+
+    def grid(META):
+        return (
+            triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"])
+            * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        )
+
+    even_Ks = (K % config["BLOCK_SIZE_K"]) == 0
+    has_zp = B_zp is not None
+    zp = B_zp if has_zp else B_scale  # stand-in; every zp load is gated by HAS_ZP
+
+    fused_moe_kernel_w4a16[grid](
+        A,
+        B_packed,
+        C,
+        B_scale,
+        zp,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        sorted_token_ids.shape[0],
+        topk_ids.numel(),
+        A.stride(0),
+        A.stride(1),
+        B_packed.stride(0),
+        B_packed.stride(2),
+        B_packed.stride(1),
+        C.stride(-2),
+        C.stride(-1),
+        B_scale.stride(0),
+        B_scale.stride(2),
+        B_scale.stride(1),
+        zp.stride(0),
+        zp.stride(2),
+        zp.stride(1),
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        top_k=top_k,
+        compute_type=compute_type,
+        group_size=group_size,
+        PACK_FACTOR=pack_factor,
+        HAS_ZP=has_zp,
+        even_Ks=even_Ks,
+        **config,
+    )
+
+
 # ---------------------------------------------------------------------------
 # Activation (SwiGLU / GeGLU) kernel
 # ---------------------------------------------------------------------------
@@ -398,23 +599,44 @@ def moe_sum_reduce_triton(
 # ---------------------------------------------------------------------------
 
 
-def get_default_config(M: int, E: int, N: int, K: int, top_k: int) -> Dict[str, int]:
+def get_default_config(
+    M: int, E: int, N: int, K: int, top_k: int, group_size: int | None = None,
+) -> Dict[str, int]:
     """Pick Triton tile sizes based on problem shape.
 
     Mirrors sglang's ``get_default_config`` for the unquantized path.
     For decode batch sizes (``M`` on the order of 1--64) we always fall
     into the ``M <= E`` branch since Qwen3-Omni has ``E == 128``.
+
+    When ``group_size`` is set for W4A16, ``BLOCK_SIZE_K`` is clamped until it is
+    divisible by both INT4 pack factor and group size. ``None`` preserves the
+    historical config.
     """
     if M <= E:
-        return {
+        config = {
             "BLOCK_SIZE_M": 16,
             "BLOCK_SIZE_N": 32,
             "BLOCK_SIZE_K": 64,
             "GROUP_SIZE_M": 1,
         }
-    return {
-        "BLOCK_SIZE_M": 64,
-        "BLOCK_SIZE_N": 64,
-        "BLOCK_SIZE_K": 32,
-        "GROUP_SIZE_M": 8,
-    }
+    else:
+        config = {
+            "BLOCK_SIZE_M": 64,
+            "BLOCK_SIZE_N": 64,
+            "BLOCK_SIZE_K": 32,
+            "GROUP_SIZE_M": 8,
+        }
+    if group_size is not None:
+        pack_factor = 8  # INT4: 32 // 4
+        bk = config["BLOCK_SIZE_K"]
+        while bk % pack_factor != 0 or bk % group_size != 0:
+            bk //= 2
+            if bk < pack_factor:
+                raise ValueError(
+                    f"cannot pick a BLOCK_SIZE_K divisible by pack_factor={pack_factor} "
+                    f"and group_size={group_size}; got down to {bk}"
+                )
+        config["BLOCK_SIZE_K"] = bk
+        assert config["BLOCK_SIZE_K"] % pack_factor == 0
+        assert config["BLOCK_SIZE_K"] % group_size == 0
+    return config
diff --git a/mstar/utils/fused_moe/runner.py b/mstar/utils/fused_moe/runner.py
index 0926fa49d..2ed68e416 100644
--- a/mstar/utils/fused_moe/runner.py
+++ b/mstar/utils/fused_moe/runner.py
@@ -16,8 +16,10 @@
     act_and_mul_triton,
     get_default_config,
     invoke_fused_moe_kernel,
+    invoke_fused_moe_kernel_w4a16,
     moe_sum_reduce_triton,
 )
+from mstar.utils.quantization import QuantizationData, QuantizationType
 
 
 def _tl_compute_type(dtype: torch.dtype) -> tl.dtype:
@@ -28,6 +30,53 @@ def _tl_compute_type(dtype: torch.dtype) -> tl.dtype:
     raise ValueError(f"fused_experts: unsupported dtype {dtype}; use bf16 or fp16")
 
 
+def _validate_expert_shapes(
+    hidden: int,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    quant: QuantizationData | None,
+) -> tuple[int, int, int]:
+    """Check the stacked expert weights against ``hidden`` and the quant scheme.
+
+    Returns ``(num_experts, two_inter, inter)``. Raising here — rather than
+    dispatching on whichever argument happened to be non-None — is what keeps a
+    newly-added scheme from silently falling through to the bf16 kernel.
+    """
+    if quant is None:
+        assert w1.is_contiguous(), "w1 must be contiguous"
+        assert w2.is_contiguous(), "w2 must be contiguous"
+        E, two_inter, k_in = w1.shape
+        assert k_in == hidden, f"w1 last dim {k_in} != hidden {hidden}"
+        _, w2_hidden, inter = w2.shape
+        assert w2_hidden == hidden, f"w2 dim[1] {w2_hidden} != hidden {hidden}"
+        assert two_inter == 2 * inter, f"w1 dim[1] {two_inter} != 2 * w2 dim[2] {2 * inter}"
+        return E, two_inter, inter
+
+    if quant.quant_type is not QuantizationType.W4A16:
+        raise ValueError(
+            f"fused_experts: no kernel for {quant.quant_type}; the Triton MoE path "
+            f"implements {QuantizationType.W4A16} only"
+        )
+
+    pack_factor = quant.pack_factor
+    assert w1.dtype == torch.int32 and w2.dtype == torch.int32, (
+        "W4A16 path expects packed int32 weights"
+    )
+    assert w1.is_contiguous(), "w1 (packed) must be contiguous"
+    assert w2.is_contiguous(), "w2 (packed) must be contiguous"
+    E, two_inter, k1_packed = w1.shape
+    assert k1_packed == hidden // pack_factor, (
+        f"w1 packed last dim {k1_packed} != hidden//pack_factor {hidden // pack_factor}"
+    )
+    _, w2_hidden, k2_packed = w2.shape
+    assert w2_hidden == hidden, f"w2 dim[1] {w2_hidden} != hidden {hidden}"
+    inter = two_inter // 2
+    assert k2_packed == inter // pack_factor, (
+        f"w2 packed last dim {k2_packed} != inter//pack_factor {inter // pack_factor}"
+    )
+    return E, two_inter, inter
+
+
 def fused_experts(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
@@ -36,6 +85,7 @@ def fused_experts(
     topk_ids: torch.Tensor,
     activation: str = "silu",
     reduce_results: bool = True,
+    quant: QuantizationData | None = None,
 ) -> torch.Tensor:
     """Grouped-GEMM Triton MoE dispatch.
 
@@ -48,10 +98,11 @@ def fused_experts(
         ``(num_experts, 2 * moe_intermediate_size, hidden)``.  Matches the
         ``experts.gate_up_proj`` parameter the WeightConverter already
         produces in :mod:`mstar.model.qwen3_omni.qwen3_omni_model`.
+        Packed int32 on the W4A16 path.
     w2 : torch.Tensor
         Down projection weights, shape
         ``(num_experts, hidden, moe_intermediate_size)``.  Matches
-        ``experts.down_proj``.
+        ``experts.down_proj``. Packed int32 on the W4A16 path.
     topk_weights : torch.Tensor
         ``(tokens, top_k)``, routing probabilities (possibly
         renormalized).  Dtype matches ``hidden_states``.
@@ -64,6 +115,11 @@ def fused_experts(
         ``(tokens, hidden)``. If False, skip the sum-reduce and return
         ``(tokens, top_k, hidden)`` — the caller is responsible for the
         reduce (e.g. after an all-reduce for TP).
+    quant : QuantizationData | None
+        Scheme descriptor plus its scales/zero points. ``None`` (default) is the
+        bf16/fp16 path. A :class:`~mstar.utils.quantization.W4A16Data` selects the
+        packed-INT4 kernel and supplies its group scales; an unimplemented scheme
+        raises rather than silently taking the bf16 path.
 
     Returns
     -------
@@ -72,25 +128,23 @@ def fused_experts(
         ``(tokens, top_k, hidden)``.
     """
     assert hidden_states.is_contiguous(), "hidden_states must be contiguous"
-    assert w1.is_contiguous(), "w1 must be contiguous"
-    assert w2.is_contiguous(), "w2 must be contiguous"
     assert hidden_states.dim() == 2
     assert topk_weights.shape == topk_ids.shape
+    # Only weights are quantized; activations stay bf16/fp16.
     assert hidden_states.dtype in (torch.bfloat16, torch.float16)
 
     num_tokens, hidden = hidden_states.shape
-    E, two_inter, k_in = w1.shape
-    assert k_in == hidden, f"w1 last dim {k_in} != hidden {hidden}"
-    _, w2_hidden, inter = w2.shape
-    assert w2_hidden == hidden, f"w2 dim[1] {w2_hidden} != hidden {hidden}"
-    assert two_inter == 2 * inter, f"w1 dim[1] {two_inter} != 2 * w2 dim[2] {2 * inter}"
+    E, two_inter, inter = _validate_expert_shapes(hidden, w1, w2, quant)
+    group_size = quant.group_size if quant is not None else None
 
     top_k = topk_ids.shape[1]
     # moe_align_block_size expects int32; torch.topk returns int64.
     topk_ids = topk_ids.to(torch.int32).contiguous()
     topk_weights = topk_weights.contiguous()
 
-    config = get_default_config(M=num_tokens, E=E, N=two_inter, K=hidden, top_k=top_k)
+    config = get_default_config(
+        M=num_tokens, E=E, N=two_inter, K=hidden, top_k=top_k, group_size=group_size,
+    )
     compute_type = _tl_compute_type(hidden_states.dtype)
 
     # 1. Token permute + per-expert block alignment.
@@ -115,20 +169,41 @@ def fused_experts(
     )
 
     # 3. Gate+up GEMM: cache1[slot] = hidden[slot // top_k] @ w1[expert].T
-    invoke_fused_moe_kernel(
-        A=hidden_states,
-        B=w1,
-        C=cache1,
-        topk_weights=topk_weights,
-        topk_ids=topk_ids,
-        sorted_token_ids=sorted_token_ids,
-        expert_ids=expert_ids,
-        num_tokens_post_padded=num_tokens_post_padded,
-        mul_routed_weight=False,
-        top_k=top_k,
-        config=config,
-        compute_type=compute_type,
-    )
+    if quant is not None:
+        invoke_fused_moe_kernel_w4a16(
+            A=hidden_states,
+            B_packed=w1,
+            C=cache1,
+            B_scale=quant.w1_scale,
+            B_zp=quant.w1_zp,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            sorted_token_ids=sorted_token_ids,
+            expert_ids=expert_ids,
+            num_tokens_post_padded=num_tokens_post_padded,
+            mul_routed_weight=False,
+            top_k=top_k,
+            config=config,
+            compute_type=compute_type,
+            K=hidden,
+            pack_factor=quant.pack_factor,
+            group_size=quant.group_size,
+        )
+    else:
+        invoke_fused_moe_kernel(
+            A=hidden_states,
+            B=w1,
+            C=cache1,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            sorted_token_ids=sorted_token_ids,
+            expert_ids=expert_ids,
+            num_tokens_post_padded=num_tokens_post_padded,
+            mul_routed_weight=False,
+            top_k=top_k,
+            config=config,
+            compute_type=compute_type,
+        )
 
     # 4. SwiGLU: cache2[slot] = silu(gate) * up
     act_and_mul_triton(cache1, cache2, activation=activation)
@@ -136,20 +211,41 @@ def fused_experts(
     # 5. Down GEMM (weighted): cache3[slot] = topk_weight[slot] * (cache2[slot] @ w2[expert].T)
     # top_k=1 for this GEMM so the kernel's offs_token // top_k is identity
     # -- it reads cache2 rows directly instead of the (slot // top_k)-th source row.
-    invoke_fused_moe_kernel(
-        A=cache2,
-        B=w2,
-        C=cache3.view(m_topk, hidden),
-        topk_weights=topk_weights,
-        topk_ids=topk_ids,
-        sorted_token_ids=sorted_token_ids,
-        expert_ids=expert_ids,
-        num_tokens_post_padded=num_tokens_post_padded,
-        mul_routed_weight=True,
-        top_k=1,
-        config=config,
-        compute_type=compute_type,
-    )
+    if quant is not None:
+        invoke_fused_moe_kernel_w4a16(
+            A=cache2,
+            B_packed=w2,
+            C=cache3.view(m_topk, hidden),
+            B_scale=quant.w2_scale,
+            B_zp=quant.w2_zp,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            sorted_token_ids=sorted_token_ids,
+            expert_ids=expert_ids,
+            num_tokens_post_padded=num_tokens_post_padded,
+            mul_routed_weight=True,
+            top_k=1,
+            config=config,
+            compute_type=compute_type,
+            K=inter,
+            pack_factor=quant.pack_factor,
+            group_size=quant.group_size,
+        )
+    else:
+        invoke_fused_moe_kernel(
+            A=cache2,
+            B=w2,
+            C=cache3.view(m_topk, hidden),
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            sorted_token_ids=sorted_token_ids,
+            expert_ids=expert_ids,
+            num_tokens_post_padded=num_tokens_post_padded,
+            mul_routed_weight=True,
+            top_k=1,
+            config=config,
+            compute_type=compute_type,
+        )
 
     # 6. Sum over the top-k slots -> (tokens, hidden).
     if not reduce_results:
diff --git a/mstar/utils/marlin/__init__.py b/mstar/utils/marlin/__init__.py
new file mode 100644
index 000000000..74a404c88
--- /dev/null
+++ b/mstar/utils/marlin/__init__.py
@@ -0,0 +1,4 @@
+"""Vendored Marlin W4A16 CUDA kernels for mstar."""
+from mstar.utils.marlin.loader import is_marlin_available
+
+__all__ = ["is_marlin_available"]
diff --git a/mstar/utils/marlin/csrc/core/scalar_type.hpp b/mstar/utils/marlin/csrc/core/scalar_type.hpp
new file mode 100644
index 000000000..b6f39ed79
--- /dev/null
+++ b/mstar/utils/marlin/csrc/core/scalar_type.hpp
@@ -0,0 +1,360 @@
+#pragma once
+
+#include <cstdint>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <variant>
+
+// For STD_TORCH_CHECK
+#include <torch/headeronly/util/Exception.h>
+
+namespace vllm {
+
+//
+//  ScalarType can represent a wide range of floating point and integer types,
+//  in particular it can be used to represent sub-byte data types (something
+//  that torch.dtype currently does not support).
+//
+//  The type definitions on the Python side can be found in: vllm/scalar_type.py
+//  these type definitions should be kept up to date with any Python API changes
+//  here.
+//
+class ScalarType {
+ public:
+  enum NanRepr : uint8_t {
+    NAN_NONE = 0,                // nans are not supported
+    NAN_IEEE_754 = 1,            // nans are: exp all 1s, mantissa not all 0s
+    NAN_EXTD_RANGE_MAX_MIN = 2,  // nans are: exp all 1s, mantissa all 1s
+
+    NAN_REPR_ID_MAX
+  };
+
+  constexpr ScalarType(uint8_t exponent, uint8_t mantissa, bool signed_,
+                       int32_t bias, bool finite_values_only = false,
+                       NanRepr nan_repr = NAN_IEEE_754)
+      : exponent(exponent),
+        mantissa(mantissa),
+        signed_(signed_),
+        bias(bias),
+        finite_values_only(finite_values_only),
+        nan_repr(nan_repr) {};
+
+  static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) {
+    return ScalarType(0, size_bits - 1, true, bias);
+  }
+
+  static constexpr ScalarType uint(uint8_t size_bits, int32_t bias = 0) {
+    return ScalarType(0, size_bits, false, bias);
+  }
+
+  // IEEE 754 compliant floating point type
+  static constexpr ScalarType float_IEEE754(uint8_t exponent,
+                                            uint8_t mantissa) {
+    STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
+    return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754);
+  }
+
+  // IEEE 754 non-compliant floating point type
+  static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa,
+                                     bool finite_values_only,
+                                     NanRepr nan_repr) {
+    STD_TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
+    STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
+    STD_TORCH_CHECK(
+        nan_repr != NAN_IEEE_754,
+        "use `float_IEEE754` constructor for floating point types that "
+        "follow IEEE 754 conventions");
+    return ScalarType(exponent, mantissa, true, 0, finite_values_only,
+                      nan_repr);
+  }
+
+  uint8_t const exponent;  // size of the exponent field (0 for integer types)
+  uint8_t const mantissa;  // size of the mantissa field (size of the integer
+                           // excluding the sign bit for integer types)
+  bool const signed_;  // flag if the type supports negative numbers (i.e. has a
+                       // sign bit)
+  int32_t const bias;  // stored values equal value + bias,
+                       // used for quantized type
+
+  // Extra Floating point info
+  bool const finite_values_only;  // i.e. no +/-inf if true
+  NanRepr const nan_repr;         // how NaNs are represented
+                                  // (not applicable for integer types)
+
+  using Id = int64_t;
+
+ private:
+  // Field size in id
+  template <typename T_>
+  static constexpr size_t member_id_field_width() {
+    using T = std::decay_t<T_>;
+    return std::is_same_v<T, bool> ? 1 : sizeof(T) * 8;
+  }
+
+  template <typename Fn, typename Init, typename Member, typename... Rest>
+  static constexpr auto reduce_members_helper(Fn f, Init val, Member member,
+                                              Rest... rest) {
+    auto new_val = f(val, member);
+    if constexpr (sizeof...(rest) > 0) {
+      return reduce_members_helper(f, new_val, rest...);
+    } else {
+      return new_val;
+    };
+  }
+
+  template <typename Fn, typename Init>
+  constexpr auto reduce_members(Fn f, Init init) const {
+    // Should be in constructor order for `from_id`
+    return reduce_members_helper(f, init, exponent, mantissa, signed_, bias,
+                                 finite_values_only, nan_repr);
+  };
+
+  template <typename Fn, typename Init>
+  static constexpr auto reduce_member_types(Fn f, Init init) {
+    constexpr auto dummy_type = ScalarType(0, 0, false, 0, false, NAN_NONE);
+    return dummy_type.reduce_members(f, init);
+  };
+
+  static constexpr auto id_size_bits() {
+    return reduce_member_types(
+        [](int acc, auto member) -> int {
+          return acc + member_id_field_width<decltype(member)>();
+        },
+        0);
+  }
+
+ public:
+  // unique id for this scalar type that can be computed at compile time for
+  //  c++17 template specialization this is not needed once we migrate to
+  //  c++20 and can pass literal classes as template parameters
+  constexpr Id id() const {
+    static_assert(id_size_bits() <= sizeof(Id) * 8,
+                  "ScalarType id is too large to be stored");
+
+    auto or_and_advance = [](std::pair<Id, uint32_t> result,
+                             auto member) -> std::pair<Id, uint32_t> {
+      auto [id, bit_offset] = result;
+      auto constexpr bits = member_id_field_width<decltype(member)>();
+      return {id | (int64_t(member) & ((uint64_t(1) << bits) - 1))
+                       << bit_offset,
+              bit_offset + bits};
+    };
+    return reduce_members(or_and_advance, std::pair<Id, uint32_t>{}).first;
+  }
+
+  // create a ScalarType from an id, for c++17 template specialization,
+  //  this is not needed once we migrate to c++20 and can pass literal
+  //  classes as template parameters
+  static constexpr ScalarType from_id(Id id) {
+    auto extract_and_advance = [id](auto result, auto member) {
+      using T = decltype(member);
+      auto [tuple, bit_offset] = result;
+      auto constexpr bits = member_id_field_width<T>();
+      auto extracted_val = static_cast<T>((int64_t(id) >> bit_offset) &
+                                          ((uint64_t(1) << bits) - 1));
+      auto new_tuple = std::tuple_cat(tuple, std::make_tuple(extracted_val));
+      return std::pair<decltype(new_tuple), int>{new_tuple, bit_offset + bits};
+    };
+
+    auto [tuple_args, _] = reduce_member_types(extract_and_advance,
+                                               std::pair<std::tuple<>, int>{});
+    return std::apply([](auto... args) { return ScalarType(args...); },
+                      tuple_args);
+  }
+
+  constexpr int64_t size_bits() const {
+    return mantissa + exponent + is_signed();
+  }
+  constexpr bool is_signed() const { return signed_; }
+  constexpr bool is_integer() const { return exponent == 0; }
+  constexpr bool is_floating_point() const { return exponent > 0; }
+  constexpr bool is_ieee_754() const {
+    return is_floating_point() && finite_values_only == false &&
+           nan_repr == NAN_IEEE_754;
+  }
+  constexpr bool has_nans() const {
+    return is_floating_point() && nan_repr != NAN_NONE;
+  }
+  constexpr bool has_infs() const {
+    return is_floating_point() && finite_values_only == false;
+  }
+  constexpr bool has_bias() const { return bias != 0; }
+
+ private:
+  double _floating_point_max() const {
+    STD_TORCH_CHECK(mantissa <= 52 && exponent <= 11,
+                    "Cannot represent max/min as a double for type ", str());
+
+    uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1;
+    if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) {
+      max_mantissa -= 1;
+    }
+
+    uint64_t max_exponent = (uint64_t(1) << exponent) - 2;
+    if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) {
+      STD_TORCH_CHECK(exponent < 11,
+                      "Cannot represent max/min as a double for type ", str());
+      max_exponent += 1;
+    }
+
+    // adjust the exponent to match that of a double
+    //  for now we assume the exponent bias is the standard 2^(e-1) -1, (where e
+    //  is the exponent bits), there is some precedent for non-standard biases,
+    //  example `float8_e4m3b11fnuz` here: https://github.com/jax-ml/ml_dtypes
+    //  but to avoid premature over complication we are just assuming the
+    //  standard exponent bias until there is a need to support non-standard
+    //  biases
+    uint64_t exponent_bias = (uint64_t(1) << (exponent - 1)) - 1;
+    uint64_t exponent_bias_double = (uint64_t(1) << 10) - 1;  // double e = 11
+
+    uint64_t max_exponent_double =
+        max_exponent - exponent_bias + exponent_bias_double;
+
+    // shift the mantissa into the position for a double and
+    // the exponent
+    uint64_t double_raw =
+        (max_mantissa << (52 - mantissa)) | (max_exponent_double << 52);
+
+    return *reinterpret_cast<double*>(&double_raw);
+  }
+
+  constexpr std::variant<int64_t, double> _raw_max() const {
+    if (is_floating_point()) {
+      return {_floating_point_max()};
+    } else {
+      STD_TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
+                      "Cannot represent max as a int64_t");
+      return {(int64_t(1) << mantissa) - 1};
+    }
+  }
+
+  constexpr std::variant<int64_t, double> _raw_min() const {
+    if (is_floating_point()) {
+      STD_TORCH_CHECK(
+          is_signed(),
+          "We currently assume all floating point types are signed");
+      constexpr uint64_t sign_bit_double = (uint64_t(1) << 63);
+
+      double max = _floating_point_max();
+      uint64_t max_raw = *reinterpret_cast<uint64_t*>(&max);
+      uint64_t min_raw = max_raw | sign_bit_double;
+      return {*reinterpret_cast<double*>(&min_raw)};
+    } else {
+      STD_TORCH_CHECK(!is_signed() || size_bits() <= 64,
+                      "Cannot represent min as a int64_t");
+      if (is_signed()) {
+        // set the top bit to 1 (i.e. INT64_MIN) and the rest to 0
+        // then perform an arithmetic shift right to set all the bits above
+        // (size_bits() - 1) to 1
+        return {INT64_MIN >> (64 - size_bits())};
+      } else {
+        return {int64_t(0)};
+      }
+    }
+  }
+
+ public:
+  // Max representable value for this scalar type.
+  // (accounting for bias if there is one)
+  constexpr std::variant<int64_t, double> max() const {
+    return std::visit(
+        [this](auto x) -> std::variant<int64_t, double> { return {x - bias}; },
+        _raw_max());
+  }
+
+  // Min representable value for this scalar type.
+  // (accounting for bias if there is one)
+  constexpr std::variant<int64_t, double> min() const {
+    return std::visit(
+        [this](auto x) -> std::variant<int64_t, double> { return {x - bias}; },
+        _raw_min());
+  }
+
+  std::string str() const {
+    /* naming generally follows: https://github.com/jax-ml/ml_dtypes
+     * for floating point types (leading f) the scheme is:
+     *  `float<size_bits>_e<exponent_bits>m<mantissa_bits>[flags]`
+     *  flags:
+     *  - no-flags: means it follows IEEE 754 conventions
+     *  - f: means finite values only (no infinities)
+     *  - n: means nans are supported (non-standard encoding)
+     * for integer types the scheme is:
+     *  `[u]int<size_bits>[b<bias>]`
+     *  - if bias is not present it means its zero
+     */
+    if (is_floating_point()) {
+      auto ret = "float" + std::to_string(size_bits()) + "_e" +
+                 std::to_string(exponent) + "m" + std::to_string(mantissa);
+      if (!is_ieee_754()) {
+        if (finite_values_only) {
+          ret += "f";
+        }
+        if (nan_repr != NAN_NONE) {
+          ret += "n";
+        }
+      }
+      return ret;
+    } else {
+      auto ret = ((is_signed()) ? "int" : "uint") + std::to_string(size_bits());
+      if (has_bias()) {
+        ret += "b" + std::to_string(bias);
+      }
+      return ret;
+    }
+  }
+
+  constexpr bool operator==(ScalarType const& other) const {
+    return mantissa == other.mantissa && exponent == other.exponent &&
+           bias == other.bias && signed_ == other.signed_ &&
+           finite_values_only == other.finite_values_only &&
+           nan_repr == other.nan_repr;
+  }
+};
+
+using ScalarTypeId = ScalarType::Id;
+
+// "rust style" names generally following:
+//   https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L60-L70
+static inline constexpr auto kS4 = ScalarType::int_(4);
+static inline constexpr auto kU4 = ScalarType::uint(4);
+static inline constexpr auto kU4B8 = ScalarType::uint(4, 8);
+static inline constexpr auto kS8 = ScalarType::int_(8);
+static inline constexpr auto kU8 = ScalarType::uint(8);
+static inline constexpr auto kU8B128 = ScalarType::uint(8, 128);
+
+static inline constexpr auto kFE2M1f =
+    ScalarType::float_(2, 1, true, ScalarType::NAN_NONE);
+static inline constexpr auto kFE3M2f =
+    ScalarType::float_(3, 2, true, ScalarType::NAN_NONE);
+static inline constexpr auto kFE4M3fn =
+    ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN);
+static inline constexpr auto kFE8M0fnu =
+    ScalarType(8, 0, false, 0, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN);
+static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2);
+static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7);
+static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10);
+
+// Fixed width style names, generally following:
+//  https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57
+static inline constexpr auto kInt4 = kS4;
+static inline constexpr auto kUint4 = kU4;
+static inline constexpr auto kUint4b8 = kU4B8;
+static inline constexpr auto kInt8 = kS8;
+static inline constexpr auto kUint8 = kU8;
+static inline constexpr auto kUint8b128 = kU8B128;
+
+static inline constexpr auto kFloat4_e2m1f = kFE2M1f;
+static inline constexpr auto kFloat6_e3m2f = kFE3M2f;
+static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn;
+static inline constexpr auto kFloat8_e5m2 = kFE5M2;
+static inline constexpr auto kFloat16_e8m7 = kFE8M7;
+static inline constexpr auto kFloat16_e5m10 = kFE5M10;
+
+// colloquial names
+static inline constexpr auto kHalf = kFE5M10;
+static inline constexpr auto kFloat16 = kHalf;
+static inline constexpr auto kBFloat16 = kFE8M7;
+
+static inline constexpr auto kFloat16Id = kFloat16.id();
+};  // namespace vllm
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py
new file mode 100644
index 000000000..3b6fc79f0
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/generate_kernels.py
@@ -0,0 +1,230 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import glob
+import itertools
+import os
+import subprocess
+import sys
+
+import jinja2
+
+ARCHS = []
+SUPPORT_FP8 = False
+SUPPORT_SM75 = False
+SUPPORT_SM80 = False
+for arch in sys.argv[1].split(","):
+    arch = arch[: arch.index(".") + 2].replace(".", "")
+    arch = int(arch)
+    # SM89 and the SM12x family (SM120 RTX 5090, SM121 DGX Spark GB10)
+    # fully support mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
+    # SM90 and SM100 can use this PTX, but it’s simulated
+    # with FP16 MMA, so it cannot achieve any acceleration.
+    if arch == 89 or arch // 10 == 12:
+        SUPPORT_FP8 = True
+    if arch >= 80:
+        SUPPORT_SM80 = True
+    if arch == 75:
+        SUPPORT_SM75 = True
+
+FILE_HEAD_COMMENT = """
+// auto generated by generate_kernels.py
+// clang-format off
+""".lstrip()
+
+FILE_HEAD = (
+    FILE_HEAD_COMMENT
+    + """
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+"""
+)
+
+TEMPLATE = (
+    "template __global__ void Marlin<"
+    "{{a_type_id}}, "
+    "{{b_type_id}}, "
+    "{{c_type_id}}, "
+    "{{s_type_id}}, "
+    "{{threads}}, "
+    "{{thread_m_blocks}}, "
+    "{{thread_n_blocks}}, "
+    "{{thread_k_blocks}}, "
+    "{{m_block_size_8}}, "
+    "{{stages}}, "
+    "{{group_blocks}}, "
+    "{{is_zp_float}}>"
+    "( MARLIN_KERNEL_PARAMS );"
+)
+
+THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)]
+
+THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
+
+# mstar: trimmed to GPTQ symmetric INT4 (kU4B8) with fp16/bf16 activation only —
+# the sole scheme Kimi-K2.7 (compressed-tensors W4A16) uses. All the AWQ / INT8 /
+# FP8 / NVFP4 / MXFP4 / W4A8 variants are dropped so the JIT compile stays small.
+# Upstream QUANT_CONFIGS kept in git history if another scheme is ever needed.
+QUANT_CONFIGS = [
+    # GPTQ-INT4
+    {
+        "b_type": "kU4B8",
+        "thread_configs": THREAD_CONFIGS,
+        "thread_m_blocks": THREAD_M_BLOCKS,
+        "group_blocks": [-1, 0, 2, 4, 8],
+    },
+]
+
+
+def remove_old_kernels():
+    for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"):
+        subprocess.call(["rm", "-f", filename])
+
+    filename = os.path.dirname(__file__) + "/kernel_selector.h"
+    subprocess.call(["rm", "-f", filename])
+
+
+def generate_new_kernels():
+    result_dict = {}
+    sm_75_result_dict = {}
+
+    for quant_config in QUANT_CONFIGS:
+        c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"])
+        a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"])
+        b_type = quant_config["b_type"]
+        all_group_blocks = quant_config["group_blocks"]
+        all_m_blocks = quant_config["thread_m_blocks"]
+        all_thread_configs = quant_config["thread_configs"]
+
+        for a_type, c_type in itertools.product(a_types, c_types):
+            if not SUPPORT_FP8 and a_type == "kFE4M3fn":
+                continue
+            if "16" in a_type and "16" in c_type and a_type != c_type:
+                continue
+            s_type = quant_config.get("s_type", c_type)
+            if (a_type, b_type, c_type) not in result_dict:
+                result_dict[(a_type, b_type, c_type)] = []
+                if a_type in ["kFloat16", "kS8"] and c_type == "kFloat16":
+                    sm_75_result_dict[(a_type, b_type, c_type)] = []
+
+            for group_blocks, m_blocks, thread_configs in itertools.product(
+                all_group_blocks, all_m_blocks, all_thread_configs
+            ):
+                thread_k, thread_n, threads = thread_configs
+
+                if threads == 256:
+                    # for small batch (m_blocks == 1),
+                    #     we only need (128, 128, 256)
+                    # for large batch (m_blocks > 1),
+                    #     we only need (64, 256, 256)
+                    if m_blocks <= 1 and (thread_k, thread_n) != (128, 128):
+                        continue
+                    if m_blocks > 1 and (thread_k, thread_n) != (64, 256):
+                        continue
+
+                config = {
+                    "threads": threads,
+                    "s_type": s_type,
+                    "thread_m_blocks": max(m_blocks, 1),
+                    "thread_k_blocks": thread_k // 16,
+                    "thread_n_blocks": thread_n // 16,
+                    "m_block_size_8": "true" if m_blocks == 0.5 else "false",
+                    "stages": 4,
+                    "group_blocks": group_blocks,
+                    "is_zp_float": "false",
+                }
+
+                if SUPPORT_SM80:
+                    result_dict[(a_type, b_type, c_type)].append(config)
+                if (a_type, b_type, c_type) in sm_75_result_dict and SUPPORT_SM75:
+                    config_sm75 = config.copy()
+                    config_sm75["stages"] = 2
+                    sm_75_result_dict[(a_type, b_type, c_type)].append(config_sm75)
+
+    kernel_selector_str = FILE_HEAD_COMMENT
+
+    for result_dict_tmp in [result_dict, sm_75_result_dict]:
+        for (a_type, b_type, c_type), config_list in result_dict_tmp.items():
+            all_template_str_list = []
+            if not config_list:
+                continue
+            for config in config_list:
+                s_type = config["s_type"]
+                template_str = jinja2.Template(TEMPLATE).render(
+                    a_type_id=f"vllm::{a_type}.id()",
+                    b_type_id=f"vllm::{b_type}.id()",
+                    c_type_id=f"vllm::{c_type}.id()",
+                    s_type_id=f"vllm::{s_type}.id()",
+                    **config,
+                )
+                all_template_str_list.append(template_str)
+
+                conditions = [
+                    f"a_type == vllm::{a_type}",
+                    f"b_type == vllm::{b_type}",
+                    f"c_type == vllm::{c_type}",
+                    f"s_type == vllm::{s_type}",
+                    f"threads == {config['threads']}",
+                    f"thread_m_blocks == {config['thread_m_blocks']}",
+                    f"thread_n_blocks == {config['thread_n_blocks']}",
+                    f"thread_k_blocks == {config['thread_k_blocks']}",
+                    f"m_block_size_8 == {config['m_block_size_8']}",
+                    f"stages == {config['stages']}",
+                    f"group_blocks == {config['group_blocks']}",
+                    f"is_zp_float == {config['is_zp_float']}",
+                ]
+                conditions = " && ".join(conditions)
+
+                if kernel_selector_str == FILE_HEAD_COMMENT:
+                    kernel_selector_str += f"if ({conditions})\n  kernel = "
+                else:
+                    kernel_selector_str += f"else if ({conditions})\n  kernel = "
+
+                kernel_template2 = (
+                    "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, "
+                    "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, "
+                    "{{thread_n_blocks}}, {{thread_k_blocks}}, "
+                    "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, "
+                    "{{is_zp_float}}>;"
+                )
+
+                kernel_selector_str += (
+                    jinja2.Template(kernel_template2).render(
+                        a_type_id=f"vllm::{a_type}.id()",
+                        b_type_id=f"vllm::{b_type}.id()",
+                        c_type_id=f"vllm::{c_type}.id()",
+                        s_type_id=f"vllm::{s_type}.id()",
+                        **config,
+                    )
+                    + "\n"
+                )
+
+            file_content = FILE_HEAD + "\n\n"
+            file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
+            if a_type == "kFE4M3fn":
+                filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
+            elif result_dict_tmp is sm_75_result_dict:
+                filename = f"sm75_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
+            else:
+                filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
+
+            filename = filename.lower()
+
+            with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
+                f.write(file_content)
+
+    if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
+        kernel_selector_str += (
+            "else if (a_type == vllm::kFE4M3fn)\n"
+            "  STD_TORCH_CHECK(false, "
+            '"marlin kernel with fp8 activation is not built.");'
+        )
+
+    with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f:
+        f.write(kernel_selector_str)
+
+
+if __name__ == "__main__":
+    remove_old_kernels()
+    generate_new_kernels()
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h
new file mode 100644
index 000000000..783736ab5
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel.h
@@ -0,0 +1,47 @@
+
+#ifndef MARLIN_NAMESPACE_NAME
+  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#endif
+
+#include "libtorch_stable/quantization/marlin/marlin.cuh"
+#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
+#include "core/scalar_type.hpp"
+
+#define MARLIN_KERNEL_PARAMS                                          \
+  const int4 *__restrict__ A, const int4 *__restrict__ B,             \
+      int4 *__restrict__ C, int4 *__restrict__ C_tmp,                 \
+      const int4 *__restrict__ b_bias_ptr,                            \
+      const float *__restrict__ a_scales_ptr,                         \
+      const int4 *__restrict__ scales_ptr,                            \
+      const float *__restrict__ global_scale_ptr,                     \
+      const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
+      const int32_t *__restrict__ sorted_token_ids_ptr,               \
+      const int32_t *__restrict__ expert_ids_ptr,                     \
+      const int32_t *__restrict__ num_tokens_past_padded_ptr,         \
+      const float *__restrict__ topk_weights_ptr, int top_k,          \
+      bool mul_topk_weights, int num_groups, int prob_m, int prob_n,  \
+      int prob_k, int *locks, bool has_bias, bool use_atomic_add,     \
+      bool use_fp32_reduce
+
+namespace MARLIN_NAMESPACE_NAME {
+template <const vllm::ScalarTypeId a_type_id,  // A ScalarType id
+          const vllm::ScalarTypeId b_type_id,  // B ScalarType id
+          const vllm::ScalarTypeId c_type_id,  // C ScalarType id
+          const vllm::ScalarTypeId s_type_id,  // B_SCALE ScalarType id
+          const int threads,          // number of threads in a threadblock
+          const int thread_m_blocks,  // number of 16x16 blocks in the m
+                                      // dimension (batchsize) of the
+                                      // threadblock
+          const int thread_n_blocks,  // same for n dimension (output)
+          const int thread_k_blocks,  // same for k dimension (reduction)
+          const bool m_block_size_8,  // whether m_block_size == 8
+                                      // only works when thread_m_blocks == 1
+          const int stages,  // number of stages for the async global->shared
+                             // fetch pipeline
+          const int group_blocks,  // number of consecutive 16x16 blocks
+                                   // with a separate quantization scale
+          const bool is_zp_float   // is zero point of float16 type?
+          >
+__global__ void Marlin(MARLIN_KERNEL_PARAMS);
+
+}
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h
new file mode 100644
index 000000000..8551dcd34
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_selector.h
@@ -0,0 +1,304 @@
+// auto generated by generate_kernels.py
+// clang-format off
+if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 0, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>;
+else if (a_type == vllm::kFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kFloat16 && s_type == vllm::kFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == -1 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 0 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 0, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 2 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 4 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == true && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 1 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 2 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 2 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 3 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 3 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 256 && thread_m_blocks == 4 && thread_n_blocks == 16 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 8 && thread_k_blocks == 4 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>;
+else if (a_type == vllm::kBFloat16 && b_type == vllm::kU4B8 && c_type == vllm::kBFloat16 && s_type == vllm::kBFloat16 && threads == 128 && thread_m_blocks == 4 && thread_n_blocks == 4 && thread_k_blocks == 8 && m_block_size_8 == false && stages == 4 && group_blocks == 8 && is_zp_float == false)
+  kernel = Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>;
+else if (a_type == vllm::kFE4M3fn)
+  STD_TORCH_CHECK(false, "marlin kernel with fp8 activation is not built.");
\ No newline at end of file
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_moe.cu b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_moe.cu
new file mode 100644
index 000000000..a7452bda1
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_moe.cu
@@ -0,0 +1,717 @@
+/*
+ * Modified by Neural Magic
+ * Copyright (C) Marlin.2024 Elias Frantar
+ *
+ * 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.
+ */
+
+/*
+ * Adapted from https://github.com/IST-DASLab/marlin
+ */
+
+#ifndef MARLIN_NAMESPACE_NAME
+  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#endif
+
+#include "kernel.h"
+
+// mstar: vendored from vLLM. The __global__ kernels + device host helpers
+// (marlin_mm etc., lines below) are kept VERBATIM; only the top includes, the
+// public host wrapper, and the op registration are ported from vLLM's
+// torch::stable ABI to the classic torch ABI mstar's vendored ops use (see
+// utils/fused_moe/csrc/moe_align_block_size.cu). STD_TORCH_CHECK, used by the
+// verbatim device code, is provided by core/scalar_type.hpp's
+// <torch/headeronly/util/Exception.h> include (reachable via kernel.h).
+//
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#include <torch/all.h>
+#include <torch/library.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t)               \
+  static_assert(std::is_same<scalar_t, half>::value ||          \
+                    std::is_same<scalar_t, nv_bfloat16>::value, \
+                "only float16 and bfloat16 is supported");
+
+namespace MARLIN_NAMESPACE_NAME {
+
+__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){};
+
+using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS);
+
+// For a given "a" of size [M,K] performs a permutation of the K columns based
+// on the given "perm" indices.
+template <int moe_block_size>
+__global__ void permute_cols_kernel(
+    int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr,
+    int4* __restrict__ out_int4_ptr,
+    const int32_t* __restrict__ sorted_token_ids_ptr,
+    const int32_t* __restrict__ expert_ids_ptr,
+    const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m,
+    int size_k, int top_k) {
+  int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
+  int num_moe_blocks = div_ceil(num_tokens_past_padded, moe_block_size);
+  int32_t block_sorted_ids[moe_block_size];
+  int block_num_valid_tokens = 0;
+  int64_t old_expert_id = 0;
+  int64_t expert_id = 0;
+  int row_stride = size_k * sizeof(half) / 16;
+
+  auto read_moe_block_data = [&](int block_id) {
+    block_num_valid_tokens = moe_block_size;
+    int4* tmp_block_sorted_ids = reinterpret_cast<int4*>(block_sorted_ids);
+    for (int i = 0; i < moe_block_size / 4; i++) {
+      tmp_block_sorted_ids[i] =
+          ((int4*)sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i];
+    }
+    for (int i = 0; i < moe_block_size; i++) {
+      if (block_sorted_ids[i] >= size_m * top_k) {
+        block_num_valid_tokens = i;
+        break;
+      };
+    }
+  };
+
+  auto permute_row = [&](int row) {
+    int iters = size_k / default_threads;
+    int rest = size_k % default_threads;
+
+    int in_offset = (row / top_k) * row_stride;
+    int out_offset = row * row_stride;
+
+    half const* a_row_half =
+        reinterpret_cast<half const*>(a_int4_ptr + in_offset);
+    half* out_half = reinterpret_cast<half*>(out_int4_ptr + out_offset);
+
+    int base_k = 0;
+
+    for (int i = 0; i < iters; i++) {
+      auto cur_k = base_k + threadIdx.x;
+      int src_pos = perm_int_ptr[cur_k];
+
+      out_half[cur_k] = a_row_half[src_pos];
+
+      base_k += default_threads;
+    }
+
+    if (rest) {
+      if (threadIdx.x < rest) {
+        auto cur_k = base_k + threadIdx.x;
+        int src_pos = perm_int_ptr[cur_k];
+
+        out_half[cur_k] = a_row_half[src_pos];
+      }
+    }
+  };
+
+  for (int index = blockIdx.x; index < num_moe_blocks; index += gridDim.x) {
+    old_expert_id = expert_id;
+    int tmp_expert_id = expert_ids_ptr[index];
+    if (tmp_expert_id == -1) continue;
+    expert_id = tmp_expert_id;
+    perm_int_ptr += (expert_id - old_expert_id) * size_k;
+    read_moe_block_data(index);
+
+    for (int i = 0; i < block_num_valid_tokens; i++)
+      permute_row(block_sorted_ids[i]);
+  }
+}
+
+typedef struct {
+  int thread_k;
+  int thread_n;
+  int num_threads;
+} thread_config_t;
+
+thread_config_t small_batch_thread_configs[] = {
+    // Ordered by priority
+
+    // thread_k, thread_n, num_threads
+    {128, 128, 256},
+    {64, 128, 128},
+    {128, 64, 128}};
+
+thread_config_t large_batch_thread_configs[] = {
+    // Ordered by priority
+
+    // thread_k, thread_n, num_threads
+    {64, 256, 256},
+    {64, 128, 128},
+    {128, 64, 128}};
+
+typedef struct {
+  int blocks_per_sm;
+  thread_config_t tb_cfg;
+} exec_config_t;
+
+int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
+                          int prob_n, int prob_k, int num_bits, int group_size,
+                          bool has_act_order, bool is_k_full, int stages) {
+  bool cache_scales_chunk = has_act_order && !is_k_full;
+
+  int tb_n = th_config.thread_n;
+  int tb_k = th_config.thread_k;
+
+  // Get max scale groups per thread-block
+  int tb_groups;
+  if (group_size == -1) {
+    tb_groups = 1;
+  } else if (group_size == 0) {
+    tb_groups = div_ceil(tb_k, 32);  // Worst case is 32 group size
+  } else {
+    tb_groups = div_ceil(tb_k, group_size);
+  }
+
+  if (cache_scales_chunk) {
+    int load_groups =
+        tb_groups * stages * 2;          // Chunk size is 2x pipeline over dim K
+    load_groups = max(load_groups, 32);  // We load at least 32 scale groups
+    return load_groups * tb_n * 2;
+  } else {
+    int tb_scales = tb_groups * tb_n * 2;
+
+    return tb_scales * stages;
+  }
+}
+
+int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
+                          int thread_m_blocks, int prob_m, int prob_n,
+                          int prob_k, int num_bits, int group_size,
+                          bool has_act_order, bool is_k_full, int has_zp,
+                          int is_zp_float, bool is_a_8bit, int stages) {
+  int pack_factor = 32 / num_bits;
+
+  // Get B size
+  int tb_k = th_config.thread_k;
+  int tb_n = th_config.thread_n;
+  int tb_m = thread_m_blocks * 16;
+
+  // shm size for block_sorted_ids/rd_block_sorted_ids/block_topk_weights
+  // both of them requires tb_m * 4 bytes (tb_m * int32 or tb_m * float32)
+  int sh_block_meta_size = tb_m * 16;
+  int sh_a_size = stages * (tb_m * tb_k) * (is_a_8bit ? 1 : 2);
+  int sh_b_size = stages * (tb_k * tb_n / pack_factor) * 4;
+  int sh_red_size = tb_m * (tb_n + 8) * 2;
+  int sh_bias_size = tb_n * 2;
+  int tmp_size =
+      (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size;
+  tmp_size = max(max(sh_b_size, sh_red_size), tmp_size);
+
+  int sh_s_size =
+      get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
+                            group_size, has_act_order, is_k_full, stages);
+  int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0;
+  int sh_zp_size = 0;
+  if (has_zp) {
+    if (is_zp_float)
+      sh_zp_size = sh_s_size;
+    else if (num_bits == 4)
+      sh_zp_size = sh_s_size / 4;
+    else if (num_bits == 8)
+      sh_zp_size = sh_s_size / 2;
+  }
+
+  int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size +
+                   sh_g_idx_size + sh_block_meta_size;
+
+  return total_size;
+}
+
+bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
+                     int thread_m_blocks, int prob_m, int prob_n, int prob_k,
+                     int num_bits, int group_size, bool has_act_order,
+                     bool is_k_full, int has_zp, int is_zp_float,
+                     bool is_a_8bit, int stages, int max_shared_mem) {
+  // Sanity
+  if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
+      th_config.num_threads == -1) {
+    return false;
+  }
+
+  // Verify K/N are divisible by thread K/N
+  if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) {
+    return false;
+  }
+
+  // Verify min for thread K/N
+  if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) {
+    return false;
+  }
+
+  // num_threads must be at least 128 (= 4 warps)
+  if (th_config.num_threads < 128) {
+    return false;
+  }
+
+  // Check that pipeline fits into cache
+  int cache_size =
+      get_kernel_cache_size(th_config, m_block_size_8, thread_m_blocks, prob_m,
+                            prob_n, prob_k, num_bits, group_size, has_act_order,
+                            is_k_full, has_zp, is_zp_float, is_a_8bit, stages);
+  return cache_size <= max_shared_mem;
+}
+
+MarlinFuncPtr get_marlin_kernel(
+    const vllm::ScalarType a_type, const vllm::ScalarType b_type,
+    const vllm::ScalarType c_type, const vllm::ScalarType s_type,
+    int thread_m_blocks, int thread_n_blocks, int thread_k_blocks,
+    bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks,
+    int threads, bool is_zp_float, int stages) {
+  int num_bits = b_type.size_bits();
+  auto kernel = MarlinDefault;
+
+#include "kernel_selector.h"
+
+  return kernel;
+}
+
+exec_config_t determine_exec_config(
+    const vllm::ScalarType& a_type, const vllm::ScalarType& b_type,
+    const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m,
+    int prob_n, int prob_k, int num_experts, int top_k, int thread_m_blocks,
+    bool m_block_size_8, int num_bits, int group_size, bool has_act_order,
+    bool is_k_full, bool has_zp, bool is_zp_float, bool is_a_8bit, int stages,
+    int max_shared_mem, int sms) {
+  exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
+  thread_config_t* thread_configs = thread_m_blocks > 1
+                                        ? large_batch_thread_configs
+                                        : small_batch_thread_configs;
+  int thread_configs_size =
+      thread_m_blocks > 1
+          ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t)
+          : sizeof(small_batch_thread_configs) / sizeof(thread_config_t);
+
+  int count = 0;
+  constexpr int device_max_reg_size = 255 * 1024;
+  for (int i = 0; i < thread_configs_size; i++) {
+    thread_config_t th_config = thread_configs[i];
+
+    if (!is_valid_config(th_config, m_block_size_8, thread_m_blocks, prob_m,
+                         prob_n, prob_k, num_bits, group_size, has_act_order,
+                         is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
+                         max_shared_mem - 512)) {
+      continue;
+    }
+
+    int cache_size = get_kernel_cache_size(
+        th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
+        num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
+        is_a_8bit, stages);
+
+    int group_blocks = 0;
+    if (!has_act_order) {
+      group_blocks = group_size == -1 ? -1 : (group_size / 16);
+    }
+
+    auto kernel =
+        get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks,
+                          th_config.thread_n / 16, th_config.thread_k / 16,
+                          m_block_size_8, has_act_order, has_zp, group_blocks,
+                          th_config.num_threads, is_zp_float, stages);
+
+    if (kernel == MarlinDefault) continue;
+
+    cudaFuncAttributes attr;
+    cudaFuncGetAttributes(&attr, kernel);
+    int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4;
+    int allow_count = min(device_max_reg_size / reg_size,
+                          max_shared_mem / (cache_size + 1536));
+    if (thread_m_blocks == 1)
+      allow_count = max(min(allow_count, 4), 1);
+    else
+      allow_count = max(min(allow_count, 2), 1);
+
+    if (prob_n / th_config.thread_n * prob_m * top_k * 4 < sms * allow_count) {
+      allow_count =
+          max(prob_n / th_config.thread_n * prob_m * top_k * 4 / sms, 1);
+    }
+
+    if (allow_count > count) {
+      count = allow_count;
+      exec_cfg = {count, th_config};
+    };
+  }
+
+  return exec_cfg;
+}
+
+void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
+               void* a_s, void* b_s, void* g_s, void* zp, void* g_idx,
+               void* perm, void* a_tmp, void* sorted_token_ids,
+               void* expert_ids, void* num_tokens_past_padded,
+               void* topk_weights, int moe_block_size, int num_experts,
+               int top_k, bool mul_topk_weights, int prob_m, int prob_n,
+               int prob_k, void* workspace, vllm::ScalarType const& a_type,
+               vllm::ScalarType const& b_type, vllm::ScalarType const& c_type,
+               vllm::ScalarType const& s_type, bool has_bias,
+               bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
+               int group_size, int dev, cudaStream_t stream, int thread_k,
+               int thread_n, int sms, int blocks_per_sm, bool use_atomic_add,
+               bool use_fp32_reduce, bool is_zp_float) {
+  int thread_m_blocks = div_ceil(moe_block_size, 16);
+  bool m_block_size_8 = moe_block_size == 8;
+  bool is_a_8bit = a_type.size_bits() == 8;
+
+  STD_TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [",
+                  prob_m, ", ", prob_n, ", ", prob_k, "]");
+
+  int group_blocks = 0;
+  if (has_act_order) {
+    if (is_k_full) {
+      STD_TORCH_CHECK(group_size != -1);
+      group_blocks = group_size / 16;
+      STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
+                      " is not divisible by group_blocks = ", group_blocks);
+    } else {
+      STD_TORCH_CHECK(group_size == 0);
+      group_blocks = 0;
+    }
+  } else {
+    if (group_size == -1) {
+      group_blocks = -1;
+    } else {
+      group_blocks = group_size / 16;
+      STD_TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
+                      " is not divisible by group_blocks = ", group_blocks);
+    }
+  }
+
+  int num_bits = b_type.size_bits();
+  const int4* A_ptr = (const int4*)A;
+  const int4* B_ptr = (const int4*)B;
+  int4* C_ptr = (int4*)C;
+  int4* C_tmp_ptr = (int4*)C_tmp;
+  const int4* bias_ptr = (const int4*)b_bias;
+  const float* a_s_ptr = (const float*)a_s;
+  const int4* b_s_ptr = (const int4*)b_s;
+  const float* g_s_ptr = (const float*)g_s;
+  const int4* zp_ptr = (const int4*)zp;
+  const int* g_idx_ptr = (const int*)g_idx;
+  const int* perm_ptr = (const int*)perm;
+  int4* a_tmp_ptr = (int4*)a_tmp;
+  const int32_t* sorted_token_ids_ptr = (const int32_t*)sorted_token_ids;
+  const int32_t* expert_ids_ptr = (const int32_t*)expert_ids;
+  const int32_t* num_tokens_past_padded_ptr =
+      (const int32_t*)num_tokens_past_padded;
+  const float* topk_weights_ptr = (const float*)topk_weights;
+  int* locks = (int*)workspace;
+
+  if (has_act_order) {
+    // Permute A columns
+    auto kernel = permute_cols_kernel<8>;
+    if (moe_block_size == 8) {
+    } else if (moe_block_size == 16)
+      kernel = permute_cols_kernel<16>;
+    else if (moe_block_size == 32)
+      kernel = permute_cols_kernel<32>;
+    else if (moe_block_size == 48)
+      kernel = permute_cols_kernel<48>;
+    else if (moe_block_size == 64)
+      kernel = permute_cols_kernel<64>;
+    else
+      STD_TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size);
+
+    // avoid ">>>" being formatted to "> > >"
+    // clang-format off
+    kernel<<<sms, default_threads, 0, stream>>>(
+        A_ptr, perm_ptr, a_tmp_ptr, sorted_token_ids_ptr, expert_ids_ptr,
+        num_tokens_past_padded_ptr, prob_m, prob_k, top_k);
+    // clang-format on
+    A_ptr = a_tmp_ptr;
+    prob_m = prob_m * top_k;
+    top_k = 1;
+
+    // If we have a full K, then we can run the non-act-order version of Marlin
+    // (since the weight rows are reordered by increasing group ids, and by
+    // having a full K, we have full original groups)
+    if (is_k_full) has_act_order = false;
+  }
+
+  int max_shared_mem = 0;
+  cudaDeviceGetAttribute(&max_shared_mem,
+                         cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
+  STD_TORCH_CHECK(max_shared_mem > 0);
+
+  int major_capability, minor_capability;
+  cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor,
+                         dev);
+  cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
+                         dev);
+  STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 75,
+                  "marlin kernel only support Turing or newer GPUs.");
+  int stages = 4;
+  if (major_capability == 7 && minor_capability == 5) {
+    stages = 2;
+    STD_TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8,
+                    "Turing only support FP16 or INT8 activation.");
+  }
+  if (a_type == vllm::kFE4M3fn) {
+    STD_TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
+                    "FP8 only support Ada Lovelace or newer GPUs.");
+    STD_TORCH_CHECK(
+        major_capability * 10 + minor_capability == 89 ||
+            major_capability == 12,
+        "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
+        "Marlin W4A16 on other devices).");
+  }
+
+  // Set thread config
+  exec_config_t exec_cfg;
+  thread_config_t thread_tfg;
+  if (thread_k != -1 && thread_n != -1) {
+    thread_tfg = thread_config_t{thread_k, thread_n, thread_k * thread_n / 64};
+    if (blocks_per_sm == -1) blocks_per_sm = 1;
+    exec_cfg = exec_config_t{blocks_per_sm, thread_tfg};
+    STD_TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
+                    " is not divisible by thread_n = ", thread_n);
+    STD_TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
+                    " is not divisible by thread_k = ", thread_k);
+  } else {
+    // Auto config
+    exec_cfg = determine_exec_config(
+        a_type, b_type, c_type, s_type, prob_m, prob_n, prob_k, num_experts,
+        top_k, thread_m_blocks, m_block_size_8, num_bits, group_size,
+        has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
+        max_shared_mem, sms);
+    thread_tfg = exec_cfg.tb_cfg;
+  }
+
+  int num_threads = thread_tfg.num_threads;
+  thread_k = thread_tfg.thread_k;
+  thread_n = thread_tfg.thread_n;
+  int blocks = sms * exec_cfg.blocks_per_sm;
+  if (exec_cfg.blocks_per_sm > 1)
+    max_shared_mem = max_shared_mem / exec_cfg.blocks_per_sm - 1024;
+
+  int thread_k_blocks = thread_k / 16;
+  int thread_n_blocks = thread_n / 16;
+
+  STD_TORCH_CHECK(
+      is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
+                      prob_n, prob_k, num_bits, group_size, has_act_order,
+                      is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
+                      max_shared_mem),
+      "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
+      ", thread_k = ", thread_tfg.thread_k,
+      ", thread_n = ", thread_tfg.thread_n,
+      ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ",
+      prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
+      ", group_size = ", group_size, ", has_act_order = ", has_act_order,
+      ", is_k_full = ", is_k_full, ", has_zp = ", has_zp,
+      ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem);
+
+  int sh_cache_size =
+      get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
+                            prob_n, prob_k, num_bits, group_size, has_act_order,
+                            is_k_full, has_zp, is_zp_float, is_a_8bit, stages);
+
+  auto kernel = get_marlin_kernel(
+      a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks,
+      thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks,
+      num_threads, is_zp_float, stages);
+
+  if (kernel == MarlinDefault) {
+    STD_TORCH_CHECK(
+        false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, ", ",
+        prob_k, "]", ", has_act_order = ", has_act_order,
+        ", num_groups = ", num_groups, ", group_size = ", group_size,
+        ", thread_m_blocks = ", thread_m_blocks,
+        ", thread_n_blocks = ", thread_n_blocks,
+        ", thread_k_blocks = ", thread_k_blocks, ", num_bits = ", num_bits);
+  }
+
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       max_shared_mem);
+  // avoid ">>>" being formatted to "> > >"
+  // clang-format off
+  kernel<<<blocks, num_threads, max_shared_mem, stream>>>(
+      A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, g_idx_ptr,
+      sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
+      topk_weights_ptr, top_k, mul_topk_weights, num_groups, prob_m,
+      prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce);
+  // clang-format on
+}
+
+}  // namespace MARLIN_NAMESPACE_NAME
+
+// ---------------------------------------------------------------------------
+// mstar classic-ABI host wrapper, trimmed to W4A16: fp16/bf16 activation,
+// symmetric GPTQ INT4 (uint4b8), no zero-point / act-order / bias / global-scale
+// / 8-bit-activation. Ports vLLM ops.cu::moe_wna16_marlin_gemm; every dropped
+// optional is handed to the verbatim device ``marlin_mm`` as a 0-element tensor,
+// exactly as upstream's "absent optional" branches did. Namespace of the device
+// symbols is ``marlin_moe_wna16`` (set by kernel.h's MARLIN_NAMESPACE_NAME), so
+// its Marlin<> instantiations never ODR-clash with the linear repack's ``marlin``.
+// ---------------------------------------------------------------------------
+torch::Tensor moe_wna16_marlin_gemm(
+    torch::Tensor a, std::optional<torch::Tensor> c_or_none,
+    torch::Tensor b_q_weight, torch::Tensor b_scales, torch::Tensor workspace,
+    torch::Tensor sorted_token_ids, torch::Tensor expert_ids,
+    torch::Tensor num_tokens_past_padded, torch::Tensor topk_weights,
+    int64_t moe_block_size, int64_t top_k, bool mul_topk_weights,
+    int64_t b_type_id, int64_t size_m, int64_t size_n, int64_t size_k,
+    bool is_k_full, bool use_atomic_add, bool use_fp32_reduce) {
+  vllm::ScalarTypeId a_type_id, c_type_id;
+  TORCH_CHECK(
+      a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
+      "moe_wna16_marlin_gemm (mstar W4A16): activation must be fp16 or bf16");
+  if (a.scalar_type() == torch::kHalf) {
+    a_type_id = vllm::kFloat16.id();
+    c_type_id = vllm::kFloat16.id();
+  } else {
+    a_type_id = vllm::kBFloat16.id();
+    c_type_id = vllm::kBFloat16.id();
+  }
+  auto c_dtype = a.scalar_type();
+  vllm::ScalarTypeId s_type_id = c_type_id;
+
+  vllm::ScalarType a_type = vllm::ScalarType::from_id(a_type_id);
+  vllm::ScalarType b_type = vllm::ScalarType::from_id(b_type_id);
+  vllm::ScalarType c_type = vllm::ScalarType::from_id(c_type_id);
+  vllm::ScalarType s_type = vllm::ScalarType::from_id(s_type_id);
+  TORCH_CHECK(b_type == vllm::kU4B8,
+              "mstar Marlin MoE supports only symmetric INT4 (uint4b8); got ",
+              b_type.str());
+
+  int pack_factor = 32 / b_type.size_bits();
+  int num_experts = b_q_weight.size(0);
+
+  if (moe_block_size != 8) {
+    TORCH_CHECK(moe_block_size % 16 == 0,
+                "unsupported moe_block_size=", moe_block_size);
+    TORCH_CHECK(moe_block_size >= 16 && moe_block_size <= 64,
+                "unsupported moe_block_size=", moe_block_size);
+  }
+
+  // Verify A
+  TORCH_CHECK(a.size(0) == size_m, "a.size(0) = ", a.size(0),
+              ", size_m = ", size_m);
+  TORCH_CHECK(a.size(1) == size_k, "a.size(1) = ", a.size(1),
+              ", size_k = ", size_k);
+
+  // Verify B (Marlin-tiled: b_q_weight is (E, size_k/tile, size_n*pack/tile))
+  TORCH_CHECK(size_k % marlin_moe_wna16::tile_size == 0, "size_k = ", size_k,
+              " not divisible by tile_size = ", marlin_moe_wna16::tile_size);
+  TORCH_CHECK((size_k / marlin_moe_wna16::tile_size) == b_q_weight.size(1),
+              "b_q_weight.size(1) = ", b_q_weight.size(1), ", size_k = ", size_k);
+  TORCH_CHECK(b_q_weight.size(2) % marlin_moe_wna16::tile_size == 0,
+              "b_q_weight.size(2) = ", b_q_weight.size(2));
+  int actual_size_n =
+      (b_q_weight.size(2) / marlin_moe_wna16::tile_size) * pack_factor;
+  TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n,
+              ", actual_size_n = ", actual_size_n);
+
+  TORCH_CHECK(a.is_cuda() && a.is_contiguous(), "A must be contiguous CUDA");
+  TORCH_CHECK(b_q_weight.is_cuda() && b_q_weight.is_contiguous(),
+              "b_q_weight must be contiguous CUDA");
+  TORCH_CHECK(b_scales.is_cuda() && b_scales.is_contiguous(),
+              "b_scales must be contiguous CUDA");
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+  int dev = a.get_device();
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  int sms = -1;
+  cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
+
+  auto opts_c = a.options().dtype(c_dtype);
+  auto opts_f = a.options().dtype(torch::kFloat);
+
+  torch::Tensor c;
+  if (c_or_none.has_value()) {
+    c = c_or_none.value();
+    TORCH_CHECK(c.is_cuda() && c.is_contiguous(), "c must be contiguous CUDA");
+    TORCH_CHECK(c.size(0) == size_m * top_k && c.size(1) == size_n,
+                "bad c shape");
+  } else {
+    c = torch::empty({size_m * top_k, size_n}, opts_c);
+  }
+
+  torch::Tensor c_tmp;
+  if (use_fp32_reduce && !use_atomic_add) {
+    long max_c_tmp_size = std::min(
+        (long)size_n * sorted_token_ids.size(0),
+        (long)sms * 4 * moe_block_size * marlin_moe_wna16::max_thread_n);
+    if (moe_block_size == 8) max_c_tmp_size *= 2;
+    c_tmp = torch::empty({max_c_tmp_size}, opts_f);
+  } else {
+    c_tmp = torch::empty({0}, opts_f);
+  }
+
+  // Grouping (no act-order): num_groups = b_scales.size(1).
+  TORCH_CHECK(b_scales.dim() == 3, "b_scales must be rank 3");
+  TORCH_CHECK(b_scales.size(2) == size_n, "b_scales dim2 != size_n");
+  int num_groups = b_scales.size(1);
+  int group_size;
+  if (num_groups > 1) {
+    TORCH_CHECK(size_k % num_groups == 0, "size_k not divisible by num_groups");
+    group_size = size_k / num_groups;
+  } else {
+    group_size = -1;
+  }
+
+  // Dropped optionals -> 0-element tensors (upstream's absent-optional branches).
+  torch::Tensor a_scales = torch::empty({0}, opts_f);
+  torch::Tensor global_scale = torch::empty({0}, opts_f);
+  torch::Tensor b_bias = torch::empty({0}, opts_c);
+  torch::Tensor b_zeros = torch::empty({0}, opts_c);
+  torch::Tensor g_idx = torch::empty({0}, opts_c);
+  torch::Tensor perm = torch::empty({0}, opts_c);
+  torch::Tensor a_tmp = torch::empty({0}, opts_c);
+
+  TORCH_CHECK(size_n % marlin_moe_wna16::min_thread_n == 0, "size_n = ", size_n,
+              " not divisible by min_thread_n");
+  int max_n_tiles = size_n / marlin_moe_wna16::min_thread_n;
+  int min_workspace_size = std::min(
+      max_n_tiles * (int)(sorted_token_ids.size(0) / moe_block_size), sms * 4);
+  TORCH_CHECK(workspace.numel() >= min_workspace_size,
+              "workspace.numel = ", workspace.numel(),
+              " < min_workspace_size = ", min_workspace_size);
+
+  marlin_moe_wna16::marlin_mm(
+      a.const_data_ptr(), b_q_weight.const_data_ptr(), c.mutable_data_ptr(),
+      c_tmp.mutable_data_ptr(), b_bias.mutable_data_ptr(),
+      a_scales.mutable_data_ptr(), b_scales.mutable_data_ptr(),
+      global_scale.mutable_data_ptr(), b_zeros.mutable_data_ptr(),
+      g_idx.mutable_data_ptr(), perm.mutable_data_ptr(), a_tmp.mutable_data_ptr(),
+      sorted_token_ids.mutable_data_ptr(), expert_ids.mutable_data_ptr(),
+      num_tokens_past_padded.mutable_data_ptr(), topk_weights.mutable_data_ptr(),
+      moe_block_size, num_experts, top_k, mul_topk_weights, size_m, size_n,
+      size_k, workspace.mutable_data_ptr(), a_type, b_type, c_type, s_type,
+      /*has_bias=*/false, /*has_act_order=*/false, is_k_full, /*has_zp=*/false,
+      num_groups, group_size, dev, stream, /*thread_k=*/-1, /*thread_n=*/-1, sms,
+      /*blocks_per_sm=*/-1, use_atomic_add, use_fp32_reduce,
+      /*is_zp_float=*/false);
+
+  return c;
+}
+
+// Single TORCH_LIBRARY block for the whole ``_mstar_marlin_C`` namespace (the
+// repack impl lives in gptq_marlin_repack.cu as a TORCH_LIBRARY_IMPL).
+TORCH_LIBRARY(_mstar_marlin_C, m) {
+  m.def(
+      "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, int size_k, "
+      "int size_n, int num_bits) -> Tensor");
+  m.def(
+      "moe_wna16_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
+      "Tensor b_scales, Tensor workspace, Tensor sorted_token_ids, "
+      "Tensor expert_ids, Tensor num_tokens_past_padded, Tensor topk_weights, "
+      "int moe_block_size, int top_k, bool mul_topk_weights, int b_type_id, "
+      "int size_m, int size_n, int size_k, bool is_k_full, bool use_atomic_add, "
+      "bool use_fp32_reduce) -> Tensor");
+}
+
+TORCH_LIBRARY_IMPL(_mstar_marlin_C, CUDA, m) {
+  m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm);
+}
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h
new file mode 100644
index 000000000..04f90101b
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/marlin_template.h
@@ -0,0 +1,2241 @@
+/*
+ * Modified by Neural Magic
+ * Copyright (C) Marlin.2024 Elias Frantar
+ *
+ * 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.
+ */
+
+/*
+ * Adapted from https://github.com/IST-DASLab/marlin
+ */
+
+#ifndef MARLIN_NAMESPACE_NAME
+  #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
+#endif
+
+#include "libtorch_stable/quantization/marlin/marlin.cuh"
+#include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh"
+#include "libtorch_stable/quantization/marlin/dequant.h"
+#include "libtorch_stable/quantization/marlin/marlin_mma.h"
+#include "core/scalar_type.hpp"
+
+#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t)               \
+  static_assert(std::is_same<scalar_t, half>::value ||          \
+                    std::is_same<scalar_t, nv_bfloat16>::value, \
+                "only float16 and bfloat16 is supported");
+
+namespace MARLIN_NAMESPACE_NAME {
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750
+
+template <typename scalar_t,  // compute dtype, half or nv_float16
+          const vllm::ScalarTypeId b_type_id,  // weight MarlinScalarType id
+          const int threads,          // number of threads in a threadblock
+          const int thread_m_blocks,  // number of 16x16 blocks in the m
+                                      // dimension (batchsize) of the
+                                      // threadblock
+          const int thread_n_blocks,  // same for n dimension (output)
+          const int thread_k_blocks,  // same for k dimension (reduction)
+          const bool m_block_size_8,  // whether m_block_size == 8
+                                      // only works when thread_m_blocks == 1
+          const int stages,  // number of stages for the async global->shared
+                             // fetch pipeline
+          const bool has_act_order,  // whether act_order is enabled
+          const int group_blocks,    // number of consecutive 16x16 blocks
+                                     // with a separate quantization scale
+          const bool is_zp_float     // is zero point of float16 type?
+          >
+__global__ void Marlin(
+    const int4* __restrict__ A,  // fp16 input matrix of shape mxk
+    const int4* __restrict__ B,  // 4bit quantized weight matrix of shape kxn
+    int4* __restrict__ C,        // fp16 output buffer of shape mxn
+    int4* __restrict__ C_tmp,    // fp32 tmp output buffer (for reduce)
+    const int4* __restrict__ scales_ptr,  // fp16 quantization scales of shape
+                                          // (k/groupsize)xn
+    const int4* __restrict__ zp_ptr,      // 4bit packed zero-points of shape
+                                          // (k/groupsize)x(n/pack_factor)
+    const int* __restrict__ g_idx,        // int32 group indices of shape k
+    const int32_t* __restrict__ sorted_token_ids_ptr,        // moe sorted_ids
+    const int32_t* __restrict__ expert_ids_ptr,              // moe expert ids
+    const int32_t* __restrict__ num_tokens_past_padded_ptr,  // moe num tokens
+    const float* __restrict__ topk_weights_ptr,              // moe top weights
+    int top_k,              // num of experts per token
+    bool mul_topk_weights,  // mul topk weights or not
+    int num_groups,         // number of scale groups per output channel
+    int prob_m,             // batch dimension m
+    int prob_n,             // output dimension n
+    int prob_k,             // reduction dimension k
+    int* locks,             // extra global storage for barrier synchronization
+    bool use_atomic_add,    // whether to use atomic add to reduce
+    bool use_fp32_reduce    // whether to use fp32 global reduce
+) {}
+
+}  // namespace MARLIN_NAMESPACE_NAME
+
+#else
+
+// Instruction for loading a full 16x16 matrix fragment of operand A from shared
+// memory, directly in tensor core layout.
+template <int count, vllm::ScalarTypeId type_id>
+__device__ inline void ldsm(typename MarlinScalarType<type_id>::FragA& frag_a,
+                            const void* smem_ptr) {
+  uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  if constexpr (count == 4) {
+    asm volatile(
+        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
+        : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3])
+        : "r"(smem));
+  } else if constexpr (count == 2) {
+    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n"
+                 : "=r"(a[0]), "=r"(a[1])
+                 : "r"(smem));
+  } else if constexpr (count == 1) {
+    asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n"
+                 : "=r"(a[0])
+                 : "r"(smem));
+  } else {
+    static_assert(count == 1 || count == 2 || count == 4, "invalid count");
+  }
+}
+
+// Multiply dequantized values by the corresponding quantization scale; used
+// only for grouped quantization.
+template <vllm::ScalarTypeId type_id>
+__device__ inline void scale(typename MarlinScalarType<type_id>::FragB& frag_b,
+                             typename MarlinScalarType<type_id>::FragS& frag_s,
+                             int i) {
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
+  scalar_t2 s = MarlinScalarType<type_id>::num2num2(
+      reinterpret_cast<scalar_t*>(&frag_s)[i]);
+  frag_b[0] = __hmul2(frag_b[0], s);
+  frag_b[1] = __hmul2(frag_b[1], s);
+}
+
+template <vllm::ScalarTypeId type_id>
+__device__ inline void scale_and_sub(
+    typename MarlinScalarType<type_id>::FragB& frag_b,
+    typename MarlinScalarType<type_id>::scalar_t s,
+    typename MarlinScalarType<type_id>::scalar_t zp) {
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
+  scalar_t2 s2 = MarlinScalarType<type_id>::num2num2(s);
+  scalar_t2 zp2 = MarlinScalarType<type_id>::num2num2(zp);
+  frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2));
+  frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2));
+}
+
+template <vllm::ScalarTypeId type_id>
+__device__ inline void sub_zp(
+    typename MarlinScalarType<type_id>::FragB& frag_b,
+    typename MarlinScalarType<type_id>::scalar_t2& frag_zp, int i) {
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
+  scalar_t2 zp = MarlinScalarType<type_id>::num2num2(
+      reinterpret_cast<scalar_t*>(&frag_zp)[i]);
+  frag_b[0] = __hsub2(frag_b[0], zp);
+  frag_b[1] = __hsub2(frag_b[1], zp);
+}
+
+// Same as above, but for act_order (each K is multiplied individually)
+template <vllm::ScalarTypeId type_id>
+__device__ inline void scale4(
+    typename MarlinScalarType<type_id>::FragB& frag_b,
+    typename MarlinScalarType<type_id>::FragS& frag_s_1,
+    typename MarlinScalarType<type_id>::FragS& frag_s_2,
+    typename MarlinScalarType<type_id>::FragS& frag_s_3,
+    typename MarlinScalarType<type_id>::FragS& frag_s_4, int i) {
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
+
+  scalar_t2 s_val_1_2;
+  s_val_1_2.x = reinterpret_cast<scalar_t*>(&frag_s_1)[i];
+  s_val_1_2.y = reinterpret_cast<scalar_t*>(&frag_s_2)[i];
+
+  scalar_t2 s_val_3_4;
+  s_val_3_4.x = reinterpret_cast<scalar_t*>(&frag_s_3)[i];
+  s_val_3_4.y = reinterpret_cast<scalar_t*>(&frag_s_4)[i];
+
+  frag_b[0] = __hmul2(frag_b[0], s_val_1_2);
+  frag_b[1] = __hmul2(frag_b[1], s_val_3_4);
+}
+
+// Given 2 floats multiply by 2 scales (halves)
+template <vllm::ScalarTypeId type_id>
+__device__ inline void scale_float(
+    float* c, typename MarlinScalarType<type_id>::FragS& s) {
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  scalar_t* s_ptr = reinterpret_cast<scalar_t*>(&s);
+  c[0] = __fmul_rn(c[0], MarlinScalarType<type_id>::num2float(s_ptr[0]));
+  c[1] = __fmul_rn(c[1], MarlinScalarType<type_id>::num2float(s_ptr[1]));
+}
+
+// Wait until barrier reaches `count`, then lock for current threadblock.
+__device__ inline void barrier_acquire(int* lock, int count) {
+  if (threadIdx.x == 0) {
+    int state = -1;
+    do
+      // Guarantee that subsequent writes by this threadblock will be visible
+      // globally.
+      asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
+                   : "=r"(state)
+                   : "l"(lock));
+    while (state != count);
+  }
+  __syncthreads();
+}
+
+// Release barrier and increment visitation count.
+__device__ inline void barrier_release(int* lock, bool reset = false) {
+  __syncthreads();
+  if (threadIdx.x == 0) {
+    if (reset) {
+      lock[0] = 0;
+      return;
+    }
+    int val = 1;
+    // Make sure that all writes since acquiring this barrier are visible
+    // globally, while releasing the barrier.
+    asm volatile("fence.acq_rel.gpu;\n");
+    asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n"
+                 :
+                 : "l"(lock), "r"(val));
+  }
+}
+
+// Wait until value of lock to be negative, and then add 1
+__device__ inline void wait_negative_and_add(int* lock) {
+  if (threadIdx.x == 0) {
+    int state = 0;
+    do
+      // Guarantee that subsequent writes by this threadblock will be visible
+      // globally.
+      asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
+                   : "=r"(state)
+                   : "l"(lock));
+    while (state >= 0);
+    atomicAdd(lock, 1);
+  }
+  __syncthreads();
+}
+
+template <const vllm::ScalarTypeId a_type_id,  // A ScalarType id
+          const vllm::ScalarTypeId b_type_id,  // B ScalarType id
+          const vllm::ScalarTypeId c_type_id,  // C ScalarType id
+          const vllm::ScalarTypeId s_type_id,  // B_SCALE ScalarType id
+          const int threads,          // number of threads in a threadblock
+          const int thread_m_blocks,  // number of 16x16 blocks in the m
+                                      // dimension (batchsize) of the
+                                      // threadblock
+          const int thread_n_blocks,  // same for n dimension (output)
+          const int thread_k_blocks,  // same for k dimension (reduction)
+          const bool m_block_size_8,  // whether m_block_size == 8
+                                      // only works when thread_m_blocks == 1
+          const int stages,  // number of stages for the async global->shared
+                             // fetch pipeline
+          const int group_blocks,  // number of consecutive 16x16 blocks
+                                   // with a separate quantization scale
+          const bool is_zp_float   // is zero point of float16 type?
+          >
+__global__ void Marlin(
+    const int4* __restrict__ A,  // fp16 input matrix of shape mxk
+    const int4* __restrict__ B,  // 4bit quantized weight matrix of shape kxn
+    int4* __restrict__ C,        // fp16 output buffer of shape mxn
+    int4* __restrict__ C_tmp,    // fp32 tmp output buffer (for reduce)
+    const int4* __restrict__ b_bias_ptr,
+    // float scales of input matrix, only used when is_a_8bit == true.
+    // shape (m,)
+    const float* __restrict__ a_scales_ptr,
+    // fp16 quantization scales. shape (k/groupsize, n)
+    const int4* __restrict__ scales_ptr,
+    // fp16 global scale (for nvfp4// only)
+    const float* __restrict__ global_scale_ptr,
+    // 4bit packed zero-points of shape
+    // (k/groupsize, n/pack_factor)
+    const int4* __restrict__ zp_ptr,
+    // int32 group indices of shape k
+    const int* __restrict__ g_idx,
+    const int32_t* __restrict__ sorted_token_ids_ptr,        // moe sorted_ids
+    const int32_t* __restrict__ expert_ids_ptr,              // moe expert ids
+    const int32_t* __restrict__ num_tokens_past_padded_ptr,  // moe num tokens
+    const float* __restrict__ topk_weights_ptr,              // moe top weights
+    int top_k,              // num of experts per token
+    bool mul_topk_weights,  // mul topk weights or not
+    int num_groups,         // number of scale groups per output channel
+    int prob_m,             // batch dimension m
+    int prob_n,             // output dimension n
+    int prob_k,             // reduction dimension k
+    int* locks,             // extra global storage for barrier synchronization
+    bool has_bias,
+    bool use_atomic_add,  // whether to use atomic add to reduce
+    bool use_fp32_reduce  // whether to use fp32 global reduce
+) {
+  // Each threadblock processes one "stripe" of the B matrix with (roughly) the
+  // same size, which might involve multiple column "slices" (of width 16 *
+  // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM
+  // example:
+  //   0 1 3
+  //   0 2 3
+  //   1 2 4
+  // While this kind of partitioning makes things somewhat more complicated, it
+  // ensures good utilization of all SMs for many kinds of shape and GPU
+  // configurations, while requiring as few slow global cross-threadblock
+  // reductions as possible.
+
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 890
+  // FP8 computation is only supported for Ada Lovelace or newer architectures.
+  if constexpr (a_type_id == vllm::kFE4M3fn.id()) return;
+  #endif
+
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+  // Turing TensorCore only supports fp16 and int8
+  if constexpr (a_type_id != vllm::kFloat16.id() && a_type_id != vllm::kS8.id())
+    return;
+  #endif
+
+  int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
+  constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
+
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+  static constexpr auto num_bits =
+      vllm::ScalarType::from_id(b_type_id).size_bits();
+  // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 &&
+  // num_bits == 4
+  constexpr bool use_fp16_accum =
+      a_type_id == vllm::kFloat16.id() &&
+      (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) &&
+       !(group_blocks == -1 && num_bits == 4));
+  #else
+  constexpr bool use_fp16_accum = false;
+  #endif
+  using Adtype = MarlinScalarType<a_type_id>;
+  using Cdtype = MarlinScalarType<c_type_id>;
+
+  using scalar_t = typename MarlinScalarType<a_type_id>::scalar_t;
+  using scalar_t2 = typename MarlinScalarType<a_type_id>::scalar_t2;
+  using scalar_32bit_t = typename MarlinScalarType<a_type_id>::scalar_32bit_t;
+
+  using c_scalar_t = typename MarlinScalarType<c_type_id>::scalar_t;
+  using c_scalar_t2 = typename MarlinScalarType<c_type_id>::scalar_t2;
+
+  using FragA = typename MarlinScalarType<a_type_id>::FragA;
+  using FragB = typename MarlinScalarType<a_type_id>::FragB;
+  using FragC = typename MarlinScalarType<a_type_id>::FragC;
+  using FragS = typename MarlinScalarType<c_type_id>::FragS;
+  using FragZP = typename MarlinScalarType<c_type_id>::FragZP;
+
+  extern __shared__ int4 sh[];
+  static constexpr auto a_type = vllm::ScalarType::from_id(a_type_id);
+  static constexpr auto b_type = vllm::ScalarType::from_id(b_type_id);
+  static constexpr auto c_type = vllm::ScalarType::from_id(c_type_id);
+  static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id);
+  if constexpr (b_type == vllm::kFE2M1f) {
+    static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 ||
+                  s_type == vllm::kFE8M0fnu && group_blocks == 2);
+  } else if constexpr (b_type == vllm::kFE4M3fn && s_type == vllm::kFE8M0fnu) {
+    static_assert(group_blocks == 2);
+  } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
+    static_assert(s_type == vllm::kBFloat16);
+  } else if constexpr (std::is_same<scalar_t, half>::value) {
+    static_assert(s_type == vllm::kFloat16);
+  }
+
+  constexpr bool is_a_8bit = a_type.size_bits() == 8;
+  if constexpr (!is_a_8bit) {
+    static_assert(std::is_same<scalar_t, c_scalar_t>::value);
+  }
+  constexpr bool has_zp = b_type == vllm::kU4 || b_type == vllm::kU8;
+  constexpr bool is_int_type = b_type == vllm::kU4 || b_type == vllm::kU8 ||
+                               b_type == vllm::kS4 || b_type == vllm::kS8 ||
+                               b_type == vllm::kU4B8 || b_type == vllm::kU8B128;
+  constexpr bool is_8bit_scale = s_type.size_bits() == 8;
+  // see comments of dequant.h for more details
+  constexpr bool dequant_skip_flop =
+      is_a_8bit || (b_type == vllm::kFE4M3fn && !(s_type == vllm::kFE8M0fnu)) ||
+      b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn ||
+      has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
+      has_zp && !is_zp_float && !(b_type == vllm::kU8);
+
+  float global_scale_f32 = 1.0f;
+
+  constexpr bool has_act_order = group_blocks == 0;
+
+  constexpr int pack_factor = 32 / b_type.size_bits();
+  static_assert(thread_m_blocks == 1 || !m_block_size_8);
+  const int group_size =
+      (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups;
+  const int scales_expert_stride =
+      prob_n * prob_k / group_size / (is_8bit_scale ? 16 : 8);
+  const int zp_expert_stride =
+      is_zp_float ? prob_n * prob_k / group_size / 8
+                  : prob_n * prob_k / group_size / (pack_factor * 4);
+  const int b_bias_expert_stride = prob_n / 8;
+
+  // parallel: num valid moe blocks
+  int parallel = num_tokens_past_padded / moe_block_size;
+
+  int k_tiles = prob_k / 16 / thread_k_blocks;
+  int n_tiles = prob_n / 16 / thread_n_blocks;
+
+  int global_mn_tiles = parallel * n_tiles;
+  int part2_mn_tiles = global_mn_tiles;
+  int part1_mn_iters = 0;
+  bool in_part2 = false;
+
+  // we use DP + two-tile SK here
+  // part1: DP
+  // part2: two-tile SK
+  // see https://github.com/vllm-project/vllm/pull/24722 for more details
+  if (global_mn_tiles > gridDim.x) {
+    part2_mn_tiles = global_mn_tiles % gridDim.x;
+    if (part2_mn_tiles * 3 <= gridDim.x) part2_mn_tiles += gridDim.x;
+    part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x;
+  }
+
+  int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x);
+
+  if constexpr (!has_act_order && group_blocks != -1) {
+    if (group_blocks >= thread_k_blocks) {
+      // Ensure that the number of tiles in each stripe is a multiple of the
+      // groupsize; this avoids an annoying special case where a stripe starts
+      // in the middle of group.
+      iters = (group_blocks / thread_k_blocks) *
+              div_ceil(iters, (group_blocks / thread_k_blocks));
+    }
+  }
+
+  int slice_row = 0;
+  int slice_col_par = blockIdx.x;
+  int slice_col;
+  int slice_iters =
+      k_tiles;  // number of threadblock tiles in the current slice
+  // total number of active threadblocks in the current slice
+  int slice_count = 1;
+  // index of threadblock in current slice; numbered bottom to top
+  int slice_idx = 0;
+
+  int par_id = 0;
+  int block_id = -1;
+  int64_t expert_id = 0;  // use int64 to avoid computation result overflow
+  int old_expert_id = 0;
+  int64_t B_expert_off = 0;
+
+  float* sh_a_s = reinterpret_cast<float*>(sh);
+  int4* sh_block_sorted_ids_int4 = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0);
+  int4* sh_rd_block_sorted_ids_int4 =
+      sh_block_sorted_ids_int4 + moe_block_size / 4;
+  int4* sh_block_topk_weights_int4 =
+      sh_rd_block_sorted_ids_int4 + moe_block_size / 4;
+  // sh_block_topk_weights_int4 only need (moe_block_size / 4);
+  // but we pad to align to 256 bytes
+  int4* sh_new = sh_block_topk_weights_int4 + moe_block_size / 2;
+  int32_t* sh_block_sorted_ids =
+      reinterpret_cast<int*>(sh_block_sorted_ids_int4);
+  int32_t* sh_rd_block_sorted_ids =
+      reinterpret_cast<int*>(sh_rd_block_sorted_ids_int4);
+  c_scalar_t2* sh_block_topk_weights =
+      reinterpret_cast<c_scalar_t2*>(sh_block_topk_weights_int4);
+
+  int32_t block_num_valid_tokens = 0;
+  int32_t locks_off = 0;
+
+  // We can easily implement parallel problem execution by just remapping
+  // indices and advancing global pointers
+  if (part2_mn_tiles >= gridDim.x) {
+    // when part2_mn_tiles >= sms
+    // then there are at most $sms$ conflict tile blocks
+    locks_off = blockIdx.x;
+  } else {
+    locks_off = (iters * blockIdx.x) / k_tiles - 1;
+  }
+
+  int prob_m_top_k = prob_m * top_k;
+  // read moe block data given block_id
+  // block_sorted_ids / block_num_valid_tokens / block_topk_weights
+  auto read_moe_block_data = [&](int block_id) {
+    block_num_valid_tokens = moe_block_size;
+
+    cp_async4_pred(sh_block_sorted_ids_int4 + threadIdx.x,
+                   reinterpret_cast<const int4*>(sorted_token_ids_ptr) +
+                       (block_id * moe_block_size / 4 + threadIdx.x),
+                   threadIdx.x < moe_block_size / 4);
+
+    cp_async_fence();
+    cp_async_wait<0>();
+
+    __syncthreads();
+
+    if (threadIdx.x >= threads - 32) {
+      constexpr int size_per_thread = div_ceil(moe_block_size, 32);
+      int lane_id = threadIdx.x - (threads - 32);
+
+      int local_count = 0;
+  #pragma unroll
+      for (int i = 0; i < size_per_thread; i++) {
+        int j = lane_id * size_per_thread + i;
+        if (j < moe_block_size) {
+          int idx = sh_block_sorted_ids[j];
+          if (idx < prob_m_top_k) local_count++;
+        }
+      }
+
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+
+      if constexpr (moe_block_size >= 16)
+        local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 16);
+      if constexpr (moe_block_size >= 8)
+        local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 8);
+      if constexpr (moe_block_size >= 4)
+        local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 4);
+      if constexpr (moe_block_size >= 2)
+        local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 2);
+
+      local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 1);
+      block_num_valid_tokens = local_count;
+  #else
+      block_num_valid_tokens = __reduce_add_sync(0xffffffff, local_count);
+  #endif
+
+      if (lane_id == 0)
+        reinterpret_cast<int*>(sh_new)[0] = block_num_valid_tokens;
+    }
+
+    if (threadIdx.x < moe_block_size) {
+      int idx = sh_block_sorted_ids[threadIdx.x];
+      sh_rd_block_sorted_ids[threadIdx.x] = idx / top_k;
+
+      if (mul_topk_weights) {
+        idx = idx < prob_m_top_k ? idx : 0;
+        float topk_weight_tmp = topk_weights_ptr[idx];
+        if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
+          topk_weight_tmp *= global_scale_f32;
+        }
+        c_scalar_t2 topk_weight_val =
+            Cdtype::num2num2(Cdtype::float2num(topk_weight_tmp));
+        sh_block_topk_weights[threadIdx.x] = topk_weight_val;
+      }
+    }
+
+    __syncthreads();
+
+    block_num_valid_tokens = reinterpret_cast<int*>(sh_new)[0];
+    __syncthreads();
+  };
+
+  // when move to next moe block, find the next block_id and expert_id
+  // and then read moe block data
+  auto update_next_moe_block_data = [&]() {
+    if (par_id >= parallel) return;
+
+    old_expert_id = expert_id;
+    block_id = par_id;
+    expert_id = expert_ids_ptr[block_id];
+
+    if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
+      global_scale_f32 = global_scale_ptr[expert_id];
+    }
+
+    B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4);
+    scales_ptr += (expert_id - old_expert_id) * scales_expert_stride;
+    if constexpr (has_zp) {
+      zp_ptr += (expert_id - old_expert_id) * zp_expert_stride;
+    }
+    if constexpr (has_act_order) {
+      g_idx += (expert_id - old_expert_id) * prob_k;
+    }
+    if (has_bias) {
+      b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride;
+    }
+
+    read_moe_block_data(block_id);
+  };
+
+  // Compute all information about the current slice which is required for
+  // synchronization.
+  bool first_init = true;
+  auto init_part2_slice = [&]() {
+    slice_iters =
+        iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);
+    if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) slice_iters = 0;
+    if (slice_iters == 0) return;
+    if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;
+    slice_count = 1;
+    slice_idx = 0;
+    int col_first = iters * div_ceil(k_tiles * slice_col_par, iters);
+    if (col_first <= k_tiles * (slice_col_par + 1)) {
+      int col_off = col_first - k_tiles * slice_col_par;
+      slice_count = div_ceil(k_tiles - col_off, iters);
+      if (col_off > 0) slice_count++;
+      int delta_first = iters * blockIdx.x - col_first;
+      if (delta_first < 0 || (col_off == 0 && delta_first == 0))
+        slice_idx = slice_count - 1;
+      else {
+        slice_idx = slice_count - 1 - delta_first / iters;
+        if (col_off > 0) slice_idx--;
+      }
+    }
+    if (part2_mn_tiles >= gridDim.x) {
+      if (slice_count > 1 && slice_idx == slice_count - 1) {
+        locks_off++;
+      }
+    } else {
+      locks_off++;
+    }
+
+    if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) {
+      constexpr int threads_per_m = 16 * thread_n_blocks / 8;
+      int m_per_thread =
+          div_ceil(block_num_valid_tokens, threads / threads_per_m);
+      for (int i = 0; i < m_per_thread; i++) {
+        int row = threads / threads_per_m * i + threadIdx.x / threads_per_m;
+        if (row < block_num_valid_tokens) {
+          int64_t sorted_row = sh_block_sorted_ids[row];
+          int col = slice_col * 16 * thread_n_blocks / 8 +
+                    threadIdx.x % threads_per_m;
+          C[sorted_row * prob_n / 8 + col] = {0, 0, 0, 0};
+        }
+      }
+      // After write zero to output, write a negative value to lock.
+      // Every SM that processes the same slice would wait for
+      // the negative value, and then atomicAdd 1 to it.
+      // After all SMs are processed, the lock value would back to 0 again.
+      __syncthreads();
+      if (threadIdx.x == 0) locks[locks_off] = 1 - slice_count;
+    }
+
+    if (slice_col == n_tiles) {
+      slice_col = 0;
+      par_id++;
+      update_next_moe_block_data();
+    }
+    if (is_a_8bit && (first_init || slice_col == 0)) {
+      __syncthreads();
+      cp_async1_ca_pred(&sh_a_s[threadIdx.x],
+                        &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]],
+                        threadIdx.x < block_num_valid_tokens);
+    }
+  };
+
+  auto init_part1_slice = [&]() {
+    if (part1_mn_iters) {
+      part1_mn_iters--;
+      par_id = slice_col_par / n_tiles;
+      slice_col = slice_col_par % n_tiles;
+      slice_iters = k_tiles;
+      update_next_moe_block_data();
+      if (is_a_8bit) {
+        __syncthreads();
+        cp_async1_ca_pred(&sh_a_s[threadIdx.x],
+                          &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]],
+                          threadIdx.x < block_num_valid_tokens);
+      }
+    }
+  };
+
+  auto init_slice = [&]() {
+    if (!in_part2 && !part1_mn_iters) {
+      in_part2 = true;
+      slice_col_par = (iters * blockIdx.x) / k_tiles;
+      slice_row = (iters * blockIdx.x) % k_tiles;
+      slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles;
+      par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles;
+      update_next_moe_block_data();
+    }
+    if (!in_part2) {
+      init_part1_slice();
+    } else {
+      init_part2_slice();
+      first_init = false;
+    }
+  };
+
+  init_slice();
+
+  // A sizes/strides
+
+  // stride of the A matrix in global memory
+  int a_gl_stride = prob_k / (is_a_8bit ? 16 : 8);
+  // stride of an A matrix tile in shared memory
+  constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8);
+  // delta between subsequent A tiles in global memory
+  constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8);
+  // between subsequent accesses within a tile
+  int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o);
+  // between shared memory writes
+  constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o);
+  // within a shared memory tile
+  constexpr int a_sh_rd_delta_i = a_sh_stride * 16;
+  // overall size of a tile
+  constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks);
+  // number of shared write iterations for a tile
+  constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta);
+
+  // B sizes/strides
+  int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4));
+  constexpr int b_sh_stride =
+      ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4);
+  constexpr int b_thread_vecs = b_type.size_bits() == 4 ? 1 : 2;
+  constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs;
+
+  int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1);
+  constexpr int b_sh_wr_delta = threads * b_thread_vecs;
+  constexpr int b_sh_stage =
+      b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1);
+  constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
+
+  // Scale sizes/strides without act_order
+  int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8);
+  constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8);
+  constexpr int s_tb_groups =
+      !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
+          ? thread_k_blocks / group_blocks
+          : 1;
+  constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
+  int s_gl_rd_delta = s_gl_stride;
+
+  // Scale size/strides with act_order
+  constexpr int tb_k = 16 * thread_k_blocks;
+  constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0;
+  // constexpr int act_s_row_stride      = 1;
+  // int           act_s_col_stride      = act_s_row_stride * num_groups;
+  constexpr int act_s_max_num_groups = 32;
+  int act_s_col_stride = 1;
+  int act_s_col_warp_stride = act_s_col_stride * 8;
+
+  constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4);
+  int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps;
+
+  // Zero-points sizes/strides
+  int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4;
+  constexpr int zp_sh_stride = is_zp_float
+                                   ? 16 * thread_n_blocks / 8
+                                   : ((16 * thread_n_blocks) / pack_factor) / 4;
+  constexpr int zp_tb_groups = s_tb_groups;
+  constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0;
+  int zp_gl_rd_delta = zp_gl_stride;
+
+  // Global A read index of current thread.
+  int a_gl_rd_row = threadIdx.x / a_gl_rd_delta_o;
+  int a_gl_rd_col = a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o;
+  // Shared write index of current thread.
+  int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +
+                (threadIdx.x % a_gl_rd_delta_o);
+  // Shared read index.
+  int a_sh_rd =
+      a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) +
+      (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1));
+  a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters;
+
+  int b_gl_rd;
+  if (threads <= b_sh_stride) {
+    b_gl_rd = threadIdx.x;
+  } else {
+    b_gl_rd =
+        b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride);
+  }
+
+  b_gl_rd += B_expert_off + b_sh_stride * slice_col;
+  b_gl_rd += b_gl_rd_delta_o * slice_row;
+  auto b_sh_rd = threadIdx.x * b_thread_vecs;
+  b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1));
+
+  // For act_order
+  int slice_k_start = tb_k * slice_row;
+  int slice_k_finish = slice_k_start + tb_k * slice_iters;
+  int slice_k_start_shared_fetch = slice_k_start;
+  int slice_n_offset = act_s_col_tb_stride * slice_col;
+
+  // No act_order
+  int s_gl_rd;
+  if constexpr (!has_act_order) {
+    if constexpr (group_blocks == -1) {
+      s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
+    } else if constexpr (group_blocks >= thread_k_blocks) {
+      s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
+                s_sh_stride * slice_col + threadIdx.x;
+    } else {
+      s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
+                               threadIdx.x / s_sh_stride) +
+                s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
+    }
+  }
+  auto s_sh_wr = threadIdx.x;
+  bool s_sh_wr_pred = threadIdx.x < s_sh_stage;
+
+  // Zero-points
+  int zp_gl_rd;
+  if constexpr (has_zp) {
+    if constexpr (group_blocks == -1) {
+      zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
+    } else if constexpr (group_blocks >= thread_k_blocks) {
+      zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
+                 zp_sh_stride * slice_col + threadIdx.x;
+    } else {
+      zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
+                                 threadIdx.x / zp_sh_stride) +
+                 zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
+    }
+  }
+  auto zp_sh_wr = threadIdx.x;
+  bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage;
+
+  // We use a different scale layout for grouped and column-wise quantization as
+  // we scale a `half2` tile in column-major layout in the former and in
+  // row-major in the latter case.
+  int s_sh_rd;
+  if constexpr (is_a_8bit) {
+    s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4);
+  } else if constexpr (group_blocks != -1)
+    s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4;
+  else if constexpr (group_blocks == -1 &&
+                     (m_block_size_8 || (has_zp && !dequant_skip_flop)))
+    s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8;
+  else
+    s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4;
+
+  int bias_sh_rd;
+  if constexpr (m_block_size_8) {
+    bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8;
+  } else {
+    bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) +
+                 (threadIdx.x % 32) % 4;
+  }
+
+  int bias_sh_wr = threadIdx.x;
+  int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
+
+  // Zero-points have the same read layout as the scales
+  // (without column-wise case)
+  constexpr int num_col_threads = 8;
+  constexpr int num_row_threads = 4;
+  constexpr int num_ints_per_thread = 8 / pack_factor;
+  int zp_sh_rd;
+  if constexpr (has_zp) {
+    if constexpr (is_zp_float) {
+      if constexpr (group_blocks != -1) {
+        zp_sh_rd =
+            8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4;
+      }
+    } else if (is_a_8bit) {
+      zp_sh_rd = num_ints_per_thread * num_col_threads *
+                     ((threadIdx.x / 32) % tb_n_warps / 2) +
+                 num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
+    } else {
+      zp_sh_rd = num_ints_per_thread * num_col_threads *
+                     ((threadIdx.x / 32) % tb_n_warps) +
+                 num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
+    }
+  }
+
+  // To ensure that writing and reading A tiles to/from shared memory, the
+  // latter in fragment format, is fully bank conflict free, we need to use a
+  // rather fancy XOR-based layout. The key here is that neither reads nor
+  // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the
+  // same shared memory banks. Further, it seems (based on NSight-Compute) that
+  // each warp must also write a consecutive memory segment?
+  auto transform_a = [&](int i) {
+    int row = i / a_gl_rd_delta_o;
+    return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8);
+  };
+  // Since the computation of this remapping is non-trivial and, due to our main
+  // loop unrolls, all shared memory accesses are static, we simply precompute
+  // both transformed reads and writes.
+  int a_sh_wr_trans[a_sh_wr_iters];
+  #pragma unroll
+  for (int i = 0; i < a_sh_wr_iters; i++)
+    a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr);
+  int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks];
+  #pragma unroll
+  for (int i = 0; i < b_sh_wr_iters; i++) {
+  #pragma unroll
+    for (int j = 0; j < thread_m_blocks; j++)
+      a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd);
+  }
+
+  // Since B-accesses have non-constant stride they have to be computed at
+  // runtime; we break dependencies between subsequent accesses with a tile by
+  // maintining multiple pointers (we have enough registers), a tiny
+  // optimization.
+
+  // Shared memory storage for global fetch pipelines.
+  constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks;
+  constexpr int sh_b_size = stages * b_sh_stage;
+  int4* sh_b = sh_new;
+  int4* sh_red = sh_new;
+
+  constexpr int sh_size_b_red_min =
+      (sh_red_size < sh_b_size ? sh_red_size : sh_b_size);
+  constexpr int sh_size_b_red_max =
+      (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
+  constexpr int sh_bias_size = (thread_n_blocks * 16 / 8);
+  constexpr int sh_b_red_bias_size =
+      sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size)
+          ? sh_size_b_red_max
+          : (sh_size_b_red_min + sh_bias_size);
+
+  int4* sh_bias = sh_new + sh_size_b_red_min;
+  int4* sh_g_idx = sh_new + sh_b_red_bias_size;
+  int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
+  constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
+                                          : (stages * s_sh_stage);
+  int4* sh_s = sh_zp + (stages * zp_sh_stage);
+  int4* sh_a = sh_s + sh_s_size;
+
+  // Register storage for double buffer of shared memory reads.
+  FragA frag_a[2][thread_m_blocks];
+  I4 frag_b_quant[2][b_thread_vecs];
+  FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2];
+  FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2];
+  FragS frag_s[2][4];  // No act-order
+  FragS frag_bias[2][4];
+  FragS act_frag_s[2][4][4];             // For act-order
+  int frag_qzp[2][num_ints_per_thread];  // Zero-points
+  FragZP frag_zp;                        // Zero-points in fp16
+  FragZP frag_zpf[2];                    // Zero-points in fp16 in HQQ
+
+  if constexpr (is_a_8bit && group_blocks != -1) {
+  #pragma unroll
+    for (int j = 0; j < 2; j++) {
+  #pragma unroll
+      for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+        for (int g = 0; g < 4; g++) {
+          frag_c_tmp[i][j][0][g] = 0.0f;
+        }
+
+  #pragma unroll
+        for (int g = 0; g < 4; g++) {
+          frag_c_tmp[i][j][1][g] = 0.0f;
+        }
+      }
+    }
+  }
+
+  // Zero accumulators.
+  auto zero_accums = [&]() {
+  #pragma unroll
+    for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++)
+      reinterpret_cast<float*>(frag_c)[i] = 0;
+  };
+
+  int sh_first_group_id = -1;
+  int sh_num_groups = -1;
+
+  auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id,
+                                              int last_group_id) {
+    sh_first_group_id = first_group_id;
+    sh_num_groups = last_group_id - first_group_id + 1;
+
+    if (sh_num_groups > act_s_max_num_groups) {
+      sh_num_groups = act_s_max_num_groups;
+    }
+
+    if (sh_first_group_id + sh_num_groups > num_groups) {
+      sh_num_groups = num_groups - sh_first_group_id;
+    }
+
+    int row_offset = first_group_id * s_gl_stride;
+
+    if (is_async) {
+      for (int i = 0; i < sh_num_groups; i++) {
+        if (threadIdx.x < s_sh_stride) {
+          cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x],
+                         &scales_ptr[row_offset + (i * s_gl_stride) +
+                                     slice_n_offset + threadIdx.x]);
+        }
+      }
+    } else {
+      for (int i = 0; i < sh_num_groups; i++) {
+        if (threadIdx.x < s_sh_stride) {
+          sh_s[(i * s_sh_stride) + threadIdx.x] =
+              scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset +
+                         threadIdx.x];
+        }
+      }
+    }
+  };
+  // Asynchronously fetch the next A, B and s tile from global to the next
+  // shared memory pipeline location.
+  auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
+    if (pred) {
+      int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe;
+  #pragma unroll
+      for (int i = 0; i < a_sh_wr_iters; i++) {
+        int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row;
+        int64_t sorted_row = 0;
+        if (!m_block_size_8 || row < 8)
+          sorted_row = sh_rd_block_sorted_ids[row];
+        int64_t true_idx =
+            sorted_row * a_gl_stride + a_gl_rd_col + a_gl_rd_delta_o * a_off;
+        cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx],
+                       row < block_num_valid_tokens);
+      }
+
+      int4* sh_b_stage = sh_b + b_sh_stage * pipe;
+  #pragma unroll
+      for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) {
+        constexpr int count = div_ceil(b_sh_stride, threads);
+        int b_gl_idx =
+            b_gl_rd + (i % count) * threads +
+            b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride);
+
+        cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]);
+      }
+
+      b_gl_rd += b_gl_rd_delta_o;
+
+      if constexpr (has_act_order) {
+        // Fetch g_idx thread-block portion
+        int full_pipe = a_off;
+        int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe;
+        if (cur_k < prob_k && cur_k < slice_k_finish) {
+          int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
+
+          int4 const* cur_g_idx_stage_ptr =
+              reinterpret_cast<int4 const*>(&g_idx[cur_k]);
+
+          if (threadIdx.x < g_idx_stage) {
+            cp_async4_pred(&sh_g_idx_stage[threadIdx.x],
+                           &cur_g_idx_stage_ptr[threadIdx.x]);
+          }
+        }
+      } else {
+        if constexpr (group_blocks != -1) {
+          int4* sh_s_stage = sh_s + s_sh_stage * pipe;
+
+          // Only fetch scales if this tile starts a new group
+          if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
+            if (s_sh_wr_pred) {
+              cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
+            }
+            s_gl_rd += s_gl_rd_delta * s_tb_groups;
+          }
+        }
+
+        if constexpr (has_zp && group_blocks != -1) {
+          int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
+
+          // Only fetch zero points if this tile starts a new group
+          if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
+            if (zp_sh_wr_pred) {
+              cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
+            }
+            zp_gl_rd += zp_gl_rd_delta * zp_tb_groups;
+          }
+        }
+      }
+    }
+    // Insert a fence even when we are winding down the pipeline to ensure that
+    // waiting is also correct at this point.
+    cp_async_fence();
+  };
+
+  auto fetch_col_zp_to_shared = [&]() {
+    if (zp_sh_wr_pred) {
+      cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]);
+    }
+  };
+
+  auto fetch_col_scale_to_shared = [&]() {
+    if (s_sh_wr_pred) {
+      cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
+    }
+  };
+
+  // Wait until the next thread tile has been loaded to shared memory.
+  auto wait_for_stage = [&]() {
+    // We only have `stages - 2` active fetches since we are double buffering
+    // and can only issue the next fetch when it is guaranteed that the previous
+    // shared memory load is fully complete (as it may otherwise be
+    // overwritten).
+    cp_async_wait<stages - 2>();
+    __syncthreads();
+  };
+
+  // Load the next sub-tile from the current location in the shared memory pipe
+  // into the current register buffer.
+  auto fetch_to_registers = [&](int k, int pipe) {
+    int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe;
+  #pragma unroll
+    for (int i = 0; i < thread_m_blocks; i++)
+      ldsm<m_block_size_8 ? 2 : 4, a_type_id>(
+          frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);
+    int4* sh_b_stage = sh_b + b_sh_stage * pipe;
+
+  #pragma unroll
+    for (int i = 0; i < b_thread_vecs; i++) {
+      frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>(
+          &sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]);
+    }
+  };
+
+  bool is_same_group[stages];
+  int same_group_id[stages];
+
+  auto init_same_group = [&](int pipe) {
+    if constexpr (!has_act_order) {
+      return;
+    }
+
+    int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
+    int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
+
+    int group_id_1 = sh_g_idx_int_ptr[0];
+    int group_id_2 = sh_g_idx_int_ptr[tb_k - 1];
+
+    is_same_group[pipe] = group_id_1 == group_id_2;
+    same_group_id[pipe] = group_id_1;
+  };
+
+  auto fetch_scales_to_registers = [&](int k, int full_pipe) {
+    int pipe = full_pipe % stages;
+    using IT1 = typename std::conditional_t<is_a_8bit, int2, int4>;
+    using IT0 = typename std::conditional_t<is_a_8bit, int, int2>;
+    constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1);
+
+    if constexpr (!has_act_order) {
+      // No act-order case
+      if constexpr (group_blocks == -1) {
+        // load only when starting a new slice
+        if (k == 0 && full_pipe == 0 && dequant_skip_flop) {
+          reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd];
+          reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
+        }
+      } else if constexpr (group_blocks != -1) {
+        if constexpr (group_blocks >= thread_k_blocks) {
+          constexpr int g = group_blocks / thread_k_blocks;
+          if (pipe % g == 0) {
+            if (k % b_sh_wr_iters == 0) {
+              int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g));
+              reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
+            } else {
+              reinterpret_cast<int4*>(&frag_s[1])[0] =
+                  reinterpret_cast<int4*>(&frag_s[0])[0];
+            }
+          }
+        } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
+          auto warp_id = threadIdx.x / 32;
+          int warp_row = warp_id / tb_n_warps;
+
+          int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
+          int cur_group_id = k_blocks / group_blocks2;
+
+          int4* sh_s_stage = sh_s + s_sh_stage * pipe;
+
+          if constexpr (!is_8bit_scale) {
+            reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
+                sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
+          } else {
+            reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
+                reinterpret_cast<int2*>(
+                    sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
+          }
+        } else if (group_blocks >= b_sh_wr_iters) {
+          if constexpr (!is_8bit_scale) {
+            reinterpret_cast<int4*>(&frag_s[1])[0] =
+                reinterpret_cast<int4*>(&frag_s[0])[0];
+          } else {
+            reinterpret_cast<int2*>(&frag_s[1])[0] =
+                reinterpret_cast<int2*>(&frag_s[0])[0];
+          }
+        }
+      }
+
+      return;
+    }
+
+    // Act-order case
+
+    // Determine K of the "current" thread-block
+    int cur_k = slice_k_start + tb_k * full_pipe;
+    if (cur_k >= prob_k || cur_k >= slice_k_finish) {
+      return;
+    }
+
+    // Reset (to current thread-block) since we read g_idx portion from the
+    // shared memory
+    cur_k = 0;
+
+    // Progress to current iteration
+    cur_k += k % b_sh_wr_iters;
+
+    // Determine "position" inside the thread-block (based on warp and
+    // thread-id)
+    auto warp_id = threadIdx.x / 32;
+    int warp_row = warp_id / tb_n_warps;
+    int warp_col = warp_id % tb_n_warps;
+
+    cur_k += warp_row * 16 * b_sh_wr_iters;
+
+    auto th_id = threadIdx.x % 32;
+    cur_k += (th_id % 4) * 2;  // Due to tensor-core layout for fp16 B matrix
+
+    int s_col_shift =
+        /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) +
+        (th_id / 4) * act_s_col_stride;
+
+    if (is_same_group[pipe]) {
+      if (k % 2 == 0) {
+        *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
+            sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride +
+                 s_col_shift];
+      } else {
+        *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
+            *(reinterpret_cast<int4*>(&(act_frag_s[(k - 1) % 2][0][0])));
+      }
+
+      for (int i = 1; i < 4; i++) {
+        *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
+            *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0])));
+      }
+      return;
+    }
+
+    int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
+    int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
+
+    constexpr int k_frag_offsets[4] = {0, 1, 8,
+                                       9};  // Tensor core offsets per thread
+
+  #pragma unroll
+    for (int i = 0; i < 4; i++) {
+      int actual_k = cur_k + k_frag_offsets[i];
+
+      int group_id = sh_g_idx_int_ptr[actual_k];
+      int rel_group_id = group_id - sh_first_group_id;
+
+      *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
+          sh_s[rel_group_id * s_sh_stride + s_col_shift];
+    }
+  };
+
+  auto fetch_zp_to_registers = [&](int k, int full_pipe) {
+    // This code does not handle group_blocks == 0,
+    // which signifies act_order.
+    // has_zp implies AWQ, which doesn't have act_order,
+    static_assert(!has_zp || group_blocks != 0);
+
+    if constexpr (has_zp && !is_zp_float) {
+      int pipe = full_pipe % stages;
+
+      if constexpr (group_blocks == -1) {
+        // load only when starting a new slice
+        if (k == 0 && full_pipe == 0 || is_a_8bit) {
+  #pragma unroll
+          for (int i = 0; i < num_ints_per_thread; i++) {
+            frag_qzp[k % 2][i] = (reinterpret_cast<int*>(sh_zp))[zp_sh_rd + i];
+          }
+        }
+      } else if constexpr (group_blocks >= thread_k_blocks) {
+        constexpr int g = group_blocks / thread_k_blocks;
+        if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) {
+          int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g));
+  #pragma unroll
+          for (int i = 0; i < num_ints_per_thread; i++) {
+            frag_qzp[k % 2][i] =
+                (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
+          }
+        }
+      } else {
+        auto warp_id = threadIdx.x / 32;
+
+        int warp_row = warp_id / tb_n_warps;
+
+        int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
+        int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1);
+
+        int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
+
+        sh_zp_stage += cur_group_id * zp_sh_stride;
+
+  #pragma unroll
+        for (int i = 0; i < num_ints_per_thread; i++) {
+          frag_qzp[k % 2][i] =
+              (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
+        }
+      }
+    }
+
+    else if constexpr (has_zp && is_zp_float) {
+      int pipe = full_pipe % stages;
+
+      if constexpr (group_blocks != -1) {
+        if constexpr (group_blocks >= thread_k_blocks) {
+          constexpr int g = group_blocks / thread_k_blocks;
+          if (pipe % g == 0 && k % b_sh_wr_iters == 0) {
+            int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g));
+            reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
+                sh_zp_stage[zp_sh_rd];
+          }
+        } else if (group_blocks < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
+          auto warp_id = threadIdx.x / 32;
+
+          int warp_row = warp_id / tb_n_warps;
+          int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
+          int cur_group_id = k_blocks / group_blocks;
+
+          int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
+
+          reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
+              sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride];
+        }
+      }
+    }
+  };
+
+  auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) {
+    if constexpr (a_type.size_bits() != b_type.size_bits()) {
+      if constexpr (is_a_8bit && has_zp) {
+        sub_zp_and_dequant<scalar_32bit_t, b_type_id, dequant_skip_flop>(
+            q, frag_b_ptr, zp);
+      } else {
+        dequant<scalar_32bit_t, b_type_id, dequant_skip_flop>(q, frag_b_ptr);
+      }
+    }
+  };
+
+  // Execute the actual tensor core matmul of a sub-tile.
+  bool is_first_matmul_in_slice = true;
+  auto matmul = [&](int k, int pipe) {
+    if (is_a_8bit) return;
+    int k2 = k % 2;
+    constexpr int g =
+        group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1;
+    const bool is_new_zp =
+        (group_blocks == 0) ||
+        ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) &&
+            (pipe % g == 0) ||
+        (group_blocks == -1 && is_first_matmul_in_slice);
+    if constexpr (has_zp && !is_zp_float) {
+      if (is_new_zp) {
+        if constexpr (group_blocks == -1) is_first_matmul_in_slice = false;
+        int zp_quant_0, zp_quant_1;
+
+        if constexpr (b_type.size_bits() == 4) {
+          zp_quant_0 = frag_qzp[k2][0];
+          zp_quant_1 = zp_quant_0 >> 8;
+        } else {
+          static_assert(b_type.size_bits() == 8);
+          zp_quant_0 = frag_qzp[k2][0];
+          zp_quant_1 = frag_qzp[k2][1];
+        }
+
+        dequant_data(zp_quant_0, reinterpret_cast<scalar_32bit_t*>(&frag_zp));
+        dequant_data(zp_quant_1,
+                     reinterpret_cast<scalar_32bit_t*>(&frag_zp) + 2);
+      }
+    }
+    if constexpr (!dequant_skip_flop && has_zp && is_zp_float) {
+      if (is_new_zp) {
+        reinterpret_cast<int4*>(&frag_zp)[0] =
+            reinterpret_cast<int4*>(&frag_zpf[k2])[0];
+      }
+    }
+
+    if constexpr (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu) {
+      int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
+      int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
+
+      dequant_fp8_scales<c_scalar_t2, s_type_id>(
+          s_quant_0, reinterpret_cast<c_scalar_t2*>(&frag_s[k2]));
+      dequant_fp8_scales<c_scalar_t2, s_type_id>(
+          s_quant_1, reinterpret_cast<c_scalar_t2*>(&frag_s[k2]) + 2);
+    }
+
+  // We have the m dimension as the inner loop in order to encourage overlapping
+  // dequantization and matmul operations.
+  #pragma unroll
+    for (int j = 0; j < 4; j++) {
+      FragB frag_b0;
+      FragB frag_b1;
+      int b_quant_0, b_quant_1;
+
+      if constexpr (b_type_id == vllm::kFE2M1f.id()) {
+        b_quant_1 = frag_b_quant[k2][0][j];
+        b_quant_0 = b_quant_1 << 8;
+      } else if constexpr (b_type.size_bits() == 4) {
+        b_quant_0 = frag_b_quant[k2][0][j];
+        b_quant_1 = b_quant_0 >> 8;
+      } else {
+        static_assert(b_type.size_bits() == 8);
+        int* frag_b_quant_ptr = reinterpret_cast<int*>(frag_b_quant[k2]);
+        b_quant_0 = frag_b_quant_ptr[j * 2 + 0];
+        b_quant_1 = frag_b_quant_ptr[j * 2 + 1];
+      }
+
+      dequant_data(b_quant_0, reinterpret_cast<scalar_32bit_t*>(&frag_b0));
+      dequant_data(b_quant_1, reinterpret_cast<scalar_32bit_t*>(&frag_b1));
+
+      if constexpr (dequant_skip_flop && has_zp && !is_zp_float && !is_a_8bit) {
+        sub_zp<a_type_id>(frag_b0, frag_zp[j], 0);
+        sub_zp<a_type_id>(frag_b1, frag_zp[j], 1);
+      }
+
+      // Apply scale to frag_b0
+      if constexpr (has_act_order && !is_a_8bit) {
+        static_assert(group_blocks != -1);
+        scale4<a_type_id>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
+                          act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
+        scale4<a_type_id>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
+                          act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
+      } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
+                           group_blocks == -1 && !is_a_8bit) {
+        int idx = (threadIdx.x / 4) % 2;
+        scalar_t2 s2 = Adtype::nums2num2(
+            reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx],
+            reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 1])[idx]);
+        if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2);
+        scale_and_sub<a_type_id>(frag_b0, s2.x, frag_zp[j].x);
+        scale_and_sub<a_type_id>(frag_b1, s2.y, frag_zp[j].y);
+      } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 &&
+                           !is_a_8bit) {
+        if (is_new_zp)
+          frag_zp[j] = __hmul2(frag_zp[j],
+                               *reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
+        scale_and_sub<a_type_id>(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x);
+        scale_and_sub<a_type_id>(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y);
+      } else if constexpr (group_blocks != -1 && !is_a_8bit) {
+        scale<a_type_id>(frag_b0, frag_s[k2][j], 0);
+        scale<a_type_id>(frag_b1, frag_s[k2][j], 1);
+      }
+
+  #pragma unroll
+      for (int i = 0; i < thread_m_blocks; i++) {
+        if constexpr (m_block_size_8) {
+          mma_trans<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b0, frag_b1,
+                                               frag_c[i][j][0]);
+        } else {
+          mma<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b0,
+                                         frag_c[i][j][0]);
+          mma<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b1,
+                                         frag_c[i][j][1]);
+        }
+      }
+    }
+  };
+
+  auto matmul_a8 = [&](int k) {
+    int k2 = k % 2;
+  #pragma unroll
+    for (int j = 0; j < 2; j++) {
+      FragB frag_b[2];
+
+      if (is_a_8bit && b_type.size_bits() == 4 && !has_zp) {
+        dequant_data(frag_b_quant[k2][0][j * 2],
+                     reinterpret_cast<scalar_32bit_t*>(&frag_b));
+        dequant_data(frag_b_quant[k2][0][j * 2 + 1],
+                     reinterpret_cast<scalar_32bit_t*>(&frag_b) + 2);
+      } else if (is_a_8bit && b_type.size_bits() == 4 && has_zp) {
+        int off = (threadIdx.x / 32) % 2 * 2 + j;
+        int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF;
+        dequant_data(frag_b_quant[k2][0][j * 2],
+                     reinterpret_cast<scalar_32bit_t*>(&frag_b), zp);
+        zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF;
+        dequant_data(frag_b_quant[k2][0][j * 2 + 1],
+                     reinterpret_cast<scalar_32bit_t*>(&frag_b) + 2, zp);
+      } else {
+        reinterpret_cast<int2*>(&frag_b)[0] =
+            reinterpret_cast<int2*>(&frag_b_quant[k2][j])[0];
+        reinterpret_cast<int2*>(&frag_b)[1] =
+            reinterpret_cast<int2*>(&frag_b_quant[k2][j])[1];
+      }
+
+  #pragma unroll
+      for (int i = 0; i < thread_m_blocks; i++) {
+        mma<a_type_id, false, 32>(
+            frag_a[k2][i], frag_b[0],
+            (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]);
+        mma<a_type_id, false, 32>(
+            frag_a[k2][i], frag_b[1],
+            (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]);
+      }
+
+      if constexpr (group_blocks != -1) {
+        if (group_blocks == 2 || k == 1) {
+          if constexpr (a_type == vllm::kS8) {
+            int2 s_vals[2];
+            s_vals[0] = {
+                (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2][0])[0],
+                (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2][0])[1]};
+            s_vals[1] = {
+                (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2 + 1][0])[0],
+                (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2 + 1][0])[1]};
+
+  #pragma unroll
+            for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+              for (int g = 0; g < 4; g++) {
+                int scale = reinterpret_cast<int*>(&s_vals[0])[g % 2];
+                *reinterpret_cast<int32_t*>(&frag_c[i][j][0][g]) +=
+                    *reinterpret_cast<int32_t*>(&frag_c_tmp[i][j][0][g]) *
+                    scale;
+                frag_c_tmp[i][j][0][g] = 0.0f;
+              }
+
+  #pragma unroll
+              for (int g = 0; g < 4; g++) {
+                int scale = reinterpret_cast<int*>(&s_vals[1])[g % 2];
+                *reinterpret_cast<int32_t*>(&frag_c[i][j][1][g]) +=
+                    *reinterpret_cast<int32_t*>(&frag_c_tmp[i][j][1][g]) *
+                    scale;
+                frag_c_tmp[i][j][1][g] = 0.0f;
+              }
+            }
+          } else {
+            float2 s_vals[2];
+            if constexpr (s_type_id != vllm::kFE8M0fnu.id()) {
+              static_assert(a_type.size_bits() == 16 ||
+                            s_type.size_bits() == 16);
+              s_vals[0] = Cdtype::num22float2(frag_s[k2][j * 2][0]);
+              s_vals[1] = Cdtype::num22float2(frag_s[k2][j * 2 + 1][0]);
+            } else {
+              int32_t* s_vals_int = reinterpret_cast<int32_t*>(&s_vals[0]);
+              int32_t s_vals_e8m0 =
+                  *reinterpret_cast<int32_t*>(&frag_s[k2][j][0]);
+
+              s_vals_int[0] = (s_vals_e8m0 & 0xFF) << 23;
+              s_vals_int[1] = (s_vals_e8m0 & 0xFF00) << 15;
+              s_vals_int[2] = (s_vals_e8m0 & 0xFF0000) << 7;
+              s_vals_int[3] = (s_vals_e8m0 & 0xFF000000) >> 1;
+            }
+
+  #pragma unroll
+            for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+              for (int g = 0; g < 4; g++) {
+                float scale = reinterpret_cast<float*>(&s_vals[0])[g % 2];
+                frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale;
+                frag_c_tmp[i][j][0][g] = 0.0f;
+              }
+
+  #pragma unroll
+              for (int g = 0; g < 4; g++) {
+                float scale = reinterpret_cast<float*>(&s_vals[1])[g % 2];
+                frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale;
+                frag_c_tmp[i][j][1][g] = 0.0f;
+              }
+            }
+          }
+        }
+      }
+    }
+  };
+
+  // Since we slice across the k dimension of a tile in order to increase the
+  // number of warps while keeping the n dimension of a tile reasonable, we have
+  // multiple warps that accumulate their partial sums of the same output
+  // location; which we have to reduce over in the end. We do in shared memory.
+  auto thread_block_reduce = [&]() {
+    constexpr int red_off = threads / b_sh_stride_threads / 2;
+    if (red_off >= 1) {
+      auto red_idx = threadIdx.x / b_sh_stride_threads;
+      constexpr int red_sh_stride =
+          b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2;
+      constexpr int red_sh_delta = b_sh_stride_threads;
+      int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
+                      (threadIdx.x % b_sh_stride_threads);
+
+      // Parallel logarithmic shared memory reduction. We make sure to avoid any
+      // unnecessary read or write iterations, e.g., for two warps we write only
+      // once by warp 1 and read only once by warp 0.
+
+  #pragma unroll
+      for (int m_block = 0; m_block < thread_m_blocks; m_block++) {
+  #pragma unroll
+        for (int i = red_off; i > 0; i /= 2) {
+          if (i <= red_idx && red_idx < 2 * i) {
+  #pragma unroll
+            for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2;
+                 j += (m_block_size_8 ? 2 : 1)) {
+              int red_sh_wr =
+                  red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
+              if (i < red_off) {
+                float* c_rd = reinterpret_cast<float*>(
+                    &sh_red[red_sh_delta * j + red_sh_rd]);
+                float* c_wr = reinterpret_cast<float*>(&sh_red[red_sh_wr]);
+  #pragma unroll
+                for (int k = 0; k < 4; k++)
+                  reinterpret_cast<FragC*>(
+                      frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] +=
+                      c_rd[k] + c_wr[k];
+              }
+              sh_red[red_sh_wr] = reinterpret_cast<int4*>(
+                  &frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j];
+            }
+          }
+          __syncthreads();
+        }
+        if (red_idx == 0) {
+  #pragma unroll
+          for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2;
+               i += (m_block_size_8 ? 2 : 1)) {
+            float* c_rd =
+                reinterpret_cast<float*>(&sh_red[red_sh_delta * i + red_sh_rd]);
+  #pragma unroll
+            for (int j = 0; j < 4; j++)
+              reinterpret_cast<FragC*>(
+                  frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j];
+          }
+        }
+        __syncthreads();
+      }
+    }
+  };
+
+  // Since multiple threadblocks may process parts of the same column slice, we
+  // finally have to globally reduce over the results. As the striped
+  // partitioning minimizes the number of such reductions and our outputs are
+  // usually rather small, we perform this reduction serially in L2 cache.
+  auto global_reduce_fp16 = [&](bool first = false, bool last = false) {
+    // We are very careful here to reduce directly in the output buffer to
+    // maximize L2 cache utilization in this step. To do this, we write out
+    // results in FP16 (but still reduce with FP32 compute).
+    constexpr int active_threads = 32 * tb_n_warps;
+    bool is_th_active = threadIdx.x < active_threads;
+    if (!is_th_active) {
+      return;
+    }
+
+    int c_gl_stride = prob_n / 8 * (is_a_8bit ? 2 : 1);
+    int c_gl_wr_delta_o = 8 * c_gl_stride;
+    int c_gl_wr_delta_i = 4 * (active_threads / 32);
+    int c_gl_wr;
+    if constexpr (m_block_size_8) {
+      c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) +
+                (threadIdx.x % 32) / 8;
+      c_gl_wr += (2 * thread_n_blocks) * slice_col;
+    } else {
+      c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) +
+                4 * (threadIdx.x / 32) + threadIdx.x % 4;
+      c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1);
+    }
+    constexpr int c_sh_wr_delta = active_threads;
+    int c_sh_wr = threadIdx.x;
+
+    if (!first) {
+
+  #pragma unroll
+      for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
+        int c_idx;
+        if constexpr (m_block_size_8)
+          c_idx = c_gl_wr + i * c_gl_stride +
+                  (threadIdx.x % 8) / 4 * c_gl_wr_delta_i;
+        else
+          c_idx =
+              c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2);
+        if (c_idx / c_gl_stride < block_num_valid_tokens) {
+          int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride];
+          int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride;
+          if constexpr (is_a_8bit) {
+            int2* sh_red_int2 = reinterpret_cast<int2*>(sh_red);
+            int2* c_int2 = reinterpret_cast<int2*>(C);
+            sh_red_int2[c_sh_wr + c_sh_wr_delta * i] = c_int2[true_idx];
+          } else {
+            sh_red[c_sh_wr + c_sh_wr_delta * i] = C[true_idx];
+          }
+        }
+      }
+    }
+
+  #pragma unroll
+    for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
+      if (!first) {
+        c_scalar_t* c_red_f16;
+        if constexpr (is_a_8bit) {
+          int2 tmp =
+              reinterpret_cast<int2*>(sh_red)[c_sh_wr + i * c_sh_wr_delta];
+          c_red_f16 = reinterpret_cast<c_scalar_t*>(&tmp);
+        } else {
+          int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta];
+          c_red_f16 = reinterpret_cast<c_scalar_t*>(&tmp);
+        }
+  #pragma unroll
+        for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) {
+          int delta = 0;
+          if constexpr (m_block_size_8) {
+            delta = j % 2 == 1 ? -2 : 0;
+          }
+          reinterpret_cast<float*>(
+              &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) +
+                       delta] += Cdtype::num2float(c_red_f16[j]);
+        }
+      }
+      if (!last) {
+        c_scalar_t c_f16[is_a_8bit ? 4 : 8];
+  #pragma unroll
+        for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) {
+          int delta = 0;
+          if constexpr (m_block_size_8) {
+            delta = j % 2 == 1 ? -2 : 0;
+          }
+          c_f16[j] = Cdtype::float2num(reinterpret_cast<float*>(
+              &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) +
+                       delta]);
+        }
+
+        int c_idx;
+        if constexpr (m_block_size_8)
+          c_idx = c_gl_wr + i * c_gl_stride +
+                  (threadIdx.x % 8) / 4 * c_gl_wr_delta_i;
+        else
+          c_idx =
+              c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2);
+        if (c_idx / c_gl_stride < block_num_valid_tokens) {
+          int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride];
+          int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride;
+          if constexpr (is_a_8bit) {
+            int2* c_int2 = reinterpret_cast<int2*>(C);
+            c_int2[true_idx] = *reinterpret_cast<int2*>(c_f16);
+          } else {
+            C[true_idx] = *reinterpret_cast<int4*>(c_f16);
+          }
+        }
+      }
+    }
+  };
+
+  // Globally reduce over threadblocks that compute the same column block.
+  // We use a tmp C buffer to reduce in full fp32 precision.
+  auto global_reduce_fp32 = [&](bool first = false, bool last = false) {
+    constexpr int tb_m = thread_m_blocks * 16;
+    constexpr int tb_n = thread_n_blocks * 16;
+
+    constexpr int c_size = tb_m * tb_n * sizeof(float) / 16;
+
+    constexpr int active_threads = 32 * tb_n_warps;
+    bool is_th_active = threadIdx.x < active_threads;
+
+    constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4;
+    constexpr int th_size = num_floats * sizeof(float) / 16;
+
+    int c_cur_offset = locks_off * c_size;
+
+    if (!is_th_active) {
+      return;
+    }
+
+    if (!first) {
+      float* frag_c_ptr = reinterpret_cast<float*>(&frag_c);
+  #pragma unroll
+      for (int k = 0; k < th_size; k++) {
+        if constexpr (m_block_size_8) {
+          if (k % 2) continue;
+        } else {
+          if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens)
+            continue;
+        }
+
+        sh_red[threadIdx.x] =
+            C_tmp[c_cur_offset + active_threads * k + threadIdx.x];
+
+        float* sh_c_ptr = reinterpret_cast<float*>(&sh_red[threadIdx.x]);
+  #pragma unroll
+        for (int f = 0; f < 4; f++) {
+          frag_c_ptr[k * 4 + f] += sh_c_ptr[f];
+        }
+      }
+    }
+
+    if (!last) {
+      int4* frag_c_ptr = reinterpret_cast<int4*>(&frag_c);
+  #pragma unroll
+      for (int k = 0; k < th_size; k++) {
+        if constexpr (m_block_size_8) {
+          if (k % 2) continue;
+        } else {
+          if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens)
+            continue;
+        }
+
+        C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k];
+      }
+    }
+  };
+
+  // Write out the reduce final result in the correct layout. We only actually
+  // reshuffle matrix fragments in this step, the reduction above is performed
+  // in fragment layout.
+  auto write_result = [&](bool last) {
+    int c_gl_stride = prob_n / 8;
+    constexpr int c_sh_stride = 2 * thread_n_blocks + 1;
+    int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks));
+    constexpr int c_sh_rd_delta =
+        c_sh_stride * (threads / (2 * thread_n_blocks));
+
+    int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) +
+                  (threadIdx.x % (2 * thread_n_blocks));
+    c_gl_wr += (2 * thread_n_blocks) * slice_col;
+    int c_sh_wr;
+    if constexpr (m_block_size_8) {
+      c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) +
+                (threadIdx.x % 32) / 4;
+      c_sh_wr += 64 * (threadIdx.x / 32);
+    } else {
+      c_sh_wr =
+          (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4;
+      c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32);
+    }
+
+    int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) +
+                  (threadIdx.x % (2 * thread_n_blocks));
+
+    // We first reorder in shared memory to guarantee the most efficient final
+    // global write patterns
+    auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
+      if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
+        if (!mul_topk_weights) {
+          c0 *= global_scale_f32;
+          c1 *= global_scale_f32;
+        }
+      }
+
+      c_scalar_t2 res =
+          Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1));
+
+      // For per-column quantization we finally apply the scale here (only for
+      // 4-bit)
+      if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit &&
+                    b_type.size_bits() == 4 &&
+                    (has_zp && dequant_skip_flop || !has_zp)) {
+        c_scalar_t2 tmp_scale = s[0];
+        if constexpr (m_block_size_8) {
+          tmp_scale = Cdtype::num2num2(
+              reinterpret_cast<scalar_t*>(&s[0])[(threadIdx.x % 8) / 4]);
+        }
+        res = __hmul2(res, tmp_scale);
+      }
+
+      if (has_bias && last) {
+        c_scalar_t2 tmp_bias = b_bias[0];
+        if constexpr (m_block_size_8) {
+          tmp_bias = Cdtype::num2num2(
+              reinterpret_cast<scalar_t*>(&b_bias[0])[(threadIdx.x % 8) / 4]);
+        }
+        res = __hadd2(res, tmp_bias);
+      }
+
+      if constexpr (m_block_size_8) {
+        ((c_scalar_t*)sh_red)[idx] = res.x;
+        ((c_scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y;
+      } else {
+        ((c_scalar_t2*)sh_red)[idx] = res;
+      }
+    };
+
+    if (threadIdx.x / 32 < tb_n_warps) {
+  #pragma unroll
+      for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+        for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) {
+          if constexpr (m_block_size_8) {
+            int wr = c_sh_wr + 16 * j;
+            write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1],
+                  frag_s[j / 2][2 * (j % 2) + 0],
+                  frag_bias[j / 2][2 * (j % 2) + 0]);
+            write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3],
+                  frag_s[j / 2][2 * (j % 2) + 1],
+                  frag_bias[j / 2][2 * (j % 2) + 1]);
+          } else {
+            int wr = c_sh_wr + 8 * j;
+            write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0],
+                  frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0],
+                  frag_bias[j / 2][2 * (j % 2) + 0]);
+            write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2],
+                  frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0],
+                  frag_bias[j / 2][2 * (j % 2) + 0]);
+            write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0],
+                  frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1],
+                  frag_bias[j / 2][2 * (j % 2) + 1]);
+            write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2],
+                  frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1],
+                  frag_bias[j / 2][2 * (j % 2) + 1]);
+          }
+        }
+        c_sh_wr += 16 * (4 * c_sh_stride);
+      }
+    }
+    __syncthreads();
+
+  #pragma unroll
+    for (int i = 0;
+         i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks));
+         i++) {
+      int row = c_gl_wr / c_gl_stride;
+      if (row < block_num_valid_tokens) {
+        int64_t sorted_row = sh_block_sorted_ids[row];
+        int64_t true_idx = sorted_row * c_gl_stride + c_gl_wr % c_gl_stride;
+        c_scalar_t2 topk_weight_score;
+        if (mul_topk_weights) topk_weight_score = sh_block_topk_weights[row];
+        if (use_atomic_add && slice_count > 1 || mul_topk_weights) {
+          c_scalar_t2* C_half2 = reinterpret_cast<c_scalar_t2*>(&C[true_idx]);
+          c_scalar_t2* sh_red_half2 =
+              reinterpret_cast<c_scalar_t2*>(&sh_red[c_sh_rd]);
+          if (mul_topk_weights) {
+  #pragma unroll
+            for (int a = 0; a < 4; a++) {
+              sh_red_half2[a] = __hmul2(sh_red_half2[a], topk_weight_score);
+            }
+          }
+
+          if (use_atomic_add && slice_count > 1) {
+  #pragma unroll
+            for (int a = 0; a < 4; a++) {
+              atomicAdd(&C_half2[a], sh_red_half2[a]);
+            }
+          } else {
+            C[true_idx] = *reinterpret_cast<int4*>(sh_red_half2);
+          }
+        } else {
+          C[true_idx] = sh_red[c_sh_rd];
+        }
+        c_gl_wr += c_gl_wr_delta;
+        c_sh_rd += c_sh_rd_delta;
+      }
+    }
+    __syncthreads();
+  };
+
+  // Start global fetch and register load pipelines.
+  auto start_pipes = [&]() {
+
+  #pragma unroll
+    for (int i = 0; i < stages - 1; i++) {
+      if (has_act_order && i == 0) {
+        int last_g_idx = slice_k_start + stages * tb_k * 2;
+        if (last_g_idx >= prob_k) {
+          last_g_idx = prob_k - 1;
+        }
+        fetch_act_order_scales_to_shared(true, g_idx[slice_k_start],
+                                         g_idx[last_g_idx]);
+      }
+
+      if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
+        if (i == 0) {
+          fetch_col_zp_to_shared();
+          if constexpr (!dequant_skip_flop) {
+            fetch_col_scale_to_shared();
+          }
+        }
+      }
+      fetch_to_shared(i, i, i < slice_iters);
+    }
+
+    zero_accums();
+    wait_for_stage();
+    init_same_group(0);
+    fetch_to_registers(0, 0);
+    fetch_scales_to_registers(0, 0);
+    fetch_zp_to_registers(0, 0);
+    a_gl_rd_col += a_gl_rd_delta_o * (stages - 1);
+    if constexpr (has_act_order) {
+      slice_k_start_shared_fetch += tb_k * (stages - 1);
+    }
+  };
+  if (slice_iters) {
+    start_pipes();
+  }
+
+  // Main loop.
+  while (slice_iters) {
+    // We unroll over both the global fetch and the register load pipeline to
+    // ensure all shared memory accesses are static. Note that both pipelines
+    // have even length meaning that the next iteration will always start at
+    // index 0.
+
+  #pragma unroll
+    for (int pipe = 0; pipe < stages;) {
+  #pragma unroll
+      for (int k = 0; k < b_sh_wr_iters; k++) {
+        fetch_to_registers(k + 1, pipe % stages);
+        fetch_scales_to_registers(k + 1, pipe);
+        fetch_zp_to_registers(k + 1, pipe);
+        if (k == b_sh_wr_iters - 2) {
+          fetch_to_shared((pipe + stages - 1) % stages, pipe,
+                          slice_iters >= stages);
+          pipe++;
+          wait_for_stage();
+          init_same_group(pipe % stages);
+        }
+
+        if constexpr (!is_a_8bit) {
+          matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0));
+        } else {
+          static_assert(group_blocks != 0 && group_blocks != 1);
+          matmul_a8(k);
+        }
+      }
+      slice_iters--;
+      if (slice_iters == 0) {
+        break;
+      }
+    }
+
+    a_gl_rd_col += a_gl_rd_delta_o * stages;
+
+    if constexpr (has_act_order) {
+      slice_k_start += tb_k * stages;
+
+      if (slice_k_start < prob_k) {
+        slice_k_start_shared_fetch += tb_k * stages;
+        int first_group_id = g_idx[slice_k_start];
+        int last_g_idx = slice_k_start + stages * tb_k * 2;
+        if (last_g_idx >= prob_k) {
+          last_g_idx = prob_k - 1;
+        }
+        int last_group_id = g_idx[last_g_idx];
+        if (last_group_id >= sh_first_group_id + sh_num_groups) {
+          fetch_act_order_scales_to_shared(false, first_group_id,
+                                           last_group_id);
+          __syncthreads();
+        }
+      }
+    }
+
+    // Process results and, if necessary, proceed to the next column slice.
+    // While this pattern may not be the most readable, other ways of writing
+    // the loop seemed to noticeably worse performance after compilation.
+    if (slice_iters == 0) {
+      // convert fp16 accum to fp32 for reduction
+      if constexpr (use_fp16_accum) {
+  #pragma unroll
+        for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) {
+          float* frag_c_part_float = reinterpret_cast<float*>(frag_c) + i * 4;
+          scalar_t* frag_c_part_half =
+              reinterpret_cast<scalar_t*>(frag_c_part_float);
+
+  #pragma unroll
+          for (int i = 3; i >= 0; i--) {
+            frag_c_part_float[i] = Cdtype::num2float(frag_c_part_half[i]);
+          }
+        }
+      }
+
+      if constexpr (is_a_8bit) {
+        float frag_a_s[2 * thread_m_blocks];
+
+        for (int i = 0; i < 2 * thread_m_blocks; i++)
+          frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4];
+
+  #pragma unroll
+        for (int j = 0; j < 2; j++) {
+  #pragma unroll
+          for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+            for (int g = 0; g < 4; g++) {
+              float c_val = frag_c[i][j][0][g];
+
+              if constexpr (a_type == vllm::kS8) {
+                c_val = __int2float_rn(*reinterpret_cast<int32_t*>(&c_val));
+              }
+              float s_val = frag_a_s[i * 2 + g / 2];
+              frag_c[i][j][0][g] = c_val * s_val;
+            }
+  #pragma unroll
+            for (int g = 0; g < 4; g++) {
+              float c_val = frag_c[i][j][1][g];
+
+              if constexpr (a_type == vllm::kS8) {
+                c_val = __int2float_rn(*reinterpret_cast<int32_t*>(&c_val));
+              }
+              float s_val = frag_a_s[i * 2 + g / 2];
+              frag_c[i][j][1][g] = c_val * s_val;
+            }
+          }
+        }
+      }
+
+      cp_async_wait<0>();
+      bool last = slice_idx == slice_count - 1;
+      // For per-column scales, we only fetch them here in the final step before
+      // write-out
+      if constexpr (!has_act_order && group_blocks == -1 &&
+                    (has_zp && dequant_skip_flop || !has_zp)) {
+        if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) {
+          if (s_sh_wr_pred) {
+            cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
+          }
+          cp_async_fence();
+        }
+      }
+
+      thread_block_reduce();
+
+      if (has_bias && last) {
+        __syncthreads();
+        cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd],
+                       threadIdx.x < 16 * thread_n_blocks / 8);
+        cp_async_fence();
+      }
+
+      if constexpr (!has_act_order && group_blocks == -1 &&
+                    (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) {
+        if constexpr (is_a_8bit) {
+          cp_async_wait<0>();
+          __syncthreads();
+          if (threadIdx.x / 32 < tb_n_warps) {
+            reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
+          }
+        } else if (b_type.size_bits() == 8 || (last || use_atomic_add)) {
+          cp_async_wait<0>();
+          __syncthreads();
+          if (threadIdx.x / 32 < tb_n_warps) {
+            reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
+            reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
+            if constexpr (m_block_size_8) {
+              int idx = (threadIdx.x / 4) % 2;
+              c_scalar_t2* frag_s_half2 =
+                  reinterpret_cast<c_scalar_t2*>(frag_s);
+  #pragma unroll
+              for (int i = 0; i < 8; i++) {
+                frag_s_half2[i] = Cdtype::num2num2(
+                    reinterpret_cast<c_scalar_t*>(&frag_s_half2[i])[idx]);
+              }
+            }
+          }
+        }
+      }
+
+      // For 8-bit channelwise, we apply the scale before the global reduction
+      // that converts the fp32 results to fp16 (so that we avoid possible
+      // overflow in fp16)
+      if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) {
+  #pragma unroll
+        for (int j = 0; j < 2; j++) {
+          float2 aa[2];
+          aa[0] = Cdtype::num22float2(frag_s[0][j * 2][0]);
+          aa[1] = Cdtype::num22float2(frag_s[0][j * 2 + 1][0]);
+
+  #pragma unroll
+          for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+            for (int g = 0; g < 4; g++) {
+              float scale = reinterpret_cast<float*>(&aa[0])[g % 2];
+              frag_c[i][j][0][g] *= scale;
+            }
+
+  #pragma unroll
+            for (int g = 0; g < 4; g++) {
+              float scale = reinterpret_cast<float*>(&aa[1])[g % 2];
+              frag_c[i][j][1][g] *= scale;
+            }
+          }
+        }
+      } else if (!has_act_order && group_blocks == -1 &&
+                 b_type.size_bits() == 8 &&
+                 (has_zp && dequant_skip_flop || !has_zp)) {
+        if (threadIdx.x / 32 < tb_n_warps) {
+  #pragma unroll
+          for (int i = 0; i < thread_m_blocks; i++) {
+  #pragma unroll
+            for (int j = 0; j < 4; j++) {
+              scale_float<c_type_id>(
+                  reinterpret_cast<float*>(&frag_c[i][j][0][0]),
+                  frag_s[j / 2][2 * (j % 2) + 0]);
+              scale_float<c_type_id>(
+                  reinterpret_cast<float*>(&frag_c[i][j][0][2]),
+                  frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]);
+
+              if constexpr (!m_block_size_8) {
+                scale_float<c_type_id>(
+                    reinterpret_cast<float*>(&frag_c[i][j][1][0]),
+                    frag_s[j / 2][2 * (j % 2) + 1]);
+                scale_float<c_type_id>(
+                    reinterpret_cast<float*>(&frag_c[i][j][1][2]),
+                    frag_s[j / 2][2 * (j % 2) + 1]);
+              }
+            }
+          }
+        }
+      }
+
+      if (slice_count > 1 && !use_atomic_add) {
+        // only globally reduce if there is more than one block in a slice
+        barrier_acquire(&locks[locks_off], slice_idx);
+        if (use_fp32_reduce) {
+          global_reduce_fp32(slice_idx == 0, last);
+        } else {
+          global_reduce_fp16(slice_idx == 0, last);
+        }
+        barrier_release(&locks[locks_off], last);
+      }
+
+      if (has_bias && last) {
+        cp_async_wait<0>();
+        __syncthreads();
+        reinterpret_cast<int4*>(&frag_bias)[0] = sh_bias[bias_sh_rd];
+        if constexpr (!is_a_8bit)
+          reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4];
+        __syncthreads();
+      }
+
+      if (use_atomic_add && slice_count > 1 && slice_idx != 0)
+        wait_negative_and_add(&locks[locks_off]);
+      if (last || use_atomic_add)
+        // only the last block in a slice actually writes the result
+        write_result(last);
+      slice_row = 0;
+      if (!in_part2) {
+        slice_col_par += gridDim.x;
+      } else {
+        slice_col_par++;
+        slice_col++;
+      }
+      is_first_matmul_in_slice = true;
+      init_slice();
+
+      if (slice_iters) {
+        a_gl_rd_col =
+            a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o;
+        b_gl_rd = B_expert_off + b_gl_stride * (threadIdx.x / b_sh_stride) +
+                  (threadIdx.x % b_sh_stride);
+        b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row;
+
+        bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
+        // Update slice k/n for scales loading
+        if constexpr (has_act_order) {
+          slice_k_start = tb_k * slice_row;
+          slice_k_finish = slice_k_start + tb_k * slice_iters;
+          slice_k_start_shared_fetch = slice_k_start;
+          slice_n_offset = act_s_col_tb_stride * slice_col;
+        } else {
+          if constexpr (group_blocks == -1) {
+            s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
+            zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
+          } else if constexpr (group_blocks >= thread_k_blocks) {
+            s_gl_rd =
+                s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
+                s_sh_stride * slice_col + threadIdx.x;
+            zp_gl_rd =
+                zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
+                zp_sh_stride * slice_col + threadIdx.x;
+          } else {
+            s_gl_rd =
+                s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
+                               threadIdx.x / s_sh_stride) +
+                s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
+            zp_gl_rd =
+                zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
+                                threadIdx.x / zp_sh_stride) +
+                zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
+          }
+        }
+        start_pipes();
+      }
+    }
+  }
+}
+
+}  // namespace MARLIN_NAMESPACE_NAME
+
+#endif
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_u4b8_bfloat16.cu b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_u4b8_bfloat16.cu
new file mode 100644
index 000000000..0948a68f8
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_bfloat16_u4b8_bfloat16.cu
@@ -0,0 +1,160 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kBFloat16.id(), vllm::kU4B8.id(), vllm::kBFloat16.id(), vllm::kBFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_u4b8_float16.cu b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_u4b8_float16.cu
new file mode 100644
index 000000000..fdace4b09
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/moe/marlin_moe_wna16/sm80_kernel_float16_u4b8_float16.cu
@@ -0,0 +1,160 @@
+// auto generated by generate_kernels.py
+// clang-format off
+
+#include "kernel.h"
+#include "marlin_template.h"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, -1, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 0, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 2, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 4, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, true, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 1, 8, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 1, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 2, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 2, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 3, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 3, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 256, 4, 16, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 8, 4, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+template __global__ void Marlin<vllm::kFloat16.id(), vllm::kU4B8.id(), vllm::kFloat16.id(), vllm::kFloat16.id(), 128, 4, 4, 8, false, 4, 8, false>( MARLIN_KERNEL_PARAMS );
+
+}
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/dequant.h b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/dequant.h
new file mode 100644
index 000000000..edd97dbfc
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/dequant.h
@@ -0,0 +1,609 @@
+/*
+Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16)
+
+The process of fast dequantization can be summarized as a combination
+of bitwise operations and floating-point computations:
+
+weight =>(bit_op / bitwise operations)=>
+f16_value =>(flop / floating-point computation)=>
+dequantized_weight
+
+Since the dequantized weights typically require subtracting the zero point and
+applying a scale factor, the floating-point computation step can be fused with
+the zero-point subtraction and scaling operations.
+
+The following are the parts that need to be modified for the fused operation
+of zero-point subtraction and scaling.
+
+## INT4 => FP16/BF16 or INT8 => FP16
+
+The floating-point computation is `__hsub2`
+
+If has zero points:
+
+    flop(bit_op(weight)) - flop(bit_op(zp))
+  = sub(bit_op(weight), bias) - sub(bit_op(zp), bias)
+  = bit_op(weight) - bit_op(zp)
+
+so we don't need additional modification.
+
+If has float zero points:
+
+    flop(bit_op(weight)) - fzp
+  = sub(bit_op(weight), bias) - fzp
+  = bit_op(weight) - (fzp + bias)
+
+where the `fzp + bias` can be computed at weight loading. But this
+may have accuracy issue, so we should not use this in most cases.
+
+If has not zero points:
+
+    scale(flop(bit_op(weight)))
+  = scale(sub(bit_op(weight), bias))
+  = scale(bit_op(weight)) - scale(bias)
+  = fma(bit_op(weight), scale_factor, scale(bias))
+
+where the `scale(bias)` can be cached. But this may have accuracy issue,
+so we should not use this in most cases.
+
+
+## INT8 => BF16
+
+INT8 => BF16 is a special case, it use byte_perm instead of flop.
+We cannot fused byte_perm with scaling.
+
+
+## FP4/FP8 => FP16/BF16
+
+    scale(flop(bit_op(weight)))
+  = scale(mul(bit_op(weight), multiplier))
+  = mul(bit_op(weight), scale_factor * multiplier)
+
+where `scale_factor * multiplier` can be computed at weight loading.
+
+*/
+
+#include "marlin_dtypes.cuh"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 750
+// Lookup-table based 3-input logical operation; explicitly used for
+// dequantization as the compiler does not seem to automatically recognize it in
+// all cases.
+template <int lut>
+__device__ inline int lop3(int a, int b, int c) {
+  int res;
+  asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
+               : "=r"(res)
+               : "r"(a), "r"(b), "r"(c), "n"(lut));
+  return res;
+}
+
+// Constructs destination register by taking bytes from 2 sources (based on
+// mask)
+template <int start_byte, int mask>
+__device__ inline uint32_t prmt(uint32_t a) {
+  uint32_t res;
+  asm volatile("prmt.b32 %0, %1, %2, %3;\n"
+               : "=r"(res)
+               : "r"(a), "n"(start_byte), "n"(mask));
+  return res;
+}
+
+template <typename scalar_t2, vllm::ScalarTypeId w_type_id,
+          bool skip_flop = false>
+__device__ inline void dequant(int q, scalar_t2* frag_b);
+
+//
+// Efficiently dequantize 4bit values packed in an int32 value into a full
+// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below,
+// with some small changes:
+// - FP16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287
+// - BF16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385
+//
+template <>
+__device__ inline void dequant<half2, vllm::kU4B8.id(), true>(int q,
+                                                              half2* frag_b) {
+  const int MASK = 0x000f000f;
+  const int EX = 0x64006400;
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+  q >>= 4;
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+
+  frag_b[0] = *reinterpret_cast<half2*>(&lo);
+  frag_b[1] = *reinterpret_cast<half2*>(&hi);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU4B8.id(), false>(int q,
+                                                               half2* frag_b) {
+  const int LO = 0x000f000f;
+  const int HI = 0x00f000f0;
+  const int EX = 0x64006400;
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  // clang-format off
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
+  // clang-format on
+  // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
+  // directly into `SUB` and `ADD`.
+  const int SUB = 0x64086408;
+  const int MUL = 0x2c002c00;
+  const int ADD = 0xd480d480;
+  frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
+                      *reinterpret_cast<const half2*>(&SUB));
+  frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
+                      *reinterpret_cast<const half2*>(&MUL),
+                      *reinterpret_cast<const half2*>(&ADD));
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU4.id(), true>(int q,
+                                                            half2* frag_b) {
+  dequant<half2, vllm::kU4B8.id(), true>(q, frag_b);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU4.id(), false>(int q,
+                                                             half2* frag_b) {
+  const int LO = 0x000f000f;
+  const int HI = 0x00f000f0;
+  const int EX = 0x64006400;
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  // clang-format off
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
+  // clang-format on
+  // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
+  // directly into `SUB` and `ADD`.
+  const int SUB = 0x64006400;
+  const int MUL = 0x2c002c00;
+  const int ADD = 0xd400d400;
+  frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
+                      *reinterpret_cast<const half2*>(&SUB));
+  frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
+                      *reinterpret_cast<const half2*>(&MUL),
+                      *reinterpret_cast<const half2*>(&ADD));
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU4B8.id(), true>(
+    int q, nv_bfloat162* frag_b) {
+  static constexpr uint32_t MASK = 0x000f000f;
+  static constexpr uint32_t EX = 0x43004300;
+
+  // Guarantee that the `(a & b) | c` operations are LOP3s.
+  // clang-format off
+  int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+  q >>= 4;
+  int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
+  // clang-format on
+
+  frag_b[0] = *reinterpret_cast<nv_bfloat162*>(&lo);
+  frag_b[1] = *reinterpret_cast<nv_bfloat162*>(&hi);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU4B8.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, vllm::kU4B8.id(), true>(q, frag_b);
+
+  static constexpr uint32_t SUB = 0x43084308;
+
+  frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+  frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU4.id(), true>(
+    int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, vllm::kU4B8.id(), true>(q, frag_b);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU4.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, vllm::kU4.id(), true>(q, frag_b);
+
+  static constexpr uint32_t SUB = 0x43004300;
+
+  frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+  frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast<const nv_bfloat162*>(&SUB));
+}
+
+//
+// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or
+// bf16 Reference:
+// - FP16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85
+// - BF16:
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175
+//
+template <>
+__device__ inline void dequant<half2, vllm::kU8B128.id(), true>(int q,
+                                                                half2* frag_b) {
+  static constexpr uint32_t mask_for_elt_01 = 0x5250;
+  static constexpr uint32_t mask_for_elt_23 = 0x5351;
+  static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
+
+  uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
+  uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
+
+  frag_b[0] = *reinterpret_cast<half2*>(&lo);
+  frag_b[1] = *reinterpret_cast<half2*>(&hi);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU8B128.id(), false>(
+    int q, half2* frag_b) {
+  dequant<half2, vllm::kU8B128.id(), true>(q, frag_b);
+
+  static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
+  frag_b[0] = __hsub2(frag_b[0],
+                      *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+  frag_b[1] = __hsub2(frag_b[1],
+                      *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU8.id(), true>(int q,
+                                                            half2* frag_b) {
+  dequant<half2, vllm::kU8B128.id(), true>(q, frag_b);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kU8.id(), false>(int q,
+                                                             half2* frag_b) {
+  dequant<half2, vllm::kU8.id(), true>(q, frag_b);
+
+  static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400;
+  frag_b[0] = __hsub2(frag_b[0],
+                      *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+  frag_b[1] = __hsub2(frag_b[1],
+                      *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU8B128.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  float fp32_intermediates[4];
+  uint32_t* fp32_intermediates_casted =
+      reinterpret_cast<uint32_t*>(fp32_intermediates);
+
+  static constexpr uint32_t fp32_base = 0x4B000000;
+  fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
+  fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
+  fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
+  fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
+
+  fp32_intermediates[0] -= 8388736.f;
+  fp32_intermediates[1] -= 8388736.f;
+  fp32_intermediates[2] -= 8388736.f;
+  fp32_intermediates[3] -= 8388736.f;
+
+  uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(frag_b);
+  bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
+                                   fp32_intermediates_casted[1], 0x7632);
+  bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
+                                   fp32_intermediates_casted[3], 0x7632);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kU8.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  float fp32_intermediates[4];
+  uint32_t* fp32_intermediates_casted =
+      reinterpret_cast<uint32_t*>(fp32_intermediates);
+
+  static constexpr uint32_t fp32_base = 0x4B000000;
+  fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
+  fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
+  fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
+  fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
+
+  fp32_intermediates[0] -= 8388608.f;
+  fp32_intermediates[1] -= 8388608.f;
+  fp32_intermediates[2] -= 8388608.f;
+  fp32_intermediates[3] -= 8388608.f;
+
+  uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(frag_b);
+  bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
+                                   fp32_intermediates_casted[1], 0x7632);
+  bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
+                                   fp32_intermediates_casted[3], 0x7632);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kFE4M3fn.id(), true>(
+    int q, half2* frag_b) {
+  // Constants for FP8 (E4M3) and FP16 formats
+  constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5;
+  constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT;
+  constexpr int MASK = 0x7F007F00;
+
+  // Extract and shift FP8 values to FP16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 8;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kFE4M3fn.id(), false>(
+    int q, half2* frag_b) {
+  dequant<half2, vllm::kFE4M3fn.id(), true>(q, frag_b);
+
+  // Constants for FP8 (E4M3) and FP16 formats
+  constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET =
+      (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1));
+  const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
+
+  // Convert to half2 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kFE4M3fn.id(), true>(
+    int q, nv_bfloat162* frag_b) {
+  // Constants for FP8 (E4M3) and BF16 formats
+  constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
+  constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
+
+  constexpr int MASK = 0x7F007F00;
+
+  // Extract and shift FP8 values to BF16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 8;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kFE4M3fn.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, vllm::kFE4M3fn.id(), true>(q, frag_b);
+
+  // Constants for FP8 (E4M3) and BF16 formats
+  constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET =
+      (1 << (BF16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1));
+  // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent
+  // position
+  constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23;
+  const nv_bfloat162 bias_reg =
+      __float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
+
+  // Convert to bfloat162 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kFE2M1f.id(), true>(int q,
+                                                                half2* frag_b) {
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5;
+  constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP4_EXPONENT;
+  constexpr int MASK = 0x70007000;
+
+  // Extract and shift FP4 values to FP16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 4;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<half2, vllm::kFE2M1f.id(), false>(
+    int q, half2* frag_b) {
+  dequant<half2, vllm::kFE2M1f.id(), true>(q, frag_b);
+
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET =
+      (1 << (FP16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1));
+  const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
+
+  // Convert to half2 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kFE2M1f.id(), true>(
+    int q, nv_bfloat162* frag_b) {
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8;
+  constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT;
+  constexpr int MASK = 0x70007000;
+
+  // Extract and shift FP4 values to FP16 format
+  int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 4;
+  int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<nv_bfloat162, vllm::kFE2M1f.id(), false>(
+    int q, nv_bfloat162* frag_b) {
+  dequant<nv_bfloat162, vllm::kFE2M1f.id(), true>(q, frag_b);
+
+  // Constants for FP4 (E2M1) and BF16 formats
+  constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8;
+
+  // Construct and apply exponent bias
+  constexpr int BIAS_OFFSET =
+      (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1));
+  // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent
+  // position
+  constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23;
+  const nv_bfloat162 bias_reg =
+      __float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
+
+  // Convert to half2 and apply bias
+  frag_b[1] = __hmul2(frag_b[1], bias_reg);
+  frag_b[0] = __hmul2(frag_b[0], bias_reg);
+}
+
+template <>
+__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kFE2M1f.id(), true>(
+    int q, __nv_fp8x4_e4m3* frag_b) {
+  // Constants for FP4 (E2M1) and FP16 formats
+  constexpr int FP4_EXPONENT = 2, FP8_EXPONENT = 4;
+  constexpr int RIGHT_SHIFT = FP8_EXPONENT - FP4_EXPONENT;
+  constexpr int MASK = 0x70707070;
+
+  // Extract and shift FP4 values to FP16 format
+  int Out1 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 4;
+  int Out2 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note1: reverse indexing is intentional because weights are permuted
+  // Note2: when dequant to 8bit type, we write to `frag_b[2]` instead of
+  //        `frag_b[1]` to fit the layout of tensorcore
+  frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant<int32_t, vllm::kU4B8.id(), true>(
+    int q, int32_t* frag_b) {
+  constexpr int repeated_zp = 0x08080808;
+  constexpr int MASK = 0x80808080;
+
+  frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
+  q >>= 4;
+  frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
+}
+
+template <>
+__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kU4B8.id(), true>(
+    int q, __nv_fp8x4_e4m3* frag_b) {
+  int s = q & 0x08080808;
+  int Out1 = ((q & 0x07070707) | (s << 4)) + (s >> 3);
+  q >>= 4;
+  s = q & 0x08080808;
+  int Out2 = ((q & 0x07070707) | (s << 4)) + (s >> 3);
+
+  frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
+  frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
+}
+
+template <typename scalar_t2, vllm::ScalarTypeId s_type_id>
+__device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b);
+
+template <>
+__device__ inline void dequant_fp8_scales<half2, vllm::kFE4M3fn.id()>(
+    int q, half2* frag_b) {
+  int Out1 = (q & 0xFF00FF00) >> 1;
+  ;
+  q <<= 8;
+  int Out2 = (q & 0xFF00FF00) >> 1;
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
+};
+
+template <>
+__device__ inline void dequant_fp8_scales<nv_bfloat162, vllm::kFE4M3fn.id()>(
+    int q, nv_bfloat162* frag_b) {
+  constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
+  constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
+  constexpr int MASK = 0x7F007F00;
+
+  // Extract and shift FP8 values to BF16 format
+  int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT);
+  q <<= 8;
+  int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT);
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+}
+
+template <>
+__device__ inline void dequant_fp8_scales<nv_bfloat162, vllm::kFE8M0fnu.id()>(
+    int q, nv_bfloat162* frag_b) {
+  // In this conversion, 2 ** -127 in FP8E8M0 would become 0 in BF16,
+  // but we assume that such a extreme value would not occur in real models.
+  int Out1 = (q & 0xFF00FF00) >> 1;
+  q <<= 7;
+  int Out2 = q & 0x7F807F80;
+
+  // Note: reverse indexing is intentional because weights are permuted
+  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
+  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
+};
+
+// subtract zero point in quanted format and then dequant
+template <typename scalar_t2, vllm::ScalarTypeId w_type_id,
+          bool skip_flop = false>
+__device__ inline void sub_zp_and_dequant(int q, scalar_t2* frag_b, int zp);
+
+template <>
+__device__ inline void sub_zp_and_dequant<int32_t, vllm::kU4.id(), true>(
+    int q, int32_t* frag_b, int zp) {
+  // INT4 with zp -> INT8
+  // see https://github.com/vllm-project/vllm/pull/24722
+  int repeated_zp = 0x01010101 * zp;
+  int MASK = 0x80808080;
+
+  frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
+  q >>= 4;
+  frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
+}
+
+template <>
+__device__ inline void sub_zp_and_dequant<__nv_fp8x4_e4m3, vllm::kU4.id(),
+                                          true>(int q, __nv_fp8x4_e4m3* frag_b,
+                                                int zp) {
+  // INT4 with zp -> FP8
+  // see https://github.com/vllm-project/vllm/pull/24722
+  uint32_t u_q = *reinterpret_cast<uint32_t*>(&q);
+  uint32_t u_zp = *reinterpret_cast<uint32_t*>(&zp);
+  uint32_t u_zp1 = u_zp + 1;
+  uint32_t repeated_zp = 0x01010101 * u_zp;
+
+  uint32_t q0, s;
+  q0 = (u_q & 0x0F0F0F0F) | 0x70707070;
+  s = (q0 + repeated_zp) & 0x80808080;
+  uint32_t Out1 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s;
+
+  u_q >>= 4;
+  q0 = (u_q & 0x0F0F0F0F) | 0x70707070;
+  s = (q0 + repeated_zp) & 0x80808080;
+  uint32_t Out2 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s;
+
+  frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
+  frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
+}
+
+#endif
+
+}  // namespace MARLIN_NAMESPACE_NAME
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu
new file mode 100644
index 000000000..d268421a8
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu
@@ -0,0 +1,373 @@
+// GPTQ-Marlin weight-repack CUDA op, vendored from vLLM (Apache-2.0).
+//
+// Source: vllm-project/vllm
+//   csrc/libtorch_stable/quantization/marlin/gptq_marlin_repack.cu
+// The __global__ ``gptq_marlin_repack_kernel`` is copied VERBATIM; only the host
+// wrapper + op registration are rewritten from vLLM's ``torch::stable`` ABI to
+// the classic ``torch::Tensor`` + ``TORCH_LIBRARY`` ABI mstar's vendored ops use
+// (see utils/fused_moe/csrc/moe_align_block_size.cu). W4A16 only: the ``is_a_8bit``
+// (W4A8) template arm is dropped (fixed to false), so no fp8/int8 activation path.
+// Registered under ``_mstar_marlin_C`` so it never collides with a real vLLM ``_C``.
+//
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+#include "marlin.cuh"
+
+#include <torch/all.h>
+#include <torch/library.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
+namespace marlin {
+
+template <int const num_threads, int const num_bits, bool const has_perm,
+          bool is_a_8bit>
+__global__ void gptq_marlin_repack_kernel(
+    uint32_t const* __restrict__ b_q_weight_ptr,
+    uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr,
+    int size_k, int size_n) {
+  constexpr int pack_factor = 32 / num_bits;
+
+  constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1);
+  constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1);
+  int k_tiles = size_k / target_tile_k_size;
+  int n_tiles = size_n / target_tile_n_size;
+  int block_k_tiles = div_ceil(k_tiles, gridDim.x);
+
+  auto start_k_tile = blockIdx.x * block_k_tiles;
+  if (start_k_tile >= k_tiles) {
+    return;
+  }
+
+  int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles);
+
+  // Wait until the next thread tile has been loaded to shared memory.
+  auto wait_for_stage = [&]() {
+    // We only have `stages - 2` active fetches since we are double buffering
+    // and can only issue the next fetch when it is guaranteed that the previous
+    // shared memory load is fully complete (as it may otherwise be
+    // overwritten).
+    cp_async_wait<repack_stages - 2>();
+    __syncthreads();
+  };
+
+  extern __shared__ int4 sh[];
+
+  constexpr int perm_size = target_tile_k_size / 4;
+
+  int4* sh_perm_ptr = sh;
+  int4* sh_pipe_ptr = sh_perm_ptr;
+  if constexpr (has_perm) {
+    sh_pipe_ptr += perm_size;
+  }
+
+  constexpr int tile_ints = target_tile_k_size / pack_factor;
+
+  constexpr int stage_n_threads = target_tile_n_size / 4;
+  constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints;
+  constexpr int stage_size = stage_k_threads * stage_n_threads;
+
+  auto load_perm_to_shared = [&](int k_tile_id) {
+    int first_k_int4 = (k_tile_id * target_tile_k_size) / 4;
+
+    int4 const* perm_int4_ptr = reinterpret_cast<int4 const*>(perm_ptr);
+
+    if (threadIdx.x < perm_size) {
+      sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x];
+    }
+    __syncthreads();
+  };
+
+  auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) {
+    if (n_tile_id >= n_tiles) {
+      cp_async_fence();
+      return;
+    }
+
+    int first_n = n_tile_id * target_tile_n_size;
+
+    int4* sh_ptr = sh_pipe_ptr + stage_size * pipe;
+
+    if constexpr (has_perm) {
+      if (threadIdx.x < stage_size) {
+        auto k_id = threadIdx.x / stage_n_threads;
+        auto n_id = threadIdx.x % stage_n_threads;
+
+        uint32_t const* sh_perm_int_ptr =
+            reinterpret_cast<uint32_t const*>(sh_perm_ptr);
+
+        int src_k = sh_perm_int_ptr[k_id];
+        int src_k_packed = src_k / pack_factor;
+
+        cp_async4(
+            &sh_ptr[k_id * stage_n_threads + n_id],
+            reinterpret_cast<int4 const*>(&(
+                b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)])));
+      }
+
+    } else {
+      if (threadIdx.x < stage_size) {
+        auto k_id = threadIdx.x / stage_n_threads;
+        auto n_id = threadIdx.x % stage_n_threads;
+
+        int first_k = k_tile_id * target_tile_k_size;
+        int first_k_packed = first_k / pack_factor;
+
+        cp_async4(&sh_ptr[k_id * stage_n_threads + n_id],
+                  reinterpret_cast<int4 const*>(
+                      &(b_q_weight_ptr[(first_k_packed + k_id) * size_n +
+                                       first_n + (n_id * 4)])));
+      }
+    }
+
+    cp_async_fence();
+  };
+
+  auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) {
+    if (n_tile_id >= n_tiles) {
+      return;
+    }
+
+    auto warp_id = threadIdx.x / 32;
+    auto th_id = threadIdx.x % 32;
+
+    if (warp_id >= 4) {
+      return;
+    }
+
+    int tc_col = th_id / 4;
+    int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2);
+
+    constexpr int tc_offsets[4] = {0, 1, 8, 9};
+
+    int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col;
+
+    constexpr int sh_stride = target_tile_n_size;
+    constexpr uint32_t mask = (1 << num_bits) - 1;
+
+    int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe;
+    uint32_t* sh_stage_int_ptr = reinterpret_cast<uint32_t*>(sh_stage_ptr);
+
+    uint32_t* sh_perm_int_ptr = reinterpret_cast<uint32_t*>(sh_perm_ptr);
+
+    uint32_t vals[8];
+
+    if constexpr (has_perm) {
+      static_assert(!is_a_8bit);
+      for (int i = 0; i < 4; i++) {
+        int k_idx = tc_row + tc_offsets[i];
+
+        uint32_t src_k = sh_perm_int_ptr[k_idx];
+        uint32_t src_k_pos = src_k % pack_factor;
+
+        uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n];
+        uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask;
+
+        uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8];
+        uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask;
+
+        vals[i] = b1_cur_val;
+        vals[4 + i] = b2_cur_val;
+      }
+
+    } else {
+      uint32_t b1_vals[tile_ints];
+      uint32_t b2_vals[tile_ints];
+
+#pragma unroll
+      for (int i = 0; i < tile_ints; i++) {
+        if constexpr (is_a_8bit) {
+          b1_vals[i] =
+              sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8];
+        } else {
+          b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
+          b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
+        }
+      }
+
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]);
+        int cur_int = cur_elem / pack_factor;
+        int cur_pos = cur_elem % pack_factor;
+
+        vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask;
+        if constexpr (is_a_8bit)
+          vals[4 + i] =
+              (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask;
+        else
+          vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
+      }
+    }
+
+    constexpr int tile_size =
+        target_tile_k_size * target_tile_n_size / pack_factor;
+    int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size;
+
+    // Result of:
+    // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
+    if constexpr (!is_a_8bit && num_bits == 4) {
+      int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
+
+      uint32_t res = 0;
+#pragma unroll
+      for (int i = 0; i < 8; i++) {
+        res |= vals[pack_idx[i]] << (i * 4);
+      }
+
+      out_ptr[out_offset + th_id * 4 + warp_id] = res;
+
+    } else if constexpr (is_a_8bit && num_bits == 4) {
+      int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7};
+
+      uint32_t res = 0;
+#pragma unroll
+      for (int i = 0; i < 8; i++) {
+        res |= vals[pack_idx[i]] << (i * 4);
+      }
+
+      out_ptr[out_offset + th_id * 4 + warp_id] = res;
+
+    } else {
+      constexpr int pack_idx[4] = {0, 2, 1, 3};
+
+      uint32_t res1 = 0;
+      uint32_t res2 = 0;
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        const int ii = is_a_8bit ? i : pack_idx[i];
+        res1 |= vals[ii] << (i * 8);
+        res2 |= vals[4 + ii] << (i * 8);
+      }
+
+      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
+      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2;
+    }
+  };
+
+  auto start_pipes = [&](int k_tile_id, int n_tile_id) {
+#pragma unroll
+    for (int pipe = 0; pipe < repack_stages - 1; pipe++) {
+      fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe);
+    }
+
+    wait_for_stage();
+  };
+#pragma unroll
+  for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) {
+    int n_tile_id = 0;
+
+    if constexpr (has_perm) {
+      load_perm_to_shared(k_tile_id);
+    }
+
+    start_pipes(k_tile_id, n_tile_id);
+
+    while (n_tile_id < n_tiles) {
+#pragma unroll
+      for (int pipe = 0; pipe < repack_stages; pipe++) {
+        fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id,
+                        n_tile_id + pipe + repack_stages - 1);
+        repack_tile(pipe, k_tile_id, n_tile_id + pipe);
+        wait_for_stage();
+      }
+      n_tile_id += repack_stages;
+    }
+  }
+}
+
+}  // namespace marlin
+
+// W4A16 only: is_a_8bit is fixed to false, so only the has_perm ∈ {false, true}
+// arms are instantiated (act-order not exercised by Kimi, but kept for reuse).
+#define CALL_IF(NUM_BITS, HAS_PERM)                                        \
+  else if (num_bits == NUM_BITS && has_perm == HAS_PERM) {                 \
+    cudaFuncSetAttribute(                                                  \
+        marlin::gptq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS,\
+                                          HAS_PERM, false>,                \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);      \
+    marlin::gptq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS,    \
+                                      HAS_PERM, false>                     \
+        <<<blocks, marlin::repack_threads, max_shared_mem, stream>>>(      \
+            b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n);            \
+  }
+
+torch::Tensor gptq_marlin_repack(torch::Tensor b_q_weight, torch::Tensor perm,
+                                 int64_t size_k, int64_t size_n,
+                                 int64_t num_bits) {
+  // Verify compatibility with marlin tile of 16x64
+  TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
+              " is not divisible by tile_k_size = ", marlin::tile_k_size);
+  TORCH_CHECK(size_n % marlin::tile_n_size == 0, "size_n = ", size_n,
+              " is not divisible by tile_n_size = ", marlin::tile_n_size);
+
+  TORCH_CHECK(num_bits == 4 || num_bits == 8,
+              "num_bits must be 4 or 8. Got = ", num_bits);
+  int const pack_factor = 32 / num_bits;
+
+  // Verify B
+  TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0),
+              "Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
+              ", size_k = ", size_k, ", pack_factor = ", pack_factor);
+  TORCH_CHECK(b_q_weight.size(1) == size_n,
+              "b_q_weight.size(1) = ", b_q_weight.size(1),
+              " is not size_n = ", size_n);
+
+  // Verify device and strides
+  TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight is not on GPU");
+  TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
+  TORCH_CHECK(b_q_weight.scalar_type() == torch::kInt,
+              "b_q_weight type is not kInt");
+
+  TORCH_CHECK(perm.is_cuda(), "perm is not on GPU");
+  TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
+  TORCH_CHECK(perm.scalar_type() == torch::kInt, "perm type is not kInt");
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const int device_index = b_q_weight.get_device();
+
+  // Alloc buffers
+  torch::Tensor out = torch::empty(
+      {size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor},
+      b_q_weight.options());
+
+  // Detect if there is act_order
+  bool has_perm = perm.size(0) != 0;
+
+  // Get ptrs
+  uint32_t const* b_q_weight_ptr =
+      reinterpret_cast<uint32_t const*>(b_q_weight.data_ptr());
+  uint32_t const* perm_ptr =
+      reinterpret_cast<uint32_t const*>(perm.data_ptr());
+  uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr());
+
+  int blocks;
+  cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_index);
+
+  int max_shared_mem = 0;
+  cudaDeviceGetAttribute(&max_shared_mem,
+                         cudaDevAttrMaxSharedMemoryPerBlockOptin, device_index);
+  TORCH_CHECK(max_shared_mem > 0);
+
+  if (false) {
+  }
+  CALL_IF(4, false)
+  CALL_IF(4, true)
+  CALL_IF(8, false)
+  CALL_IF(8, true)
+  else {
+    TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
+                ", has_perm = ", has_perm);
+  }
+
+  return out;
+}
+
+// The op schemas for the whole ``_mstar_marlin_C`` namespace are declared by the
+// single TORCH_LIBRARY block in marlin_moe.cu (one per namespace per extension);
+// this file only provides the repack impl.
+TORCH_LIBRARY_IMPL(_mstar_marlin_C, CUDA, m) {
+  m.impl("gptq_marlin_repack", &gptq_marlin_repack);
+}
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin.cuh b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin.cuh
new file mode 100644
index 000000000..da8f39f53
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin.cuh
@@ -0,0 +1,182 @@
+// Marlin common device header, vendored VERBATIM from vLLM (Apache-2.0).
+//
+// Source: vllm-project/vllm  csrc/libtorch_stable/quantization/marlin/marlin.cuh
+// Self-contained: includes only <cuda*.h> + <iostream>. Provides the Marlin
+// tile/thread constants and the cp.async helpers used by the repack and GEMM
+// kernels. Kept byte-identical to upstream so future syncs are a clean diff.
+//
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#pragma once
+
+#ifndef _marlin_cuh
+  #define _marlin_cuh
+  #include <cuda.h>
+  #include <cuda_fp16.h>
+  #include <cuda_runtime.h>
+  #include <iostream>
+
+  #ifndef MARLIN_NAMESPACE_NAME
+    #define MARLIN_NAMESPACE_NAME marlin
+  #endif
+
+namespace MARLIN_NAMESPACE_NAME {
+
+// Marlin params
+
+// 8 warps are a good choice since every SM has 4 schedulers and having more
+// than 1 warp per schedule allows some more latency hiding. At the same time,
+// we want relatively few warps to have many registers per warp and small tiles.
+static constexpr int default_threads = 256;
+
+static constexpr int pipe_stages =
+    4;  // 4 pipeline stages fit into shared memory
+
+static constexpr int min_thread_n = 64;
+static constexpr int min_thread_k = 64;
+static constexpr int max_thread_n = 256;
+
+static constexpr int tile_size = 16;
+static constexpr int max_par = 16;
+
+// Repack params
+static constexpr int repack_stages = 8;
+
+static constexpr int repack_threads = 256;
+
+static constexpr int tile_k_size = tile_size;
+static constexpr int tile_n_size = tile_k_size * 4;
+
+// Helpers
+template <typename T, int n>
+struct Vec {
+  T elems[n];
+  __device__ T& operator[](int i) { return elems[i]; }
+};
+
+using I4 = Vec<int, 4>;
+
+constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
+
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+
+__device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  if (pred) {
+    reinterpret_cast<int32_t*>(smem_ptr)[0] =
+        reinterpret_cast<const int32_t*>(glob_ptr)[0];
+  }
+}
+
+__device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  if (pred) {
+    reinterpret_cast<int64_t*>(smem_ptr)[0] =
+        reinterpret_cast<const int64_t*>(glob_ptr)[0];
+  }
+}
+
+__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  if (pred) {
+    reinterpret_cast<int4*>(smem_ptr)[0] =
+        reinterpret_cast<const int4*>(glob_ptr)[0];
+  }
+}
+
+__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
+                                      bool pred = true) {
+  if (pred) {
+    reinterpret_cast<int4*>(smem_ptr)[0] =
+        reinterpret_cast<const int4*>(glob_ptr)[0];
+  }
+}
+
+__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
+  reinterpret_cast<int4*>(smem_ptr)[0] =
+      reinterpret_cast<const int4*>(glob_ptr)[0];
+}
+
+__device__ inline void cp_async_fence() {}
+
+template <int n>
+__device__ inline void cp_async_wait() {}
+
+  #else
+
+__device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  const int BYTES = 4;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   .reg .pred p;\n"
+      "   setp.ne.b32 p, %0, 0;\n"
+      "   @p cp.async.ca.shared.global [%1], [%2], %3;\n"
+      "}\n" ::"r"((int)pred),
+      "r"(smem), "l"(glob_ptr), "n"(BYTES));
+}
+
+__device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  const int BYTES = 8;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   .reg .pred p;\n"
+      "   setp.ne.b32 p, %0, 0;\n"
+      "   @p cp.async.ca.shared.global [%1], [%2], %3;\n"
+      "}\n" ::"r"((int)pred),
+      "r"(smem), "l"(glob_ptr), "n"(BYTES));
+}
+
+__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
+                                         bool pred = true) {
+  const int BYTES = 16;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   .reg .pred p;\n"
+      "   setp.ne.b32 p, %0, 0;\n"
+      "   @p cp.async.ca.shared.global [%1], [%2], %3;\n"
+      "}\n" ::"r"((int)pred),
+      "r"(smem), "l"(glob_ptr), "n"(BYTES));
+}
+
+__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
+                                      bool pred = true) {
+  const int BYTES = 16;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   .reg .pred p;\n"
+      "   setp.ne.b32 p, %0, 0;\n"
+      "   @p cp.async.cg.shared.global [%1], [%2], %3;\n"
+      "}\n" ::"r"((int)pred),
+      "r"(smem), "l"(glob_ptr), "n"(BYTES));
+}
+
+__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
+  const int BYTES = 16;
+  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+  asm volatile(
+      "{\n"
+      "   cp.async.cg.shared.global [%0], [%1], %2;\n"
+      "}\n" ::"r"(smem),
+      "l"(glob_ptr), "n"(BYTES));
+}
+
+__device__ inline void cp_async_fence() {
+  asm volatile("cp.async.commit_group;\n" ::);
+}
+
+template <int n>
+__device__ inline void cp_async_wait() {
+  asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
+}
+
+  #endif
+
+}  // namespace MARLIN_NAMESPACE_NAME
+
+#endif
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh
new file mode 100644
index 000000000..a4807a688
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_dtypes.cuh
@@ -0,0 +1,149 @@
+
+#ifndef _data_types_cuh
+#define _data_types_cuh
+#include "marlin.cuh"
+#include "core/scalar_type.hpp"
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+#ifndef MARLIN_NAMESPACE_NAME
+  #define MARLIN_NAMESPACE_NAME marlin
+#endif
+
+namespace MARLIN_NAMESPACE_NAME {
+
+template <long scalar_type_id>
+class MarlinScalarType {};
+
+template <>
+class MarlinScalarType<vllm::kFloat16.id()> {
+ public:
+  using scalar_t = half;
+  using scalar_t2 = half2;
+  using scalar_t4 = half2;
+  using scalar_32bit_t = half2;
+
+  // Matrix fragments for tensor core instructions; their precise layout is
+  // documented here:
+  // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type
+  using FragA = Vec<half2, 4>;
+  using FragB = Vec<half2, 2>;
+  using FragC = Vec<float, 4>;
+  using FragS = Vec<half2, 1>;
+  using FragS0 = Vec<__nv_fp8x2_e4m3, 1>;
+  using FragZP = Vec<half2, 4>;
+
+  static __device__ float inline num2float(const half x) {
+    return __half2float(x);
+  }
+
+  static __device__ half2 inline num2num2(const half x) {
+    return __half2half2(x);
+  }
+
+  static __device__ half2 inline nums2num2(const half x1, const half x2) {
+    return __halves2half2(x1, x2);
+  }
+
+  static __host__ __device__ half inline float2num(const float x) {
+    return __float2half(x);
+  }
+
+  static __host__ __device__ float2 inline num22float2(const half2 x) {
+    return __half22float2(x);
+  }
+};
+
+template <>
+class MarlinScalarType<vllm::kBFloat16.id()> {
+ public:
+  using scalar_t = nv_bfloat16;
+  using scalar_t2 = nv_bfloat162;
+  using scalar_t4 = nv_bfloat162;
+  using scalar_32bit_t = nv_bfloat162;
+
+  using FragA = Vec<nv_bfloat162, 4>;
+  using FragB = Vec<nv_bfloat162, 2>;
+  using FragC = Vec<float, 4>;
+  using FragS = Vec<nv_bfloat162, 1>;
+  using FragS0 = Vec<__nv_fp8x2_e4m3, 1>;
+  using FragZP = Vec<nv_bfloat162, 4>;
+
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
+  static __device__ float inline num2float(const nv_bfloat16 x) {
+    return __bfloat162float(x);
+  }
+
+  static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
+    return __bfloat162bfloat162(x);
+  }
+
+  static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1,
+                                                  const nv_bfloat16 x2) {
+    return __halves2bfloat162(x1, x2);
+  }
+
+  static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
+    return __float2bfloat16(x);
+  }
+
+  static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) {
+    return __bfloat1622float2(x);
+  }
+#endif
+};
+
+template <>
+class MarlinScalarType<vllm::kFE4M3fn.id()> {
+ public:
+  using scalar_t = __nv_fp8_e4m3;
+  using scalar_t2 = __nv_fp8x2_e4m3;
+  using scalar_t4 = __nv_fp8x4_e4m3;
+  using scalar_32bit_t = __nv_fp8x4_e4m3;
+
+  using FragA = Vec<__nv_fp8x4_e4m3, 4>;
+  using FragB = Vec<__nv_fp8x4_e4m3, 2>;
+  using FragC = Vec<float, 4>;
+  using FragZP = Vec<__nv_fp8x2_e4m3, 4>;
+
+  static __host__ __device__
+      float2 inline num22float2(const __nv_fp8x2_e4m3 x) {
+    return (float2)x;
+  }
+};
+
+template <>
+class MarlinScalarType<vllm::kS8.id()> {
+ public:
+  using scalar_t = int8_t;
+  using scalar_t2 = int16_t;
+  using scalar_t4 = int32_t;
+  using scalar_32bit_t = int32_t;
+
+  using FragA = Vec<int32_t, 4>;
+  using FragB = Vec<int32_t, 2>;
+  using FragC = Vec<float, 4>;
+  using FragZP = Vec<int16_t, 4>;
+};
+
+template <typename scalar_t>
+class MarlinScalarType2 {};
+
+template <>
+class MarlinScalarType2<half> : public MarlinScalarType<vllm::kFloat16.id()> {};
+
+template <>
+class MarlinScalarType2<nv_bfloat16>
+    : public MarlinScalarType<vllm::kBFloat16.id()> {};
+
+template <>
+class MarlinScalarType2<__nv_fp8_e4m3>
+    : public MarlinScalarType<vllm::kFE4M3fn.id()> {};
+
+template <>
+class MarlinScalarType2<int8_t> : public MarlinScalarType<vllm::kS8.id()> {};
+
+}  // namespace MARLIN_NAMESPACE_NAME
+
+#endif
diff --git a/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_mma.h b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_mma.h
new file mode 100644
index 000000000..6ec2aaafc
--- /dev/null
+++ b/mstar/utils/marlin/csrc/libtorch_stable/quantization/marlin/marlin_mma.h
@@ -0,0 +1,269 @@
+
+#include "marlin_dtypes.cuh"
+
+namespace MARLIN_NAMESPACE_NAME {
+
+// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
+// output/accumulation.
+template <vllm::ScalarTypeId type_id, bool use_fp16_accum, int k_size = 16>
+__device__ inline void mma(
+    const typename MarlinScalarType<type_id>::FragA& a_frag,
+    const typename MarlinScalarType<type_id>::FragB& frag_b,
+    typename MarlinScalarType<type_id>::FragC& frag_c, int idx = 0) {
+  const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
+  const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  if constexpr (!std::is_same<scalar_t, half>::value || k_size != 16) {
+    static_assert(!use_fp16_accum);
+  }
+
+  if constexpr (k_size == 16) {
+    if constexpr (std::is_same<scalar_t, half>::value && !use_fp16_accum) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[0]), "r"(a[1]), "r"(b[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
+            "f"(c[3]));
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[2]), "r"(a[3]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
+            "f"(c[3]));
+#else
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+#endif
+    } else if constexpr (std::is_same<scalar_t, half>::value &&
+                         use_fp16_accum) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(a[0]), "r"(a[1]), "r"(b[0]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(a[2]), "r"(a[3]), "r"(b[1]), "r"(c[0]), "r"(c[1]));
+#else
+      uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+            "r"(c[0]), "r"(c[1]));
+#endif
+    } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "f"(c[0]),
+            "f"(c[1]), "f"(c[2]), "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
+      int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
+          : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "r"(c[0]),
+            "r"(c[1]), "r"(c[2]), "r"(c[3]));
+    }
+  } else if (k_size == 32) {
+    if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
+      int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(a[0]), "r"(b[0]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[2]), "=r"(c[3])
+          : "r"(a[1]), "r"(b[0]), "r"(c[2]), "r"(c[3]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(a[2]), "r"(b[1]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[2]), "=r"(c[3])
+          : "r"(a[3]), "r"(b[1]), "r"(c[2]), "r"(c[3]));
+#else
+      asm volatile(
+          "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
+          : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
+            "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
+#endif
+    }
+  }
+}
+
+template <vllm::ScalarTypeId type_id, bool use_fp16_accum, int k_size = 16>
+__device__ inline void mma_trans(
+    const typename MarlinScalarType<type_id>::FragA& a_frag,
+    const typename MarlinScalarType<type_id>::FragB& frag_b,
+    const typename MarlinScalarType<type_id>::FragB& frag_b2,
+    typename MarlinScalarType<type_id>::FragC& frag_c) {
+  const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
+  const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
+  const uint32_t* b2 = reinterpret_cast<const uint32_t*>(&frag_b2);
+  float* c = reinterpret_cast<float*>(&frag_c);
+  using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
+  if constexpr (!std::is_same<scalar_t, half>::value || k_size != 16) {
+    static_assert(!use_fp16_accum);
+  }
+
+  if constexpr (k_size == 16) {
+    if constexpr (std::is_same<scalar_t, half>::value && !use_fp16_accum) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
+            "f"(c[3]));
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
+            "f"(c[3]));
+#else
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+#endif
+    } else if constexpr (std::is_same<scalar_t, half>::value &&
+                         use_fp16_accum) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "r"(c[0]), "r"(c[1]));
+#else
+      uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+          "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
+            "r"(c[0]), "r"(c[1]));
+#endif
+    } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
+            "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
+      int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
+          : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]), "r"(c[2]),
+            "r"(c[3]));
+    }
+  } else {
+    if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
+      float* c = reinterpret_cast<float*>(&frag_c);
+      asm volatile(
+          "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
+            "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
+    } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
+      int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(b[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[2]), "=r"(c[3])
+          : "r"(b2[1]), "r"(a[0]), "r"(c[2]), "r"(c[3]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[0]), "=r"(c[1])
+          : "r"(b[0]), "r"(a[1]), "r"(c[0]), "r"(c[1]));
+      asm volatile(
+          "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
+          : "=r"(c[2]), "=r"(c[3])
+          : "r"(b2[1]), "r"(a[1]), "r"(c[2]), "r"(c[3]));
+#else
+      asm volatile(
+          "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite "
+          "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
+          : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
+          : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
+            "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
+#endif
+    }
+  }
+}
+
+}  // namespace MARLIN_NAMESPACE_NAME
\ No newline at end of file
diff --git a/mstar/utils/marlin/loader.py b/mstar/utils/marlin/loader.py
new file mode 100644
index 000000000..db72a5321
--- /dev/null
+++ b/mstar/utils/marlin/loader.py
@@ -0,0 +1,57 @@
+"""JIT build + load of the vendored Marlin W4A16 CUDA ops."""
+from __future__ import annotations
+
+import functools
+import logging
+import os
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+_CSRC = os.path.join(os.path.dirname(__file__), "csrc")
+_MARLIN = os.path.join(_CSRC, "libtorch_stable", "quantization", "marlin")
+_MARLIN_MOE = os.path.join(_CSRC, "libtorch_stable", "moe", "marlin_moe_wna16")
+
+_SOURCES = [
+    os.path.join(_MARLIN, "gptq_marlin_repack.cu"),
+    os.path.join(_MARLIN_MOE, "marlin_moe.cu"),
+    os.path.join(_MARLIN_MOE, "sm80_kernel_bfloat16_u4b8_bfloat16.cu"),
+    os.path.join(_MARLIN_MOE, "sm80_kernel_float16_u4b8_float16.cu"),
+]
+
+_MIN_CAPABILITY = (8, 0)
+
+
+@functools.lru_cache(maxsize=1)
+def is_marlin_available() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    capability = torch.cuda.get_device_capability()
+    if capability < _MIN_CAPABILITY:
+        logger.warning(
+            "Marlin W4A16 needs sm%d%d+ (device is sm%d%d); using the Triton "
+            "in-kernel-dequant fallback.",
+            _MIN_CAPABILITY[0], _MIN_CAPABILITY[1], capability[0], capability[1],
+        )
+        return False
+    try:
+        from torch.utils.cpp_extension import load
+
+        load(
+            name="_mstar_marlin_C",
+            sources=list(_SOURCES),
+            is_python_module=False,
+            extra_include_paths=[_CSRC],
+            extra_cuda_cflags=["-O3", "-std=c++17", "--expt-relaxed-constexpr"],
+            verbose=False,
+        )
+        _ = torch.ops._mstar_marlin_C.moe_wna16_marlin_gemm
+        return True
+    except Exception as e:  # pragma: no cover -- depends on the build toolchain
+        logger.warning(
+            "Marlin W4A16: could not build the CUDA ops (%s); using the Triton "
+            "in-kernel-dequant fallback.",
+            e,
+        )
+        return False
diff --git a/mstar/utils/marlin/ops.py b/mstar/utils/marlin/ops.py
new file mode 100644
index 000000000..814e0a72c
--- /dev/null
+++ b/mstar/utils/marlin/ops.py
@@ -0,0 +1,136 @@
+"""Python launchers over the vendored Marlin torch ops."""
+from __future__ import annotations
+
+import torch
+
+from mstar.utils.fused_moe.align import moe_align_block_size
+from mstar.utils.fused_moe.kernels import act_and_mul_triton, moe_sum_reduce_triton
+from mstar.utils.marlin.scalar_type import UINT4B8_ID
+
+_MARLIN_TILE = 16
+
+
+def gptq_marlin_repack(
+    b_q_weight: torch.Tensor, size_k: int, size_n: int, num_bits: int = 4
+) -> torch.Tensor:
+    perm = torch.empty(0, dtype=torch.int32, device=b_q_weight.device)
+    return torch.ops._mstar_marlin_C.gptq_marlin_repack(
+        b_q_weight, perm, size_k, size_n, num_bits
+    )
+
+
+def gptq_marlin_moe_repack(
+    b_q_weight: torch.Tensor, size_k: int, size_n: int, num_bits: int = 4
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    perm = torch.empty(0, dtype=torch.int32, device=b_q_weight.device)
+    output = torch.empty(
+        (num_experts, size_k // _MARLIN_TILE, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = torch.ops._mstar_marlin_C.gptq_marlin_repack(
+            b_q_weight[e], perm, size_k, size_n, num_bits
+        )
+    return output
+
+
+def _get_scale_perms() -> tuple[list[int], list[int]]:
+    scale_perm: list[int] = []
+    for i in range(8):
+        scale_perm.extend([i + 8 * j for j in range(8)])
+    scale_perm_single: list[int] = []
+    for i in range(4):
+        scale_perm_single.extend([2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]])
+    return scale_perm, scale_perm_single
+
+
+def marlin_permute_scales(
+    s: torch.Tensor, size_k: int, size_n: int, group_size: int
+) -> torch.Tensor:
+    scale_perm, scale_perm_single = _get_scale_perms()
+    if group_size < size_k and group_size != -1:
+        s = s.reshape((-1, len(scale_perm)))[:, scale_perm]
+    else:
+        s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single]
+    return s.reshape((-1, size_n)).contiguous()
+
+
+def marlin_moe_permute_scales(
+    s: torch.Tensor, size_k: int, size_n: int, group_size: int
+) -> torch.Tensor:
+    num_experts = s.shape[0]
+    output = torch.empty_like(s)
+    for e in range(num_experts):
+        output[e] = marlin_permute_scales(s[e], size_k, size_n, group_size)
+    return output
+
+
+def marlin_make_workspace(device: torch.device, max_blocks_per_sm: int = 4) -> torch.Tensor:
+    sms = torch.cuda.get_device_properties(device).multi_processor_count
+    return torch.zeros(sms * max_blocks_per_sm, dtype=torch.int, device=device)
+
+
+def fused_marlin_moe(
+    hidden_states: torch.Tensor,
+    w1_marlin: torch.Tensor,
+    w2_marlin: torch.Tensor,
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    workspace: torch.Tensor,
+    *,
+    activation: str = "silu",
+    reduce_results: bool = True,
+) -> torch.Tensor:
+    assert hidden_states.is_contiguous() and hidden_states.dim() == 2
+    assert hidden_states.dtype in (torch.bfloat16, torch.float16)
+    M, K = hidden_states.shape
+    E = w1_marlin.shape[0]
+    top_k = topk_ids.shape[1]
+    # w2 marlin is (E, N//16, K*num_bits//8) -> intermediate size N = shape[1]*16.
+    N = w2_marlin.shape[1] * _MARLIN_TILE
+
+    # Block-size selection (vLLM heuristic): smallest block that keeps ~>=0.9
+    # expert occupancy.
+    block_size_m = 8
+    for block_size_m in (8, 16, 32, 48, 64):
+        if M * top_k / E / block_size_m < 0.9:
+            break
+
+    # Marlin requires fp32 combine weights and int32 expert ids.
+    topk_weights = topk_weights.to(torch.float32).contiguous()
+    topk_ids_i32 = topk_ids.to(torch.int32).contiguous()
+    sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+        topk_ids_i32, block_size_m, E
+    )
+
+    # Gate+up GEMM: (M, K) x experts -> (M*top_k, 2N).
+    gate_up = torch.ops._mstar_marlin_C.moe_wna16_marlin_gemm(
+        hidden_states, None, w1_marlin, w1_scale, workspace,
+        sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights,
+        block_size_m, top_k, False, UINT4B8_ID, M, 2 * N, K,
+        True, False, True,
+    )
+
+    # SwiGLU: silu(gate) * up -> (M*top_k, N).
+    down_in = torch.empty((M * top_k, N), device=hidden_states.device, dtype=hidden_states.dtype)
+    act_and_mul_triton(gate_up, down_in, activation=activation)
+
+    # Down GEMM (weighted): (M*top_k, N) x experts -> (M*top_k, K), routing
+    # weight folded in (mul_topk_weights=True).
+    down = torch.ops._mstar_marlin_C.moe_wna16_marlin_gemm(
+        down_in, None, w2_marlin, w2_scale, workspace,
+        sorted_token_ids, expert_ids, num_tokens_post_padded, topk_weights,
+        block_size_m, 1, True, UINT4B8_ID, M * top_k, K, N,
+        True, False, True,
+    )
+
+    cache3 = down.view(M, top_k, K)
+    if not reduce_results:
+        return cache3
+    output = torch.empty_like(hidden_states)
+    moe_sum_reduce_triton(cache3, output, routed_scaling_factor=1.0)
+    return output
diff --git a/mstar/utils/marlin/scalar_type.py b/mstar/utils/marlin/scalar_type.py
new file mode 100644
index 000000000..c1edc41cd
--- /dev/null
+++ b/mstar/utils/marlin/scalar_type.py
@@ -0,0 +1,7 @@
+"""vLLM ``ScalarType`` ids mirrored in Python for the Marlin ops."""
+from __future__ import annotations
+
+# vllm::kU4B8.id(): (mantissa=4 << 8) | (bias=8 << 17) | (nan_repr=1 << 50).
+UINT4B8_ID = (4 << 8) | (8 << 17) | (1 << 50)  # == 1125899907892224
+
+assert UINT4B8_ID == 1125899907892224, "uint4b8 id drifted from the vendored header"
diff --git a/mstar/utils/quantization.py b/mstar/utils/quantization.py
new file mode 100644
index 000000000..7fc0435f0
--- /dev/null
+++ b/mstar/utils/quantization.py
@@ -0,0 +1,90 @@
+"""Quantization scheme descriptors shared by the MoE kernels and model code.
+
+These are pure data — no ``nn.Module``, no kernel imports — so both
+:mod:`mstar.utils.fused_moe` (which consumes them at runtime) and
+:mod:`mstar.model.components.quantization` (which produces them) can depend on
+this module without a cycle. ``model -> utils`` is the established direction;
+defining them model-side would invert it.
+
+A quantized GEMM entry point takes one :class:`QuantizationData` instead of a
+widening list of optional per-scheme kwargs, and branches on
+:attr:`QuantizationData.quant_type` rather than sniffing arguments for ``None``.
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import ClassVar
+
+import torch
+
+
+class QuantizationType(Enum):
+    """Weight/activation precision scheme of a quantized GEMM.
+
+    Members are added alongside a kernel path that implements them — a member
+    with no dispatch branch is a promise the kernels cannot keep.
+    """
+
+    W4A16 = "w4a16"
+
+
+class QuantizationData(ABC):
+    """Per-call companion data for a quantized GEMM (scales, zero points, layout).
+
+    One subclass per :class:`QuantizationType`. Kernel entry points branch on
+    :attr:`quant_type` and read the subclass's fields, so supporting a new scheme
+    is a subclass plus one dispatch branch — not another six optional kwargs whose
+    legal combinations exist only in a maintainer's head.
+
+    Distinct from a *checkpoint* descriptor such as
+    ``CompressedTensorsQuantConfig``: that says how the weights were written to
+    disk (format, ignore list, bit width); this carries the live tensors a
+    specific kernel call needs. The checkpoint descriptor builds one of these.
+    """
+
+    @property
+    @abstractmethod
+    def quant_type(self) -> QuantizationType:
+        """Which scheme this data describes. Kernels dispatch on it."""
+
+
+@dataclass(frozen=True)
+class W4A16Data(QuantizationData):
+    """Group-wise INT4 weights with bf16/fp16 activations (compressed-tensors layout).
+
+    Weights reach the kernel as low-order-first int32 containers holding
+    :attr:`pack_factor` nibbles each; ``w1_scale`` / ``w2_scale`` are
+    ``(num_experts, N, K // group_size)``. ``w1_zp`` / ``w2_zp`` are ``None`` for
+    symmetric quantization (Kimi-K2.7's case) — that ``None`` *is* the
+    symmetric/asymmetric discriminant, read off one object rather than re-derived
+    at each call depth.
+
+    ``frozen=True`` blocks rebinding a field, not mutation of the tensors a field
+    points at; it marks the object as a description rather than a workspace.
+    """
+
+    # The Triton W4A16 kernel extracts nibbles with ``& 0xF`` and subtracts the
+    # 4-bit offset-binary bias, so the width is fixed by the class, not a knob.
+    NUM_BITS: ClassVar[int] = 4
+
+    w1_scale: torch.Tensor
+    w2_scale: torch.Tensor
+    group_size: int
+    w1_zp: torch.Tensor | None = None
+    w2_zp: torch.Tensor | None = None
+
+    @property
+    def quant_type(self) -> QuantizationType:
+        return QuantizationType.W4A16
+
+    @property
+    def pack_factor(self) -> int:
+        """INT4 values per int32 container (8). Derived, so it cannot disagree
+        with the bit width the kernel actually implements."""
+        return 32 // self.NUM_BITS
+
+    @property
+    def symmetric(self) -> bool:
+        return self.w1_zp is None and self.w2_zp is None
diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py
index 97be3ed33..204edc912 100644
--- a/mstar/worker/worker.py
+++ b/mstar/worker/worker.py
@@ -985,7 +985,14 @@ def _send_outputs(
                 request_id, outputs.persist
             )
 
-        if outputs.new_token_outputs:
+        # Only the tp-leader computes new-token counts: the conductor consumes
+        # ``new_token_counts`` exclusively from ``is_first_tp_rank`` messages
+        # (see conductor.py — ``if body.is_first_tp_rank``), and for a replicated
+        # (non-persisted) EMIT_TO_CLIENT + conductor_new_token edge the token
+        # tensor is only kept alive on rank 0. Without this gate the 7 non-leader
+        # ranks call get_tensor() on an already-dereferenced/GC'd uuid and crash
+        # with KeyError.
+        if outputs.new_token_outputs and outputs.is_first_tp_rank:
             name_to_count: dict[str, int] = {}
             for signal in outputs.new_token_outputs:
                 if signal.name in name_to_count:
diff --git a/pyproject.toml b/pyproject.toml
index 273a1c10f..1467f82fd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -205,6 +205,10 @@ exclude = ["tests*", "benchmark*", "examples*"]
 # Ship the vendored CUDA source so it can be JIT-compiled on first use.
 [tool.setuptools.package-data]
 "mstar.utils.fused_moe" = ["csrc/*.cu", "csrc/*.cuh", "csrc/*.h", "csrc/*.cpp"]
+"mstar.utils.marlin" = [
+    "csrc/**/*.cu", "csrc/**/*.cuh", "csrc/**/*.h", "csrc/**/*.hpp",
+    "csrc/**/*.cpp", "csrc/**/*.py",
+]
 
 [tool.ruff]
 exclude = [".git", ".ruff_cache", ".venv", "ref"]
diff --git a/test/integration/test_kimi_components.py b/test/integration/test_kimi_components.py
new file mode 100644
index 000000000..4a29a3d69
--- /dev/null
+++ b/test/integration/test_kimi_components.py
@@ -0,0 +1,121 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from mstar.model.kimi_k2_7.components.language_model import (
+    build_dense_mlp,
+    build_embedding,
+    build_lm_head,
+    build_rmsnorm,
+)
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="M1 golden tests need a GPU (mstar RMSNorm uses a FlashInfer kernel)",
+)
+
+DEVICE = "cuda"
+
+
+def _cfg() -> KimiK2Config:
+    return KimiK2Config.reduced()
+
+
+def _ref_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    # DeepSeek/Llama RMSNorm: normalize in fp32, scale in input dtype.
+    orig_dtype = x.dtype
+    x32 = x.float()
+    var = x32.pow(2).mean(dim=-1, keepdim=True)
+    x32 = x32 * torch.rsqrt(var + eps)
+    return weight * x32.to(orig_dtype)
+
+
+def _ref_swiglu(
+    x: torch.Tensor, gate_w: torch.Tensor, up_w: torch.Tensor, down_w: torch.Tensor
+) -> torch.Tensor:
+    # DeepseekV2MLP uses bias-free down_proj(silu(gate) * up).
+    gate = F.linear(x, gate_w)
+    up = F.linear(x, up_w)
+    return F.linear(F.silu(gate) * up, down_w)
+
+
+def _ref_embedding(ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    return F.embedding(ids, weight)
+
+
+def _ref_lm_head(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    return F.linear(x, weight)
+
+
+def test_rmsnorm_matches_reference():
+    torch.manual_seed(0)
+    cfg = _cfg()
+    dtype = torch.bfloat16  # FlashInfer rmsnorm runs in half precision
+    n_tokens = 7
+
+    norm = build_rmsnorm(cfg).to(device=DEVICE, dtype=dtype)
+    weight = torch.randn(cfg.hidden_size, device=DEVICE, dtype=dtype)
+    norm.weight.data.copy_(weight)
+
+    x = torch.randn(n_tokens, cfg.hidden_size, device=DEVICE, dtype=dtype)
+
+    got = norm(x)
+    expected = _ref_rmsnorm(x, weight, cfg.rms_norm_eps)
+    torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2)
+
+
+def test_dense_mlp_matches_reference():
+    torch.manual_seed(1)
+    cfg = _cfg()
+    dtype = torch.float32
+    n_tokens = 5
+    h, i = cfg.hidden_size, cfg.intermediate_size
+
+    mlp = build_dense_mlp(cfg).to(device=DEVICE, dtype=dtype)
+    gate_w = torch.randn(i, h, device=DEVICE, dtype=dtype) * 0.05
+    up_w = torch.randn(i, h, device=DEVICE, dtype=dtype) * 0.05
+    down_w = torch.randn(h, i, device=DEVICE, dtype=dtype) * 0.05
+    mlp.gate_up_proj.weight_loader(mlp.gate_up_proj.weight, gate_w, loaded_shard_id=0)
+    mlp.gate_up_proj.weight_loader(mlp.gate_up_proj.weight, up_w, loaded_shard_id=1)
+    mlp.down_proj.weight_loader(mlp.down_proj.weight, down_w)
+
+    x = torch.randn(n_tokens, h, device=DEVICE, dtype=dtype)
+
+    got = mlp(x)
+    expected = _ref_swiglu(x, gate_w, up_w, down_w)
+    torch.testing.assert_close(got, expected, rtol=1e-4, atol=1e-4)
+
+
+def test_embedding_matches_reference():
+    torch.manual_seed(2)
+    cfg = _cfg()
+    dtype = torch.float32
+
+    emb = build_embedding(cfg).to(device=DEVICE, dtype=dtype)
+    weight = torch.randn(cfg.vocab_size, cfg.hidden_size, device=DEVICE, dtype=dtype)
+    emb.weight_loader(emb.weight, weight)
+
+    ids = torch.randint(0, cfg.vocab_size, (9,), device=DEVICE)
+
+    got = emb(ids)
+    expected = _ref_embedding(ids, weight)
+    torch.testing.assert_close(got, expected, rtol=0, atol=0)
+
+
+def test_lm_head_matches_reference():
+    torch.manual_seed(3)
+    cfg = _cfg()
+    dtype = torch.float32
+    n_tokens = 4
+
+    head = build_lm_head(cfg).to(device=DEVICE, dtype=dtype)
+    weight = torch.randn(cfg.vocab_size, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.02
+    head.weight_loader(head.weight, weight)
+
+    x = torch.randn(n_tokens, cfg.hidden_size, device=DEVICE, dtype=dtype)
+
+    got = head(x)
+    expected = _ref_lm_head(x, weight)
+    assert got.shape == (n_tokens, cfg.vocab_size)
+    torch.testing.assert_close(got, expected, rtol=1e-4, atol=1e-4)
diff --git a/test/integration/test_kimi_decoder_layer.py b/test/integration/test_kimi_decoder_layer.py
new file mode 100644
index 000000000..f7b75f750
--- /dev/null
+++ b/test/integration/test_kimi_decoder_layer.py
@@ -0,0 +1,225 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from mstar.model.kimi_k2_7.components.decoder_layer import KimiDecoderLayer
+from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+from mstar.model.kimi_k2_7.components.rope import (
+    _yarn_find_correction_range,
+    _yarn_linear_ramp_mask,
+    rotate_gptj,
+    yarn_get_mscale,
+)
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="M4 golden tests need a GPU (RMSNorm + fused expert GEMM are CUDA-only)",
+)
+
+DEVICE = "cuda"
+
+
+def _ref_rmsnorm(x, weight, eps):
+    x32 = x.float()
+    x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps)
+    return weight * x32.to(x.dtype)
+
+
+def _ref_yarn_rope(pos, q_pe, k_pe, rotary_dim, base, factor, max_pos,
+                   beta_fast, beta_slow, mscale, mscale_all_dim):
+    pos_freqs = base ** (torch.arange(0, rotary_dim, 2, device=q_pe.device).float() / rotary_dim)
+    ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs)
+    low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos)
+    mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float).to(q_pe.device)
+    inv_freq = interp * (1 - mask) + ext * mask
+    amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim)
+    freqs = torch.outer(pos.float(), inv_freq)
+    cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin
+    kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin
+    return qr.to(q_pe.dtype), kr.to(k_pe.dtype)
+
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))  # (H,T,D)
+    T = q.shape[0]
+    causal = torch.triu(
+        torch.full((T, T), float("-inf"), device=q.device), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+def _ref_attn_forward(attn, cfg, h_normed, pos):
+    T, H = h_normed.shape[0], attn.num_heads
+    Dnope, Drope, Dv, L = (
+        cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank)
+    eps = cfg.rms_norm_eps
+    q = _ref_rmsnorm(F.linear(h_normed, attn.q_a_proj.weight), attn.q_a_layernorm.weight, eps)
+    q = F.linear(q, attn.q_b_proj.weight).view(T, H, cfg.qk_head_dim)
+    q_nope, q_pe = q.split([Dnope, Drope], dim=-1)
+    latent = F.linear(h_normed, attn.kv_a_proj_with_mqa.weight)
+    kv_a, k_pe = latent.split([L, Drope], dim=-1)
+    kv = F.linear(_ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps),
+                  attn.kv_b_proj.weight).view(T, H, Dnope + Dv)
+    k_nope, v = kv.split([Dnope, Dv], dim=-1)
+    k_pe = k_pe.view(T, 1, Drope)
+    r = cfg.rope_scaling
+    q_pe, k_pe = _ref_yarn_rope(
+        pos, q_pe, k_pe, Drope, cfg.rope_theta, r["factor"],
+        r["original_max_position_embeddings"], r.get("beta_fast", 32),
+        r.get("beta_slow", 1), r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0))
+    q = torch.cat([q_nope, q_pe], dim=-1) * attn.softmax_scale_boost
+    k = torch.cat([k_nope, k_pe.expand(T, H, Drope)], dim=-1)
+    v = F.pad(v, [0, cfg.qk_head_dim - Dv])
+    out = _sdpa_causal(q, k, v, cfg.qk_head_dim ** -0.5)
+    out = out[..., :Dv].reshape(T, H * Dv)
+    return F.linear(out, attn.o_proj.weight)
+
+
+def _ref_grouped_topk(logits, bias, n_group, topk_group, top_k,
+                      norm_topk_prob, routed_scaling_factor):
+    scores = logits.float().sigmoid()
+    T = scores.shape[0]
+    original = scores
+    scores = scores + bias.unsqueeze(0)
+    group_scores = scores.view(T, n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1]
+    group_mask = torch.zeros_like(group_scores)
+    group_mask.scatter_(1, group_idx, 1)
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(T, n_group, scores.shape[-1] // n_group)
+        .reshape(T, -1)
+    )
+    masked = scores.masked_fill(~score_mask.bool(), float("-inf"))
+    ids = torch.topk(masked, k=top_k, dim=-1, sorted=False)[1]
+    weights = original.gather(1, ids)
+    if norm_topk_prob:
+        weights = weights / weights.sum(dim=-1, keepdim=True)
+    weights = weights * routed_scaling_factor
+    return weights, ids
+
+
+def _ref_routed_experts(h, gate_up, down, weights, ids):
+    T, H = h.shape
+    inter = down.shape[-1]
+    out = torch.zeros(T, H, dtype=h.dtype, device=h.device)
+    for t in range(T):
+        for j in range(ids.shape[1]):
+            e = int(ids[t, j])
+            gu = gate_up[e] @ h[t]
+            g, u = gu[:inter], gu[inter:]
+            out[t] += weights[t, j] * (down[e] @ (F.silu(g) * u))
+    return out
+
+
+def _ref_swiglu(x, gate_w, up_w, down_w):
+    return F.linear(F.silu(F.linear(x, gate_w)) * F.linear(x, up_w), down_w)
+
+
+def _ref_mlp_forward(mlp, cfg, h_normed):
+    if isinstance(mlp, KimiSparseMoeBlock):
+        I = cfg.moe_intermediate_size
+        si = cfg.moe_intermediate_size * cfg.n_shared_experts
+        logits = F.linear(h_normed.float(), mlp.gate.weight.float())
+        weights, ids = _ref_grouped_topk(
+            logits, mlp.gate.e_score_correction_bias, cfg.n_group, cfg.topk_group,
+            cfg.num_experts_per_tok, cfg.norm_topk_prob, cfg.routed_scaling_factor)
+        routed = _ref_routed_experts(
+            h_normed, mlp.experts.gate_up_proj, mlp.experts.down_proj,
+            weights.to(h_normed.dtype), ids)
+        sh = mlp.shared_expert
+        shared = _ref_swiglu(
+            h_normed, sh.gate_up_proj.weight[:si], sh.gate_up_proj.weight[si:],
+            sh.down_proj.weight)
+        return routed + shared
+    i = cfg.intermediate_size
+    return _ref_swiglu(
+        h_normed, mlp.gate_up_proj.weight[:i], mlp.gate_up_proj.weight[i:],
+        mlp.down_proj.weight)
+
+
+def _ref_decoder_layer(layer, cfg, h, pos):
+    eps = cfg.rms_norm_eps
+    attn_in = _ref_rmsnorm(h, layer.input_layernorm.weight, eps)
+    h1 = h + _ref_attn_forward(layer.self_attn, cfg, attn_in, pos)
+    mlp_in = _ref_rmsnorm(h1, layer.post_attention_layernorm.weight, eps)
+    return h1 + _ref_mlp_forward(layer.mlp, cfg, mlp_in)
+
+
+class _MockMLACache:
+    def __init__(self, head_dim):
+        self.scale = head_dim ** -0.5
+
+    def set_layer_idx(self, _i):
+        pass
+
+    def advance_seq_lens(self, *_a, **_k):
+        pass
+
+    def run_attention(self, q, k, v):
+        return _sdpa_causal(q, k, v, self.scale)
+
+
+def _build_layer(cfg, layer_idx, dtype):
+    layer = KimiDecoderLayer(cfg, layer_idx).to(device=DEVICE, dtype=dtype)
+    a = layer.self_attn
+    for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (a.q_a_layernorm, a.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    layer.input_layernorm.weight.data.normal_(1.0, 0.02)
+    layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02)
+    mlp = layer.mlp
+    if isinstance(mlp, KimiSparseMoeBlock):
+        # Keep the router fp32 for deterministic selection.
+        mlp.gate.weight.data = torch.randn(
+            cfg.n_routed_experts, cfg.hidden_size, device=DEVICE)
+        mlp.gate.e_score_correction_bias.data = torch.randn(
+            cfg.n_routed_experts, device=DEVICE)
+        mlp.experts.gate_up_proj.data.normal_(0, 0.05)
+        mlp.experts.down_proj.data.normal_(0, 0.05)
+        mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05)
+    else:
+        mlp.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.down_proj.weight.data.normal_(0, 0.05)
+    return layer
+
+
+def test_dense_decoder_layer_matches_reference():
+    torch.manual_seed(0)
+    cfg = KimiK2Config.reduced()
+    dtype = torch.bfloat16
+    layer = _build_layer(cfg, layer_idx=0, dtype=dtype)  # dense (< first_k_dense_replace)
+    assert not isinstance(layer.mlp, KimiSparseMoeBlock)
+
+    T = 6
+    h = torch.randn(T, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1
+    pos = torch.arange(T, device=DEVICE)
+
+    got = layer(h, _MockMLACache(cfg.qk_head_dim), pos)
+    expected = _ref_decoder_layer(layer, cfg, h, pos)
+
+    assert got.shape == (T, cfg.hidden_size)
+    torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2)
+
+
+def test_moe_decoder_layer_matches_reference():
+    torch.manual_seed(1)
+    cfg = KimiK2Config.reduced()
+    dtype = torch.bfloat16
+    layer = _build_layer(cfg, layer_idx=1, dtype=dtype)  # MoE (>= first_k_dense_replace)
+    assert isinstance(layer.mlp, KimiSparseMoeBlock)
+
+    T = 7
+    h = torch.randn(T, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1
+    pos = torch.arange(T, device=DEVICE)
+
+    got = layer(h, _MockMLACache(cfg.qk_head_dim), pos)
+    expected = _ref_decoder_layer(layer, cfg, h, pos)
+
+    assert got.shape == (T, cfg.hidden_size)
+    torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2)
diff --git a/test/integration/test_kimi_flashinfer_attention.py b/test/integration/test_kimi_flashinfer_attention.py
new file mode 100644
index 000000000..aed3e20c7
--- /dev/null
+++ b/test/integration/test_kimi_flashinfer_attention.py
@@ -0,0 +1,112 @@
+import os
+
+import pytest
+import torch
+
+from mstar.communication.tensors import LocalTransferEngine
+from mstar.engine.cache_manager import WorkspaceBufferManager, create_cache_manager
+from mstar.engine.kv_store import (
+    KVCacheConfig,
+    PagedAllocationManager,
+    TransferEngineInfo,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="real FlashInfer paged attention needs a GPU",
+)
+
+DEVICE = torch.device("cuda")
+
+
+def _make_real_cache_manager(num_heads, head_dim, dtype, page_size=128, max_num_pages=8):
+    kv_cache = torch.zeros(
+        2, max_num_pages, 2, page_size, num_heads, head_dim,
+        dtype=dtype, device=DEVICE,
+    ).contiguous()
+    kv_cfg = KVCacheConfig(
+        num_layers=2, num_kv_heads=num_heads, head_dim=head_dim,
+        max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages,
+        page_size=page_size, num_qo_heads=num_heads,
+    )
+    transfer_info = TransferEngineInfo(
+        my_entity_id="kimi_flashinfer_test",
+        my_session_id="kimi_session",
+        transfer_engine=LocalTransferEngine("localhost"),
+    )
+    alloc = PagedAllocationManager(
+        config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info,
+    )
+    alloc.add_request("r0", ["main"])
+    buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE)
+    cm = create_cache_manager(
+        request_ids=["r0"],
+        active_labels_per_request={"r0": "main"},
+        kv_cache=kv_cache,
+        alloc_manager=alloc,
+        buffer_manager=buffers,
+        kv_cache_config=kv_cfg,
+        device=DEVICE,
+    )
+    return cm, alloc
+
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))  # (H,T,D)
+    T = q.shape[0]
+    causal = torch.triu(
+        torch.full((T, T), float("-inf"), device=q.device), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+@pytest.mark.parametrize("head_dim", [128, 256])
+def test_real_paged_run_attention_matches_sdpa(head_dim):
+    torch.manual_seed(0)
+    num_heads, T = 4, 6
+    dtype = torch.bfloat16
+    cm, alloc = _make_real_cache_manager(num_heads, head_dim, dtype)
+    try:
+        q = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1
+        k = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1
+        v = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1
+
+        cm.set_active_label("main")
+        cm.plan_attention(seq_lens=[T], is_causal=True, dtype=dtype)
+        cm.set_layer_idx(0)
+        got = cm.run_attention(q=q, k=k, v=v)
+        torch.cuda.synchronize()
+
+        expected = _sdpa_causal(q, k, v, head_dim ** -0.5)
+        assert got.shape == (T, num_heads, head_dim)
+        torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2)
+    finally:
+        alloc.cleanup()
+
+
+@pytest.mark.skipif(
+    os.environ.get("KIMI_TEST_FLASHINFER_192") != "1",
+    reason="opt-in (~60s failing JIT): set KIMI_TEST_FLASHINFER_192=1 to record "
+           "the head_dim=192 SM90 static_assert rejection",
+)
+def test_flashinfer_rejects_head_dim_192():
+    torch.manual_seed(0)
+    num_heads, T, head_dim = 4, 6, 192
+    dtype = torch.bfloat16
+    cm, alloc = _make_real_cache_manager(num_heads, head_dim, dtype)
+    try:
+        q = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1
+        k = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1
+        v = torch.randn(T, num_heads, head_dim, device=DEVICE, dtype=dtype) * 0.1
+        cm.set_active_label("main")
+        # FlashInfer may JIT the failing kernel in plan_attention or run_attention.
+        with pytest.raises(Exception) as exc_info:
+            cm.plan_attention(seq_lens=[T], is_causal=True, dtype=dtype)
+            cm.set_layer_idx(0)
+            cm.run_attention(q=q, k=k, v=v)
+            torch.cuda.synchronize()
+        # Guard against catching an unrelated error.
+        msg = str(exc_info.value).lower()
+        assert "ninja" in msg or "build" in msg or "192" in msg
+    finally:
+        alloc.cleanup()
diff --git a/test/integration/test_kimi_forward.py b/test/integration/test_kimi_forward.py
new file mode 100644
index 000000000..98de088aa
--- /dev/null
+++ b/test/integration/test_kimi_forward.py
@@ -0,0 +1,230 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM
+from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+from mstar.model.kimi_k2_7.components.rope import (
+    _yarn_find_correction_range,
+    _yarn_linear_ramp_mask,
+    rotate_gptj,
+    yarn_get_mscale,
+)
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="M4 full-forward golden needs a GPU (RMSNorm + fused expert GEMM)",
+)
+
+DEVICE = "cuda"
+
+
+def _ref_rmsnorm(x, weight, eps):
+    x32 = x.float()
+    x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps)
+    return weight * x32.to(x.dtype)
+
+
+def _ref_yarn_rope(pos, q_pe, k_pe, rotary_dim, base, factor, max_pos,
+                   beta_fast, beta_slow, mscale, mscale_all_dim):
+    pos_freqs = base ** (torch.arange(0, rotary_dim, 2, device=q_pe.device).float() / rotary_dim)
+    ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs)
+    low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos)
+    mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float).to(q_pe.device)
+    inv_freq = interp * (1 - mask) + ext * mask
+    amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim)
+    freqs = torch.outer(pos.float(), inv_freq)
+    cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin
+    kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin
+    return qr.to(q_pe.dtype), kr.to(k_pe.dtype)
+
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))
+    T = q.shape[0]
+    causal = torch.triu(
+        torch.full((T, T), float("-inf"), device=q.device), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+def _ref_attn_forward(attn, cfg, h_normed, pos):
+    T, H = h_normed.shape[0], attn.num_heads
+    Dnope, Drope, Dv, L = (
+        cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank)
+    eps = cfg.rms_norm_eps
+    q = _ref_rmsnorm(F.linear(h_normed, attn.q_a_proj.weight), attn.q_a_layernorm.weight, eps)
+    q = F.linear(q, attn.q_b_proj.weight).view(T, H, cfg.qk_head_dim)
+    q_nope, q_pe = q.split([Dnope, Drope], dim=-1)
+    latent = F.linear(h_normed, attn.kv_a_proj_with_mqa.weight)
+    kv_a, k_pe = latent.split([L, Drope], dim=-1)
+    kv = F.linear(_ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps),
+                  attn.kv_b_proj.weight).view(T, H, Dnope + Dv)
+    k_nope, v = kv.split([Dnope, Dv], dim=-1)
+    k_pe = k_pe.view(T, 1, Drope)
+    r = cfg.rope_scaling
+    q_pe, k_pe = _ref_yarn_rope(
+        pos, q_pe, k_pe, Drope, cfg.rope_theta, r["factor"],
+        r["original_max_position_embeddings"], r.get("beta_fast", 32),
+        r.get("beta_slow", 1), r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0))
+    # softmax_scale_boost restores DeepSeek's scale after padded-head attention.
+    pad = cfg.padded_head_dim
+    q = F.pad(torch.cat([q_nope, q_pe], dim=-1), [0, pad - cfg.qk_head_dim]) * attn.softmax_scale_boost
+    k = F.pad(torch.cat([k_nope, k_pe.expand(T, H, Drope)], dim=-1), [0, pad - cfg.qk_head_dim])
+    v = F.pad(v, [0, pad - Dv])
+    out = _sdpa_causal(q, k, v, cfg.padded_head_dim ** -0.5)
+    out = out[..., :Dv].reshape(T, H * Dv)
+    return F.linear(out, attn.o_proj.weight)
+
+
+def _ref_grouped_topk(logits, bias, n_group, topk_group, top_k,
+                      norm_topk_prob, routed_scaling_factor):
+    scores = logits.float().sigmoid()
+    T = scores.shape[0]
+    original = scores
+    scores = scores + bias.unsqueeze(0)
+    group_scores = scores.view(T, n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1]
+    group_mask = torch.zeros_like(group_scores)
+    group_mask.scatter_(1, group_idx, 1)
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(T, n_group, scores.shape[-1] // n_group)
+        .reshape(T, -1)
+    )
+    masked = scores.masked_fill(~score_mask.bool(), float("-inf"))
+    ids = torch.topk(masked, k=top_k, dim=-1, sorted=False)[1]
+    weights = original.gather(1, ids)
+    if norm_topk_prob:
+        weights = weights / weights.sum(dim=-1, keepdim=True)
+    weights = weights * routed_scaling_factor
+    return weights, ids
+
+
+def _ref_routed_experts(h, gate_up, down, weights, ids):
+    T, H = h.shape
+    inter = down.shape[-1]
+    out = torch.zeros(T, H, dtype=h.dtype, device=h.device)
+    for t in range(T):
+        for j in range(ids.shape[1]):
+            e = int(ids[t, j])
+            gu = gate_up[e] @ h[t]
+            g, u = gu[:inter], gu[inter:]
+            out[t] += weights[t, j] * (down[e] @ (F.silu(g) * u))
+    return out
+
+
+def _ref_swiglu(x, gate_w, up_w, down_w):
+    return F.linear(F.silu(F.linear(x, gate_w)) * F.linear(x, up_w), down_w)
+
+
+def _ref_mlp_forward(mlp, cfg, h_normed):
+    if isinstance(mlp, KimiSparseMoeBlock):
+        si = cfg.moe_intermediate_size * cfg.n_shared_experts
+        logits = F.linear(h_normed.float(), mlp.gate.weight.float())
+        weights, ids = _ref_grouped_topk(
+            logits, mlp.gate.e_score_correction_bias, cfg.n_group, cfg.topk_group,
+            cfg.num_experts_per_tok, cfg.norm_topk_prob, cfg.routed_scaling_factor)
+        routed = _ref_routed_experts(
+            h_normed, mlp.experts.gate_up_proj, mlp.experts.down_proj,
+            weights.to(h_normed.dtype), ids)
+        sh = mlp.shared_expert
+        shared = _ref_swiglu(
+            h_normed, sh.gate_up_proj.weight[:si], sh.gate_up_proj.weight[si:],
+            sh.down_proj.weight)
+        return routed + shared
+    i = cfg.intermediate_size
+    return _ref_swiglu(
+        h_normed, mlp.gate_up_proj.weight[:i], mlp.gate_up_proj.weight[i:],
+        mlp.down_proj.weight)
+
+
+def _ref_decoder_layer(layer, cfg, h, pos):
+    eps = cfg.rms_norm_eps
+    attn_in = _ref_rmsnorm(h, layer.input_layernorm.weight, eps)
+    h1 = h + _ref_attn_forward(layer.self_attn, cfg, attn_in, pos)
+    mlp_in = _ref_rmsnorm(h1, layer.post_attention_layernorm.weight, eps)
+    return h1 + _ref_mlp_forward(layer.mlp, cfg, mlp_in)
+
+
+def _ref_forward(model, cfg, ids, pos):
+    h = F.embedding(ids, model.model.embed_tokens.weight)
+    for layer in model.model.layers:
+        h = _ref_decoder_layer(layer, cfg, h, pos)
+    h = _ref_rmsnorm(h, model.model.norm.weight, cfg.rms_norm_eps)
+    return F.linear(h, model.lm_head.weight)
+
+
+class _MockMLACache:
+    def __init__(self, head_dim):
+        self.scale = head_dim ** -0.5
+        self.layer_idx = 0
+        self.advance_calls = 0
+
+    def set_layer_idx(self, i):
+        self.layer_idx = i
+
+    def advance_seq_lens(self, *_a, **_k):
+        self.advance_calls += 1
+
+    def run_attention(self, q, k, v):
+        return _sdpa_causal(q, k, v, self.scale)
+
+
+def _fill_layer(layer, cfg):
+    a = layer.self_attn
+    for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (a.q_a_layernorm, a.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    layer.input_layernorm.weight.data.normal_(1.0, 0.02)
+    layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02)
+    mlp = layer.mlp
+    if isinstance(mlp, KimiSparseMoeBlock):
+        mlp.gate.weight.data = torch.randn(
+            cfg.n_routed_experts, cfg.hidden_size, device=DEVICE)
+        mlp.gate.e_score_correction_bias.data = torch.randn(
+            cfg.n_routed_experts, device=DEVICE)
+        mlp.experts.gate_up_proj.data.normal_(0, 0.05)
+        mlp.experts.down_proj.data.normal_(0, 0.05)
+        mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05)
+    else:
+        mlp.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.down_proj.weight.data.normal_(0, 0.05)
+
+
+def _build_model(cfg, dtype):
+    model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=dtype)
+    model.model.embed_tokens.weight.data.normal_(0, 0.05)
+    model.model.norm.weight.data.normal_(1.0, 0.02)
+    model.lm_head.weight.data.normal_(0, 0.02)
+    for layer in model.model.layers:
+        _fill_layer(layer, cfg)
+    return model.eval()
+
+
+def test_full_forward_logits_match_reference():
+    torch.manual_seed(0)
+    cfg = KimiK2Config.reduced()
+    dtype = torch.bfloat16
+    model = _build_model(cfg, dtype)
+    assert not isinstance(model.model.layers[0].mlp, KimiSparseMoeBlock)
+    assert isinstance(model.model.layers[1].mlp, KimiSparseMoeBlock)
+
+    T = 8
+    ids = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE)
+    pos = torch.arange(T, device=DEVICE)
+
+    cache = _MockMLACache(cfg.padded_head_dim)
+    with torch.no_grad():
+        got = model(ids, cache, pos)
+    expected = _ref_forward(model, cfg, ids, pos)
+
+    # advance_seq_lens belongs after the layer loop, not once per layer.
+    assert cache.advance_calls == 1
+    assert got.shape == (T, cfg.vocab_size)
+    torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2)
diff --git a/test/integration/test_kimi_mla.py b/test/integration/test_kimi_mla.py
new file mode 100644
index 000000000..347796584
--- /dev/null
+++ b/test/integration/test_kimi_mla.py
@@ -0,0 +1,167 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention
+from mstar.model.kimi_k2_7.components.rope import (
+    KimiYarnRotaryEmbedding,
+    _yarn_find_correction_range,
+    _yarn_linear_ramp_mask,
+    rotate_gptj,
+    yarn_get_mscale,
+)
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="M3 golden tests need a GPU (MLA RMSNorm uses a FlashInfer kernel)",
+)
+
+DEVICE = "cuda"
+
+
+def _ref_rmsnorm(x, weight, eps):
+    x32 = x.float()
+    x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps)
+    return weight * x32.to(x.dtype)
+
+
+def _ref_yarn_rope(pos, q_pe, k_pe, rotary_dim, base, factor, max_pos,
+                   beta_fast, beta_slow, mscale, mscale_all_dim):
+    pos_freqs = base ** (torch.arange(0, rotary_dim, 2, device=q_pe.device).float() / rotary_dim)
+    ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs)
+    low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos)
+    mask = (1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float).to(q_pe.device))
+    inv_freq = interp * (1 - mask) + ext * mask
+    amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim)
+    freqs = torch.outer(pos.float(), inv_freq)
+    cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin
+    kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin
+    return qr.to(q_pe.dtype), kr.to(k_pe.dtype)
+
+
+class _MockMLACache:
+
+    def __init__(self, head_dim: int):
+        self.scale = head_dim ** -0.5
+        self.captured: dict = {}
+
+    def set_layer_idx(self, _i):  # noqa: D401
+        pass
+
+    def set_active_label(self, _l):
+        pass
+
+    def advance_seq_lens(self, *_a, **_k):
+        pass
+
+    def run_attention(self, q, k, v):
+        self.captured = {"q": q.clone(), "k": k.clone(), "v": v.clone()}
+        qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))  # (H,T,D)
+        scores = torch.einsum("hqd,hkd->hqk", qt, kt) * self.scale
+        num_tokens = q.shape[0]
+        causal = torch.triu(
+            torch.full((num_tokens, num_tokens), float("-inf"), device=q.device), diagonal=1)
+        attn = (scores + causal).softmax(-1)
+        return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+def _rope_kwargs(cfg):
+    r = cfg.rope_scaling
+    return dict(
+        rotary_dim=cfg.qk_rope_head_dim, base=cfg.rope_theta, factor=r["factor"],
+        max_pos=r["original_max_position_embeddings"],
+        beta_fast=r.get("beta_fast", 32), beta_slow=r.get("beta_slow", 1),
+        mscale=r.get("mscale", 1.0), mscale_all_dim=r.get("mscale_all_dim", 0.0),
+    )
+
+
+def _ref_mla(attn: KimiMLAAttention, cfg, h, pos, scale, boost):
+    T, H = h.shape[0], attn.num_heads
+    Dnope, Drope, Dv, L = (
+        cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank)
+    eps = cfg.rms_norm_eps
+    q = F.linear(h, attn.q_a_proj.weight)
+    q = _ref_rmsnorm(q, attn.q_a_layernorm.weight, eps)
+    q = F.linear(q, attn.q_b_proj.weight).view(T, H, cfg.qk_head_dim)
+    q_nope, q_pe = q.split([Dnope, Drope], dim=-1)
+    latent = F.linear(h, attn.kv_a_proj_with_mqa.weight)
+    kv_a, k_pe = latent.split([L, Drope], dim=-1)
+    kv_a = _ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps)
+    kv = F.linear(kv_a, attn.kv_b_proj.weight).view(T, H, Dnope + Dv)
+    k_nope, v = kv.split([Dnope, Dv], dim=-1)
+    k_pe = k_pe.view(T, 1, Drope)
+    rk = _rope_kwargs(cfg)
+    q_pe, k_pe = _ref_yarn_rope(pos, q_pe, k_pe, **rk)
+    # Reference mirrors q/k/v zero-padding to padded_head_dim.
+    pad = cfg.padded_head_dim
+    q = F.pad(torch.cat([q_nope, q_pe], dim=-1), [0, pad - cfg.qk_head_dim]) * boost
+    k = F.pad(torch.cat([k_nope, k_pe.expand(T, H, Drope)], dim=-1), [0, pad - cfg.qk_head_dim])
+    v = F.pad(v, [0, pad - Dv])
+    return q, k, v
+
+
+def test_yarn_rope_matches_reference():
+    torch.manual_seed(0)
+    # mscale != mscale_all_dim so the cos/sin amplitude factor is non-trivial.
+    rd, base, factor, max_pos = 8, 50000.0, 32.0, 4096
+    ms, msad = 1.0, 0.5
+    rope = KimiYarnRotaryEmbedding(rd, base, factor, max_pos, 32, 1, ms, msad).to(DEVICE)
+    pos = torch.arange(6, device=DEVICE)
+    q = torch.randn(6, 4, rd, device=DEVICE)
+    k = torch.randn(6, 1, rd, device=DEVICE)
+
+    gq, gk = rope(pos, q, k)
+    rq, rk = _ref_yarn_rope(pos, q, k, rd, base, factor, max_pos, 32, 1, ms, msad)
+    torch.testing.assert_close(gq, rq, rtol=1e-4, atol=1e-4)
+    torch.testing.assert_close(gk, rk, rtol=1e-4, atol=1e-4)
+    assert abs(rope.mscale - 1.0) > 1e-3  # amplitude path is exercised
+
+
+def _build_attention(cfg, dtype):
+    attn = KimiMLAAttention(cfg).to(device=DEVICE, dtype=dtype)
+    for lin in (attn.q_a_proj, attn.q_b_proj, attn.kv_a_proj_with_mqa,
+                attn.kv_b_proj, attn.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (attn.q_a_layernorm, attn.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    return attn
+
+
+def test_mla_qkv_assembly_matches_reference():
+    torch.manual_seed(1)
+    cfg = KimiK2Config.reduced()
+    dtype = torch.bfloat16
+    attn = _build_attention(cfg, dtype)
+    cache = _MockMLACache(cfg.padded_head_dim)
+    h = torch.randn(5, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1
+    pos = torch.arange(5, device=DEVICE)
+
+    attn(h, cache, pos)  # populates cache.captured with the assembled q/k/v
+    ref_q, ref_k, ref_v = _ref_mla(attn, cfg, h, pos, cache.scale, attn.softmax_scale_boost)
+
+    torch.testing.assert_close(cache.captured["q"], ref_q, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(cache.captured["k"], ref_k, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(cache.captured["v"], ref_v, rtol=2e-2, atol=2e-2)
+
+
+def test_mla_attention_forward_matches_reference():
+    torch.manual_seed(2)
+    cfg = KimiK2Config.reduced()
+    dtype = torch.bfloat16
+    attn = _build_attention(cfg, dtype)
+    cache = _MockMLACache(cfg.padded_head_dim)
+    h = torch.randn(7, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1
+    pos = torch.arange(7, device=DEVICE)
+
+    got = attn(h, cache, pos)
+
+    ref_q, ref_k, ref_v = _ref_mla(attn, cfg, h, pos, cache.scale, attn.softmax_scale_boost)
+    ref_attn = _MockMLACache(cfg.padded_head_dim).run_attention(ref_q, ref_k, ref_v)
+    ref_out = ref_attn[..., : cfg.v_head_dim].reshape(7, attn.num_heads * cfg.v_head_dim)
+    expected = F.linear(ref_out, attn.o_proj.weight)
+
+    assert got.shape == (7, cfg.hidden_size)
+    torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2)
diff --git a/test/integration/test_kimi_mla_absorb_forward.py b/test/integration/test_kimi_mla_absorb_forward.py
new file mode 100644
index 000000000..659e517f6
--- /dev/null
+++ b/test/integration/test_kimi_mla_absorb_forward.py
@@ -0,0 +1,135 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention
+from mstar.model.kimi_k2_7.components.rope import (
+    _yarn_find_correction_range,
+    _yarn_linear_ramp_mask,
+    rotate_gptj,
+    yarn_get_mscale,
+)
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="the absorbed forward runs MLA RMSNorm (a FlashInfer kernel) on GPU",
+)
+
+DEVICE = "cuda"
+
+
+def _ref_rmsnorm(x, weight, eps):
+    x32 = x.float()
+    x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps)
+    return weight * x32.to(x.dtype)
+
+
+def _ref_yarn_rope(pos, q_pe, k_pe, cfg):
+    r = cfg.rope_scaling
+    rotary_dim, base, factor = cfg.qk_rope_head_dim, cfg.rope_theta, r["factor"]
+    max_pos = r["original_max_position_embeddings"]
+    beta_fast, beta_slow = r.get("beta_fast", 32), r.get("beta_slow", 1)
+    mscale, mscale_all_dim = r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0)
+    pos_freqs = base ** (torch.arange(0, rotary_dim, 2, device=q_pe.device).float() / rotary_dim)
+    ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs)
+    low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos)
+    mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float).to(q_pe.device)
+    inv_freq = interp * (1 - mask) + ext * mask
+    amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim)
+    freqs = torch.outer(pos.float(), inv_freq)
+    cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin
+    kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin
+    return qr.to(q_pe.dtype), kr.to(k_pe.dtype)
+
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))  # (H,T,D)
+    t = q.shape[0]
+    causal = torch.triu(torch.full((t, t), float("-inf"), device=q.device), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+def _deepseek_scale(cfg):
+    r = cfg.rope_scaling
+    mscale = yarn_get_mscale(r["factor"], r.get("mscale_all_dim", 0.0))
+    return cfg.qk_head_dim ** -0.5 * mscale * mscale
+
+
+def _ref_deepseek_mla(attn, cfg, h, pos):
+    t, heads = h.shape[0], attn.num_heads
+    d_nope, d_rope, d_v, latent = (
+        cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank)
+    eps = cfg.rms_norm_eps
+    q = _ref_rmsnorm(F.linear(h, attn.q_a_proj.weight), attn.q_a_layernorm.weight, eps)
+    q = F.linear(q, attn.q_b_proj.weight).view(t, heads, cfg.qk_head_dim)
+    q_nope, q_pe = q.split([d_nope, d_rope], dim=-1)
+    lat = F.linear(h, attn.kv_a_proj_with_mqa.weight)
+    kv_a, k_pe = lat.split([latent, d_rope], dim=-1)
+    kv = F.linear(_ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps),
+                  attn.kv_b_proj.weight).view(t, heads, d_nope + d_v)
+    k_nope, v = kv.split([d_nope, d_v], dim=-1)
+    k_pe = k_pe.view(t, 1, d_rope)
+    q_pe, k_pe = _ref_yarn_rope(pos, q_pe, k_pe, cfg)
+    q = torch.cat([q_nope, q_pe], dim=-1)
+    k = torch.cat([k_nope, k_pe.expand(t, heads, d_rope)], dim=-1)
+    out = _sdpa_causal(q, k, v, _deepseek_scale(cfg)).reshape(t, heads * d_v)
+    return F.linear(out, attn.o_proj.weight)
+
+
+class _MockMLALatentCache:
+
+    def __init__(self, sm_scale):
+        self.sm_scale = sm_scale
+
+    def set_layer_idx(self, _i):
+        pass
+
+    def set_active_label(self, _l):
+        pass
+
+    def advance_seq_lens(self, *_a, **_k):
+        pass
+
+    def run_attention_mla(self, q_nope, q_pe, kv_c, k_pe):
+        t, heads, latent = q_nope.shape
+        d_rope = q_pe.shape[-1]
+        query = torch.cat([q_nope, q_pe], dim=-1)                 # (T,H,L+Drope)
+        kv_c_h = kv_c.expand(t, heads, latent)                    # MQA broadcast
+        key = torch.cat([kv_c_h, k_pe.expand(t, heads, d_rope)], dim=-1)
+        return _sdpa_causal(query, key, kv_c_h, self.sm_scale)    # (T,H,L)
+
+
+def _build_attention(cfg, dtype):
+    attn = KimiMLAAttention(cfg).to(device=DEVICE, dtype=dtype)
+    for lin in (attn.q_a_proj, attn.q_b_proj, attn.kv_a_proj_with_mqa,
+                attn.kv_b_proj, attn.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (attn.q_a_layernorm, attn.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    attn.process_weights_after_loading()  # build w_kc / w_vc on-device
+    return attn
+
+
+def test_absorbed_forward_matches_deepseek():
+    torch.manual_seed(0)
+    cfg = KimiK2Config.reduced()
+    cfg.mla_absorb = True
+    dtype = torch.bfloat16
+    attn = _build_attention(cfg, dtype)
+    assert attn.w_kc is not None and attn.w_vc is not None
+
+    t = 7
+    h = torch.randn(t, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1
+    pos = torch.arange(t, device=DEVICE)
+
+    cache = _MockMLALatentCache(_deepseek_scale(cfg))
+    with torch.no_grad():
+        got = attn(h, cache, pos)
+
+    expected = _ref_deepseek_mla(attn, cfg, h, pos)
+    assert got.shape == (t, cfg.hidden_size)
+    torch.testing.assert_close(got, expected, rtol=3e-2, atol=3e-2)
diff --git a/test/integration/test_kimi_mla_absorb_kernel.py b/test/integration/test_kimi_mla_absorb_kernel.py
new file mode 100644
index 000000000..a699b37ce
--- /dev/null
+++ b/test/integration/test_kimi_mla_absorb_kernel.py
@@ -0,0 +1,279 @@
+import pytest
+import torch
+
+from mstar.communication.tensors import LocalTransferEngine
+from mstar.engine.cache_manager import (
+    MlaAbsorbCacheManager,
+    _mla_kernel_available,
+    create_cache_manager,
+)
+from mstar.engine.kv_store import (
+    KVCacheConfig,
+    PagedAllocationManager,
+    TransferEngineInfo,
+)
+
+_IS_SM90 = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 9
+
+pytestmark = pytest.mark.skipif(
+    not _IS_SM90,
+    reason="the FlashInfer MLA kernel fast path requires a Hopper (sm90) GPU",
+)
+
+DEVICE = torch.device("cuda")
+L, DROPE = 512, 64
+
+
+def _make_cm(softmax_scale, ckv_dim, request_ids, page_size=4, max_num_pages=64):
+    latent_width = L + DROPE
+    kv_cache = torch.zeros(
+        2, max_num_pages, page_size, latent_width, dtype=torch.bfloat16, device=DEVICE,
+    ).contiguous()
+    kv_cfg = KVCacheConfig(
+        num_layers=2, num_kv_heads=1, head_dim=latent_width,
+        max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages,
+        page_size=page_size, num_qo_heads=1,
+        attention_backend="mla_absorb", softmax_scale=softmax_scale,
+        mla_ckv_dim=ckv_dim,
+    )
+    transfer_info = TransferEngineInfo(
+        my_entity_id="kimi_mla_kernel_test",
+        my_session_id="kimi_session",
+        transfer_engine=LocalTransferEngine("localhost"),
+    )
+    alloc = PagedAllocationManager(
+        config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info,
+    )
+    for rid in request_ids:
+        alloc.add_request(rid, ["main"])
+    from mstar.engine.cache_manager import WorkspaceBufferManager
+    buffers = WorkspaceBufferManager(128 * 1024 * 1024, device=DEVICE)
+    cm = create_cache_manager(
+        request_ids=list(request_ids),
+        active_labels_per_request={rid: "main" for rid in request_ids},
+        kv_cache=kv_cache,
+        alloc_manager=alloc,
+        buffer_manager=buffers,
+        kv_cache_config=kv_cfg,
+        device=DEVICE,
+    )
+    assert isinstance(cm, MlaAbsorbCacheManager)
+    cm.set_active_label("main")
+    cm.set_layer_idx(0)
+    return cm, alloc
+
+def _ref_mla_step(q_nope_new, q_pe_new, kv_c_all, k_pe_all, scale):
+    sl, _H, _L = q_nope_new.shape
+    total = kv_c_all.shape[0]
+    old_len = total - sl
+    query = torch.cat([q_nope_new, q_pe_new], dim=-1)
+    key = torch.cat([kv_c_all.squeeze(1), k_pe_all.squeeze(1)], dim=-1)  # [total,L+Drope]
+    value = kv_c_all.squeeze(1)
+    qt = query.transpose(0, 1).float()
+    scores = torch.einsum("hqd,kd->hqk", qt, key.float()) * scale
+    q_pos = old_len + torch.arange(sl, device=DEVICE)
+    k_pos = torch.arange(total, device=DEVICE)
+    mask = torch.where(k_pos[None, :] <= q_pos[:, None], 0.0,
+                       torch.tensor(float("-inf"), device=DEVICE))
+    attn = (scores + mask).softmax(-1)
+    out = torch.einsum("hqk,kd->hqd", attn, value.float())
+    return out.transpose(0, 1).to(q_nope_new.dtype)
+
+
+def _rand_step(sl, H, dtype=torch.bfloat16):
+    return (
+        torch.randn(sl, H, L, device=DEVICE, dtype=dtype) * 0.1,
+        torch.randn(sl, H, DROPE, device=DEVICE, dtype=dtype) * 0.1,
+        torch.randn(sl, 1, L, device=DEVICE, dtype=dtype) * 0.1,
+        torch.randn(sl, 1, DROPE, device=DEVICE, dtype=dtype) * 0.1,
+    )
+
+
+def _run_single_req(cm, alloc, T, H, scale):
+    q_nope, q_pe, kv_c, k_pe = _rand_step(T, H)
+    cm.plan_attention(seq_lens=[T], is_causal=True, dtype=torch.bfloat16)
+    with torch.no_grad():
+        got_prefill = cm.run_attention_mla(q_nope, q_pe, kv_c, k_pe)
+    torch.cuda.synchronize()
+    ref_prefill = _ref_mla_step(q_nope, q_pe, kv_c, k_pe, scale)
+    cm.advance_seq_lens()
+
+    q_nope1, q_pe1, kv_c1, k_pe1 = _rand_step(1, H)
+    cm.plan_attention(seq_lens=[1], is_causal=True, dtype=torch.bfloat16)
+    with torch.no_grad():
+        got_decode = cm.run_attention_mla(q_nope1, q_pe1, kv_c1, k_pe1)
+    torch.cuda.synchronize()
+    ref_decode = _ref_mla_step(
+        q_nope1, q_pe1,
+        torch.cat([kv_c, kv_c1], 0), torch.cat([k_pe, k_pe1], 0), scale,
+    )
+    return (got_prefill, ref_prefill), (got_decode, ref_decode)
+
+
+def test_kernel_matches_sdpa_and_reference():
+    H, T, ps = 2, 6, 4
+    scale = (L + DROPE) ** -0.5 * 1.3
+    assert _mla_kernel_available(L, DROPE, 9) is True
+
+    torch.manual_seed(0)
+    cm_k, alloc_k = _make_cm(scale, ckv_dim=L, request_ids=["r0"], page_size=ps)
+    try:
+        (kp, refp), (kd, refd) = _run_single_req(cm_k, alloc_k, T, H, scale)
+    finally:
+        alloc_k.cleanup()
+
+    torch.manual_seed(0)
+    cm_s, alloc_s = _make_cm(scale, ckv_dim=None, request_ids=["r0"], page_size=ps)
+    try:
+        (sp, _refp), (sd, _refd) = _run_single_req(cm_s, alloc_s, T, H, scale)
+    finally:
+        alloc_s.cleanup()
+
+    assert cm_k._plan_states["main"].wrapper is not None
+    assert cm_s._plan_states["main"].wrapper is None
+
+    assert kp.shape == (T, H, L) and kd.shape == (1, H, L)
+    torch.testing.assert_close(kp, refp, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(kd, refd, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(sp, refp, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(sd, refd, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(kp, sp, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(kd, sd, rtol=2e-2, atol=2e-2)
+
+
+def test_kernel_batched_decode():
+    H, ps = 2, 4
+    scale = (L + DROPE) ** -0.5
+    lens = [5, 9]
+    rids = ["r0", "r1"]
+
+    torch.manual_seed(1)
+    cm, alloc = _make_cm(scale, ckv_dim=L, request_ids=rids, page_size=ps)
+    try:
+        pref = [_rand_step(sl, H) for sl in lens]
+        q_nope = torch.cat([p[0] for p in pref], 0)
+        q_pe = torch.cat([p[1] for p in pref], 0)
+        kv_c = torch.cat([p[2] for p in pref], 0)
+        k_pe = torch.cat([p[3] for p in pref], 0)
+        cm.plan_attention(seq_lens=lens, is_causal=True, dtype=torch.bfloat16)
+        with torch.no_grad():
+            cm.run_attention_mla(q_nope, q_pe, kv_c, k_pe)
+        torch.cuda.synchronize()
+        cm.advance_seq_lens()
+        assert cm._plan_states["main"].wrapper is not None
+
+        dec = [_rand_step(1, H) for _ in lens]
+        dq_nope = torch.cat([d[0] for d in dec], 0)
+        dq_pe = torch.cat([d[1] for d in dec], 0)
+        dkv_c = torch.cat([d[2] for d in dec], 0)
+        dk_pe = torch.cat([d[3] for d in dec], 0)
+        cm.plan_attention(seq_lens=[1, 1], is_causal=True, dtype=torch.bfloat16)
+        with torch.no_grad():
+            got = cm.run_attention_mla(dq_nope, dq_pe, dkv_c, dk_pe)
+        torch.cuda.synchronize()
+        assert got.shape == (2, H, L)
+
+        for i in range(2):
+            kv_c_all = torch.cat([pref[i][2], dec[i][2]], 0)
+            k_pe_all = torch.cat([pref[i][3], dec[i][3]], 0)
+            ref = _ref_mla_step(dec[i][0], dec[i][1], kv_c_all, k_pe_all, scale)
+            torch.testing.assert_close(got[i:i + 1], ref, rtol=2e-2, atol=2e-2)
+    finally:
+        alloc.cleanup()
+
+
+def test_mla_wrapper_cuda_graph_capture_replay():
+    from mstar.utils.flashinfer_utils import FlashInferMLAWrapper
+
+    bs, H, ps, max_pages = 2, 2, 4, 64
+    latent_w = L + DROPE
+    scale = latent_w ** -0.5 * 1.1
+    dtype = torch.bfloat16
+
+    cache = torch.zeros(max_pages, ps, latent_w, device=DEVICE, dtype=dtype)
+    ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=DEVICE)
+    wrapper = FlashInferMLAWrapper(
+        ws, num_heads=H, head_dim_ckv=L, head_dim_kpe=DROPE, page_size=ps,
+        sm_scale=scale, batch_size=bs, max_num_pages=max_pages,
+        max_total_tokens=bs, device=DEVICE, use_cuda_graph=True,
+    )
+
+    # Decode steps stay within fixed pages; only kv_len and scatter offsets advance.
+    req_pages = [[0, 1], [2, 3, 4]]
+    prefill = [6, 9]
+    torch.manual_seed(7)
+    prefix = []  # [prefill_r, latent_w] latents already resident per request
+    for r in range(bs):
+        lat = torch.randn(prefill[r], latent_w, device=DEVICE, dtype=dtype) * 0.1
+        for t in range(prefill[r]):
+            cache[req_pages[r][t // ps], t % ps] = lat[t]
+        prefix.append(lat)
+
+    kv_indptr = torch.tensor(
+        [0, len(req_pages[0]), len(req_pages[0]) + len(req_pages[1])],
+        device=DEVICE, dtype=torch.int32,
+    )
+    kv_indices = torch.tensor(req_pages[0] + req_pages[1], device=DEVICE, dtype=torch.int32)
+    qo_indptr = torch.tensor([0, 1, 2], device=DEVICE, dtype=torch.int32)
+
+    q_nope_s = torch.zeros(bs, H, L, device=DEVICE, dtype=dtype)
+    q_pe_s = torch.zeros(bs, H, DROPE, device=DEVICE, dtype=dtype)
+    latent_s = torch.zeros(bs, latent_w, device=DEVICE, dtype=dtype)
+
+    def plan_step(step):  # step 1 -> position prefill_r, step 2 -> prefill_r+1
+        kv_len_arr = torch.tensor(
+            [prefill[r] + step for r in range(bs)], device=DEVICE, dtype=torch.int32,
+        )
+        wrapper.plan(qo_indptr, kv_indptr, kv_indices, kv_len_arr, causal=True, dtype=dtype)
+
+    def fill_inputs(seed):
+        torch.manual_seed(seed)
+        q_nope_s.copy_(torch.randn(bs, H, L, device=DEVICE, dtype=dtype) * 0.1)
+        q_pe_s.copy_(torch.randn(bs, H, DROPE, device=DEVICE, dtype=dtype) * 0.1)
+        latent_s.copy_(torch.randn(bs, latent_w, device=DEVICE, dtype=dtype) * 0.1)
+
+    plan_step(1)
+    fill_inputs(100)
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(3):
+            wrapper.set_latent(cache, latent_s)
+            wrapper.run(q_nope_s, q_pe_s, cache[..., :L], cache[..., L:])
+    torch.cuda.current_stream().wait_stream(s)
+
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        wrapper.set_latent(cache, latent_s)
+        out_static = wrapper.run(q_nope_s, q_pe_s, cache[..., :L], cache[..., L:])
+
+    decode_hist = [[] for _ in range(bs)]  # decode latents scattered so far, per req
+    prev = None
+    for step in (1, 2):
+        fill_inputs(step)
+        plan_step(step)
+        g.replay()
+        torch.cuda.synchronize()
+        for r in range(bs):
+            decode_hist[r].append(latent_s[r:r + 1].clone())
+        got = out_static.clone()
+
+        ref = torch.empty(bs, H, L, device=DEVICE, dtype=dtype)
+        for r in range(bs):
+            all_lat = torch.cat([prefix[r]] + decode_hist[r], dim=0)  # [prefill+step, w]
+            ref[r:r + 1] = _ref_mla_step(
+                q_nope_s[r:r + 1], q_pe_s[r:r + 1],
+                all_lat[:, :L].unsqueeze(1), all_lat[:, L:].unsqueeze(1), scale,
+            )
+        torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2)
+        if prev is not None:
+            # Distinct inputs prove replay reads static buffers, not baked values.
+            assert not torch.allclose(got, prev, atol=1e-3)
+        prev = got
+
+
+def test_probe_declines_reduced_dims():
+    assert _mla_kernel_available(L, DROPE, 9) is True     # real dims, Hopper
+    assert _mla_kernel_available(32, 8, 9) is False        # reduced dims
+    assert _mla_kernel_available(L, DROPE, 8) is False     # pre-sm90
+    assert _mla_kernel_available(L, DROPE, 10) is False    # Blackwell (wants trtllm path)
diff --git a/test/integration/test_kimi_mla_absorb_marlin_merge.py b/test/integration/test_kimi_mla_absorb_marlin_merge.py
new file mode 100644
index 000000000..f98da5815
--- /dev/null
+++ b/test/integration/test_kimi_mla_absorb_marlin_merge.py
@@ -0,0 +1,90 @@
+import pytest
+import torch
+from torch import nn
+
+from mstar.model.components.quantization import process_weights_after_loading
+from mstar.model.kimi_k2_7._testing import fake_quantize_weight
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+pytestmark = pytest.mark.skipif(
+    not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8),
+    reason="the Marlin backend needs a CUDA GPU with sm80+",
+)
+
+DEVICE = "cuda"
+GROUP_SIZE = 32
+PACK_FACTOR = 8
+
+
+def _quantize_stack(weight):
+    E, N, K = weight.shape
+    packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE)
+    scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE)
+    for e in range(E):
+        p, s, _d = fake_quantize_weight(
+            weight[e], num_bits=4, group_size=GROUP_SIZE, symmetric=True,
+            scale_dtype=torch.bfloat16,
+        )
+        packed[e], scale[e] = p.to(DEVICE), s.to(DEVICE)
+    return packed, scale
+
+
+def _build_marlin_moe(cfg):
+    from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+
+    with torch.device("meta"):
+        block = KimiSparseMoeBlock(cfg)
+    block = block.to(torch.bfloat16)
+    block.to_empty(device=DEVICE)
+    for p in block.parameters():
+        if p.dtype.is_floating_point:
+            with torch.no_grad():
+                p.copy_(torch.randn_like(p) * 0.1)
+    E, H, I = cfg.n_routed_experts, cfg.hidden_size, cfg.moe_intermediate_size
+    w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w1_packed, w1_scale = _quantize_stack(w1)
+    w2_packed, w2_scale = _quantize_stack(w2)
+    with torch.no_grad():
+        block.experts.gate_up_proj_packed.data = w1_packed
+        block.experts.gate_up_proj_scale.data = w1_scale
+        block.experts.down_proj_packed.data = w2_packed
+        block.experts.down_proj_scale.data = w2_scale
+    return block
+
+
+def _build_absorbed_attn(cfg):
+    from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention
+
+    torch.manual_seed(0)
+    return KimiMLAAttention(cfg).to(device=DEVICE, dtype=torch.bfloat16)
+
+
+def test_generic_walker_composes_absorbed_mla_and_marlin_moe():
+    cfg = KimiK2Config.reduced_marlin()
+    cfg.mla_absorb = True  # absorbed attention + Marlin experts in one config
+
+    attn = _build_absorbed_attn(cfg)
+    moe = _build_marlin_moe(cfg)
+    root = nn.Module()
+    root.add_module("attn", attn)
+    root.add_module("moe", moe)
+
+    assert attn.mla_absorb
+    assert attn.w_kc is None and attn.w_vc is None and attn.fused_qkv_a_proj_weight is None
+    assert moe.experts.gate_up_proj_packed.numel() > 0
+    assert not getattr(moe, "_use_marlin", False)
+
+    process_weights_after_loading(root, torch.device(DEVICE))
+
+    assert attn.w_kc is not None, "walker did not build the absorbed w_kc"
+    assert attn.w_vc is not None, "walker did not build the absorbed w_vc"
+    assert attn.fused_qkv_a_proj_weight is not None, "walker did not build fused_qkv_a_proj"
+
+    assert moe._use_marlin, "walker did not select the Marlin backend"
+    assert moe.experts.gate_up_proj_packed.numel() == 0, "packed experts not freed after repack"
+
+    x = torch.randn(4, cfg.hidden_size, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    with torch.no_grad():
+        out = moe(x)
+    assert out.shape == x.shape and torch.isfinite(out).all()
diff --git a/test/integration/test_kimi_mla_absorb_paged.py b/test/integration/test_kimi_mla_absorb_paged.py
new file mode 100644
index 000000000..483479a7f
--- /dev/null
+++ b/test/integration/test_kimi_mla_absorb_paged.py
@@ -0,0 +1,163 @@
+import pytest
+import torch
+
+from mstar.communication.tensors import LocalTransferEngine
+from mstar.engine.cache_manager import (
+    MlaAbsorbCacheManager,
+    MlaSdpaPlan,
+    WorkspaceBufferManager,
+    _PlanState,
+    create_cache_manager,
+)
+from mstar.engine.kv_store import (
+    KVCacheConfig,
+    PagedAllocationManager,
+    TransferEngineInfo,
+)
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="the paged compressed-latent MLA backend runs on GPU",
+)
+
+DEVICE = torch.device("cuda")
+
+
+def _make_latent_cache_manager(
+    latent_width, dtype, softmax_scale, page_size=4, max_num_pages=64
+):
+    # 4D mla_absorb cache; small page_size forces multi-page coverage.
+    kv_cache = torch.zeros(
+        2, max_num_pages, page_size, latent_width,
+        dtype=dtype, device=DEVICE,
+    ).contiguous()
+    kv_cfg = KVCacheConfig(
+        num_layers=2, num_kv_heads=1, head_dim=latent_width,
+        max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages,
+        page_size=page_size, num_qo_heads=1,
+        attention_backend="mla_absorb", softmax_scale=softmax_scale,
+    )
+    transfer_info = TransferEngineInfo(
+        my_entity_id="kimi_mla_absorb_test",
+        my_session_id="kimi_session",
+        transfer_engine=LocalTransferEngine("localhost"),
+    )
+    alloc = PagedAllocationManager(
+        config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info,
+    )
+    alloc.add_request("r0", ["main"])
+    buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE)
+    cm = create_cache_manager(
+        request_ids=["r0"],
+        active_labels_per_request={"r0": "main"},
+        kv_cache=kv_cache,
+        alloc_manager=alloc,
+        buffer_manager=buffers,
+        kv_cache_config=kv_cfg,
+        device=DEVICE,
+    )
+    assert isinstance(cm, MlaAbsorbCacheManager)
+    return cm, alloc
+
+
+def _ref_mla_step(q_nope_new, q_pe_new, kv_c_all, k_pe_all, scale):
+    sl, _H, L = q_nope_new.shape
+    total = kv_c_all.shape[0]
+    old_len = total - sl
+
+    query = torch.cat([q_nope_new, q_pe_new], dim=-1)          # [sl,H,L+Drope]
+    kv_c_h = kv_c_all.squeeze(1)                               # [total,L]
+    k_pe_h = k_pe_all.squeeze(1)                               # [total,Drope]
+    key = torch.cat([kv_c_h, k_pe_h], dim=-1)                 # [total,L+Drope]
+    value = kv_c_h                                             # [total,L]
+
+    qt = query.transpose(0, 1).float()                        # [H,sl,L+Drope]
+    scores = torch.einsum("hqd,kd->hqk", qt, key.float()) * scale  # [H,sl,total]
+    q_pos = old_len + torch.arange(sl, device=DEVICE)
+    k_pos = torch.arange(total, device=DEVICE)
+    mask = torch.where(
+        k_pos[None, :] <= q_pos[:, None],
+        0.0,
+        torch.tensor(float("-inf"), device=DEVICE),
+    )
+    attn = (scores + mask).softmax(-1)
+    out = torch.einsum("hqk,kd->hqd", attn, value.float())    # [H,sl,L]
+    return out.transpose(0, 1).to(q_nope_new.dtype)           # [sl,H,L]
+
+
+def _rand_step(sl, H, L, Drope, dtype):
+    return (
+        torch.randn(sl, H, L, device=DEVICE, dtype=dtype) * 0.1,
+        torch.randn(sl, H, Drope, device=DEVICE, dtype=dtype) * 0.1,
+        torch.randn(sl, 1, L, device=DEVICE, dtype=dtype) * 0.1,
+        torch.randn(sl, 1, Drope, device=DEVICE, dtype=dtype) * 0.1,
+    )
+
+
+def _run_prefill_then_decode(L, Drope, H, T, page_size):
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    # Arbitrary MLA-style scale (the backend must apply exactly this value).
+    scale = (L + Drope) ** -0.5 * 1.3
+
+    cm, alloc = _make_latent_cache_manager(L + Drope, dtype, scale, page_size=page_size)
+    try:
+        cm.set_active_label("main")
+        cm.set_layer_idx(0)
+
+        assert T > page_size, "T must span >1 page to exercise page boundaries"
+        q_nope, q_pe, kv_c, k_pe = _rand_step(T, H, L, Drope, dtype)
+        cm.plan_attention(seq_lens=[T], is_causal=True, dtype=dtype)
+        with torch.no_grad():
+            got_prefill = cm.run_attention_mla(q_nope, q_pe, kv_c, k_pe)
+        torch.cuda.synchronize()
+
+        ref_prefill = _ref_mla_step(q_nope, q_pe, kv_c, k_pe, scale)
+        assert got_prefill.shape == (T, H, L)
+        torch.testing.assert_close(got_prefill, ref_prefill, rtol=2e-2, atol=2e-2)
+
+        cm.advance_seq_lens()
+
+        q_nope1, q_pe1, kv_c1, k_pe1 = _rand_step(1, H, L, Drope, dtype)
+        cm.plan_attention(seq_lens=[1], is_causal=True, dtype=dtype)
+        with torch.no_grad():
+            got_decode = cm.run_attention_mla(q_nope1, q_pe1, kv_c1, k_pe1)
+        torch.cuda.synchronize()
+
+        kv_c_all = torch.cat([kv_c, kv_c1], dim=0)     # [T+1,1,L]
+        k_pe_all = torch.cat([k_pe, k_pe1], dim=0)     # [T+1,1,Drope]
+        ref_decode = _ref_mla_step(q_nope1, q_pe1, kv_c_all, k_pe_all, scale)
+        assert got_decode.shape == (1, H, L)
+        torch.testing.assert_close(got_decode, ref_decode, rtol=2e-2, atol=2e-2)
+    finally:
+        alloc.cleanup()
+
+
+def test_paged_latent_mla_real_dims():
+    _run_prefill_then_decode(L=512, Drope=64, H=2, T=6, page_size=4)
+
+
+def test_paged_latent_mla_reduced_dims():
+    _run_prefill_then_decode(L=32, Drope=8, H=4, T=6, page_size=4)
+
+
+def test_sdpa_plan_clears_any_injected_wrapper():
+    """``run_attention_mla`` picks its path with ``ps.wrapper is not None``, so the
+    SDPA branch must clear the slot — otherwise a wrapper injected by CUDA-graph
+    capture routes latent attention into a paged kernel that cannot read the 4-D
+    cache. Mirrors how the paged path clears ``dense_gen``."""
+    L, Drope = 32, 8  # reduced dims: no MLA kernel, so plan_attention takes SDPA
+    cm, alloc = _make_latent_cache_manager(L + Drope, torch.bfloat16, 0.1)
+    try:
+        cm.set_active_label("main")
+        # Stand in for a persistent wrapper injected via cuda_graph_plan_states.
+        cm._plan_states["main"] = _PlanState(wrapper=object())
+
+        cm.plan_attention(seq_lens=[4], is_causal=True, dtype=torch.bfloat16)
+
+        ps = cm._plan_states["main"]
+        assert ps.wrapper is None, "SDPA plan left a stale wrapper in the plan state"
+        assert isinstance(ps.mla, MlaSdpaPlan)
+        assert [r.seq_len for r in ps.mla.requests] == [4]
+    finally:
+        alloc.cleanup()
diff --git a/test/integration/test_kimi_mla_absorb_serve.py b/test/integration/test_kimi_mla_absorb_serve.py
new file mode 100644
index 000000000..d444b1c83
--- /dev/null
+++ b/test/integration/test_kimi_mla_absorb_serve.py
@@ -0,0 +1,216 @@
+import pytest
+import torch
+
+from mstar.communication.tensors import LocalTransferEngine
+from mstar.engine.cache_manager import WorkspaceBufferManager, create_cache_manager
+from mstar.engine.kv_store import (
+    KVCacheConfig,
+    PagedAllocationManager,
+    TransferEngineInfo,
+)
+from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM
+from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+from mstar.model.kimi_k2_7.components.rope import yarn_get_mscale
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.kimi_k2_7.kimi_model import KimiK2Model
+from mstar.model.kimi_k2_7.submodules import KimiLLMSubmodule
+from mstar.model.submodule_base import ModelInputsFromEngine
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="absorbed serve smoke needs a GPU (real paged latent cache)",
+)
+
+DEVICE = torch.device("cuda")
+
+def _fill_layer(layer, cfg):
+    a = layer.self_attn
+    for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (a.q_a_layernorm, a.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    layer.input_layernorm.weight.data.normal_(1.0, 0.02)
+    layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02)
+    mlp = layer.mlp
+    if isinstance(mlp, KimiSparseMoeBlock):
+        mlp.gate.weight.data.normal_(0, 1)
+        mlp.gate.e_score_correction_bias.data = torch.randn(
+            cfg.n_routed_experts, device=DEVICE, dtype=torch.float32)
+        mlp.experts.gate_up_proj.data.normal_(0, 0.05)
+        mlp.experts.down_proj.data.normal_(0, 0.05)
+        mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05)
+    else:
+        mlp.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.down_proj.weight.data.normal_(0, 0.05)
+
+
+def _build_reference(cfg):
+    model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16)
+    model.model.embed_tokens.weight.data.normal_(0, 0.05)
+    model.model.norm.weight.data.normal_(1.0, 0.02)
+    model.lm_head.weight.data.normal_(0, 0.02)
+    for layer in model.model.layers:
+        _fill_layer(layer, cfg)
+    return model.eval()
+
+
+def _hf_checkpoint(model, cfg):
+    inter = cfg.intermediate_size
+    moe_inter = cfg.moe_intermediate_size
+    shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts
+    m = model.model
+    sd = {"model.embed_tokens.weight": m.embed_tokens.weight}
+    for i, layer in enumerate(m.layers):
+        p = f"model.layers.{i}."
+        a = layer.self_attn
+        sd[p + "self_attn.q_a_proj.weight"] = a.q_a_proj.weight
+        sd[p + "self_attn.q_a_layernorm.weight"] = a.q_a_layernorm.weight
+        sd[p + "self_attn.q_b_proj.weight"] = a.q_b_proj.weight
+        sd[p + "self_attn.kv_a_proj_with_mqa.weight"] = a.kv_a_proj_with_mqa.weight
+        sd[p + "self_attn.kv_a_layernorm.weight"] = a.kv_a_layernorm.weight
+        sd[p + "self_attn.kv_b_proj.weight"] = a.kv_b_proj.weight
+        sd[p + "self_attn.o_proj.weight"] = a.o_proj.weight
+        sd[p + "input_layernorm.weight"] = layer.input_layernorm.weight
+        sd[p + "post_attention_layernorm.weight"] = layer.post_attention_layernorm.weight
+        mlp = layer.mlp
+        if isinstance(mlp, KimiSparseMoeBlock):
+            sd[p + "mlp.gate.weight"] = mlp.gate.weight
+            sd[p + "mlp.gate.e_score_correction_bias"] = mlp.gate.e_score_correction_bias
+            gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj
+            for e in range(cfg.n_routed_experts):
+                sd[p + f"mlp.experts.{e}.gate_proj.weight"] = gup[e, :moe_inter, :]
+                sd[p + f"mlp.experts.{e}.up_proj.weight"] = gup[e, moe_inter:, :]
+                sd[p + f"mlp.experts.{e}.down_proj.weight"] = dwn[e]
+            sh = mlp.shared_expert
+            sd[p + "mlp.shared_experts.gate_proj.weight"] = sh.gate_up_proj.weight[:shared_inter]
+            sd[p + "mlp.shared_experts.up_proj.weight"] = sh.gate_up_proj.weight[shared_inter:]
+            sd[p + "mlp.shared_experts.down_proj.weight"] = sh.down_proj.weight
+        else:
+            sd[p + "mlp.gate_proj.weight"] = mlp.gate_up_proj.weight[:inter]
+            sd[p + "mlp.up_proj.weight"] = mlp.gate_up_proj.weight[inter:]
+            sd[p + "mlp.down_proj.weight"] = mlp.down_proj.weight
+    sd["model.norm.weight"] = m.norm.weight
+    sd["lm_head.weight"] = model.lm_head.weight
+    return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()}
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))
+    t = q.shape[0]
+    causal = torch.triu(torch.full((t, t), float("-inf"), device=q.device), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+class _MockMLACache:
+    def __init__(self, head_dim):
+        self.scale = head_dim ** -0.5
+
+    def set_layer_idx(self, _i):
+        pass
+
+    def advance_seq_lens(self, *_a, **_k):
+        pass
+
+    def run_attention(self, q, k, v):
+        return _sdpa_causal(q, k, v, self.scale)
+
+def _absorbed_softmax_scale(cfg):
+    r = cfg.rope_scaling
+    mscale = yarn_get_mscale(r["factor"], r.get("mscale_all_dim", 0.0))
+    return cfg.qk_head_dim ** -0.5 * mscale * mscale
+
+
+def _make_latent_cache_manager(cfg, dtype, page_size=128, max_num_pages=8):
+    latent_dim = cfg.kv_lora_rank + cfg.qk_rope_head_dim
+    kv_cache = torch.zeros(
+        cfg.num_hidden_layers, max_num_pages, page_size, latent_dim,
+        dtype=dtype, device=DEVICE,
+    ).contiguous()  # 4D latent cache (matches KVCacheEngine's mla_absorb branch)
+    kv_cfg = KVCacheConfig(
+        num_layers=cfg.num_hidden_layers, num_kv_heads=1, head_dim=latent_dim,
+        max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages,
+        page_size=page_size, num_qo_heads=cfg.num_attention_heads,
+        attention_backend="mla_absorb", softmax_scale=_absorbed_softmax_scale(cfg),
+    )
+    transfer_info = TransferEngineInfo(
+        my_entity_id="kimi_absorb_serve", my_session_id="kimi_session",
+        transfer_engine=LocalTransferEngine("localhost"),
+    )
+    alloc = PagedAllocationManager(
+        config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info)
+    alloc.add_request("r0", ["main"])
+    buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE)
+    cm = create_cache_manager(
+        request_ids=["r0"], active_labels_per_request={"r0": "main"},
+        kv_cache=kv_cache, alloc_manager=alloc, buffer_manager=buffers,
+        kv_cache_config=kv_cfg, device=DEVICE,
+    )
+    return cm, alloc
+
+
+def _make_model(cfg, checkpoint_dir) -> KimiK2Model:
+    model = object.__new__(KimiK2Model)  # skip __init__ (tokenizer/full config)
+    model.config = cfg
+    model.model_path_hf = str(checkpoint_dir)
+    model.cache_dir = None
+    model._submodule_cache = {}
+    return model
+
+
+def _step(submodule, cm, graph_walk, token_ids):
+    engine_inputs = ModelInputsFromEngine(
+        request_ids=["r0"], per_request_info={}, cache_manager=cm)
+    ar_in = submodule.prepare_inputs(
+        graph_walk=graph_walk, fwd_info=None, inputs={"text_inputs": [token_ids]})
+    packed = submodule.preprocess(graph_walk, engine_inputs, [ar_in])
+    with torch.no_grad():
+        out = submodule.forward(graph_walk, engine_inputs, **packed)
+    return out["logits"][0]
+
+
+def test_absorbed_serve_matches_naive_reference(tmp_path):
+    from safetensors.torch import save_file
+
+    torch.manual_seed(0)
+    cfg_naive = KimiK2Config.reduced()            # mla_absorb=False (reduced default)
+    assert cfg_naive.mla_absorb is False
+    ref = _build_reference(cfg_naive)
+    save_file(_hf_checkpoint(ref, cfg_naive), str(tmp_path / "model.safetensors"))
+
+    T = 6
+    prompt = torch.randint(0, cfg_naive.vocab_size, (T,), device=DEVICE)
+    pos = torch.arange(T, device=DEVICE)
+    with torch.no_grad():
+        ref_hidden = ref.model(prompt, _MockMLACache(cfg_naive.padded_head_dim), pos)
+        ref_logits = ref.lm_head(ref_hidden[-1:])  # (1, vocab)
+
+    cfg_absorb = KimiK2Config.reduced()
+    cfg_absorb.mla_absorb = True
+    model = _make_model(cfg_absorb, tmp_path)
+    submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16)
+    assert isinstance(submodule, KimiLLMSubmodule)
+    buf_names = {n for n, _ in submodule.language_model.named_buffers()}
+    assert any("w_kc" in n for n in buf_names) and any("fused_qkv_a_proj" in n for n in buf_names)
+
+    cm, alloc = _make_latent_cache_manager(cfg_absorb, torch.bfloat16)
+    try:
+        prefill_logits = _step(submodule, cm, "prefill", prompt)
+        assert prefill_logits.shape == (1, cfg_absorb.vocab_size)
+        assert torch.isfinite(prefill_logits).all()
+        torch.testing.assert_close(prefill_logits, ref_logits, rtol=5e-2, atol=5e-2)
+
+        next_token = prefill_logits.argmax(-1)
+        generated = [int(next_token.item())]
+        for _ in range(4):
+            logits = _step(submodule, cm, "decode", next_token)
+            assert logits.shape == (1, cfg_absorb.vocab_size)
+            assert torch.isfinite(logits).all()
+            next_token = logits.argmax(-1)
+            tok = int(next_token.item())
+            assert 0 <= tok < cfg_absorb.vocab_size
+            generated.append(tok)
+    finally:
+        alloc.cleanup()
+
+    assert len(generated) == 5
diff --git a/test/integration/test_kimi_mla_paged.py b/test/integration/test_kimi_mla_paged.py
new file mode 100644
index 000000000..101f246f8
--- /dev/null
+++ b/test/integration/test_kimi_mla_paged.py
@@ -0,0 +1,155 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from mstar.communication.tensors import LocalTransferEngine
+from mstar.engine.cache_manager import WorkspaceBufferManager, create_cache_manager
+from mstar.engine.kv_store import (
+    KVCacheConfig,
+    PagedAllocationManager,
+    TransferEngineInfo,
+)
+from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention
+from mstar.model.kimi_k2_7.components.rope import (
+    _yarn_find_correction_range,
+    _yarn_linear_ramp_mask,
+    rotate_gptj,
+    yarn_get_mscale,
+)
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="real FlashInfer paged MLA needs a GPU",
+)
+
+DEVICE = torch.device("cuda")
+
+
+def _make_real_cache_manager(num_heads, head_dim, dtype, page_size=128, max_num_pages=8):
+    kv_cache = torch.zeros(
+        2, max_num_pages, 2, page_size, num_heads, head_dim,
+        dtype=dtype, device=DEVICE,
+    ).contiguous()
+    kv_cfg = KVCacheConfig(
+        num_layers=2, num_kv_heads=num_heads, head_dim=head_dim,
+        max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages,
+        page_size=page_size, num_qo_heads=num_heads,
+    )
+    transfer_info = TransferEngineInfo(
+        my_entity_id="kimi_mla_paged_test",
+        my_session_id="kimi_session",
+        transfer_engine=LocalTransferEngine("localhost"),
+    )
+    alloc = PagedAllocationManager(
+        config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info,
+    )
+    alloc.add_request("r0", ["main"])
+    buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE)
+    cm = create_cache_manager(
+        request_ids=["r0"],
+        active_labels_per_request={"r0": "main"},
+        kv_cache=kv_cache,
+        alloc_manager=alloc,
+        buffer_manager=buffers,
+        kv_cache_config=kv_cfg,
+        device=DEVICE,
+    )
+    return cm, alloc
+
+
+def _ref_rmsnorm(x, weight, eps):
+    x32 = x.float()
+    x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps)
+    return weight * x32.to(x.dtype)
+
+
+def _ref_yarn_rope(pos, q_pe, k_pe, rotary_dim, base, factor, max_pos,
+                   beta_fast, beta_slow, mscale, mscale_all_dim):
+    pos_freqs = base ** (torch.arange(0, rotary_dim, 2, device=q_pe.device).float() / rotary_dim)
+    ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs)
+    low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos)
+    mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float).to(q_pe.device)
+    inv_freq = interp * (1 - mask) + ext * mask
+    amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim)
+    freqs = torch.outer(pos.float(), inv_freq)
+    cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin
+    kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin
+    return qr.to(q_pe.dtype), kr.to(k_pe.dtype)
+
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))  # (H,T,D)
+    T = q.shape[0]
+    causal = torch.triu(
+        torch.full((T, T), float("-inf"), device=q.device), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+def _ref_deepseek_mla(attn: KimiMLAAttention, cfg, h, pos):
+    T, H = h.shape[0], attn.num_heads
+    Dnope, Drope, Dv, L = (
+        cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank)
+    eps = cfg.rms_norm_eps
+    q = _ref_rmsnorm(F.linear(h, attn.q_a_proj.weight), attn.q_a_layernorm.weight, eps)
+    q = F.linear(q, attn.q_b_proj.weight).view(T, H, cfg.qk_head_dim)
+    q_nope, q_pe = q.split([Dnope, Drope], dim=-1)
+    latent = F.linear(h, attn.kv_a_proj_with_mqa.weight)
+    kv_a, k_pe = latent.split([L, Drope], dim=-1)
+    kv = F.linear(_ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps),
+                  attn.kv_b_proj.weight).view(T, H, Dnope + Dv)
+    k_nope, v = kv.split([Dnope, Dv], dim=-1)
+    k_pe = k_pe.view(T, 1, Drope)
+    r = cfg.rope_scaling
+    q_pe, k_pe = _ref_yarn_rope(
+        pos, q_pe, k_pe, Drope, cfg.rope_theta, r["factor"],
+        r["original_max_position_embeddings"], r.get("beta_fast", 32),
+        r.get("beta_slow", 1), r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0))
+    q = torch.cat([q_nope, q_pe], dim=-1)  # (T, H, Dqk) — NOT padded
+    k = torch.cat([k_nope, k_pe.expand(T, H, Drope)], dim=-1)  # (T, H, Dqk)
+    mscale = yarn_get_mscale(r["factor"], r.get("mscale_all_dim", 0.0))
+    deepseek_scale = cfg.qk_head_dim ** -0.5 * mscale * mscale
+    out = _sdpa_causal(q, k, v, deepseek_scale)  # v is Dv-wide, output Dv-wide
+    out = out.reshape(T, H * Dv)
+    return F.linear(out, attn.o_proj.weight)
+
+
+def _build_attention(cfg, dtype):
+    attn = KimiMLAAttention(cfg).to(device=DEVICE, dtype=dtype)
+    for lin in (attn.q_a_proj, attn.q_b_proj, attn.kv_a_proj_with_mqa,
+                attn.kv_b_proj, attn.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (attn.q_a_layernorm, attn.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    return attn
+
+
+def test_paged_mla_matches_deepseek_sdpa():
+    torch.manual_seed(0)
+    cfg = KimiK2Config.reduced()
+    assert cfg.qk_head_dim == 24 and cfg.padded_head_dim == 64  # the mitigation
+    dtype = torch.bfloat16
+    attn = _build_attention(cfg, dtype)
+
+    T = 6
+    h = torch.randn(T, cfg.hidden_size, device=DEVICE, dtype=dtype) * 0.1
+    pos = torch.arange(T, device=DEVICE)
+
+    cm, alloc = _make_real_cache_manager(cfg.num_attention_heads, cfg.padded_head_dim, dtype)
+    try:
+        cm.set_active_label("main")
+        cm.plan_attention(seq_lens=[T], is_causal=True, dtype=dtype)
+        cm.set_layer_idx(0)
+        with torch.no_grad():
+            got = attn(h, cm, pos)
+        torch.cuda.synchronize()
+    finally:
+        alloc.cleanup()
+
+    expected = _ref_deepseek_mla(attn, cfg, h, pos)
+    assert got.shape == (T, cfg.hidden_size)
+    # Any residual after exact scale compensation is bf16 FlashInfer rounding.
+    torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2)
diff --git a/test/integration/test_kimi_moe.py b/test/integration/test_kimi_moe.py
new file mode 100644
index 000000000..e067f68d9
--- /dev/null
+++ b/test/integration/test_kimi_moe.py
@@ -0,0 +1,169 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+from mstar.model.components.moe import _dispatch
+from mstar.model.kimi_k2_7.components.language_model import build_moe_block
+from mstar.model.kimi_k2_7.components.moe import KimiMoEGate
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="M2 golden tests need a GPU (fused expert GEMM is CUDA/bf16-only)",
+)
+
+DEVICE = "cuda"
+
+def _ref_grouped_topk(
+    logits: torch.Tensor, bias: torch.Tensor, n_group, topk_group, top_k,
+    norm_topk_prob, routed_scaling_factor,
+):
+    scores = logits.float().sigmoid()
+    T = scores.shape[0]
+    original = scores
+    scores = scores + bias.unsqueeze(0)
+    group_scores = scores.view(T, n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1]
+    group_mask = torch.zeros_like(group_scores)
+    group_mask.scatter_(1, group_idx, 1)
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(T, n_group, scores.shape[-1] // n_group)
+        .reshape(T, -1)
+    )
+    masked = scores.masked_fill(~score_mask.bool(), float("-inf"))
+    ids = torch.topk(masked, k=top_k, dim=-1, sorted=False)[1]
+    weights = original.gather(1, ids)
+    if norm_topk_prob:
+        weights = weights / weights.sum(dim=-1, keepdim=True)
+    weights = weights * routed_scaling_factor
+    return weights, ids
+
+
+def _ref_routed_experts(h, gate_up, down, weights, ids):
+    T, H = h.shape
+    inter = down.shape[-1]
+    out = torch.zeros(T, H, dtype=h.dtype, device=h.device)
+    for t in range(T):
+        for j in range(ids.shape[1]):
+            e = int(ids[t, j])
+            gu = gate_up[e] @ h[t]  # (2*inter,)
+            g, u = gu[:inter], gu[inter:]
+            expert_out = down[e] @ (F.silu(g) * u)  # (H,)
+            out[t] += weights[t, j] * expert_out
+    return out
+
+
+def _ref_swiglu(x, gate_w, up_w, down_w):
+    return F.linear(F.silu(F.linear(x, gate_w)) * F.linear(x, up_w), down_w)
+
+
+def _dense_combine(ids, weights, num_experts):
+    dense = torch.zeros(ids.shape[0], num_experts, device=ids.device)
+    dense.scatter_(1, ids, weights.float())
+    return dense
+
+def test_moe_gate_matches_reference():
+    torch.manual_seed(0)
+    cfg = KimiK2Config.reduced()
+    gate = KimiMoEGate(
+        cfg.hidden_size, cfg.n_routed_experts, cfg.num_experts_per_tok,
+        cfg.n_group, cfg.topk_group, cfg.routed_scaling_factor,
+        cfg.scoring_func, cfg.topk_method, cfg.norm_topk_prob,
+    ).to(DEVICE)
+    W = torch.randn(cfg.n_routed_experts, cfg.hidden_size, device=DEVICE)
+    bias = torch.randn(cfg.n_routed_experts, device=DEVICE)
+    gate.weight.data.copy_(W)
+    gate.e_score_correction_bias.data.copy_(bias)
+
+    h = torch.randn(6, cfg.hidden_size, device=DEVICE)
+    got_w, got_ids = gate(h)
+
+    logits = F.linear(h.float(), W.float())
+    exp_w, exp_ids = _ref_grouped_topk(
+        logits, bias, cfg.n_group, cfg.topk_group, cfg.num_experts_per_tok,
+        cfg.norm_topk_prob, cfg.routed_scaling_factor,
+    )
+    torch.testing.assert_close(
+        _dense_combine(got_ids, got_w, cfg.n_routed_experts),
+        _dense_combine(exp_ids, exp_w, cfg.n_routed_experts),
+        rtol=1e-5, atol=1e-5,
+    )
+
+
+def test_moe_gate_group_limited_routing():
+    torch.manual_seed(1)
+    n_experts, n_group, topk_group, top_k = 8, 2, 1, 2
+    experts_per_group = n_experts // n_group
+    gate = KimiMoEGate(
+        hidden_size=16, n_routed_experts=n_experts, num_experts_per_tok=top_k,
+        n_group=n_group, topk_group=topk_group, routed_scaling_factor=1.0,
+    ).to(DEVICE)
+    gate.weight.data.copy_(torch.randn(n_experts, 16, device=DEVICE))
+    gate.e_score_correction_bias.data.copy_(torch.randn(n_experts, device=DEVICE))
+
+    h = torch.randn(20, 16, device=DEVICE)
+    _, ids = gate(h)
+
+    # Each token's chosen experts share one group index.
+    groups = ids // experts_per_group
+    assert (groups == groups[:, :1]).all(), "experts crossed group boundary"
+
+def test_expert_dispatch_matches_naive():
+    torch.manual_seed(2)
+    cfg = KimiK2Config.reduced()
+    dtype = torch.bfloat16
+    T, H, I, E = 5, cfg.hidden_size, cfg.moe_intermediate_size, cfg.n_routed_experts
+
+    h = torch.randn(T, H, device=DEVICE, dtype=dtype) * 0.1
+    gate_up = torch.randn(E, 2 * I, H, device=DEVICE, dtype=dtype) * 0.05
+    down = torch.randn(E, H, I, device=DEVICE, dtype=dtype) * 0.05
+    ids = torch.tensor([[0, 1]] * T, device=DEVICE)
+    weights = torch.full((T, 2), 0.5, device=DEVICE, dtype=dtype)
+
+    got = _dispatch(h, gate_up, down, E, ids, weights)
+    expected = _ref_routed_experts(h, gate_up, down, weights, ids)
+    torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2)
+
+def test_moe_block_matches_reference():
+    torch.manual_seed(3)
+    cfg = KimiK2Config.reduced()
+    dtype = torch.bfloat16
+    T, H, I, E = 7, cfg.hidden_size, cfg.moe_intermediate_size, cfg.n_routed_experts
+    shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts
+
+    block = build_moe_block(cfg).to(device=DEVICE, dtype=dtype)
+
+    gate_w = torch.randn(E, H, device=DEVICE)  # fp32 router
+    bias = torch.randn(E, device=DEVICE)
+    expert_gate_up = torch.randn(E, 2 * I, H, device=DEVICE, dtype=dtype) * 0.05
+    expert_down = torch.randn(E, H, I, device=DEVICE, dtype=dtype) * 0.05
+    sh_gate = torch.randn(shared_inter, H, device=DEVICE, dtype=dtype) * 0.05
+    sh_up = torch.randn(shared_inter, H, device=DEVICE, dtype=dtype) * 0.05
+    sh_down = torch.randn(H, shared_inter, device=DEVICE, dtype=dtype) * 0.05
+
+    # Keep the router fp32 for deterministic selection.
+    block.gate.weight.data = gate_w
+    block.gate.e_score_correction_bias.data = bias
+    block.experts.gate_up_proj.data.copy_(expert_gate_up)
+    block.experts.down_proj.data.copy_(expert_down)
+    block.shared_expert.gate_up_proj.weight_loader(
+        block.shared_expert.gate_up_proj.weight, sh_gate, loaded_shard_id=0)
+    block.shared_expert.gate_up_proj.weight_loader(
+        block.shared_expert.gate_up_proj.weight, sh_up, loaded_shard_id=1)
+    block.shared_expert.down_proj.weight_loader(
+        block.shared_expert.down_proj.weight, sh_down)
+
+    h = torch.randn(T, H, device=DEVICE, dtype=dtype) * 0.1
+    got = block(h)
+
+    logits = F.linear(h.float(), gate_w.float())
+    weights, ids = _ref_grouped_topk(
+        logits, bias, cfg.n_group, cfg.topk_group, cfg.num_experts_per_tok,
+        cfg.norm_topk_prob, cfg.routed_scaling_factor,
+    )
+    routed = _ref_routed_experts(
+        h, expert_gate_up, expert_down, weights.to(dtype), ids)
+    shared = _ref_swiglu(h, sh_gate, sh_up, sh_down)
+    expected = routed + shared
+    torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2)
diff --git a/test/integration/test_kimi_moe_inkernel_dequant.py b/test/integration/test_kimi_moe_inkernel_dequant.py
new file mode 100644
index 000000000..7e4a246f5
--- /dev/null
+++ b/test/integration/test_kimi_moe_inkernel_dequant.py
@@ -0,0 +1,89 @@
+import pytest
+import torch
+
+from mstar.model.components.quantization import W4A16Data, unpack_int32
+from mstar.model.kimi_k2_7._testing import fake_quantize_weight
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="W4A16 fused-expert kernel golden needs a GPU",
+)
+
+DEVICE = "cuda"
+GROUP_SIZE = 32
+PACK_FACTOR = 8
+
+
+def _quantize_stack(weight):
+    E, N, K = weight.shape
+    packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE)
+    scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE)
+    deq = torch.empty((E, N, K), dtype=torch.bfloat16, device=DEVICE)
+    for e in range(E):
+        p, s, d = fake_quantize_weight(
+            weight[e], num_bits=4, group_size=GROUP_SIZE, symmetric=True,
+            scale_dtype=torch.bfloat16,
+        )
+        packed[e], scale[e], deq[e] = p.to(DEVICE), s.to(DEVICE), d.to(DEVICE)
+    return packed, scale, deq
+
+
+def _random_topk(num_tokens, E, top_k):
+    logits = torch.randn(num_tokens, E, device=DEVICE)
+    weights, ids = torch.topk(logits.softmax(-1), top_k, dim=-1)
+    weights = weights / weights.sum(-1, keepdim=True)
+    return weights.to(torch.bfloat16), ids
+
+
+@pytest.mark.parametrize("num_tokens", [8, 3])  # M > E and M <= E branches
+def test_w4a16_matches_bf16_on_same_dequant(num_tokens):
+    from mstar.utils.fused_moe.runner import fused_experts
+
+    torch.manual_seed(0)
+    E, H, I, top_k = 4, 128, 64, 2
+    # Wide init exercises top-nibble sign masking.
+    w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16)
+
+    w1_packed, w1_scale, w1_deq = _quantize_stack(w1)
+    w2_packed, w2_scale, w2_deq = _quantize_stack(w2)
+
+    # Guard that the negative-container path is actually covered.
+    assert (w1_packed < 0).any(), "no int32 with bit 31 set — top-nibble path untested"
+    top_nibbles = unpack_int32(w1_packed.cpu(), num_bits=4)[..., PACK_FACTOR - 1 :: PACK_FACTOR]
+    assert (top_nibbles >= 8).any()
+
+    x = (torch.randn(num_tokens, H, device=DEVICE) * 0.5).to(torch.bfloat16)
+    topk_weights, topk_ids = _random_topk(num_tokens, E, top_k)
+
+    out_quant = fused_experts(
+        x, w1_packed, w2_packed, topk_weights, topk_ids,
+        quant=W4A16Data(w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE),
+    )
+    out_bf16 = fused_experts(x, w1_deq, w2_deq, topk_weights, topk_ids)
+
+    assert out_quant.shape == (num_tokens, H)
+    assert out_quant.dtype == torch.bfloat16
+    torch.testing.assert_close(out_quant, out_bf16, rtol=1e-2, atol=1e-2)
+
+
+def test_w4a16_reduce_results_false_shape():
+    from mstar.utils.fused_moe.runner import fused_experts
+
+    torch.manual_seed(1)
+    E, H, I, top_k, num_tokens = 4, 128, 64, 2, 6
+    w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w1_packed, w1_scale, w1_deq = _quantize_stack(w1)
+    w2_packed, w2_scale, w2_deq = _quantize_stack(w2)
+    x = (torch.randn(num_tokens, H, device=DEVICE) * 0.5).to(torch.bfloat16)
+    topk_weights, topk_ids = _random_topk(num_tokens, E, top_k)
+
+    got = fused_experts(
+        x, w1_packed, w2_packed, topk_weights, topk_ids,
+        quant=W4A16Data(w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE),
+        reduce_results=False,
+    )
+    exp = fused_experts(x, w1_deq, w2_deq, topk_weights, topk_ids, reduce_results=False)
+    assert got.shape == (num_tokens, top_k, H)
+    torch.testing.assert_close(got, exp, rtol=1e-2, atol=1e-2)
diff --git a/test/integration/test_kimi_moe_marlin.py b/test/integration/test_kimi_moe_marlin.py
new file mode 100644
index 000000000..34ce2909f
--- /dev/null
+++ b/test/integration/test_kimi_moe_marlin.py
@@ -0,0 +1,115 @@
+import pytest
+import torch
+
+from mstar.model.kimi_k2_7._testing import fake_quantize_weight
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 0),
+    reason="Marlin-backed KimiSparseMoeBlock golden needs a CUDA GPU with sm80+",
+)
+
+DEVICE = "cuda"
+GROUP_SIZE = 32
+PACK_FACTOR = 8
+
+
+def _quantize_stack(weight):
+    E, N, K = weight.shape
+    packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE)
+    scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE)
+    deq = torch.empty((E, N, K), dtype=torch.bfloat16, device=DEVICE)
+    for e in range(E):
+        p, s, d = fake_quantize_weight(
+            weight[e], num_bits=4, group_size=GROUP_SIZE, symmetric=True,
+            scale_dtype=torch.bfloat16,
+        )
+        packed[e], scale[e], deq[e] = p.to(DEVICE), s.to(DEVICE), d.to(DEVICE)
+    return packed, scale, deq
+
+
+def _build_block():
+    from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+    from mstar.model.kimi_k2_7.config import KimiK2Config
+
+    torch.manual_seed(0)
+    cfg = KimiK2Config.reduced_marlin()
+    with torch.device("meta"):
+        block = KimiSparseMoeBlock(cfg)
+    block = block.to(torch.bfloat16)
+    block.to_empty(device=DEVICE)
+
+    for p in block.parameters():
+        if p.dtype.is_floating_point:
+            with torch.no_grad():
+                p.copy_(torch.randn_like(p) * 0.1)
+
+    E, H, I = cfg.n_routed_experts, cfg.hidden_size, cfg.moe_intermediate_size
+    w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w1_packed, w1_scale, w1_deq = _quantize_stack(w1)
+    w2_packed, w2_scale, w2_deq = _quantize_stack(w2)
+    with torch.no_grad():
+        block.experts.gate_up_proj_packed.data = w1_packed
+        block.experts.gate_up_proj_scale.data = w1_scale
+        block.experts.down_proj_packed.data = w2_packed
+        block.experts.down_proj_scale.data = w2_scale
+    return cfg, block, (w1_deq, w2_deq)
+
+
+def test_marlin_block_matches_bf16_reference():
+    from mstar.utils.fused_moe.runner import fused_experts
+
+    cfg, block, (w1_deq, w2_deq) = _build_block()
+    H = cfg.hidden_size
+    x = (torch.randn(5, H, device=DEVICE) * 0.5).to(torch.bfloat16)
+
+    # Build the bf16 reference before Marlin repack frees packed params.
+    with torch.no_grad():
+        topk_w, topk_ids = block.gate(x)
+        routed_ref = fused_experts(x, w1_deq, w2_deq, topk_w.to(x.dtype), topk_ids)
+        ref = (routed_ref + block.shared_expert(x)).view(x.shape)
+
+    block.process_weights_after_loading(torch.device(DEVICE))
+    assert block._use_marlin, "reduced_marlin config should select the Marlin backend"
+    assert block.experts.gate_up_proj_packed.numel() == 0
+
+    with torch.no_grad():
+        out = block(x)
+
+    assert out.shape == x.shape and out.dtype == torch.bfloat16
+    cos = torch.nn.functional.cosine_similarity(
+        out.flatten().float(), ref.flatten().float(), dim=0
+    ).item()
+    rel_l2 = ((out - ref).float().norm() / ref.float().norm()).item()
+    assert cos > 0.999, f"cosine vs bf16 reference too low: {cos}"
+    assert rel_l2 < 0.02, f"relative-L2 vs bf16 reference too high: {rel_l2}"
+
+
+def test_forced_marlin_raises_on_illegal_shapes():
+    from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+    from mstar.model.kimi_k2_7.config import KimiK2Config
+
+    cfg = KimiK2Config.reduced_quantized_inkernel()
+    cfg.quant_kernel = "marlin"
+    with torch.device("meta"):
+        block = KimiSparseMoeBlock(cfg)
+    block = block.to(torch.bfloat16)
+    block.to_empty(device=DEVICE)
+    with pytest.raises(RuntimeError, match="Marlin is ineligible"):
+        block.process_weights_after_loading(torch.device(DEVICE))
+
+
+def test_triton_backend_still_selected_when_forced():
+    from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+    from mstar.model.kimi_k2_7.config import KimiK2Config
+
+    cfg = KimiK2Config.reduced_marlin()
+    cfg.quant_kernel = "triton"
+    with torch.device("meta"):
+        block = KimiSparseMoeBlock(cfg)
+    block = block.to(torch.bfloat16)
+    block.to_empty(device=DEVICE)
+    block.process_weights_after_loading(torch.device(DEVICE))
+    assert not block._use_marlin
+    # Packed params are retained for the Triton kernel (not freed).
+    assert block.experts.gate_up_proj_packed.numel() > 0
diff --git a/test/integration/test_kimi_quant_inkernel_weight_loading.py b/test/integration/test_kimi_quant_inkernel_weight_loading.py
new file mode 100644
index 000000000..52c5ec4c8
--- /dev/null
+++ b/test/integration/test_kimi_quant_inkernel_weight_loading.py
@@ -0,0 +1,306 @@
+import pytest
+import torch
+
+from mstar.distributed.communication import CommGroup
+from mstar.model.kimi_k2_7._testing import fake_quantize_weight
+from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM
+from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.loader import load_weights as driver_load_weights
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="Kimi packed-expert weight-loading golden needs a GPU (W4A16 fused expert GEMM)",
+)
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))
+    T = q.shape[0]
+    causal = torch.triu(torch.full((T, T), float("-inf"), device=q.device), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+class _MockMLACache:
+    def __init__(self, head_dim):
+        self.scale = head_dim ** -0.5
+
+    def set_layer_idx(self, _i):
+        pass
+
+    def advance_seq_lens(self, *_a, **_k):
+        pass
+
+    def run_attention(self, q, k, v):
+        return _sdpa_causal(q, k, v, self.scale)
+
+def _fill_layer(layer, cfg):
+    a = layer.self_attn
+    for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (a.q_a_layernorm, a.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    layer.input_layernorm.weight.data.normal_(1.0, 0.02)
+    layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02)
+    mlp = layer.mlp
+    if isinstance(mlp, KimiSparseMoeBlock):
+        mlp.gate.weight.data.normal_(0, 1)
+        mlp.gate.e_score_correction_bias.data = torch.randn(
+            cfg.n_routed_experts, device=DEVICE, dtype=torch.float32)
+        # Wide init exercises the top-nibble sign path.
+        mlp.experts.gate_up_proj.data.normal_(0, 0.2)
+        mlp.experts.down_proj.data.normal_(0, 0.2)
+        mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05)
+    else:
+        mlp.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.down_proj.weight.data.normal_(0, 0.05)
+
+
+def _build_reference(cfg):
+    model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=DTYPE)
+    model.model.embed_tokens.weight.data.normal_(0, 0.05)
+    model.model.norm.weight.data.normal_(1.0, 0.02)
+    model.lm_head.weight.data.normal_(0, 0.02)
+    for layer in model.model.layers:
+        _fill_layer(layer, cfg)
+    model.requires_grad_(False)
+    return model.eval()
+
+
+def _keep_bf16(key):
+    if key == "lm_head.weight":
+        return True
+    return (
+        key.endswith("norm.weight")
+        or key.endswith("mlp.gate.weight")
+        or key.endswith("e_score_correction_bias")
+        or key.endswith("embed_tokens.weight")
+    )
+
+
+def _emit(sd, key, view, quant_cfg):
+    eligible = (
+        not _keep_bf16(key)
+        and view.dim() == 2
+        and view.shape[-1] % quant_cfg.group_size == 0
+    )
+    if not eligible:
+        sd[key] = view
+        return
+    packed, scale, deq = fake_quantize_weight(
+        view, num_bits=quant_cfg.num_bits, group_size=quant_cfg.group_size,
+        symmetric=quant_cfg.symmetric, scale_dtype=DTYPE,
+    )
+    base = key[: -len(".weight")]
+    sd[base + ".weight_packed"] = packed
+    sd[base + ".weight_scale"] = scale
+    view.copy_(deq.to(view.dtype))
+
+
+def _hf_quant_checkpoint(model, cfg, quant_cfg):
+    inter = cfg.intermediate_size
+    moe_inter = cfg.moe_intermediate_size
+    shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts
+    m = model.model
+    sd = {}
+    _emit(sd, "model.embed_tokens.weight", m.embed_tokens.weight, quant_cfg)
+    for i, layer in enumerate(m.layers):
+        p = f"model.layers.{i}."
+        a = layer.self_attn
+        _emit(sd, p + "self_attn.q_a_proj.weight", a.q_a_proj.weight, quant_cfg)
+        _emit(sd, p + "self_attn.q_a_layernorm.weight", a.q_a_layernorm.weight, quant_cfg)
+        _emit(sd, p + "self_attn.q_b_proj.weight", a.q_b_proj.weight, quant_cfg)
+        _emit(sd, p + "self_attn.kv_a_proj_with_mqa.weight",
+              a.kv_a_proj_with_mqa.weight, quant_cfg)
+        _emit(sd, p + "self_attn.kv_a_layernorm.weight", a.kv_a_layernorm.weight, quant_cfg)
+        _emit(sd, p + "self_attn.kv_b_proj.weight", a.kv_b_proj.weight, quant_cfg)
+        _emit(sd, p + "self_attn.o_proj.weight", a.o_proj.weight, quant_cfg)
+        _emit(sd, p + "input_layernorm.weight", layer.input_layernorm.weight, quant_cfg)
+        _emit(sd, p + "post_attention_layernorm.weight",
+              layer.post_attention_layernorm.weight, quant_cfg)
+        mlp = layer.mlp
+        if isinstance(mlp, KimiSparseMoeBlock):
+            _emit(sd, p + "mlp.gate.weight", mlp.gate.weight, quant_cfg)
+            _emit(sd, p + "mlp.gate.e_score_correction_bias",
+                  mlp.gate.e_score_correction_bias, quant_cfg)
+            gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj
+            for e in range(cfg.n_routed_experts):
+                _emit(sd, p + f"mlp.experts.{e}.gate_proj.weight",
+                      gup[e, :moe_inter, :], quant_cfg)
+                _emit(sd, p + f"mlp.experts.{e}.up_proj.weight",
+                      gup[e, moe_inter:, :], quant_cfg)
+                _emit(sd, p + f"mlp.experts.{e}.down_proj.weight", dwn[e], quant_cfg)
+            sh = mlp.shared_expert
+            _emit(sd, p + "mlp.shared_experts.gate_proj.weight",
+                  sh.gate_up_proj.weight[:shared_inter], quant_cfg)
+            _emit(sd, p + "mlp.shared_experts.up_proj.weight",
+                  sh.gate_up_proj.weight[shared_inter:], quant_cfg)
+            _emit(sd, p + "mlp.shared_experts.down_proj.weight", sh.down_proj.weight, quant_cfg)
+        else:
+            _emit(sd, p + "mlp.gate_proj.weight", mlp.gate_up_proj.weight[:inter], quant_cfg)
+            _emit(sd, p + "mlp.up_proj.weight", mlp.gate_up_proj.weight[inter:], quant_cfg)
+            _emit(sd, p + "mlp.down_proj.weight", mlp.down_proj.weight, quant_cfg)
+    _emit(sd, "model.norm.weight", m.norm.weight, quant_cfg)
+    _emit(sd, "lm_head.weight", model.lm_head.weight, quant_cfg)
+    return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()}
+
+
+def _build_loaded(cfg, checkpoint_dir):
+    with torch.device("meta"):
+        model = KimiForCausalLM(cfg)
+    model = model.to(DTYPE)
+    model.to_empty(device=DEVICE)
+    loaded = driver_load_weights(model, checkpoint_dir, device=DEVICE)
+    return model.eval(), loaded
+
+def test_inkernel_weight_loading_and_forward_vs_dequant_on_load(tmp_path):
+    from safetensors.torch import save_file
+
+    torch.manual_seed(0)
+    cfg_a = KimiK2Config.reduced_quantized()
+    cfg_b = KimiK2Config.reduced_quantized_inkernel()
+    assert cfg_b.moe_in_kernel_dequant and cfg_b.quantization_config is not None
+
+    ref = _build_reference(cfg_a)
+    assert not isinstance(ref.model.layers[0].mlp, KimiSparseMoeBlock)
+    assert isinstance(ref.model.layers[1].mlp, KimiSparseMoeBlock)
+
+    sd = _hf_quant_checkpoint(ref, cfg_a, cfg_a.quantization_config)
+    assert any(k.endswith("mlp.experts.0.gate_proj.weight_packed") for k in sd)
+    assert any(k.endswith("mlp.experts.0.down_proj.weight_packed") for k in sd)
+    save_file(sd, str(tmp_path / "model.safetensors"))
+
+    model_a, _ = _build_loaded(cfg_a, tmp_path)
+    model_b, loaded_b = _build_loaded(cfg_b, tmp_path)
+
+    all_params_b = set(dict(model_b.named_parameters()).keys())
+    assert loaded_b == all_params_b, (
+        f"unloaded: {all_params_b - loaded_b}; spurious: {loaded_b - all_params_b}")
+    moe_prefix = "model.layers.1.mlp.experts."
+    for suffix in ("gate_up_proj_packed", "gate_up_proj_scale",
+                   "down_proj_packed", "down_proj_scale"):
+        assert moe_prefix + suffix in all_params_b, f"missing {suffix}"
+    assert moe_prefix + "gate_up_proj" not in all_params_b  # bf16 fused param gone
+    assert moe_prefix + "down_proj" not in all_params_b
+
+    # PyTorch .to(dtype) skips integer tensors; packed params must stay int32.
+    experts_b = model_b.model.layers[1].mlp.experts
+    assert experts_b.gate_up_proj_packed.dtype == torch.int32
+    assert experts_b.down_proj_packed.dtype == torch.int32
+    assert experts_b.gate_up_proj_scale.dtype == DTYPE
+    assert experts_b.down_proj_scale.dtype == DTYPE
+
+    assert model_b.model.layers[1].mlp.gate.e_score_correction_bias.dtype == torch.float32
+    assert {n for n, _ in model_b.named_buffers()} == set()
+
+    a_sd = dict(model_a.named_parameters())
+    b_sd = dict(model_b.named_parameters())
+    shared_keys = set(a_sd) & set(b_sd)
+    assert "model.layers.1.mlp.gate.weight" in shared_keys
+    for name in shared_keys:
+        assert torch.equal(a_sd[name], b_sd[name]), f"shared param mismatch at {name}"
+
+    T = 8
+    ids = torch.randint(0, cfg_b.vocab_size, (T,), device=DEVICE)
+    pos = torch.arange(T, device=DEVICE)
+    with torch.no_grad():
+        got = model_b(ids, _MockMLACache(cfg_b.padded_head_dim), pos)
+        expected = model_a(ids, _MockMLACache(cfg_a.padded_head_dim), pos)
+    assert got.shape == (T, cfg_b.vocab_size)
+    torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2)
+
+class _NoCommGroup(CommGroup):
+
+    def __init__(self, rank: int) -> None:
+        super().__init__(my_global_rank=rank, my_group_rank=rank, group_members=[0, 1])
+
+    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
+        return input_
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        return input_
+
+
+def _packed_source(cfg, seed):
+    g = torch.Generator().manual_seed(seed)
+    E, H, I = cfg.n_routed_experts, cfg.hidden_size, cfg.moe_intermediate_size
+    sh = I * cfg.n_shared_experts
+    qc = cfg.quantization_config
+
+    def rn(*shape, std=0.05, mean=0.0):
+        return torch.randn(*shape, generator=g) * std + mean
+
+    def quant_stack(w):
+        packs, scales = [], []
+        for e in range(w.shape[0]):
+            pk, sc, _ = fake_quantize_weight(
+                w[e], num_bits=qc.num_bits, group_size=qc.group_size,
+                symmetric=qc.symmetric, scale_dtype=DTYPE)
+            packs.append(pk)
+            scales.append(sc)
+        return torch.stack(packs), torch.stack(scales)
+
+    gate = rn(E, I, H, std=0.2)
+    up = rn(E, I, H, std=0.2)
+    down = rn(E, H, I, std=0.2)
+    return {
+        "router_w": torch.randn(E, H, generator=g),
+        "router_b": torch.randn(E, generator=g),
+        "gate_packed": quant_stack(gate),
+        "up_packed": quant_stack(up),
+        "down_packed": quant_stack(down),
+        "sh_gate": rn(sh, H), "sh_up": rn(sh, H), "sh_down": rn(H, sh),
+    }
+
+
+def _load_moe_packed(block, src):
+    block.gate.weight.data = src["router_w"].to(DEVICE)
+    block.gate.e_score_correction_bias.data = src["router_b"].to(DEVICE)
+    gup_p, gup_s = block.experts.gate_up_proj_packed, block.experts.gate_up_proj_scale
+    dwn_p, dwn_s = block.experts.down_proj_packed, block.experts.down_proj_scale
+    gate_pk, gate_sc = src["gate_packed"]
+    up_pk, up_sc = src["up_packed"]
+    down_pk, down_sc = src["down_packed"]
+    for e in range(block.num_experts):
+        gup_p.weight_loader(gup_p, gate_pk[e].to(DEVICE), loaded_shard_id=f"gate:{e}")
+        gup_p.weight_loader(gup_p, up_pk[e].to(DEVICE), loaded_shard_id=f"up:{e}")
+        gup_s.weight_loader(gup_s, gate_sc[e].to(DEVICE), loaded_shard_id=f"gate:{e}")
+        gup_s.weight_loader(gup_s, up_sc[e].to(DEVICE), loaded_shard_id=f"up:{e}")
+        dwn_p.weight_loader(dwn_p, down_pk[e].to(DEVICE), loaded_shard_id=f"down:{e}")
+        dwn_s.weight_loader(dwn_s, down_sc[e].to(DEVICE), loaded_shard_id=f"down:{e}")
+    s = block.shared_expert
+    s.gate_up_proj.weight.weight_loader(
+        s.gate_up_proj.weight, src["sh_gate"].to(DEVICE, DTYPE), loaded_shard_id=0)
+    s.gate_up_proj.weight.weight_loader(
+        s.gate_up_proj.weight, src["sh_up"].to(DEVICE, DTYPE), loaded_shard_id=1)
+    s.down_proj.weight.weight_loader(s.down_proj.weight, src["sh_down"].to(DEVICE, DTYPE))
+
+
+def test_packed_moe_block_tp2_sim_matches_tp1():
+    cfg = KimiK2Config.reduced_quantized_inkernel()
+    src = _packed_source(cfg, seed=707)
+    g = torch.Generator().manual_seed(808)
+    h = (torch.randn(7, cfg.hidden_size, generator=g) * 0.1).to(DEVICE, DTYPE)
+
+    ref = KimiSparseMoeBlock(cfg, CommGroup.trivial()).to(DEVICE, DTYPE)
+    _load_moe_packed(ref, src)
+    full_inter = ref.experts.gate_up_proj_packed.shape[1]  # 2*moe_inter
+    out_ref = ref(h)
+
+    partials = []
+    for rank in range(2):
+        block = KimiSparseMoeBlock(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE)
+        assert block.experts.gate_up_proj_packed.shape[1] == full_inter // 2
+        assert block.experts.down_proj_packed.dtype == torch.int32
+        _load_moe_packed(block, src)
+        partials.append(block(h))
+
+    out_tp2 = partials[0] + partials[1]
+    max_abs = (out_tp2 - out_ref).abs().max().item()
+    assert max_abs < 5e-2, f"packed MoE tp2 vs tp1 max abs diff {max_abs}"
+    torch.testing.assert_close(out_tp2, out_ref, rtol=2e-2, atol=2e-2)
diff --git a/test/integration/test_kimi_quant_weight_loading.py b/test/integration/test_kimi_quant_weight_loading.py
new file mode 100644
index 000000000..f0bad8058
--- /dev/null
+++ b/test/integration/test_kimi_quant_weight_loading.py
@@ -0,0 +1,201 @@
+import pytest
+import torch
+
+from mstar.model.kimi_k2_7._testing import fake_quantize_weight
+from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM
+from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.loader import load_weights as driver_load_weights
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="Kimi quant weight-loading golden needs a GPU (RMSNorm + fused expert GEMM)",
+)
+
+DEVICE = "cuda"
+
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))
+    T = q.shape[0]
+    causal = torch.triu(
+        torch.full((T, T), float("-inf"), device=q.device), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+class _MockMLACache:
+    def __init__(self, head_dim):
+        self.scale = head_dim ** -0.5
+
+    def set_layer_idx(self, _i):
+        pass
+
+    def advance_seq_lens(self, *_a, **_k):
+        pass
+
+    def run_attention(self, q, k, v):
+        return _sdpa_causal(q, k, v, self.scale)
+
+def _fill_layer(layer, cfg):
+    a = layer.self_attn
+    for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (a.q_a_layernorm, a.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    layer.input_layernorm.weight.data.normal_(1.0, 0.02)
+    layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02)
+    mlp = layer.mlp
+    if isinstance(mlp, KimiSparseMoeBlock):
+        mlp.gate.weight.data.normal_(0, 1)
+        mlp.gate.e_score_correction_bias.data = torch.randn(
+            cfg.n_routed_experts, device=DEVICE, dtype=torch.float32)
+        mlp.experts.gate_up_proj.data.normal_(0, 0.05)
+        mlp.experts.down_proj.data.normal_(0, 0.05)
+        mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05)
+    else:
+        mlp.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.down_proj.weight.data.normal_(0, 0.05)
+
+
+def _build_reference(cfg):
+    model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16)
+    model.model.embed_tokens.weight.data.normal_(0, 0.05)
+    model.model.norm.weight.data.normal_(1.0, 0.02)
+    model.lm_head.weight.data.normal_(0, 0.02)
+    for layer in model.model.layers:
+        _fill_layer(layer, cfg)
+    # Disable grad so in-place dequant write-back into leaf params is allowed.
+    model.requires_grad_(False)
+    return model.eval()
+
+def _keep_bf16(key):
+    if key == "lm_head.weight":
+        return True
+    return (
+        key.endswith("norm.weight")            # all RMSNorms incl. *_layernorm
+        or key.endswith("mlp.gate.weight")     # MoE router — never quantized
+        or key.endswith("e_score_correction_bias")
+        or key.endswith("embed_tokens.weight")
+    )
+
+
+def _emit(sd, key, view, quant_cfg):
+    eligible = (
+        not _keep_bf16(key)
+        and view.dim() == 2
+        and view.shape[-1] % quant_cfg.group_size == 0
+    )
+    if not eligible:
+        sd[key] = view
+        return
+    # Match compressed-tensors checkpoints: bf16 scale drives the reference dequant.
+    packed, scale, deq = fake_quantize_weight(
+        view, num_bits=quant_cfg.num_bits, group_size=quant_cfg.group_size,
+        symmetric=quant_cfg.symmetric, scale_dtype=torch.bfloat16,
+    )
+    base = key[: -len(".weight")]
+    sd[base + ".weight_packed"] = packed
+    sd[base + ".weight_scale"] = scale
+    view.copy_(deq.to(view.dtype))  # reference now holds the dequantized weight
+
+
+def _hf_quant_checkpoint(model, cfg, quant_cfg):
+    inter = cfg.intermediate_size
+    moe_inter = cfg.moe_intermediate_size
+    shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts
+    m = model.model
+    sd = {}
+    _emit(sd, "model.embed_tokens.weight", m.embed_tokens.weight, quant_cfg)
+    for i, layer in enumerate(m.layers):
+        p = f"model.layers.{i}."
+        a = layer.self_attn
+        _emit(sd, p + "self_attn.q_a_proj.weight", a.q_a_proj.weight, quant_cfg)
+        _emit(sd, p + "self_attn.q_a_layernorm.weight", a.q_a_layernorm.weight, quant_cfg)
+        _emit(sd, p + "self_attn.q_b_proj.weight", a.q_b_proj.weight, quant_cfg)
+        _emit(sd, p + "self_attn.kv_a_proj_with_mqa.weight",
+              a.kv_a_proj_with_mqa.weight, quant_cfg)
+        _emit(sd, p + "self_attn.kv_a_layernorm.weight", a.kv_a_layernorm.weight, quant_cfg)
+        _emit(sd, p + "self_attn.kv_b_proj.weight", a.kv_b_proj.weight, quant_cfg)
+        _emit(sd, p + "self_attn.o_proj.weight", a.o_proj.weight, quant_cfg)
+        _emit(sd, p + "input_layernorm.weight", layer.input_layernorm.weight, quant_cfg)
+        _emit(sd, p + "post_attention_layernorm.weight",
+              layer.post_attention_layernorm.weight, quant_cfg)
+        mlp = layer.mlp
+        if isinstance(mlp, KimiSparseMoeBlock):
+            _emit(sd, p + "mlp.gate.weight", mlp.gate.weight, quant_cfg)
+            _emit(sd, p + "mlp.gate.e_score_correction_bias",
+                  mlp.gate.e_score_correction_bias, quant_cfg)
+            gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj
+            for e in range(cfg.n_routed_experts):
+                _emit(sd, p + f"mlp.experts.{e}.gate_proj.weight",
+                      gup[e, :moe_inter, :], quant_cfg)
+                _emit(sd, p + f"mlp.experts.{e}.up_proj.weight",
+                      gup[e, moe_inter:, :], quant_cfg)
+                _emit(sd, p + f"mlp.experts.{e}.down_proj.weight", dwn[e], quant_cfg)
+            sh = mlp.shared_expert
+            _emit(sd, p + "mlp.shared_experts.gate_proj.weight",
+                  sh.gate_up_proj.weight[:shared_inter], quant_cfg)
+            _emit(sd, p + "mlp.shared_experts.up_proj.weight",
+                  sh.gate_up_proj.weight[shared_inter:], quant_cfg)
+            _emit(sd, p + "mlp.shared_experts.down_proj.weight", sh.down_proj.weight, quant_cfg)
+        else:
+            _emit(sd, p + "mlp.gate_proj.weight", mlp.gate_up_proj.weight[:inter], quant_cfg)
+            _emit(sd, p + "mlp.up_proj.weight", mlp.gate_up_proj.weight[inter:], quant_cfg)
+            _emit(sd, p + "mlp.down_proj.weight", mlp.down_proj.weight, quant_cfg)
+    _emit(sd, "model.norm.weight", m.norm.weight, quant_cfg)
+    _emit(sd, "lm_head.weight", model.lm_head.weight, quant_cfg)
+    # Clone to cpu + break storage aliasing (safetensors rejects shared storage).
+    return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()}
+
+
+def _build_loaded(cfg, checkpoint_dir):
+    with torch.device("meta"):
+        model = KimiForCausalLM(cfg)
+    model = model.to(torch.bfloat16)
+    model.to_empty(device=DEVICE)
+    loaded = driver_load_weights(model, checkpoint_dir, device=DEVICE)
+    return model.eval(), loaded
+
+
+def test_quant_weight_loading_roundtrip_and_forward(tmp_path):
+    from safetensors.torch import save_file
+
+    torch.manual_seed(0)
+    cfg = KimiK2Config.reduced_quantized()
+    assert cfg.quantization_config is not None
+    ref = _build_reference(cfg)
+    assert not isinstance(ref.model.layers[0].mlp, KimiSparseMoeBlock)
+    assert isinstance(ref.model.layers[1].mlp, KimiSparseMoeBlock)
+
+    sd = _hf_quant_checkpoint(ref, cfg, cfg.quantization_config)
+
+    assert any(k.endswith("mlp.experts.0.gate_proj.weight_packed") for k in sd)
+    assert any(k.endswith("self_attn.o_proj.weight_packed") for k in sd)
+    assert "lm_head.weight" in sd and "lm_head.weight_packed" not in sd
+    assert "model.embed_tokens.weight" in sd
+
+    save_file(sd, str(tmp_path / "model.safetensors"))
+    model, loaded = _build_loaded(cfg, tmp_path)
+
+    all_params = set(dict(model.named_parameters()).keys())
+    assert loaded == all_params, (
+        f"unloaded: {all_params - loaded}; spurious: {loaded - all_params}")
+
+    ref_sd = dict(ref.named_parameters())
+    for name, param in model.named_parameters():
+        assert torch.equal(param, ref_sd[name]), f"mismatch at {name}"
+
+    bias = model.model.layers[1].mlp.gate.e_score_correction_bias
+    assert bias.dtype == torch.float32
+    assert {n for n, _ in model.named_buffers()} == set()
+
+    T = 8
+    ids = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE)
+    pos = torch.arange(T, device=DEVICE)
+    with torch.no_grad():
+        got = model(ids, _MockMLACache(cfg.padded_head_dim), pos)
+        expected = ref(ids, _MockMLACache(cfg.padded_head_dim), pos)
+    assert got.shape == (T, cfg.vocab_size)
+    torch.testing.assert_close(got, expected, rtol=1e-3, atol=1e-3)
diff --git a/test/integration/test_kimi_serve_e2e.py b/test/integration/test_kimi_serve_e2e.py
new file mode 100644
index 000000000..f0bd9edbd
--- /dev/null
+++ b/test/integration/test_kimi_serve_e2e.py
@@ -0,0 +1,242 @@
+import pytest
+import torch
+
+from mstar.communication.tensors import LocalTransferEngine
+from mstar.conductor.request_info import CurrentForwardPassInfo
+from mstar.engine.cache_manager import WorkspaceBufferManager, create_cache_manager
+from mstar.engine.kv_store import (
+    KVCacheConfig,
+    PagedAllocationManager,
+    TransferEngineInfo,
+)
+from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM
+from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.kimi_k2_7.kimi_model import KimiK2Model
+from mstar.model.kimi_k2_7.submodules import KimiLLMSubmodule
+from mstar.model.submodule_base import ModelInputsFromEngine
+from mstar.utils.sampling import Sampler, SamplingConfig
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="serve e2e needs a GPU (real FlashInfer paged cache)",
+)
+
+DEVICE = torch.device("cuda")
+
+
+def _fill_layer(layer, cfg):
+    a = layer.self_attn
+    for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (a.q_a_layernorm, a.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    layer.input_layernorm.weight.data.normal_(1.0, 0.02)
+    layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02)
+    mlp = layer.mlp
+    if isinstance(mlp, KimiSparseMoeBlock):
+        mlp.gate.weight.data.normal_(0, 1)
+        mlp.gate.e_score_correction_bias.data = torch.randn(
+            cfg.n_routed_experts, device=DEVICE, dtype=torch.float32)
+        mlp.experts.gate_up_proj.data.normal_(0, 0.05)
+        mlp.experts.down_proj.data.normal_(0, 0.05)
+        mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05)
+    else:
+        mlp.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.down_proj.weight.data.normal_(0, 0.05)
+
+
+def _build_reference(cfg):
+    model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16)
+    model.model.embed_tokens.weight.data.normal_(0, 0.05)
+    model.model.norm.weight.data.normal_(1.0, 0.02)
+    model.lm_head.weight.data.normal_(0, 0.02)
+    for layer in model.model.layers:
+        _fill_layer(layer, cfg)
+    return model.eval()
+
+
+def _hf_checkpoint(model, cfg):
+    inter = cfg.intermediate_size
+    moe_inter = cfg.moe_intermediate_size
+    shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts
+    m = model.model
+    sd = {"model.embed_tokens.weight": m.embed_tokens.weight}
+    for i, layer in enumerate(m.layers):
+        p = f"model.layers.{i}."
+        a = layer.self_attn
+        sd[p + "self_attn.q_a_proj.weight"] = a.q_a_proj.weight
+        sd[p + "self_attn.q_a_layernorm.weight"] = a.q_a_layernorm.weight
+        sd[p + "self_attn.q_b_proj.weight"] = a.q_b_proj.weight
+        sd[p + "self_attn.kv_a_proj_with_mqa.weight"] = a.kv_a_proj_with_mqa.weight
+        sd[p + "self_attn.kv_a_layernorm.weight"] = a.kv_a_layernorm.weight
+        sd[p + "self_attn.kv_b_proj.weight"] = a.kv_b_proj.weight
+        sd[p + "self_attn.o_proj.weight"] = a.o_proj.weight
+        sd[p + "input_layernorm.weight"] = layer.input_layernorm.weight
+        sd[p + "post_attention_layernorm.weight"] = layer.post_attention_layernorm.weight
+        mlp = layer.mlp
+        if isinstance(mlp, KimiSparseMoeBlock):
+            sd[p + "mlp.gate.weight"] = mlp.gate.weight
+            sd[p + "mlp.gate.e_score_correction_bias"] = mlp.gate.e_score_correction_bias
+            gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj
+            for e in range(cfg.n_routed_experts):
+                sd[p + f"mlp.experts.{e}.gate_proj.weight"] = gup[e, :moe_inter, :]
+                sd[p + f"mlp.experts.{e}.up_proj.weight"] = gup[e, moe_inter:, :]
+                sd[p + f"mlp.experts.{e}.down_proj.weight"] = dwn[e]
+            sh = mlp.shared_expert
+            sd[p + "mlp.shared_experts.gate_proj.weight"] = sh.gate_up_proj.weight[:shared_inter]
+            sd[p + "mlp.shared_experts.up_proj.weight"] = sh.gate_up_proj.weight[shared_inter:]
+            sd[p + "mlp.shared_experts.down_proj.weight"] = sh.down_proj.weight
+        else:
+            sd[p + "mlp.gate_proj.weight"] = mlp.gate_up_proj.weight[:inter]
+            sd[p + "mlp.up_proj.weight"] = mlp.gate_up_proj.weight[inter:]
+            sd[p + "mlp.down_proj.weight"] = mlp.down_proj.weight
+    sd["model.norm.weight"] = m.norm.weight
+    sd["lm_head.weight"] = model.lm_head.weight
+    return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()}
+
+
+def _write_checkpoint(tmp_path, seed=0):
+    from safetensors.torch import save_file
+
+    torch.manual_seed(seed)
+    cfg = KimiK2Config.reduced()
+    ref = _build_reference(cfg)
+    save_file(_hf_checkpoint(ref, cfg), str(tmp_path / "model.safetensors"))
+    return cfg
+
+
+def _make_real_cache_manager(cfg, dtype, page_size=128, max_num_pages=8):
+    num_heads = cfg.num_attention_heads
+    head_dim = cfg.padded_head_dim
+    kv_cache = torch.zeros(
+        cfg.num_hidden_layers, max_num_pages, 2, page_size, num_heads, head_dim,
+        dtype=dtype, device=DEVICE,
+    ).contiguous()
+    kv_cfg = KVCacheConfig(
+        num_layers=cfg.num_hidden_layers, num_kv_heads=num_heads, head_dim=head_dim,
+        max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages,
+        page_size=page_size, num_qo_heads=num_heads,
+    )
+    transfer_info = TransferEngineInfo(
+        my_entity_id="kimi_serve_e2e", my_session_id="kimi_serve_e2e",
+        transfer_engine=LocalTransferEngine("localhost"),
+    )
+    alloc = PagedAllocationManager(
+        config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info)
+    alloc.add_request("r0", ["main"])
+    buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE)
+    cm = create_cache_manager(
+        request_ids=["r0"], active_labels_per_request={"r0": "main"},
+        kv_cache=kv_cache, alloc_manager=alloc, buffer_manager=buffers,
+        kv_cache_config=kv_cfg, device=DEVICE,
+    )
+    return cm, alloc
+
+
+def _greedy_sampler(cfg):
+    sampler = Sampler(device=DEVICE)
+    sampler.add_request("r0")
+    sampler.set_config("r0", vocab_size=cfg.vocab_size, temperature=0.0,
+                       top_k=0, top_p=1.0, repetition_penalty=1.0)
+    return sampler
+
+
+def _fwd_info(max_tokens, cfg):
+    return CurrentForwardPassInfo(
+        request_id="r0", graph_walk="decode", requires_cfg=False, fwd_index=0,
+        random_seed=0, max_tokens=max_tokens,
+        sampling_config={"LLM": SamplingConfig(vocab_size=cfg.vocab_size,
+                                               ignore_eos=cfg.ignore_eos)},
+        dynamic_loop_iter_counts={},
+    )
+
+
+def _run_generation(model, submodule, cfg, prompt_ids, max_tokens):
+    cm, alloc = _make_real_cache_manager(cfg, torch.bfloat16)
+    sampler = _greedy_sampler(cfg)
+    engine_inputs = ModelInputsFromEngine(
+        request_ids=["r0"], per_request_info={}, cache_manager=cm, sampler=sampler,
+    )
+    info = _fwd_info(max_tokens, cfg)
+    generated: list[int] = []
+    stopped = False
+    try:
+        ar = submodule.prepare_inputs("prefill", None, {"text_inputs": [prompt_ids]})
+        packed = submodule.preprocess("prefill", engine_inputs, [ar])
+        with torch.no_grad():
+            logits = submodule.forward("prefill", engine_inputs, **packed)["logits"][0]
+        assert logits.shape == (1, cfg.vocab_size)
+        assert torch.isfinite(logits).all()
+        next_token = sampler.sample(["r0"], logits).clone()  # (1,)
+        generated.append(int(next_token.item()))
+
+        for step in range(max_tokens + 4):  # +slack; check_stop must break first
+            ar = submodule.prepare_inputs("decode", None, {"text_inputs": [next_token]})
+            packed = submodule.preprocess("decode", engine_inputs, [ar])
+            with torch.no_grad():
+                out = submodule.forward_batched("decode", engine_inputs, **packed)
+            new_token = out["r0"]["new_token"][0]
+            outputs = {"new_token": [new_token]}
+            submodule.postprocess("r0", info, outputs)  # rebinds text_inputs
+            assert outputs["text_inputs"] is outputs["new_token"]
+            info.dynamic_loop_iter_counts["decode_loop"] = step
+            stop = submodule.check_stop("r0", info, outputs)
+            generated.append(int(new_token.item()))
+            next_token = new_token
+            if stop:
+                stopped = True
+                break
+    finally:
+        alloc.cleanup()
+        sampler.remove_request("r0")
+    return generated, stopped
+
+
+def test_serve_path_prefill_decode_loop(tmp_path):
+    cfg = _write_checkpoint(tmp_path, seed=0)
+
+    model = KimiK2Model(
+        model_path_hf=str(tmp_path), config_variant="reduced",
+        tokenizer_mode="byte",
+    )
+    assert model.config.vocab_size == 256
+
+    prompt_tensors = model.process_prompt("hello kimi", ["text"], ["text"])
+    prompt_ids = prompt_tensors["text_inputs"][0].to(DEVICE)
+    assert prompt_ids.tolist() == list("hello kimi".encode("utf-8"))
+    assert prompt_ids.max().item() < cfg.vocab_size
+
+    submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16)
+    assert isinstance(submodule, KimiLLMSubmodule)
+    assert list(submodule.language_model.named_buffers()) == []
+
+    MAX_TOKENS = 6
+    generated, stopped = _run_generation(model, submodule, cfg, prompt_ids, MAX_TOKENS)
+
+    assert stopped, "decode loop did not terminate via check_stop"
+    # check_stop fires after exactly MAX_TOKENS decode steps.
+    assert len(generated) == 1 + MAX_TOKENS, generated
+    assert all(0 <= t < cfg.vocab_size for t in generated), generated
+
+    out_bytes = model.postprocess(torch.tensor(generated), "text")
+    assert isinstance(out_bytes, bytes)
+    assert len(out_bytes) == len(generated)
+
+
+def test_serve_path_is_deterministic(tmp_path):
+    cfg = _write_checkpoint(tmp_path, seed=1)
+    model = KimiK2Model(
+        model_path_hf=str(tmp_path), config_variant="reduced",
+        tokenizer_mode="byte",
+    )
+    submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16)
+    prompt_ids = model.process_prompt("serve", ["text"], ["text"])["text_inputs"][0].to(DEVICE)
+
+    runs = [
+        _run_generation(model, submodule, cfg, prompt_ids, max_tokens=5)[0]
+        for _ in range(2)
+    ]
+    assert runs[0] == runs[1], runs
+    assert len(runs[0]) == 1 + 5
diff --git a/test/integration/test_kimi_submodule.py b/test/integration/test_kimi_submodule.py
new file mode 100644
index 000000000..0f6fdec0d
--- /dev/null
+++ b/test/integration/test_kimi_submodule.py
@@ -0,0 +1,243 @@
+import pytest
+import torch
+
+from mstar.communication.tensors import LocalTransferEngine
+from mstar.engine.cache_manager import WorkspaceBufferManager, create_cache_manager
+from mstar.engine.kv_store import (
+    KVCacheConfig,
+    PagedAllocationManager,
+    TransferEngineInfo,
+)
+from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM
+from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.kimi_k2_7.kimi_model import KimiK2Model
+from mstar.model.kimi_k2_7.submodules import KimiLLMSubmodule
+from mstar.model.submodule_base import ModelInputsFromEngine
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="submodule e2e needs a GPU (real FlashInfer paged cache)",
+)
+
+DEVICE = torch.device("cuda")
+
+def _fill_layer(layer, cfg):
+    a = layer.self_attn
+    for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (a.q_a_layernorm, a.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    layer.input_layernorm.weight.data.normal_(1.0, 0.02)
+    layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02)
+    mlp = layer.mlp
+    if isinstance(mlp, KimiSparseMoeBlock):
+        mlp.gate.weight.data.normal_(0, 1)
+        mlp.gate.e_score_correction_bias.data = torch.randn(
+            cfg.n_routed_experts, device=DEVICE, dtype=torch.float32)
+        mlp.experts.gate_up_proj.data.normal_(0, 0.05)
+        mlp.experts.down_proj.data.normal_(0, 0.05)
+        mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05)
+    else:
+        mlp.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.down_proj.weight.data.normal_(0, 0.05)
+
+
+def _build_reference(cfg):
+    model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16)
+    model.model.embed_tokens.weight.data.normal_(0, 0.05)
+    model.model.norm.weight.data.normal_(1.0, 0.02)
+    model.lm_head.weight.data.normal_(0, 0.02)
+    for layer in model.model.layers:
+        _fill_layer(layer, cfg)
+    return model.eval()
+
+
+def _hf_checkpoint(model, cfg):
+    inter = cfg.intermediate_size
+    moe_inter = cfg.moe_intermediate_size
+    shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts
+    m = model.model
+    sd = {"model.embed_tokens.weight": m.embed_tokens.weight}
+    for i, layer in enumerate(m.layers):
+        p = f"model.layers.{i}."
+        a = layer.self_attn
+        sd[p + "self_attn.q_a_proj.weight"] = a.q_a_proj.weight
+        sd[p + "self_attn.q_a_layernorm.weight"] = a.q_a_layernorm.weight
+        sd[p + "self_attn.q_b_proj.weight"] = a.q_b_proj.weight
+        sd[p + "self_attn.kv_a_proj_with_mqa.weight"] = a.kv_a_proj_with_mqa.weight
+        sd[p + "self_attn.kv_a_layernorm.weight"] = a.kv_a_layernorm.weight
+        sd[p + "self_attn.kv_b_proj.weight"] = a.kv_b_proj.weight
+        sd[p + "self_attn.o_proj.weight"] = a.o_proj.weight
+        sd[p + "input_layernorm.weight"] = layer.input_layernorm.weight
+        sd[p + "post_attention_layernorm.weight"] = layer.post_attention_layernorm.weight
+        mlp = layer.mlp
+        if isinstance(mlp, KimiSparseMoeBlock):
+            sd[p + "mlp.gate.weight"] = mlp.gate.weight
+            sd[p + "mlp.gate.e_score_correction_bias"] = mlp.gate.e_score_correction_bias
+            gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj
+            for e in range(cfg.n_routed_experts):
+                sd[p + f"mlp.experts.{e}.gate_proj.weight"] = gup[e, :moe_inter, :]
+                sd[p + f"mlp.experts.{e}.up_proj.weight"] = gup[e, moe_inter:, :]
+                sd[p + f"mlp.experts.{e}.down_proj.weight"] = dwn[e]
+            sh = mlp.shared_expert
+            sd[p + "mlp.shared_experts.gate_proj.weight"] = sh.gate_up_proj.weight[:shared_inter]
+            sd[p + "mlp.shared_experts.up_proj.weight"] = sh.gate_up_proj.weight[shared_inter:]
+            sd[p + "mlp.shared_experts.down_proj.weight"] = sh.down_proj.weight
+        else:
+            sd[p + "mlp.gate_proj.weight"] = mlp.gate_up_proj.weight[:inter]
+            sd[p + "mlp.up_proj.weight"] = mlp.gate_up_proj.weight[inter:]
+            sd[p + "mlp.down_proj.weight"] = mlp.down_proj.weight
+    sd["model.norm.weight"] = m.norm.weight
+    sd["lm_head.weight"] = model.lm_head.weight
+    return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()}
+
+def _make_real_cache_manager(cfg, dtype, page_size=128, max_num_pages=8):
+    num_heads = cfg.num_attention_heads
+    head_dim = cfg.padded_head_dim
+    kv_cache = torch.zeros(
+        cfg.num_hidden_layers, max_num_pages, 2, page_size, num_heads, head_dim,
+        dtype=dtype, device=DEVICE,
+    ).contiguous()
+    kv_cfg = KVCacheConfig(
+        num_layers=cfg.num_hidden_layers, num_kv_heads=num_heads, head_dim=head_dim,
+        max_seq_len=page_size * max_num_pages, max_num_pages=max_num_pages,
+        page_size=page_size, num_qo_heads=num_heads,
+    )
+    transfer_info = TransferEngineInfo(
+        my_entity_id="kimi_submodule_test", my_session_id="kimi_session",
+        transfer_engine=LocalTransferEngine("localhost"),
+    )
+    alloc = PagedAllocationManager(
+        config=kv_cfg, kv_cache=kv_cache, transfer_engine_info=transfer_info)
+    alloc.add_request("r0", ["main"])
+    buffers = WorkspaceBufferManager(64 * 1024 * 1024, device=DEVICE)
+    cm = create_cache_manager(
+        request_ids=["r0"], active_labels_per_request={"r0": "main"},
+        kv_cache=kv_cache, alloc_manager=alloc, buffer_manager=buffers,
+        kv_cache_config=kv_cfg, device=DEVICE,
+    )
+    return cm, alloc
+
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))
+    T = q.shape[0]
+    causal = torch.triu(torch.full((T, T), float("-inf"), device=q.device), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+class _MockMLACache:
+    def __init__(self, head_dim):
+        self.scale = head_dim ** -0.5
+
+    def set_layer_idx(self, _i):
+        pass
+
+    def advance_seq_lens(self, *_a, **_k):
+        pass
+
+    def run_attention(self, q, k, v):
+        return _sdpa_causal(q, k, v, self.scale)
+
+
+def _make_model(cfg, checkpoint_dir) -> KimiK2Model:
+    model = object.__new__(KimiK2Model)
+    model.config = cfg
+    model.model_path_hf = str(checkpoint_dir)
+    model.cache_dir = None
+    model._submodule_cache = {}
+    return model
+
+def _engine_inputs(cm):
+    return ModelInputsFromEngine(
+        request_ids=["r0"], per_request_info={}, cache_manager=cm,
+    )
+
+
+def _step(submodule, cm, graph_walk, token_ids):
+    engine_inputs = _engine_inputs(cm)
+    ar_in = submodule.prepare_inputs(
+        graph_walk=graph_walk, fwd_info=None,
+        inputs={"text_inputs": [token_ids]},
+    )
+    packed = submodule.preprocess(graph_walk, engine_inputs, [ar_in])
+    with torch.no_grad():
+        out = submodule.forward(graph_walk, engine_inputs, **packed)
+    return out["logits"][0], packed  # (1, vocab), packed dict
+
+
+def test_submodule_prefill_decode_over_real_paged_cache(tmp_path):
+    from safetensors.torch import save_file
+
+    torch.manual_seed(0)
+    cfg = KimiK2Config.reduced()
+    ref = _build_reference(cfg)
+    save_file(_hf_checkpoint(ref, cfg), str(tmp_path / "model.safetensors"))
+
+    model = _make_model(cfg, tmp_path)
+    submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16)
+    assert isinstance(submodule, KimiLLMSubmodule)
+    assert model.get_submodule("LLM") is submodule
+    # Derived buffers must not survive meta -> to_empty as garbage.
+    assert list(submodule.language_model.named_buffers()) == []
+    p = next(submodule.language_model.parameters())
+    assert p.device.type == "cuda" and p.dtype == torch.bfloat16
+
+    T = 6
+    prompt = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE)
+    cm, alloc = _make_real_cache_manager(cfg, torch.bfloat16)
+    try:
+        prefill_logits, _ = _step(submodule, cm, "prefill", prompt)
+        assert prefill_logits.shape == (1, cfg.vocab_size)
+        assert torch.isfinite(prefill_logits).all()
+
+        # Same loaded weights through the mock cache at the DeepSeek-correct scale.
+        pos = torch.arange(T, device=DEVICE)
+        with torch.no_grad():
+            ref_hidden = submodule.language_model.model(
+                prompt, _MockMLACache(cfg.padded_head_dim), pos)
+        ref_logits = submodule.lm_head(ref_hidden[-1:])
+        torch.testing.assert_close(prefill_logits, ref_logits, rtol=5e-2, atol=5e-2)
+
+        next_token = prefill_logits.argmax(-1)  # (1,)
+        generated = [int(next_token.item())]
+        assert 0 <= generated[-1] < cfg.vocab_size
+        for _ in range(4):
+            logits, _ = _step(submodule, cm, "decode", next_token)
+            assert logits.shape == (1, cfg.vocab_size)
+            assert torch.isfinite(logits).all()
+            next_token = logits.argmax(-1)
+            tok = int(next_token.item())
+            assert 0 <= tok < cfg.vocab_size
+            generated.append(tok)
+    finally:
+        alloc.cleanup()
+
+    assert len(generated) == 5
+    assert all(0 <= t < cfg.vocab_size for t in generated)
+
+
+def test_submodule_paged_decode_is_deterministic(tmp_path):
+    from safetensors.torch import save_file
+
+    torch.manual_seed(1)
+    cfg = KimiK2Config.reduced()
+    ref = _build_reference(cfg)
+    save_file(_hf_checkpoint(ref, cfg), str(tmp_path / "model.safetensors"))
+    model = _make_model(cfg, tmp_path)
+    submodule = model.get_submodule("LLM", device="cuda", autocast_dtype=torch.bfloat16)
+
+    T = 5
+    prompt = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE)
+    tokens = []
+    for _ in range(2):
+        cm, alloc = _make_real_cache_manager(cfg, torch.bfloat16)
+        try:
+            logits, _ = _step(submodule, cm, "prefill", prompt)
+            tokens.append(int(logits.argmax(-1).item()))
+        finally:
+            alloc.cleanup()
+    assert tokens[0] == tokens[1]
diff --git a/test/integration/test_kimi_tp.py b/test/integration/test_kimi_tp.py
new file mode 100644
index 000000000..35e763012
--- /dev/null
+++ b/test/integration/test_kimi_tp.py
@@ -0,0 +1,273 @@
+from __future__ import annotations
+
+import os
+import socket
+
+import pytest
+import torch
+
+from mstar.distributed.communication import CommGroup
+from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention
+from mstar.model.kimi_k2_7.components.decoder_layer import KimiDecoderLayer
+from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+from mstar.model.kimi_k2_7.config import KimiK2Config
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="Kimi TP goldens need a GPU (fused expert GEMM + MLA RMSNorm are CUDA-only)",
+)
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+TP = 2
+
+
+class _NoCommGroup(CommGroup):
+    def __init__(self, rank: int) -> None:
+        super().__init__(my_global_rank=rank, my_group_rank=rank, group_members=[0, 1])
+
+    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:  # noqa: D401
+        return input_
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        return input_
+
+
+class _MockMLACache:
+
+    def __init__(self, head_dim: int) -> None:
+        self.scale = head_dim ** -0.5
+
+    def set_layer_idx(self, _i):
+        pass
+
+    def set_active_label(self, _l):
+        pass
+
+    def advance_seq_lens(self, *_a, **_k):
+        pass
+
+    def run_attention(self, q, k, v):
+        qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))  # (H,T,D)
+        scores = torch.einsum("hqd,hkd->hqk", qt, kt) * self.scale
+        num_tokens = q.shape[0]
+        causal = torch.triu(
+            torch.full((num_tokens, num_tokens), float("-inf"), device=q.device),
+            diagonal=1,
+        )
+        attn = (scores + causal).softmax(-1)
+        return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+def _source_weights(cfg: KimiK2Config, seed: int) -> dict:
+    g = torch.Generator().manual_seed(seed)
+    H = cfg.num_attention_heads
+    E, Hd, I = cfg.n_routed_experts, cfg.hidden_size, cfg.moe_intermediate_size
+    sh = I * cfg.n_shared_experts
+
+    def rn(*shape, std=0.03, mean=0.0):
+        return torch.randn(*shape, generator=g) * std + mean
+
+    return {
+        "q_a": rn(cfg.q_lora_rank, cfg.hidden_size),
+        "q_a_norm": rn(cfg.q_lora_rank, std=0.02, mean=1.0),
+        "q_b": rn(H * cfg.qk_head_dim, cfg.q_lora_rank),
+        "kv_a": rn(cfg.kv_lora_rank + cfg.qk_rope_head_dim, cfg.hidden_size),
+        "kv_a_norm": rn(cfg.kv_lora_rank, std=0.02, mean=1.0),
+        "kv_b": rn(H * (cfg.qk_nope_head_dim + cfg.v_head_dim), cfg.kv_lora_rank),
+        "o": rn(cfg.hidden_size, H * cfg.v_head_dim),
+        "router_w": torch.randn(E, Hd, generator=g),
+        "router_b": torch.randn(E, generator=g),
+        "gate": rn(E, I, Hd, std=0.05),
+        "up": rn(E, I, Hd, std=0.05),
+        "down": rn(E, Hd, I, std=0.05),
+        "sh_gate": rn(sh, Hd, std=0.05),
+        "sh_up": rn(sh, Hd, std=0.05),
+        "sh_down": rn(Hd, sh, std=0.05),
+        "in_ln": rn(Hd, std=0.02, mean=1.0),
+        "post_ln": rn(Hd, std=0.02, mean=1.0),
+    }
+
+
+def _load_attention(attn: KimiMLAAttention, src: dict) -> None:
+    attn.q_a_proj.weight.data.copy_(src["q_a"].to(DEVICE, DTYPE))
+    attn.q_a_layernorm.weight.data.copy_(src["q_a_norm"].to(DEVICE, DTYPE))
+    attn.q_b_proj.weight.weight_loader(attn.q_b_proj.weight, src["q_b"].to(DEVICE, DTYPE))
+    attn.kv_a_proj_with_mqa.weight.data.copy_(src["kv_a"].to(DEVICE, DTYPE))
+    attn.kv_a_layernorm.weight.data.copy_(src["kv_a_norm"].to(DEVICE, DTYPE))
+    attn.kv_b_proj.weight.weight_loader(attn.kv_b_proj.weight, src["kv_b"].to(DEVICE, DTYPE))
+    attn.o_proj.weight.weight_loader(attn.o_proj.weight, src["o"].to(DEVICE, DTYPE))
+
+
+def _load_moe(block: KimiSparseMoeBlock, src: dict) -> None:
+    block.gate.weight.data = src["router_w"].to(DEVICE)  # keep router fp32
+    block.gate.e_score_correction_bias.data = src["router_b"].to(DEVICE)
+    gu, dp = block.experts.gate_up_proj, block.experts.down_proj
+    for e in range(block.num_experts):
+        gu.weight_loader(gu, src["gate"][e].to(DEVICE, DTYPE), loaded_shard_id=f"gate:{e}")
+        gu.weight_loader(gu, src["up"][e].to(DEVICE, DTYPE), loaded_shard_id=f"up:{e}")
+        dp.weight_loader(dp, src["down"][e].to(DEVICE, DTYPE), loaded_shard_id=f"down:{e}")
+    s = block.shared_expert
+    s.gate_up_proj.weight.weight_loader(
+        s.gate_up_proj.weight, src["sh_gate"].to(DEVICE, DTYPE), loaded_shard_id=0)
+    s.gate_up_proj.weight.weight_loader(
+        s.gate_up_proj.weight, src["sh_up"].to(DEVICE, DTYPE), loaded_shard_id=1)
+    s.down_proj.weight.weight_loader(s.down_proj.weight, src["sh_down"].to(DEVICE, DTYPE))
+
+
+def _load_decoder(layer: KimiDecoderLayer, src: dict) -> None:
+    _load_attention(layer.self_attn, src)
+    _load_moe(layer.mlp, src)
+    layer.input_layernorm.weight.data.copy_(src["in_ln"].to(DEVICE, DTYPE))
+    layer.post_attention_layernorm.weight.data.copy_(src["post_ln"].to(DEVICE, DTYPE))
+
+
+def _inputs(cfg: KimiK2Config, num_tokens: int, seed: int):
+    g = torch.Generator().manual_seed(seed)
+    h = (torch.randn(num_tokens, cfg.hidden_size, generator=g) * 0.1).to(DEVICE, DTYPE)
+    pos = torch.arange(num_tokens, device=DEVICE)
+    return h, pos
+
+def test_mla_attention_tp2_sim_matches_tp1():
+    cfg = KimiK2Config.reduced()
+    src = _source_weights(cfg, seed=101)
+    h, pos = _inputs(cfg, num_tokens=6, seed=202)
+
+    ref = KimiMLAAttention(cfg, CommGroup.trivial()).to(DEVICE, DTYPE)
+    _load_attention(ref, src)
+    assert ref.num_heads == cfg.num_attention_heads
+    out_ref = ref(h, _MockMLACache(cfg.padded_head_dim), pos)
+
+    partials = []
+    for rank in range(TP):
+        attn = KimiMLAAttention(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE)
+        assert attn.num_heads == cfg.num_attention_heads // TP
+        _load_attention(attn, src)
+        partials.append(attn(h, _MockMLACache(cfg.padded_head_dim), pos))
+
+    out_tp2 = partials[0] + partials[1]
+    max_abs = (out_tp2 - out_ref).abs().max().item()
+    assert max_abs < 5e-2, f"MLA tp2 vs tp1 max abs diff {max_abs}"
+    torch.testing.assert_close(out_tp2, out_ref, rtol=2e-2, atol=2e-2)
+
+
+def test_moe_block_tp2_sim_matches_tp1():
+    cfg = KimiK2Config.reduced()
+    src = _source_weights(cfg, seed=303)
+    h, _ = _inputs(cfg, num_tokens=7, seed=404)
+
+    ref = KimiSparseMoeBlock(cfg, CommGroup.trivial()).to(DEVICE, DTYPE)
+    _load_moe(ref, src)
+    full_inter = ref.experts.gate_up_proj.shape[1]
+    out_ref = ref(h)
+
+    partials = []
+    for rank in range(TP):
+        block = KimiSparseMoeBlock(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE)
+        assert block.experts.gate_up_proj.shape[1] == full_inter // TP
+        _load_moe(block, src)
+        partials.append(block(h))
+
+    out_tp2 = partials[0] + partials[1]
+    max_abs = (out_tp2 - out_ref).abs().max().item()
+    assert max_abs < 5e-2, f"MoE tp2 vs tp1 max abs diff {max_abs}"
+    torch.testing.assert_close(out_tp2, out_ref, rtol=2e-2, atol=2e-2)
+
+
+def test_tp2_sim_is_stable_across_repeats():
+    cfg = KimiK2Config.reduced()
+
+    def attn_diff():
+        src = _source_weights(cfg, seed=101)
+        h, pos = _inputs(cfg, num_tokens=6, seed=202)
+        ref = KimiMLAAttention(cfg, CommGroup.trivial()).to(DEVICE, DTYPE)
+        _load_attention(ref, src)
+        o_ref = ref(h, _MockMLACache(cfg.padded_head_dim), pos)
+        parts = []
+        for rank in range(TP):
+            a = KimiMLAAttention(cfg, _NoCommGroup(rank)).to(DEVICE, DTYPE)
+            _load_attention(a, src)
+            parts.append(a(h, _MockMLACache(cfg.padded_head_dim), pos))
+        return (parts[0] + parts[1] - o_ref).abs().max().item()
+
+    diffs = [attn_diff() for _ in range(3)]
+    assert diffs[0] == diffs[1] == diffs[2], f"unstable tp2 sim diffs: {diffs}"
+
+def _free_port() -> int:
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.bind(("127.0.0.1", 0))
+    port = s.getsockname()[1]
+    s.close()
+    return port
+
+
+def _nccl_worker(rank: int, world_size: int, port: int, result_path: str) -> None:
+    import torch.distributed as dist
+
+    os.environ.setdefault("NCCL_IB_DISABLE", "1")
+    torch.cuda.set_device(rank)
+    dist.init_process_group(
+        backend="nccl",
+        init_method=f"tcp://127.0.0.1:{port}",
+        world_size=world_size,
+        rank=rank,
+    )
+    try:
+        cfg = KimiK2Config.reduced()
+        src = _source_weights(cfg, seed=505)
+        h, pos = _inputs(cfg, num_tokens=6, seed=606)
+
+        cg = CommGroup(my_global_rank=rank, my_group_rank=rank, group_members=[0, 1])
+        cg.device_group = None
+        cg.initialized = True
+
+        attn = KimiMLAAttention(cfg, cg).to(DEVICE, DTYPE)
+        _load_attention(attn, src)
+        moe = KimiSparseMoeBlock(cfg, cg).to(DEVICE, DTYPE)
+        _load_moe(moe, src)
+        dec = KimiDecoderLayer(cfg, layer_idx=1, comm_group=cg).to(DEVICE, DTYPE)
+        _load_decoder(dec, src)
+        assert isinstance(dec.mlp, KimiSparseMoeBlock)
+
+        attn_tp2 = attn(h, _MockMLACache(cfg.padded_head_dim), pos)
+        moe_tp2 = moe(h)
+        dec_tp2 = dec(h, _MockMLACache(cfg.padded_head_dim), pos)
+
+        attn_ref = KimiMLAAttention(cfg, CommGroup.trivial()).to(DEVICE, DTYPE)
+        _load_attention(attn_ref, src)
+        moe_ref = KimiSparseMoeBlock(cfg, CommGroup.trivial()).to(DEVICE, DTYPE)
+        _load_moe(moe_ref, src)
+        dec_ref = KimiDecoderLayer(cfg, layer_idx=1, comm_group=CommGroup.trivial()).to(DEVICE, DTYPE)
+        _load_decoder(dec_ref, src)
+
+        o_attn = attn_ref(h, _MockMLACache(cfg.padded_head_dim), pos)
+        o_moe = moe_ref(h)
+        o_dec = dec_ref(h, _MockMLACache(cfg.padded_head_dim), pos)
+
+        diffs = {
+            "attn": (attn_tp2 - o_attn).abs().max().item(),
+            "moe": (moe_tp2 - o_moe).abs().max().item(),
+            "decoder": (dec_tp2 - o_dec).abs().max().item(),
+        }
+        dist.barrier()
+        if rank == 0:
+            torch.save(diffs, result_path)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.skipif(
+    torch.cuda.device_count() < 2,
+    reason="real NCCL tp=2 golden needs >= 2 CUDA devices",
+)
+def test_tp2_nccl_matches_tp1(tmp_path):
+    import torch.multiprocessing as mp
+
+    result_path = str(tmp_path / "tp2_diffs.pt")
+    port = _free_port()
+    mp.spawn(_nccl_worker, args=(TP, port, result_path), nprocs=TP, join=True)
+
+    diffs = torch.load(result_path)
+    assert diffs["attn"] < 5e-2, diffs
+    assert diffs["moe"] < 5e-2, diffs
+    assert diffs["decoder"] < 5e-2, diffs
diff --git a/test/integration/test_kimi_weight_loading.py b/test/integration/test_kimi_weight_loading.py
new file mode 100644
index 000000000..26eb631f4
--- /dev/null
+++ b/test/integration/test_kimi_weight_loading.py
@@ -0,0 +1,170 @@
+import pytest
+import torch
+
+from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM
+from mstar.model.kimi_k2_7.components.moe import KimiSparseMoeBlock
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.loader import load_weights as driver_load_weights
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="M5 weight-loading golden needs a GPU (RMSNorm + fused expert GEMM)",
+)
+
+DEVICE = "cuda"
+
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))
+    T = q.shape[0]
+    causal = torch.triu(
+        torch.full((T, T), float("-inf"), device=q.device), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+class _MockMLACache:
+    def __init__(self, head_dim):
+        self.scale = head_dim ** -0.5
+
+    def set_layer_idx(self, _i):
+        pass
+
+    def advance_seq_lens(self, *_a, **_k):
+        pass
+
+    def run_attention(self, q, k, v):
+        return _sdpa_causal(q, k, v, self.scale)
+
+def _fill_layer(layer, cfg):
+    a = layer.self_attn
+    for lin in (a.q_a_proj, a.q_b_proj, a.kv_a_proj_with_mqa, a.kv_b_proj, a.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (a.q_a_layernorm, a.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    layer.input_layernorm.weight.data.normal_(1.0, 0.02)
+    layer.post_attention_layernorm.weight.data.normal_(1.0, 0.02)
+    mlp = layer.mlp
+    if isinstance(mlp, KimiSparseMoeBlock):
+        # In place so the router weight keeps the model dtype (bf16).
+        mlp.gate.weight.data.normal_(0, 1)
+        mlp.gate.e_score_correction_bias.data = torch.randn(
+            cfg.n_routed_experts, device=DEVICE, dtype=torch.float32)
+        mlp.experts.gate_up_proj.data.normal_(0, 0.05)
+        mlp.experts.down_proj.data.normal_(0, 0.05)
+        mlp.shared_expert.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.shared_expert.down_proj.weight.data.normal_(0, 0.05)
+    else:
+        mlp.gate_up_proj.weight.data.normal_(0, 0.05)
+        mlp.down_proj.weight.data.normal_(0, 0.05)
+
+
+def _build_reference(cfg):
+    model = KimiForCausalLM(cfg).to(device=DEVICE, dtype=torch.bfloat16)
+    model.model.embed_tokens.weight.data.normal_(0, 0.05)
+    model.model.norm.weight.data.normal_(1.0, 0.02)
+    model.lm_head.weight.data.normal_(0, 0.02)
+    for layer in model.model.layers:
+        _fill_layer(layer, cfg)
+    return model.eval()
+
+
+def _hf_checkpoint(model, cfg):
+    inter = cfg.intermediate_size
+    moe_inter = cfg.moe_intermediate_size
+    shared_inter = cfg.moe_intermediate_size * cfg.n_shared_experts
+    m = model.model
+    sd = {"model.embed_tokens.weight": m.embed_tokens.weight}
+    for i, layer in enumerate(m.layers):
+        p = f"model.layers.{i}."
+        a = layer.self_attn
+        sd[p + "self_attn.q_a_proj.weight"] = a.q_a_proj.weight
+        sd[p + "self_attn.q_a_layernorm.weight"] = a.q_a_layernorm.weight
+        sd[p + "self_attn.q_b_proj.weight"] = a.q_b_proj.weight
+        sd[p + "self_attn.kv_a_proj_with_mqa.weight"] = a.kv_a_proj_with_mqa.weight
+        sd[p + "self_attn.kv_a_layernorm.weight"] = a.kv_a_layernorm.weight
+        sd[p + "self_attn.kv_b_proj.weight"] = a.kv_b_proj.weight
+        sd[p + "self_attn.o_proj.weight"] = a.o_proj.weight
+        sd[p + "input_layernorm.weight"] = layer.input_layernorm.weight
+        sd[p + "post_attention_layernorm.weight"] = layer.post_attention_layernorm.weight
+        mlp = layer.mlp
+        if isinstance(mlp, KimiSparseMoeBlock):
+            sd[p + "mlp.gate.weight"] = mlp.gate.weight
+            sd[p + "mlp.gate.e_score_correction_bias"] = mlp.gate.e_score_correction_bias
+            gup, dwn = mlp.experts.gate_up_proj, mlp.experts.down_proj
+            for e in range(cfg.n_routed_experts):
+                sd[p + f"mlp.experts.{e}.gate_proj.weight"] = gup[e, :moe_inter, :]
+                sd[p + f"mlp.experts.{e}.up_proj.weight"] = gup[e, moe_inter:, :]
+                sd[p + f"mlp.experts.{e}.down_proj.weight"] = dwn[e]
+            sh = mlp.shared_expert
+            sd[p + "mlp.shared_experts.gate_proj.weight"] = sh.gate_up_proj.weight[:shared_inter]
+            sd[p + "mlp.shared_experts.up_proj.weight"] = sh.gate_up_proj.weight[shared_inter:]
+            sd[p + "mlp.shared_experts.down_proj.weight"] = sh.down_proj.weight
+        else:
+            sd[p + "mlp.gate_proj.weight"] = mlp.gate_up_proj.weight[:inter]
+            sd[p + "mlp.up_proj.weight"] = mlp.gate_up_proj.weight[inter:]
+            sd[p + "mlp.down_proj.weight"] = mlp.down_proj.weight
+    sd["model.norm.weight"] = m.norm.weight
+    sd["lm_head.weight"] = model.lm_head.weight
+    # Clone to cpu + break storage aliasing (safetensors rejects shared storage).
+    return {k: v.detach().cpu().clone().contiguous() for k, v in sd.items()}
+
+
+def _build_loaded(cfg, checkpoint_dir):
+    with torch.device("meta"):
+        model = KimiForCausalLM(cfg)
+    model = model.to(torch.bfloat16)
+    model.to_empty(device=DEVICE)
+    loaded = driver_load_weights(model, checkpoint_dir, device=DEVICE)
+    return model.eval(), loaded
+
+
+def test_weight_loading_roundtrip_and_forward(tmp_path):
+    from safetensors.torch import save_file
+
+    torch.manual_seed(0)
+    cfg = KimiK2Config.reduced()
+    ref = _build_reference(cfg)
+    assert not isinstance(ref.model.layers[0].mlp, KimiSparseMoeBlock)
+    assert isinstance(ref.model.layers[1].mlp, KimiSparseMoeBlock)
+
+    save_file(_hf_checkpoint(ref, cfg), str(tmp_path / "model.safetensors"))
+    model, loaded = _build_loaded(cfg, tmp_path)
+
+    all_params = set(dict(model.named_parameters()).keys())
+    assert loaded == all_params, (
+        f"unloaded: {all_params - loaded}; spurious: {loaded - all_params}")
+
+    ref_sd = dict(ref.named_parameters())
+    for name, param in model.named_parameters():
+        assert torch.equal(param, ref_sd[name]), f"mismatch at {name}"
+
+    bias = model.model.layers[1].mlp.gate.e_score_correction_bias
+    assert bias.dtype == torch.float32
+
+    # Derived buffers must not survive meta -> to_empty as uninitialized memory.
+    buffer_names = {n for n, _ in model.named_buffers()}
+    assert buffer_names == set(), f"unexpected buffers survived the load path: {buffer_names}"
+
+    inter = cfg.intermediate_size
+    d_gup = model.model.layers[0].mlp.gate_up_proj.weight
+    r_gup = ref.model.layers[0].mlp.gate_up_proj.weight
+    assert torch.equal(d_gup[:inter], r_gup[:inter])   # gate half
+    assert torch.equal(d_gup[inter:], r_gup[inter:])   # up half
+    mi = cfg.moe_intermediate_size
+    l_gup = model.model.layers[1].mlp.experts.gate_up_proj
+    r_egup = ref.model.layers[1].mlp.experts.gate_up_proj
+    r_edwn = ref.model.layers[1].mlp.experts.down_proj
+    for e in range(cfg.n_routed_experts):
+        assert torch.equal(l_gup[e, :mi], r_egup[e, :mi])   # gate:e
+        assert torch.equal(l_gup[e, mi:], r_egup[e, mi:])   # up:e
+    assert torch.equal(model.model.layers[1].mlp.experts.down_proj, r_edwn)
+
+    T = 8
+    ids = torch.randint(0, cfg.vocab_size, (T,), device=DEVICE)
+    pos = torch.arange(T, device=DEVICE)
+    with torch.no_grad():
+        got = model(ids, _MockMLACache(cfg.padded_head_dim), pos)
+        expected = ref(ids, _MockMLACache(cfg.padded_head_dim), pos)
+    assert got.shape == (T, cfg.vocab_size)
+    torch.testing.assert_close(got, expected, rtol=1e-3, atol=1e-3)
diff --git a/test/integration/test_marlin_kernels.py b/test/integration/test_marlin_kernels.py
new file mode 100644
index 000000000..764ac1b9b
--- /dev/null
+++ b/test/integration/test_marlin_kernels.py
@@ -0,0 +1,114 @@
+"""GPU golden tests for the vendored Marlin W4A16 kernels."""
+import pytest
+import torch
+
+from mstar.model.components.quantization import W4A16Data
+from mstar.model.kimi_k2_7._testing import fake_quantize_weight
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 0),
+    reason="Marlin W4A16 kernels need a CUDA GPU with sm80+ (Ampere/Hopper)",
+)
+
+DEVICE = "cuda"
+GROUP_SIZE = 32
+PACK_FACTOR = 8
+
+
+def _quantize_stack(weight):
+    E, N, K = weight.shape
+    packed = torch.empty((E, N, K // PACK_FACTOR), dtype=torch.int32, device=DEVICE)
+    scale = torch.empty((E, N, K // GROUP_SIZE), dtype=torch.bfloat16, device=DEVICE)
+    deq = torch.empty((E, N, K), dtype=torch.bfloat16, device=DEVICE)
+    for e in range(E):
+        p, s, d = fake_quantize_weight(
+            weight[e], num_bits=4, group_size=GROUP_SIZE, symmetric=True,
+            scale_dtype=torch.bfloat16,
+        )
+        packed[e], scale[e], deq[e] = p.to(DEVICE), s.to(DEVICE), d.to(DEVICE)
+    return packed, scale, deq
+
+
+def _random_topk(num_tokens, E, top_k):
+    logits = torch.randn(num_tokens, E, device=DEVICE)
+    weights, ids = torch.topk(logits.softmax(-1), top_k, dim=-1)
+    weights = weights / weights.sum(-1, keepdim=True)
+    return weights.to(torch.bfloat16), ids
+
+
+def _rel_l2(a, b):
+    return ((a - b).float().norm() / b.float().norm()).item()
+
+
+def test_marlin_builds_and_registers():
+    from mstar.utils.marlin import is_marlin_available
+
+    assert is_marlin_available(), "the vendored Marlin CUDA extension failed to build"
+    assert hasattr(torch.ops._mstar_marlin_C, "gptq_marlin_repack")
+    assert hasattr(torch.ops._mstar_marlin_C, "moe_wna16_marlin_gemm")
+
+
+def test_repack_shape_and_deterministic():
+    from mstar.utils.marlin import is_marlin_available, ops
+
+    assert is_marlin_available()
+    size_k, size_n = 256, 128  # k%16==0, n%64==0
+    b = torch.randint(
+        -(2**31), 2**31 - 1, (size_k // PACK_FACTOR, size_n), dtype=torch.int32, device=DEVICE
+    )
+    out = ops.gptq_marlin_repack(b, size_k, size_n, num_bits=4)
+    assert out.shape == (size_k // 16, size_n * 16 // PACK_FACTOR)
+    assert torch.equal(out, ops.gptq_marlin_repack(b, size_k, size_n, num_bits=4))
+
+
+@pytest.mark.parametrize("num_tokens", [8, 3, 1])  # M > E, M <= E, single-token decode
+def test_marlin_moe_matches_bf16_and_triton(num_tokens):
+    from mstar.model.components.quantization import MarlinMoEMethod
+    from mstar.utils.fused_moe.runner import fused_experts
+
+    torch.manual_seed(0)
+    # Marlin-legal shapes: hidden % 128, moe_inter % 128, 2*inter % 64, hidden % 64.
+    E, H, I, top_k = 4, 256, 256, 2
+    w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w1_packed, w1_scale, w1_deq = _quantize_stack(w1)
+    w2_packed, w2_scale, w2_deq = _quantize_stack(w2)
+
+    x = (torch.randn(num_tokens, H, device=DEVICE) * 0.5).to(torch.bfloat16)
+    topk_weights, topk_ids = _random_topk(num_tokens, E, top_k)
+
+    method = MarlinMoEMethod(num_bits=4, group_size=GROUP_SIZE)
+    method.prepare(w1_packed, w1_scale, w2_packed, w2_scale, torch.device(DEVICE))
+    out_marlin = method.apply(x, topk_weights, topk_ids)
+
+    out_bf16 = fused_experts(x, w1_deq, w2_deq, topk_weights, topk_ids)  # ground truth
+    out_triton = fused_experts(  # same packed nibbles, Triton W4A16 kernel
+        x, w1_packed, w2_packed, topk_weights, topk_ids,
+        quant=W4A16Data(w1_scale=w1_scale, w2_scale=w2_scale, group_size=GROUP_SIZE),
+    )
+
+    assert out_marlin.shape == (num_tokens, H) and out_marlin.dtype == torch.bfloat16
+    cos = torch.nn.functional.cosine_similarity(
+        out_marlin.flatten().float(), out_bf16.flatten().float(), dim=0
+    ).item()
+    assert cos > 0.999, f"cosine vs bf16 too low: {cos}"
+    assert _rel_l2(out_marlin, out_bf16) < 0.02, "relative-L2 vs bf16 too high"
+    assert _rel_l2(out_marlin, out_triton) < 0.02, "relative-L2 vs Triton W4A16 too high"
+
+
+def test_marlin_moe_reduce_results_false_shape():
+    from mstar.model.components.quantization import MarlinMoEMethod
+
+    torch.manual_seed(1)
+    E, H, I, top_k, num_tokens = 4, 256, 256, 2, 6
+    w1 = (torch.randn(E, 2 * I, H, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w2 = (torch.randn(E, H, I, device=DEVICE) * 0.3).to(torch.bfloat16)
+    w1_packed, w1_scale, _ = _quantize_stack(w1)
+    w2_packed, w2_scale, _ = _quantize_stack(w2)
+    x = (torch.randn(num_tokens, H, device=DEVICE) * 0.5).to(torch.bfloat16)
+    topk_weights, topk_ids = _random_topk(num_tokens, E, top_k)
+
+    method = MarlinMoEMethod(num_bits=4, group_size=GROUP_SIZE)
+    method.prepare(w1_packed, w1_scale, w2_packed, w2_scale, torch.device(DEVICE))
+    got = method.apply(x, topk_weights, topk_ids, reduce_results=False)
+    assert got.shape == (num_tokens, top_k, H)
diff --git a/test/modular/test_cache_manager_backends.py b/test/modular/test_cache_manager_backends.py
index faa39f093..99d962030 100644
--- a/test/modular/test_cache_manager_backends.py
+++ b/test/modular/test_cache_manager_backends.py
@@ -1,8 +1,13 @@
 """KV-cache attention-backend selection: ``KVCacheConfig.attention_backend``
 names the ``BatchedCacheManager`` subclass ``create_cache_manager``
-instantiates, unknown names are rejected, and the base class is abstract."""
+instantiates, unknown names are rejected, and the base class is abstract.
+
+Also covers the ``mla_absorb`` kernel predicate and the CUDA-graph capture
+decision it drives — the two must agree, or capture builds a paged wrapper that
+cannot read the 4-D latent cache."""
 
 import pytest
+import torch
 
 from mstar.engine.cache_manager import (
     ATTENTION_BACKENDS,
@@ -10,7 +15,9 @@
     DenseGenCacheManager,
     FlashInferCacheManager,
     create_cache_manager,
+    mla_kernel_available_for,
 )
+from mstar.engine.cuda_graph_runner import mla_absorb_capture_blocked
 from mstar.engine.kv_store import KVCacheConfig
 
 
@@ -77,4 +84,80 @@ def test_base_class_is_abstract():
 
 
 def test_registry_names():
-    assert set(ATTENTION_BACKENDS) == {"flashinfer", "dense_gen"}
+    assert set(ATTENTION_BACKENDS) == {"flashinfer", "dense_gen", "mla_absorb"}
+
+
+# --- mla_absorb: kernel predicate + capture decision ----------------------
+
+# Real Kimi latent dims; flashinfer's MLA wrapper is hard-locked to these.
+_REAL_CKV, _REAL_KPE = 512, 64
+_CUDA = torch.device("cuda:0")  # constructible without a GPU present
+
+
+def _mla_cfg(mla_ckv_dim: int | None, head_dim: int) -> KVCacheConfig:
+    return KVCacheConfig(
+        num_layers=2, num_kv_heads=1, head_dim=head_dim, max_seq_len=64,
+        attention_backend="mla_absorb", mla_ckv_dim=mla_ckv_dim,
+    )
+
+
+@pytest.fixture
+def sm90(monkeypatch):
+    """Pretend we are on a Hopper GPU, without needing one."""
+    import mstar.engine.cache_manager as cm_mod
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda d: (9, 0))
+    # _mla_kernel_available is functools.cache'd and imports flashinfer; stub the
+    # module attribute so the predicate's own logic is what's under test.
+    monkeypatch.setattr(
+        cm_mod, "_mla_kernel_available",
+        lambda ckv, kpe, sm: (ckv, kpe, sm) == (_REAL_CKV, _REAL_KPE, 9),
+    )
+
+
+def test_kernel_predicate_true_only_for_real_dims_on_sm90(sm90):
+    cfg = _mla_cfg(_REAL_CKV, _REAL_CKV + _REAL_KPE)
+    assert mla_kernel_available_for(cfg, _CUDA) is True
+
+
+def test_kernel_predicate_false_for_reduced_dims(sm90):
+    # The reduced test config: right backend, wrong latent width.
+    assert mla_kernel_available_for(_mla_cfg(32, 40), _CUDA) is False
+
+
+def test_kernel_predicate_false_without_ckv_dim(sm90):
+    assert mla_kernel_available_for(_mla_cfg(None, 40), _CUDA) is False
+
+
+def test_kernel_predicate_is_total_on_cpu(sm90):
+    """Must return False, not raise: torch.cuda.get_device_capability rejects a
+    CPU device, and the absorbed-SDPA fallback has to stay reachable there."""
+    cfg = _mla_cfg(_REAL_CKV, _REAL_CKV + _REAL_KPE)
+    assert mla_kernel_available_for(cfg, torch.device("cpu")) is False
+
+
+def test_capture_not_blocked_for_other_backends(sm90):
+    for backend in ("flashinfer", "dense_gen"):
+        assert mla_absorb_capture_blocked(_cfg(backend), _CUDA) is None
+
+
+def test_capture_not_blocked_when_kernel_serves_the_dims(sm90):
+    cfg = _mla_cfg(_REAL_CKV, _REAL_CKV + _REAL_KPE)
+    assert mla_absorb_capture_blocked(cfg, _CUDA) is None
+
+
+@pytest.mark.parametrize(
+    "ckv, head_dim", [(32, 40), (None, 40), (_REAL_CKV, _REAL_CKV + 128)]
+)
+def test_capture_blocked_when_absorbed_falls_back_to_sdpa(sm90, ckv, head_dim):
+    """No wrapper is planned on the SDPA path, so a captured graph could only
+    hold a paged wrapper — capture must be cancelled, not attempted."""
+    reason = mla_absorb_capture_blocked(_mla_cfg(ckv, head_dim), _CUDA)
+    assert reason is not None
+    assert "mla_absorb" in reason and "eager" in reason
+
+
+def test_capture_blocked_on_cpu_device(sm90):
+    cfg = _mla_cfg(_REAL_CKV, _REAL_CKV + _REAL_KPE)
+    assert mla_absorb_capture_blocked(cfg, torch.device("cpu")) is not None
diff --git a/test/modular/test_compressed_tensors.py b/test/modular/test_compressed_tensors.py
new file mode 100644
index 000000000..1e9a7c098
--- /dev/null
+++ b/test/modular/test_compressed_tensors.py
@@ -0,0 +1,198 @@
+import pytest
+import torch
+
+from mstar.model.components.quantization import (
+    CompressedTensorsQuantConfig,
+    dequant_compressed_tensors_stream,
+    dequantize_weight,
+    pack_int32,
+    unpack_int32,
+)
+from mstar.model.kimi_k2_7._testing import fake_quantize_weight
+
+
+def test_pack_known_answer():
+    nibbles = torch.arange(8, dtype=torch.int64).reshape(1, 8)
+    packed = pack_int32(nibbles, num_bits=4)
+    assert packed.dtype == torch.int32
+    assert packed.shape == (1, 1)
+    assert packed.item() == 0x76543210
+
+    # Top nibble >= 8 makes the int32 container negative; unpack reads it unsigned.
+    top = torch.tensor([[0, 0, 0, 0, 0, 0, 0, 8]], dtype=torch.int64)
+    packed_top = pack_int32(top, num_bits=4)
+    assert packed_top.item() == -(2**31)  # 0x80000000 as signed int32
+    back = unpack_int32(packed_top, num_bits=4)
+    assert torch.equal(back, top)
+
+
+def test_pack_unpack_roundtrip():
+    torch.manual_seed(0)
+    nibbles = torch.randint(0, 16, (5, 32), dtype=torch.int64)  # in=32 -> packed 4
+    packed = pack_int32(nibbles, num_bits=4)
+    assert packed.shape == (5, 4)
+    assert torch.equal(unpack_int32(packed, num_bits=4), nibbles)
+
+def test_dequantize_symmetric_known_answer():
+    unsigned = torch.tensor([[8, 9, 7, 8, 10, 6, 8, 8]], dtype=torch.int64)
+    signed = torch.tensor([[0, 1, -1, 0, 2, -2, 0, 0]], dtype=torch.float32)
+    packed = pack_int32(unsigned, num_bits=4)
+    scale = torch.tensor([[2.0]])  # (out=1, groups=1)
+    got = dequantize_weight(
+        packed, scale, num_bits=4, group_size=8, symmetric=True,
+        out_dtype=torch.float32,
+    )
+    assert torch.equal(got, signed * 2.0)
+
+
+def test_dequantize_two_groups_broadcast():
+    unsigned = torch.full((1, 16), 8, dtype=torch.int64)
+    unsigned[0, 0] = 9   # +1 in group 0
+    unsigned[0, 8] = 9   # +1 in group 1
+    packed = pack_int32(unsigned, num_bits=4)
+    scale = torch.tensor([[3.0, 5.0]])  # group0=3, group1=5
+    got = dequantize_weight(
+        packed, scale, num_bits=4, group_size=8, symmetric=True,
+        out_dtype=torch.float32,
+    )
+    assert got[0, 0] == 3.0
+    assert got[0, 8] == 5.0
+    assert got[0, 1] == 0.0
+
+@pytest.mark.parametrize("group_size", [8, 16, -1])
+def test_fake_quantize_dequantize_exact(group_size):
+    torch.manual_seed(1)
+    w = torch.randn(12, 32) * 0.1
+    packed, scale, deq = fake_quantize_weight(
+        w, num_bits=4, group_size=group_size, symmetric=True,
+    )
+    # Both paths compute q*scale in fp32 then cast to bf16, so this is bit-exact.
+    got = dequantize_weight(
+        packed, scale, num_bits=4, group_size=group_size, symmetric=True,
+    )
+    assert got.dtype == torch.bfloat16
+    assert torch.equal(got, deq)
+    assert (got.float() - w).abs().max() < 0.05
+
+
+def _quant_components(base, w, cfg):
+    # Match checkpoint bf16 scales so stream reconstruction is bit-exact.
+    packed, scale, deq = fake_quantize_weight(
+        w, num_bits=cfg.num_bits, group_size=cfg.group_size,
+        symmetric=cfg.symmetric, scale_dtype=torch.bfloat16,
+    )
+    return {
+        f"{base}.weight_packed": packed,
+        f"{base}.weight_scale": scale,
+        f"{base}.weight_shape": torch.tensor(list(w.shape), dtype=torch.int64),
+    }, deq
+
+
+def test_stream_dequantizes_and_passes_through():
+    torch.manual_seed(2)
+    cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=16, symmetric=True)
+    w_a = torch.randn(8, 32) * 0.1
+    w_b = torch.randn(4, 16) * 0.1
+    comp_a, deq_a = _quant_components("layer.0.a_proj", w_a, cfg)
+    comp_b, deq_b = _quant_components("layer.0.b_proj", w_b, cfg)
+    norm = torch.randn(8)
+
+    # Reassembly is by base name, independent of stream order.
+    stream = [
+        ("layer.0.a_proj.weight_scale", comp_a["layer.0.a_proj.weight_scale"]),
+        ("layer.0.norm.weight", norm),
+        ("layer.0.b_proj.weight_packed", comp_b["layer.0.b_proj.weight_packed"]),
+        ("layer.0.a_proj.weight_shape", comp_a["layer.0.a_proj.weight_shape"]),
+        ("layer.0.a_proj.weight_packed", comp_a["layer.0.a_proj.weight_packed"]),
+        ("layer.0.b_proj.weight_scale", comp_b["layer.0.b_proj.weight_scale"]),
+    ]
+    out = dict(dequant_compressed_tensors_stream(iter(stream), cfg))
+
+    assert set(out) == {"layer.0.a_proj.weight", "layer.0.b_proj.weight", "layer.0.norm.weight"}
+    assert torch.equal(out["layer.0.a_proj.weight"], deq_a)
+    assert torch.equal(out["layer.0.b_proj.weight"], deq_b)
+    assert torch.equal(out["layer.0.norm.weight"], norm)
+
+
+def test_stream_incomplete_group_raises():
+    cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=16, symmetric=True)
+    w = torch.randn(4, 16) * 0.1
+    comp, _ = _quant_components("x.proj", w, cfg)
+    stream = [("x.proj.weight_packed", comp["x.proj.weight_packed"])]
+    with pytest.raises(ValueError, match="incomplete"):
+        list(dequant_compressed_tensors_stream(iter(stream), cfg))
+
+def test_stream_keep_packed_passthrough():
+    cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=16, symmetric=True)
+    exp_base = "model.layers.1.mlp.experts.3.gate_proj"
+    mla_base = "model.layers.1.self_attn.o_proj"
+    comp_exp, _ = _quant_components(exp_base, torch.randn(8, 32) * 0.1, cfg)
+    comp_mla, deq_mla = _quant_components(mla_base, torch.randn(4, 16) * 0.1, cfg)
+
+    def keep_packed(base):
+        return ".experts.3.gate_proj" in base
+
+    # No dangling quant buffers should remain after mixed expert/MLA dequant.
+    stream = [
+        (f"{exp_base}.weight_packed", comp_exp[f"{exp_base}.weight_packed"]),
+        (f"{exp_base}.weight_scale", comp_exp[f"{exp_base}.weight_scale"]),
+        (f"{exp_base}.weight_shape", comp_exp[f"{exp_base}.weight_shape"]),
+        (f"{mla_base}.weight_packed", comp_mla[f"{mla_base}.weight_packed"]),
+        (f"{mla_base}.weight_scale", comp_mla[f"{mla_base}.weight_scale"]),
+    ]
+    out = dict(dequant_compressed_tensors_stream(iter(stream), cfg, keep_packed=keep_packed))
+
+    assert out[f"{exp_base}.weight_packed"].dtype == torch.int32
+    assert f"{exp_base}.weight_scale" in out
+    assert f"{exp_base}.weight_shape" in out
+    assert f"{exp_base}.weight" not in out
+    assert torch.equal(out[f"{mla_base}.weight"], deq_mla)
+    assert f"{mla_base}.weight_packed" not in out
+
+def test_grouped_gemm_reference_math_and_top_nibble():
+    torch.manual_seed(7)
+    N, K, gs = 6, 32, 16  # two groups along the packed K axis
+    w = torch.randn(N, K) * 0.5
+    packed, scale, deq = fake_quantize_weight(
+        w, num_bits=4, group_size=gs, symmetric=True, scale_dtype=torch.bfloat16,
+    )
+    assert packed.shape == (N, K // 8)  # packed along the last (input/K) axis
+    assert (packed < 0).any(), "no negative container — top-nibble sign path untested"
+
+    nibbles = unpack_int32(packed, num_bits=4).to(torch.float32)  # (N, K) unsigned
+    scale_bc = scale.to(torch.float32).repeat_interleave(gs, dim=-1)  # (N, K)
+    manual_deq = ((nibbles - 8.0) * scale_bc).to(torch.bfloat16)
+    assert torch.equal(manual_deq, deq)  # kernel math == dequantize_weight
+
+    x = torch.randn(4, K)
+    y_manual = torch.einsum("tk,nk->tn", x, manual_deq.float())
+    y_deq = torch.einsum("tk,nk->tn", x, deq.float())
+    assert torch.equal(y_manual, y_deq)
+
+def test_quant_config_from_hf_dict():
+    raw = {
+        "format": "pack-quantized",
+        "quant_method": "compressed-tensors",
+        "ignore": ["lm_head", "re:.*gate$"],
+        "config_groups": {
+            "group_0": {
+                "weights": {
+                    "num_bits": 4,
+                    "group_size": 32,
+                    "symmetric": True,
+                    "strategy": "group",
+                    "type": "int",
+                },
+                "targets": ["Linear"],
+            }
+        },
+    }
+    cfg = CompressedTensorsQuantConfig.from_hf_config_dict(raw)
+    assert cfg is not None
+    assert cfg.num_bits == 4
+    assert cfg.group_size == 32
+    assert cfg.symmetric is True
+    assert cfg.pack_factor == 8
+    assert cfg.ignore == ("lm_head", "re:.*gate$")
+    assert CompressedTensorsQuantConfig.from_hf_config_dict(None) is None
+    assert CompressedTensorsQuantConfig.from_hf_config_dict({}) is None
diff --git a/test/modular/test_kimi_k27_code_wiring.py b/test/modular/test_kimi_k27_code_wiring.py
new file mode 100644
index 000000000..684b33085
--- /dev/null
+++ b/test/modular/test_kimi_k27_code_wiring.py
@@ -0,0 +1,172 @@
+import json
+
+import torch
+
+from mstar.model.kimi_k2_7.components.causal_lm import KimiForCausalLM
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.kimi_k2_7.kimi_model import KimiK2Model
+from mstar.model.kimi_k2_7.weight_loader import (
+    build_kimi_stacked_params,
+    kimi_name_remapper,
+)
+from mstar.model.loader.base import _apply_stacked
+
+
+def test_k27_code_config_full_dims_packed_and_beta_fast():
+    cfg = KimiK2Config.k27_code()
+
+    assert cfg.moe_in_kernel_dequant is True
+    assert cfg.quantization_config is None
+
+    assert cfg.rope_scaling["beta_fast"] == 32.0
+    assert cfg.rope_scaling["factor"] == 64.0
+    assert cfg.rope_scaling["rope_type"] == "deepseek_yarn"
+
+    assert cfg.num_hidden_layers == 61
+    assert cfg.n_routed_experts == 384
+    assert cfg.hidden_size == 7168
+    assert cfg.q_lora_rank == 1536
+    assert cfg.kv_lora_rank == 512
+    assert cfg.moe_intermediate_size == 2048
+    assert cfg.routed_scaling_factor == 2.827
+    assert cfg.qk_nope_head_dim == 128
+    assert cfg.qk_rope_head_dim == 64
+    assert cfg.v_head_dim == 128
+
+    base = KimiK2Config()
+    assert cfg.num_hidden_layers == base.num_hidden_layers
+    assert cfg.rope_scaling == base.rope_scaling  # NO beta_fast override
+    assert base.moe_in_kernel_dequant is False and cfg.moe_in_kernel_dequant is True
+
+def _route(name, stacked):
+    mapped = kimi_name_remapper(name)
+    if mapped is None:
+        return None, None
+    return _apply_stacked(mapped, stacked)
+
+
+def test_remapper_language_model_prefix_and_packed_experts():
+    cfg = KimiK2Config.reduced_quantized_inkernel()
+    with torch.device("meta"):
+        model = KimiForCausalLM(cfg)
+    params = set(dict(model.named_parameters()).keys())
+    stacked = build_kimi_stacked_params(cfg.n_routed_experts, packed_experts=True)
+
+    assert kimi_name_remapper(
+        "language_model.model.layers.0.self_attn.q_a_proj.weight"
+    ) == "model.layers.0.self_attn.q_a_proj.weight"
+    assert (
+        kimi_name_remapper("language_model.model.embed_tokens.weight")
+        == "model.embed_tokens.weight"
+    )
+    assert kimi_name_remapper("language_model.lm_head.weight") == "lm_head.weight"
+    for landed in (
+        "model.layers.0.self_attn.q_a_proj.weight",
+        "model.embed_tokens.weight",
+        "lm_head.weight",
+    ):
+        assert landed in params
+
+    shared = kimi_name_remapper(
+        "language_model.model.layers.1.mlp.shared_experts.down_proj.weight"
+    )
+    assert shared == "model.layers.1.mlp.shared_expert.down_proj.weight"
+    assert shared in params
+
+    gate_p, gate_sid = _route(
+        "language_model.model.layers.1.mlp.experts.0.gate_proj.weight_packed", stacked
+    )
+    assert gate_p == "model.layers.1.mlp.experts.gate_up_proj_packed"
+    assert gate_sid == "gate:0"
+    assert gate_p in params
+
+    scale_p, scale_sid = _route(
+        "language_model.model.layers.1.mlp.experts.0.gate_proj.weight_scale", stacked
+    )
+    assert scale_p == "model.layers.1.mlp.experts.gate_up_proj_scale"
+    assert scale_sid == "gate:0"
+    assert scale_p in params
+
+    down_p, down_sid = _route(
+        "language_model.model.layers.1.mlp.experts.0.down_proj.weight_packed", stacked
+    )
+    assert down_p == "model.layers.1.mlp.experts.down_proj_packed"
+    assert down_sid == "down:0"
+    assert down_p in params
+
+    for vkey in (
+        "vision_tower.encoder.blocks.0.wqkv.weight",
+        "mm_projector.proj.0.weight",
+    ):
+        assert kimi_name_remapper(vkey) == vkey  # identity — no surgery
+        target, _ = _route(vkey, stacked)
+        assert target not in params  # dropped
+
+    ws_target, _ = _route(
+        "language_model.model.layers.1.mlp.experts.0.gate_proj.weight_shape", stacked
+    )
+    assert ws_target not in params
+
+    assert (
+        kimi_name_remapper("model.layers.0.self_attn.q_a_proj.weight")
+        == "model.layers.0.self_attn.q_a_proj.weight"
+    )
+
+_QUANT_BLOCK = {
+    "format": "pack-quantized",
+    "quant_method": "compressed-tensors",
+    "ignore": ["lm_head", "re:.*self_attn.*", "re:.*shared_experts.*"],
+    "config_groups": {
+        "group_0": {
+            "weights": {
+                "num_bits": 4,
+                "group_size": 32,
+                "symmetric": True,
+                "strategy": "group",
+                "type": "int",
+            },
+            "targets": ["Linear"],
+        }
+    },
+}
+
+
+def _make_model_with_config_json(tmp_dir, config_dict):
+    (tmp_dir / "config.json").write_text(json.dumps(config_dict))
+    model = object.__new__(KimiK2Model)
+    model.config = KimiK2Config()
+    return model
+
+
+def test_nested_quant_config_read(tmp_path):
+    d = tmp_path / "nested"
+    d.mkdir()
+    model = _make_model_with_config_json(
+        d, {"text_config": {"num_hidden_layers": 61, "quantization_config": _QUANT_BLOCK}}
+    )
+    model._maybe_apply_checkpoint_quant_config(str(d))
+    qc = model.config.quantization_config
+    assert qc is not None
+    assert qc.num_bits == 4
+    assert qc.group_size == 32
+    assert qc.symmetric is True
+    assert qc.quant_format == "pack-quantized"
+
+
+def test_flat_quant_config_read_backward_compat(tmp_path):
+    d = tmp_path / "flat"
+    d.mkdir()
+    model = _make_model_with_config_json(d, {"quantization_config": _QUANT_BLOCK})
+    model._maybe_apply_checkpoint_quant_config(str(d))
+    qc = model.config.quantization_config
+    assert qc is not None
+    assert qc.num_bits == 4
+    assert qc.group_size == 32
+
+
+def test_plain_bf16_config_stays_none(tmp_path):
+    d = tmp_path / "bf16"
+    d.mkdir()
+    model = _make_model_with_config_json(d, {"text_config": {"num_hidden_layers": 61}})
+    model._maybe_apply_checkpoint_quant_config(str(d))
+    assert model.config.quantization_config is None
diff --git a/test/modular/test_kimi_mla_absorb.py b/test/modular/test_kimi_mla_absorb.py
new file mode 100644
index 000000000..3b93aacaf
--- /dev/null
+++ b/test/modular/test_kimi_mla_absorb.py
@@ -0,0 +1,174 @@
+import torch
+import torch.nn.functional as F
+
+from mstar.model.kimi_k2_7.components.attention import KimiMLAAttention
+from mstar.model.kimi_k2_7.components.rope import (
+    _yarn_find_correction_range,
+    _yarn_linear_ramp_mask,
+    rotate_gptj,
+    yarn_get_mscale,
+)
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.kimi_k2_7.kimi_model import KimiK2Model
+
+
+def _ref_rmsnorm(x, weight, eps):
+    x32 = x.float()
+    x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps)
+    return weight * x32.to(x.dtype)
+
+
+def _ref_yarn_rope(pos, q_pe, k_pe, cfg):
+    r = cfg.rope_scaling
+    rotary_dim, base, factor = cfg.qk_rope_head_dim, cfg.rope_theta, r["factor"]
+    max_pos = r["original_max_position_embeddings"]
+    beta_fast, beta_slow = r.get("beta_fast", 32), r.get("beta_slow", 1)
+    mscale, mscale_all_dim = r.get("mscale", 1.0), r.get("mscale_all_dim", 0.0)
+    pos_freqs = base ** (torch.arange(0, rotary_dim, 2).float() / rotary_dim)
+    ext, interp = 1.0 / pos_freqs, 1.0 / (factor * pos_freqs)
+    low, high = _yarn_find_correction_range(beta_fast, beta_slow, rotary_dim, base, max_pos)
+    mask = 1 - _yarn_linear_ramp_mask(low, high, rotary_dim // 2, torch.float)
+    inv_freq = interp * (1 - mask) + ext * mask
+    amp = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim)
+    freqs = torch.outer(pos.float(), inv_freq)
+    cos = (freqs.cos() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    sin = (freqs.sin() * amp).repeat_interleave(2, -1).unsqueeze(-2)
+    qr = q_pe.float() * cos + rotate_gptj(q_pe.float()) * sin
+    kr = k_pe.float() * cos + rotate_gptj(k_pe.float()) * sin
+    return qr.to(q_pe.dtype), kr.to(k_pe.dtype)
+
+
+def _sdpa_causal(q, k, v, scale):
+    qt, kt, vt = (t.transpose(0, 1).float() for t in (q, k, v))  # (H,T,D)
+    t = q.shape[0]
+    causal = torch.triu(torch.full((t, t), float("-inf")), diagonal=1)
+    attn = (torch.einsum("hqd,hkd->hqk", qt, kt) * scale + causal).softmax(-1)
+    return torch.einsum("hqk,hkd->hqd", attn, vt).transpose(0, 1).to(q.dtype)
+
+
+def _q_and_latent(attn, cfg, h, pos):
+    t, heads = h.shape[0], attn.num_heads
+    d_nope, d_rope, latent = cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.kv_lora_rank
+    eps = cfg.rms_norm_eps
+    q = _ref_rmsnorm(F.linear(h, attn.q_a_proj.weight), attn.q_a_layernorm.weight, eps)
+    q = F.linear(q, attn.q_b_proj.weight).view(t, heads, cfg.qk_head_dim)
+    q_nope, q_pe = q.split([d_nope, d_rope], dim=-1)
+    lat = F.linear(h, attn.kv_a_proj_with_mqa.weight)
+    kv_a, k_pe = lat.split([latent, d_rope], dim=-1)
+    kv_c = _ref_rmsnorm(kv_a, attn.kv_a_layernorm.weight, eps)  # (T, L)
+    k_pe = k_pe.view(t, 1, d_rope)
+    q_pe, k_pe = _ref_yarn_rope(pos, q_pe, k_pe, cfg)
+    return q_nope, q_pe, kv_c, k_pe
+
+
+def _deepseek_scale(cfg):
+    r = cfg.rope_scaling
+    mscale = yarn_get_mscale(r["factor"], r.get("mscale_all_dim", 0.0))
+    return cfg.qk_head_dim ** -0.5 * mscale * mscale
+
+
+def _ref_deepseek_mla(attn, cfg, h, pos):
+    t, heads = h.shape[0], attn.num_heads
+    d_nope, d_rope, d_v = cfg.qk_nope_head_dim, cfg.qk_rope_head_dim, cfg.v_head_dim
+    q_nope, q_pe, kv_c, k_pe = _q_and_latent(attn, cfg, h, pos)
+    kv = F.linear(kv_c, attn.kv_b_proj.weight).view(t, heads, d_nope + d_v)
+    k_nope, v = kv.split([d_nope, d_v], dim=-1)
+    q = torch.cat([q_nope, q_pe], dim=-1)
+    k = torch.cat([k_nope, k_pe.expand(t, heads, d_rope)], dim=-1)
+    out = _sdpa_causal(q, k, v, _deepseek_scale(cfg)).reshape(t, heads * d_v)
+    return F.linear(out, attn.o_proj.weight)
+
+
+def _absorbed_mla(attn, cfg, h, pos):
+    t, heads = h.shape[0], attn.num_heads
+    d_rope, d_v, latent = cfg.qk_rope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank
+    q_nope, q_pe, kv_c, k_pe = _q_and_latent(attn, cfg, h, pos)
+    q_nope = torch.einsum("thd,hdl->thl", q_nope, attn.w_kc)  # (T,H,L)
+    query = torch.cat([q_nope, q_pe], dim=-1)                 # (T,H,L+Drope)
+    kv_c_h = kv_c.unsqueeze(1).expand(t, heads, latent)       # MQA: shared over heads
+    key = torch.cat([kv_c_h, k_pe.expand(t, heads, d_rope)], dim=-1)
+    attn_latent = _sdpa_causal(query, key, kv_c_h, _deepseek_scale(cfg))  # (T,H,L)
+    out = torch.einsum("thl,hdl->thd", attn_latent, attn.w_vc).reshape(t, heads * d_v)
+    return F.linear(out, attn.o_proj.weight)
+
+
+def _build_attention_cpu(seed=0):
+    torch.manual_seed(seed)
+    cfg = KimiK2Config.reduced()
+    cfg.mla_absorb = True
+    attn = KimiMLAAttention(cfg)  # CPU, float32
+    for lin in (attn.q_a_proj, attn.q_b_proj, attn.kv_a_proj_with_mqa,
+                attn.kv_b_proj, attn.o_proj):
+        lin.weight.data.normal_(0, 0.03)
+    for norm in (attn.q_a_layernorm, attn.kv_a_layernorm):
+        norm.weight.data.normal_(1.0, 0.02)
+    attn.process_weights_after_loading()  # split kv_b_proj -> w_kc / w_vc
+    return attn, cfg
+
+
+def test_absorb_reconstructs_kv_b_proj():
+    attn, cfg = _build_attention_cpu(seed=1)
+    heads, d_nope, d_v, latent = (
+        attn.num_heads, cfg.qk_nope_head_dim, cfg.v_head_dim, cfg.kv_lora_rank)
+    assert tuple(attn.w_kc.shape) == (heads, d_nope, latent)
+    assert tuple(attn.w_vc.shape) == (heads, d_v, latent)
+
+    kv_c = torch.randn(5, latent)
+    kv = F.linear(kv_c, attn.kv_b_proj.weight).view(5, heads, d_nope + d_v)
+    k_nope_ref, v_ref = kv.split([d_nope, d_v], dim=-1)
+
+    k_nope_abs = torch.einsum("tl,hdl->thd", kv_c, attn.w_kc)
+    v_abs = torch.einsum("tl,hdl->thd", kv_c, attn.w_vc)
+    torch.testing.assert_close(k_nope_abs, k_nope_ref, rtol=1e-5, atol=1e-5)
+    torch.testing.assert_close(v_abs, v_ref, rtol=1e-5, atol=1e-5)
+
+
+def test_fused_qkv_a_proj():
+    attn, cfg = _build_attention_cpu(seed=3)
+    expected = torch.cat(
+        [attn.q_a_proj.weight, attn.kv_a_proj_with_mqa.weight], dim=0)
+    assert tuple(attn.fused_qkv_a_proj_weight.shape) == (
+        cfg.q_lora_rank + cfg.kv_lora_rank + cfg.qk_rope_head_dim, cfg.hidden_size)
+    torch.testing.assert_close(
+        attn.fused_qkv_a_proj_weight, expected, rtol=0, atol=0)
+
+
+def test_absorbed_math_matches_deepseek():
+    attn, cfg = _build_attention_cpu(seed=2)
+    t = 7
+    h = torch.randn(t, cfg.hidden_size) * 0.1
+    pos = torch.arange(t)
+
+    absorbed = _absorbed_mla(attn, cfg, h, pos)
+    reference = _ref_deepseek_mla(attn, cfg, h, pos)
+
+    assert absorbed.shape == (t, cfg.hidden_size)
+    # Pure fp32 algebra; residual comes only from op ordering.
+    torch.testing.assert_close(absorbed, reference, rtol=1e-4, atol=1e-4)
+
+
+def _kv_cfg(mla_absorb):
+    m = object.__new__(KimiK2Model)          # skip __init__ (tokenizer/weights/GPU)
+    m.config = KimiK2Config.reduced()
+    m.config.mla_absorb = mla_absorb
+    return m.get_kv_cache_config()[0]
+
+
+def test_kv_cache_config_absorbed_shrinks_latent():
+    cfg = KimiK2Config.reduced()
+    kv = _kv_cfg(mla_absorb=True)
+    assert kv.num_kv_heads == 1
+    assert kv.head_dim == cfg.kv_lora_rank + cfg.qk_rope_head_dim  # 32 + 8 = 40
+    assert kv.num_qo_heads == cfg.num_attention_heads              # 4 (q still sharded)
+    assert kv.attention_backend == "mla_absorb"
+    naive_elems = 2 * cfg.num_attention_heads * cfg.padded_head_dim
+    absorbed_elems = kv.num_kv_heads * kv.head_dim
+    assert naive_elems == 512 and absorbed_elems == 40
+
+
+def test_kv_cache_config_flag_off_is_naive():
+    cfg = KimiK2Config.reduced()
+    kv = _kv_cfg(mla_absorb=False)
+    assert kv.num_kv_heads == cfg.num_attention_heads   # 4
+    assert kv.head_dim == cfg.padded_head_dim           # 64
+    assert kv.attention_backend == "flashinfer"         # default naive backend
diff --git a/test/modular/test_kimi_model.py b/test/modular/test_kimi_model.py
new file mode 100644
index 000000000..baef9eac9
--- /dev/null
+++ b/test/modular/test_kimi_model.py
@@ -0,0 +1,90 @@
+import sys
+
+sys.path.insert(0, ".")
+
+from mstar.conductor.request_info import CurrentForwardConductorMetadata
+from mstar.engine.base import EngineType
+from mstar.graph.base import Loop
+from mstar.model.kimi_k2_7.config import KimiK2Config
+from mstar.model.kimi_k2_7.kimi_model import KimiK2Model
+
+
+def _make_model() -> KimiK2Model:
+    model = object.__new__(KimiK2Model)
+    model.config = KimiK2Config.reduced()
+    model._submodule_cache = {}
+    return model
+
+
+def test_kimi_graph_walks_and_engine_types():
+    model = _make_model()
+
+    walks = model.get_graph_walk_graphs()
+    assert set(walks) == {"prefill", "decode"}
+    assert isinstance(walks["decode"], Loop)
+    assert walks["decode"].name == "decode_loop"
+
+    assert model.get_node_engine_types() == {"LLM": EngineType.KV_CACHE}
+
+
+def test_kimi_kv_cache_config_matches_reduced_mla_dims():
+    model = _make_model()
+    cfg = model.config
+
+    kv = model.get_kv_cache_config()
+    assert len(kv) == 1
+    (kv,) = kv
+
+    assert kv.num_layers == cfg.num_hidden_layers == 2
+    assert kv.num_kv_heads == cfg.num_attention_heads == 4
+    assert kv.num_qo_heads == cfg.num_attention_heads == 4
+    # FlashInfer-SM90 requires padded_head_dim, not raw qk_head_dim.
+    assert cfg.qk_head_dim == cfg.qk_nope_head_dim + cfg.qk_rope_head_dim == 24
+    assert kv.head_dim == cfg.padded_head_dim == 64
+    assert kv.max_seq_len == cfg.max_position_embeddings
+
+
+def test_kimi_prefill_transitions_to_decode():
+    model = _make_model()
+    metadata = CurrentForwardConductorMetadata(
+        input_modalities=["text"],
+        output_modalities=["text"],
+        graph_walk="prefill",
+        is_prefill=True,
+    )
+
+    result = model.get_partition_forward_pass_args(
+        partition_name="default",
+        partition_metadata=metadata,
+        persist_signals={"new_token": []},
+    )
+
+    assert result.full_metadata.graph_walk == "decode"
+    assert result.full_metadata.is_prefill is False
+    assert result.step_metadata["is_prefill"] is False
+    assert result.request_done is False
+
+
+def test_kimi_decode_completion_marks_done():
+    model = _make_model()
+    metadata = CurrentForwardConductorMetadata(
+        input_modalities=["text"],
+        output_modalities=["text"],
+        graph_walk="decode",
+        is_prefill=False,
+    )
+
+    result = model.get_partition_forward_pass_args(
+        partition_name="default",
+        partition_metadata=metadata,
+        persist_signals={},
+    )
+
+    assert result.request_done is True
+    assert result.full_metadata.kwargs["decode_finished"] is True
+
+
+def test_kimi_get_submodule_is_dummy_mode():
+    model = _make_model()
+    assert getattr(model, "model_path_hf", None) is None
+    assert model.get_submodule("LLM") is None
diff --git a/test/modular/test_quantization_types.py b/test/modular/test_quantization_types.py
new file mode 100644
index 000000000..a90840179
--- /dev/null
+++ b/test/modular/test_quantization_types.py
@@ -0,0 +1,136 @@
+"""Quantization descriptors: the tagged union that replaced ``fused_experts``'
+per-scheme kwargs, and the checkpoint-config -> kernel-data seam."""
+
+import dataclasses
+
+import pytest
+import torch
+
+from mstar.model.components.quantization import (
+    CompressedTensorsQuantConfig,
+    MarlinMoEMethod,
+    QuantizationData,
+    QuantizationType,
+    W4A16Data,
+)
+
+
+def _data(**overrides) -> W4A16Data:
+    kwargs = {
+        "w1_scale": torch.ones(2, 4, 2),
+        "w2_scale": torch.ones(2, 4, 2),
+        "group_size": 32,
+    }
+    kwargs.update(overrides)
+    return W4A16Data(**kwargs)
+
+
+# --- the union ------------------------------------------------------------
+
+
+def test_quantization_data_is_abstract():
+    with pytest.raises(TypeError, match="abstract"):
+        QuantizationData()
+
+
+def test_w4a16_tags_itself_and_derives_pack_factor():
+    d = _data()
+    assert d.quant_type is QuantizationType.W4A16
+    # Derived from the class's NUM_BITS, so it cannot disagree with the kernel's
+    # hardcoded 4-bit nibble extraction.
+    assert d.pack_factor == 8
+    assert "pack_factor" not in {f.name for f in dataclasses.fields(d)}
+    assert "NUM_BITS" not in {f.name for f in dataclasses.fields(d)}
+
+
+def test_zero_points_are_the_symmetry_discriminant():
+    assert _data().symmetric is True
+    assert _data(w1_zp=torch.zeros(2, 4, 2), w2_zp=torch.zeros(2, 4, 2)).symmetric is False
+
+
+def test_frozen_blocks_rebinding_but_not_tensor_mutation():
+    d = _data()
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        d.group_size = 64
+    # Documented limit: frozen freezes the binding, not the pointee.
+    d.w1_scale[0, 0, 0] = 7.0
+    assert d.w1_scale[0, 0, 0] == 7.0
+
+
+# --- the checkpoint-config -> kernel-data seam ----------------------------
+
+
+def test_config_maps_num_bits_to_scheme():
+    assert CompressedTensorsQuantConfig(num_bits=4).quant_type is QuantizationType.W4A16
+
+
+@pytest.mark.parametrize("num_bits", [2, 8, 16])
+def test_config_rejects_unimplemented_width_at_load(num_bits):
+    # The whole point: an INT8 checkpoint fails here, naming the checkpoint,
+    # instead of reaching a kernel that would mask its 8-bit values to nibbles.
+    cfg = CompressedTensorsQuantConfig(num_bits=num_bits)
+    assert cfg.quant_type is None  # the query answers; only the ensure raises
+    with pytest.raises(ValueError, match="mstar does not implement"):
+        cfg.ensure_kernel_support()
+
+
+def test_ensure_kernel_support_returns_the_scheme():
+    cfg = CompressedTensorsQuantConfig(num_bits=4)
+    assert cfg.ensure_kernel_support() is QuantizationType.W4A16
+
+
+def test_moe_quant_data_round_trips_group_size_and_scales():
+    cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=64)
+    w1s, w2s = torch.ones(2, 4, 2), torch.zeros(2, 4, 2)
+    data = cfg.moe_quant_data(w1s, w2s)
+
+    assert isinstance(data, W4A16Data)
+    assert data.quant_type is QuantizationType.W4A16
+    assert data.group_size == 64
+    assert data.pack_factor == cfg.pack_factor
+    assert data.w1_scale is w1s and data.w2_scale is w2s
+
+
+def test_symmetric_config_drops_zero_points():
+    sym = CompressedTensorsQuantConfig(num_bits=4, symmetric=True)
+    zp = torch.zeros(2, 4, 2)
+    data = sym.moe_quant_data(torch.ones(2, 4, 2), torch.ones(2, 4, 2), w1_zp=zp, w2_zp=zp)
+    assert data.w1_zp is None and data.w2_zp is None and data.symmetric
+
+    asym = CompressedTensorsQuantConfig(num_bits=4, symmetric=False)
+    data = asym.moe_quant_data(torch.ones(2, 4, 2), torch.ones(2, 4, 2), w1_zp=zp, w2_zp=zp)
+    assert data.w1_zp is zp and not data.symmetric
+
+
+def test_marlin_reads_num_bits_forward_from_the_config():
+    cfg = CompressedTensorsQuantConfig(num_bits=4, group_size=128)
+    method = MarlinMoEMethod.from_quant_config(cfg)
+    assert method.num_bits == 4
+    assert method.group_size == 128
+    assert method.quant_type is QuantizationType.W4A16
+
+    with pytest.raises(ValueError, match="mstar does not implement"):
+        MarlinMoEMethod.from_quant_config(CompressedTensorsQuantConfig(num_bits=8))
+
+
+# --- kernel dispatch ------------------------------------------------------
+
+
+def test_fused_experts_rejects_an_unhandled_scheme():
+    """A future QuantizationType with no branch must raise, not silently take
+    the bf16 path — that is what the sentinel-based dispatch could not do."""
+    pytest.importorskip("triton")
+    from mstar.utils.fused_moe.runner import _validate_expert_shapes
+
+    class _FutureData(QuantizationData):
+        @property
+        def quant_type(self):
+            return "int8-someday"
+
+    with pytest.raises(ValueError, match="no kernel for"):
+        _validate_expert_shapes(
+            hidden=128,
+            w1=torch.zeros(2, 8, 16, dtype=torch.int32),
+            w2=torch.zeros(2, 128, 1, dtype=torch.int32),
+            quant=_FutureData(),
+        )